Add new csvgenerator method from spz

hitier
1 parent 57278073
Showing 1 changed file with 47 additions and 3 deletions Show diff stats
web/run.py
@@ -11,8 +11,6 @@ import time
 import urllib.request as urllib_request
 import requests
 import re  # regex
-import numpy as np
-import sys
  
 from csv import writer as csv_writer, DictWriter as csv_dict_writer
 from math import sqrt, isnan
@@ -30,6 +28,9 @@ from jinja2 import Environment, FileSystemLoader
 from jinja2.utils import markupsafe
 from yaml import load as yaml_load, dump
  
+from speasy import amda
+import pandas as pd
+
 try:
     from yaml import CLoader as Loader
 except ImportError:
@@ -874,6 +875,44 @@ def generate_csv_contents(target_slug, input_slug, started_at, stopped_at):
     return si.getvalue()
  
  
+def generate_csv_contents_spz(target_slug, input_slug, started_at, stopped_at):
+    target_config = get_target_config(target_slug)
+    params_dict = target_config['models'][input_slug][0]['parameters']
+
+    list_df = []
+    for _name, _id in params_dict.items():
+        _df = amda.get_data(_id, started_at, stopped_at).to_dataframe()
+        if _name == 'xy_v':
+            _df = _df.rename(columns={_df.columns[0]: 'vrad', _df.columns[1]: 'vtan'})
+        elif _name == 'xy_hee':
+            _df = _df.drop(_df.columns[2], axis=1)
+            _df = _df.rename(columns={_df.columns[0]: 'xhee', _df.columns[1]: 'yhee'})
+        else:
+            _df = _df.rename(columns={_df.columns[0]: _name})
+
+        # _df = _df[~_df.index.duplicated()]
+
+        # resample to frequency, for later concatenation
+        _df = _df.resample('1H').mean()
+
+        # _df = _df.loc[_df.first_valid_index():_df.last_valid_index()]
+
+        list_df.append(_df)
+
+    from math import sqrt
+    final_df = pd.concat(list_df, axis=1)
+    # Is ther a vtot param ? else build it
+    if 'vtot' not in final_df.columns:
+        final_df['vtot'] = final_df.apply(lambda x: sqrt(x.vtan * x.vtan + x.vrad * x.vrad), axis=1)
+
+    cols_ordered = ['vrad', 'vtan', 'vtot', 'btan', 'brad', 'temp', 'pdyn', 'dens', 'atse', 'xhee', 'yhee']
+    final_df = final_df[cols_ordered]
+    final_df.index.name = 'time'
+    return final_df.to_csv(date_format='%Y-%m-%dT%H:%M:%S+00:00', float_format="%.5f",
+                           header=True,
+                           sep=",")
+
+
 def generate_csv_file_if_needed(target_slug, input_slug,
                                 started_at, stopped_at):
     filename = "%s_%s_%s_%s.csv" % (target_slug, input_slug,
@@ -892,11 +931,16 @@ def generate_csv_file_if_needed(target_slug, input_slug,
                     generate = False
                     break
  
+    if target_slug == 'mars':
+        csv_generator = generate_csv_contents_spz
+    else:
+        csv_generator = generate_csv_contents
+
     if generate:
         log.info("Generating CSV '%s'..." % local_csv_file)
         try:
             with open(local_csv_file, mode="w+") as f:
-                f.write(generate_csv_contents(
+                f.write(csv_generator(
                     target_slug=target_slug, input_slug=input_slug,
                     started_at=started_at, stopped_at=stopped_at
                 ))