Skip to content

Commit

Permalink
Add vertical profile code
Browse files Browse the repository at this point in the history
  • Loading branch information
@awarde96
awarde96 committed Aug 7, 2024
1 parent ed73aff commit dbb21a4
Showing 1 changed file with 118 additions and 64 deletions.
182 changes: 118 additions & 64 deletions covjsonkit/encoder/VerticalProfile.py
View file
Original file line number Diff line number Diff line change
@@ -1,13 +1,11 @@
import json

import pandas as pd

from .encoder import Encoder


class VerticalProfile(Encoder):
def __init__(self, type, domaintype):
super().__init__(type, domaintype)
self.covjson["domainType"] = "VerticalProfile"
self.covjson["coverages"] = []

def add_coverage(self, mars_metadata, coords, values):
new_coverage = {}
Expand All @@ -34,12 +32,13 @@ def add_domain(self, coverage, coords):

def add_range(self, coverage, values):
for parameter in self.parameters:
coverage["ranges"][parameter] = {}
coverage["ranges"][parameter]["type"] = "NdArray"
coverage["ranges"][parameter]["dataType"] = "float"
coverage["ranges"][parameter]["shape"] = [len(values[parameter])]
coverage["ranges"][parameter]["axisNames"] = ["z"]
coverage["ranges"][parameter]["values"] = values[parameter] # [values[parameter]]
param = self.convert_param_id_to_param(parameter)
coverage["ranges"][param] = {}
coverage["ranges"][param]["type"] = "NdArray"
coverage["ranges"][param]["dataType"] = "float"
coverage["ranges"][param]["shape"] = [len(values[parameter])]
coverage["ranges"][param]["axisNames"] = ["z"]
coverage["ranges"][param]["values"] = values[parameter] # [values[parameter]]

def add_mars_metadata(self, coverage, metadata):
coverage["mars:metadata"] = metadata
Expand Down Expand Up @@ -84,28 +83,33 @@ def from_xarray(self, dataset):
return self.covjson

def from_polytope(self, result):
ancestors = [val.get_ancestors() for val in result.leaves]
values = [val.result for val in result.leaves]

columns = []
df_dict = {}
# Create empty dataframe
for feature in ancestors[0]:
columns.append(str(feature).split("=")[0])
df_dict[str(feature).split("=")[0]] = []

# populate dataframe
for ancestor in ancestors:
for feature in ancestor:
df_dict[str(feature).split("=")[0]].append(str(feature).split("=")[1])
values = [val.result for val in result.leaves]
df_dict["values"] = values
df = pd.DataFrame(df_dict)

params = df["param"].unique()

for param in params:
self.add_parameter(param)
coords = {}
# coords['x'] = []
# coords['y'] = []
# coords['z'] = []
# coords['t'] = []
mars_metadata = {}
range_dict = {}
lat = 0
param = 0
# number = 0
step = 0
long = 0
levels = 0
dates = 0

self.func(
result,
lat,
long,
coords,
mars_metadata,
param,
range_dict,
step,
levels,
dates,
)

self.add_reference(
{
Expand All @@ -117,37 +121,87 @@ def from_polytope(self, result):
}
)

mars_metadata = {}
mars_metadata["class"] = df["class"].unique()[0]
mars_metadata["expver"] = df["expver"].unique()[0]
mars_metadata["levtype"] = df["levtype"].unique()[0]
mars_metadata["type"] = df["type"].unique()[0]
mars_metadata["date"] = df["date"].unique()[0]
mars_metadata["domain"] = df["domain"].unique()[0]
mars_metadata["stream"] = df["stream"].unique()[0]
for date in range_dict.keys():
for param in range_dict[date].keys():
# self.coord_length = len(range_dict[date][param])
self.add_parameter(param)
break

coords = {}
coords["x"] = list(df["latitude"].unique())
coords["y"] = list(df["longitude"].unique())
coords["z"] = list(df["level"].unique())
coords["t"] = list(df["date"].unique())

if "number" not in df.columns:
new_metadata = mars_metadata.copy()
range_dict = {}
for param in params:
df_param = df[df["param"] == param]
range_dict[param] = df_param["values"].values.tolist()
self.add_coverage(new_metadata, coords, range_dict)
for date in range_dict.keys():
self.add_coverage(mars_metadata, coords[date], range_dict[date])

# return json.loads(self.get_json())
return self.covjson

def func(
self,
tree,
lat,
long,
coords,
mars_metadata,
param,
range_dict,
step,
levels,
dates,
):
if len(tree.children) != 0:
# recurse while we are not a leaf
for c in tree.children:
if (
c.axis.name != "latitude"
and c.axis.name != "longitude"
and c.axis.name != "param"
and c.axis.name != "step"
and c.axis.name != "date"
and c.axis.name != "levelist"
):
mars_metadata[c.axis.name] = c.values[0]
if c.axis.name == "latitude":
lat = c.values[0]
if c.axis.name == "param":
param = c.values
for date in dates:
for para in param:
range_dict[date][para] = []
if c.axis.name == "date":
dates = [str(date) + "Z" for date in c.values]
for date in dates:
coords[date] = {}
range_dict[date] = {}
mars_metadata[c.axis.name] = str(c.values[0]) + "Z"
if c.axis.name == "levelist":
levels = c.values

self.func(
c,
lat,
long,
coords,
mars_metadata,
param,
range_dict,
step,
levels,
dates,
)
else:
for number in df["number"].unique():
new_metadata = mars_metadata.copy()
new_metadata["number"] = number
df_number = df[df["number"] == number]
range_dict = {}
for param in params:
df_param = df_number[df_number["param"] == param]
range_dict[param] = df_param["values"].values.tolist()
self.add_coverage(new_metadata, coords, range_dict)

return json.loads(self.get_json())
tree.values = [float(val) for val in tree.values]
tree.result = [float(val) for val in tree.result]
# num_intervals = int(len(tree.result)/len(number))
# para_intervals = int(num_intervals/len(param))
len_paras = len(param) * len(levels)
len_levels = len(param)

for date in dates:

coords[date]["x"] = [lat]
coords[date]["y"] = [long]
coords[date]["z"] = list(levels)
coords[date]["t"] = date

for i, date in enumerate(dates):
for j, level in enumerate(list(levels)):
for k, para in enumerate(param):
range_dict[date][para].append(tree.result[i * len_paras + j * len_levels + k])

0 comments on commit dbb21a4

Please sign in to comment.