Skip to content

Commit

Permalink
Merge pull request #54 from ecmwf/develop
Browse files Browse the repository at this point in the history 
Develop
  • Loading branch information
@awarde96
awarde96 authored Nov 13, 2024
2 parents c17e5a6 + a82622c commit cd6e0ee
Show file tree
Hide file tree
Showing 3 changed files with 155 additions and 140 deletions.
97 changes: 69 additions & 28 deletions covjsonkit/decoder/VerticalProfile.py
View file
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
Expand Down Expand Up @@ -53,39 +54,79 @@ def to_geopandas(self):
pass

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

# Get coordinates
coords = self.get_domains()
x = coords[0]["axes"]["x"]["values"]
y = coords[0]["axes"]["y"]["values"]
z = coords[0]["axes"]["z"]["values"]
steps = coords[0]["axes"]["t"]["values"]
steps = [step.replace("Z", "") for step in steps]
steps = pd.to_datetime(steps)
# steps = list(range(len(steps)))

num = []
datetime = []
steps = []
for coverage in self.covjson["coverages"]:
num.append(coverage["mars:metadata"]["number"])
datetime.append(coverage["mars:metadata"]["Forecast date"])
steps.append(coverage["mars:metadata"]["step"])

nums = list(set(num))
datetime = list(set(datetime))
steps = list(set(steps))

param_values = {}

for parameter in self.parameters:
param_values = [[[self.get_values()[parameter]]]]
for ind, value in enumerate(self.get_values()[parameter]):
coords = self.get_coordinates()[parameter]
x = [coords[ind][0][0]]
y = [coords[ind][0][1]]
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,
"t": t,
"number": num,
"z": z,
}
dataarray = xr.DataArray(
param_values,
dims=dims,
coords=param_coords,
name=parameter,
)
dataarray.attrs["type"] = self.get_parameter_metadata(parameter)["type"]
dataarray.attrs["units"] = self.get_parameter_metadata(parameter)["unit"]["symbol"]
dataarray.attrs["long_name"] = self.get_parameter_metadata(parameter)["observedProperty"]["id"]
dataarraydict[dataarray.attrs["long_name"]] = dataarray
param_values[parameter] = []
for i, num in enumerate(nums):
param_values[parameter].append([])
for j, date in enumerate(datetime):
param_values[parameter][i].append([])
for k, step in enumerate(steps):
param_values[parameter][i][j].append([])
for coverage in self.covjson["coverages"]:
if (
coverage["mars:metadata"]["number"] == num
and coverage["mars:metadata"]["Forecast date"] == date
and coverage["mars:metadata"]["step"] == step
):
param_values[parameter][i][j][k] = coverage["ranges"][parameter]["values"]

ds = xr.Dataset(dataarraydict)
for parameter in self.parameters:
param_coords = {
"x": x,
"y": y,
"number": nums,
"datetime": datetime,
"t": steps,
"z": z,
}

dataarray = xr.DataArray(
[[param_values[parameter]]],
dims=dims,
coords=param_coords,
name=parameter,
)

dataarray.attrs["type"] = self.get_parameter_metadata(parameter)["type"]
dataarray.attrs["units"] = self.get_parameter_metadata(parameter)["unit"]["symbol"]
dataarray.attrs["long_name"] = self.get_parameter_metadata(parameter)["observedProperty"]["id"]
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]
Expand Down
196 changes: 85 additions & 111 deletions covjsonkit/encoder/VerticalProfile.py
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,8 @@
import logging
import time
from datetime import datetime, timedelta

import pandas as pd

from .encoder import Encoder

Expand Down Expand Up @@ -85,27 +89,24 @@ def from_polytope(self, result):
coords = {}
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,
number,
)
fields = {}
fields["lat"] = 0
fields["param"] = 0
fields["number"] = [0]
fields["step"] = 0
fields["dates"] = []
fields["levels"] = [0]

start = time.time()
logging.debug("Tree walking starts at: %s", start) # noqa: E501
self.walk_tree(result, fields, coords, mars_metadata, range_dict)
end = time.time()
delta = end - start
logging.debug("Tree walking ends at: %s", end) # noqa: E501
logging.debug("Tree walking takes: %s", delta) # noqa: E501

