Skip to content

Commit

Permalink
Vettical profile now works for ensemble data
Browse files Browse the repository at this point in the history
awarde96 committed Sep 11, 2024
1 parent 4885cd8 commit 621dbec
Showing 3 changed files with 39 additions and 34 deletions.
4 changes: 3 additions & 1 deletion covjsonkit/decoder/TimeSeries.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import pandas as pd
import xarray as xr

from .decoder import Decoder
@@ -67,7 +68,8 @@ def to_xarray(self):
y = coords[0]["axes"]["y"]["values"]
z = coords[0]["axes"]["z"]["values"]
steps = coords[0]["axes"]["t"]["values"]
steps = list(range(len(steps)))
steps = pd.to_datetime(steps)
# steps = list(range(len(steps)))

num = []
datetime = []
10 changes: 5 additions & 5 deletions covjsonkit/decoder/VerticalProfile.py
Original file line number Diff line number Diff line change
@@ -53,7 +53,7 @@ def to_geopandas(self):
pass

def to_xarray(self):
dims = ["x", "y","number", "t", "z"]
dims = ["x", "y", "t", "number", "z"]
dataarraydict = {}

for parameter in self.parameters:
@@ -62,15 +62,15 @@ def to_xarray(self):
coords = self.get_coordinates()[parameter]
x = [coords[ind][0][0]]
y = [coords[ind][0][1]]
t = [coord[0][4] for coord in coords]
num = [coords[ind][0][3]]
t = [coords[ind][0][4]]
num = [coord[0][3] for coord in coords]
coords_z = coords[ind]
z = [int(coord[2]) for coord in coords_z]
param_coords = {
"x": x,
"y": y,
"number": num,
"t": t,
"number": num,
"z": z,
}
dataarray = xr.DataArray(
@@ -85,7 +85,7 @@ def to_xarray(self):
dataarraydict[dataarray.attrs["long_name"]] = dataarray

ds = xr.Dataset(dataarraydict)

for mars_metadata in self.mars_metadata[0]:
if mars_metadata != "date" and mars_metadata != "step":
ds.attrs[mars_metadata] = self.mars_metadata[0][mars_metadata]
59 changes: 31 additions & 28 deletions covjsonkit/encoder/VerticalProfile.py
Original file line number Diff line number Diff line change
@@ -94,7 +94,7 @@ def from_polytope(self, result):
range_dict = {}
lat = 0
param = 0
# number = 0
number = 0
step = 0
long = 0
levels = 0
@@ -111,6 +111,7 @@ def from_polytope(self, result):
step,
levels,
dates,
number,
)

self.add_reference(
@@ -127,17 +128,20 @@ def from_polytope(self, result):
logging.debug("The coordinates returned from walking tree: %s", coords) # noqa: E501

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)
for num in range_dict[date].keys():
for param in range_dict[date][num].keys():
# self.coord_length = len(range_dict[date][param])
self.add_parameter(param)
break
break

for date in range_dict.keys():
mm = mars_metadata.copy()
# mm["number"] = num
# mm["Forecast date"] = date
del mm["date"]
self.add_coverage(mm, coords[date], range_dict[date])
for num in range_dict[date].keys():
mm = mars_metadata.copy()
mm["number"] = num
# mm["Forecast date"] = date
del mm["date"]
self.add_coverage(mm, coords[date], range_dict[date][num])

# return json.loads(self.get_json())
return self.covjson
@@ -154,6 +158,7 @@ def func(
step,
levels,
dates,
number,
):
if len(tree.children) != 0:
# recurse while we are not a leaf
@@ -172,8 +177,14 @@ def func(
if c.axis.name == "param":
param = c.values
for date in dates:
for para in param:
range_dict[date][para] = []
for num in number:
for para in param:
range_dict[date][num][para] = []
if c.axis.name == "number":
number = c.values
for date in dates:
for num in number:
range_dict[date][num] = {}
if c.axis.name == "date":
dates = [str(date) + "Z" for date in c.values]
for date in dates:
@@ -183,25 +194,14 @@ def func(
if c.axis.name == "levelist":
levels = c.values

self.func(
c,
lat,
long,
coords,
mars_metadata,
param,
range_dict,
step,
levels,
dates,
)
self.func(c, lat, long, coords, mars_metadata, param, range_dict, step, levels, dates, number)
else:
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_paras = len(param)
len_levels = len(param)
len_nums = len_paras * len(levels)

for date in dates:

@@ -211,6 +211,9 @@ def func(
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])
for j, num in enumerate(number):
for l, level in enumerate(list(levels)): # noqa: E741
for k, para in enumerate(param):
range_dict[date][num][para].append(
tree.result[i * len_paras + l * len_levels + j * len_nums + k]
)

0 comments on commit 621dbec

Please sign in to comment.