start = time.time()
logging.debug("Coords creation: %s", start) # noqa: E501

self.add_reference(
{
Expand All @@ -117,99 +118,72 @@ def from_polytope(self, result):
}
)

logging.debug("The values returned from walking tree: %s", range_dict) # noqa: E501
logging.debug("The coordinates returned from walking tree: %s", coords) # noqa: E501
coordinates = {}

for date in range_dict.keys():
for num in range_dict[date].keys():
for param in range_dict[date][num].keys():
self.add_parameter(param)
break
break
levels = fields["levels"]
if fields["param"] == 0:
raise ValueError("No parameters were returned, date requested may be out of range")
for para in fields["param"]:
self.add_parameter(para)

logging.debug("The parameters added were: %s", self.parameters) # noqa: E501

for date in range_dict.keys():
for num in range_dict[date].keys():
mm = mars_metadata.copy()
mm["number"] = num
del mm["date"]
self.add_coverage(mm, coords[date], range_dict[date][num])
for date in fields["dates"]:
coordinates[date] = {}
for level in fields["levels"]:
for num in fields["number"]:
for para in fields["param"]:
for step in fields["step"]:
date_format = "%Y%m%dT%H%M%S"
new_date = pd.Timestamp(date).strftime(date_format)
start_time = datetime.strptime(new_date, date_format)
# add current date to list by converting it to iso format
stamp = start_time + timedelta(hours=int(step))
coordinates[date][step] = {
"x": [coords[date]["composite"][0][0]],
"y": [coords[date]["composite"][0][1]],
"z": list(levels),
}
coordinates[date][step]["t"] = [stamp.isoformat() + "Z"]
# coordinates[date]["t"].append(stamp.isoformat() + "Z")
break
break
break

# return json.loads(self.get_json())
return self.covjson
end = time.time()
delta = end - start
logging.debug("Coords creation: %s", end) # noqa: E501
logging.debug("Coords creation: %s", delta) # noqa: E501

# logging.debug("The values returned from walking tree: %s", range_dict) # noqa: E501
# logging.debug("The coordinates returned from walking tree: %s", coordinates) # noqa: E501

start = time.time()
logging.debug("Coverage creation: %s", start) # noqa: E501

for date in fields["dates"]:
for num in fields["number"]:
val_dict = {}
for step in fields["step"]:
val_dict[step] = {}
for para in fields["param"]:
val_dict[step][para] = []
for level in fields["levels"]:
key = (date, level, num, para, step)
# for k, v in range_dict.items():
# if k == key:
# val_dict[para].append(v[0])
val_dict[step][para].append(range_dict[key][0])
mm = mars_metadata.copy()
mm["number"] = num
mm["Forecast date"] = date
mm["step"] = step
# del mm["step"]
self.add_coverage(mm, coordinates[date][step], val_dict[step])

end = time.time()
delta = end - start
logging.debug("Coverage creation: %s", end) # noqa: E501
logging.debug("Coverage creation: %s", delta) # noqa: E501

def func(
self,
tree,
lat,
long,
coords,
mars_metadata,
param,
range_dict,
step,
levels,
dates,
number,
):
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 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:
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, number)
else:
tree.values = [float(val) for val in tree.values]
tree.result = [float(val) for val in tree.result]
# para_intervals = int(num_intervals/len(param))
try:
len(param)
except TypeError:
raise ValueError("No parameters were returned, date requested may be out of range")
len_paras = len(param)
len_levels = len(param)
len_nums = len_paras * len(levels)
for date in dates:

coords[date]["x"] = [lat]
coords[date]["y"] = [tree.values[0]]
coords[date]["z"] = list(levels)
coords[date]["t"] = date

for i, date in enumerate(dates):
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]
)
return self.covjson
2 changes: 1 addition & 1 deletion covjsonkit/version.py
View file
Original file line number Diff line number Diff line change
@@ -1 +1 @@
__version__ = "0.0.27"
__version__ = "0.0.28"

0 comments on commit cd6e0ee

Please sign in to comment.