Skip to content

Commit

Permalink
Merge pull request #64 from ecmwf/feature/climate_dt_fixes
Browse files Browse the repository at this point in the history 
Change walk tree for climate dt
  • Loading branch information
@awarde96
awarde96 authored Feb 3, 2025
2 parents 2209ccf + d0ccc55 commit 66c829e
Show file tree
Hide file tree
Showing 4 changed files with 159 additions and 26 deletions.
62 changes: 38 additions & 24 deletions covjsonkit/encoder/TimeSeries.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -144,6 +144,10 @@ def from_polytope(self, result):
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
try:
int(step)
except ValueError:
step = step[0]
stamp = start_time + timedelta(hours=int(step))
coordinates[date]["t"].append(stamp.isoformat() + "Z")
break
Expand Down Expand Up @@ -197,10 +201,11 @@ def from_polytope_step(self, result):
fields["step"] = [0]
fields["dates"] = []
fields["levels"] = [0]
fields["times"] = []

start = time.time()
logging.debug("Tree walking starts at: %s", start) # noqa: E501
self.walk_tree(result, fields, coords, mars_metadata, range_dict)
self.walk_tree_step(result, fields, coords, mars_metadata, range_dict)
end = time.time()
delta = end - start
logging.debug("Tree walking ends at: %s", end) # noqa: E501
Expand All @@ -219,6 +224,12 @@ def from_polytope_step(self, result):
}
)

# print("Fields: ", fields)
# print("Coords: ", coords)
# print("Mars Metadata: ", mars_metadata)
# print("Range Dict: ", range_dict)
# print("Range Dict keys: ", range_dict.keys())

coordinates = {}

levels = fields["levels"]
Expand All @@ -240,12 +251,14 @@ def from_polytope_step(self, result):
for num in fields["number"]:
for para in fields["param"]:
for date in fields["dates"]:
# 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[fields["dates"][0]]["t"].append(date)
for times in fields["times"]:
# 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))
datetime = pd.Timestamp(date) + times
coordinates[fields["dates"][0]]["t"].append(str(datetime).split("+")[0] + "Z")
break
break
break
Expand All @@ -255,29 +268,30 @@ def from_polytope_step(self, result):
logging.debug("Coords creation: %s", end) # noqa: E501
logging.debug("Coords creation: %s", delta) # noqa: E501

# print(coordinates)

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

end = time.time()
delta = end - start
Expand Down
2 changes: 1 addition & 1 deletion covjsonkit/encoder/Wkt.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -89,7 +89,7 @@ def from_polytope(self, result):
fields["lat"] = 0
fields["param"] = 0
fields["number"] = [0]
fields["step"] = 0
fields["step"] = [0]
fields["dates"] = []
fields["levels"] = [0]

Expand Down
119 changes: 119 additions & 0 deletions covjsonkit/encoder/encoder.py
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
from abc import ABC, abstractmethod

import orjson
import pandas as pd
from covjson_pydantic.coverage import CoverageCollection
from covjson_pydantic.domain import DomainType

Expand Down Expand Up @@ -189,6 +190,124 @@ def append_composite_coords(dates, tree_values, lat, coords):
range_dict[key] = []
range_dict[key].extend(tree.result[start_index:end_index])

def walk_tree_step(self, tree, fields, coords, mars_metadata, range_dict):
def create_composite_key_step(date, level, num, para):
return (date, level, num, para)

def handle_non_leaf_node_step(child):
non_leaf_axes = ["latitude", "longitude", "param", "date", "time"]
if child.axis.name not in non_leaf_axes:
mars_metadata[child.axis.name] = child.values[0]

def handle_specific_axes_step(child):
if child.axis.name == "latitude":
return child.values[0]
if child.axis.name == "levelist":
return child.values
if child.axis.name == "param":
return child.values
if child.axis.name in ["date"]:
dates = [f"{date}Z" for date in child.values]
# mars_metadata["Forecast date"] = str(child.values[0])
# for date in dates:
# coords[date] = {}
# coords[date]["composite"] = []
# coords[date]["t"] = [date]
return dates
if child.axis.name == "number":
return child.values
if child.axis.name == "step":
return child.values
if child.axis.name == "time":
for date in fields["dates"]:
coords[date] = {}
coords[date]["composite"] = []
coords[date]["t"] = []
for time in child.values:
datetime = pd.Timestamp(date) + time
coords[date]["t"].append(str(datetime).split("+")[0] + "Z")
return child.values
return None

def calculate_index_bounds_step(level_len, num_len, para_len, step_len, l, i, j, k): # noqa: E741
start_index = int(l * level_len) + int(i * num_len) + int(j * para_len) + int(k * step_len)
end_index = start_index + int(step_len)
return start_index, end_index

def append_composite_coords_step(dates, tree_values, lat, coords):
# for date in dates:
for value in tree_values:
coords[dates]["composite"].append([lat, value])

if len(tree.children) != 0:
for child in tree.children:
handle_non_leaf_node_step(child)
result = handle_specific_axes_step(child)
if result is not None:
if child.axis.name == "latitude":
fields["lat"] = result
elif child.axis.name == "levelist":
fields["levels"] = result
if "l" in fields:
fields["l"].extend(result)
elif child.axis.name == "param":
fields["param"] = result
elif child.axis.name in ["date"]:
fields["dates"].extend(result)
elif child.axis.name == "number":
fields["number"] = result
elif child.axis.name == "step":
fields["step"] = result
if "s" in fields:
fields["s"].extend(result)
elif child.axis.name == "time":
fields["times"].extend(result)

self.walk_tree_step(child, fields, coords, mars_metadata, range_dict)
else:
tree.values = [float(val) for val in tree.values]
if all(val is None for val in tree.result):
fields["dates"] = fields["dates"][:-1]
for date in fields["dates"]:
for level in fields["levels"]:
for num in fields["number"]:
for para in fields["param"]:
for s in fields["step"]:
key = create_composite_key_step(date, level, num, para)
if key in range_dict:
del range_dict[key]
else:
tree.result = [float(val) if val is not None else val for val in tree.result]
date_len = len(tree.result) / len(fields["dates"])
para_len = date_len / len(fields["param"])

for date in fields["dates"]:
append_composite_coords_step(date, tree.values, fields["lat"], coords)

for d, date in enumerate(fields["dates"]):
for l, level in enumerate(fields["levels"]): # noqa: E741
for i, num in enumerate(fields["number"]):
for j, para in enumerate(fields["param"]):
# for k, t in enumerate(fields["times"]):
# start_index, end_index = calculate_index_bounds_step(
# level_len, num_len, para_len, time_len, l, i, j, k
# )
key = create_composite_key_step(date, level, num, para)
if key not in range_dict:
range_dict[key] = []
# range_dict[key].extend(tree.result[start_index:end_index])
# print(d, date_len,j, para_len)
# print(d*date_len+j*para_len)
# print(int(d*date_len+j*para_len+len(fields["times"])))
# print(tree.result[int(d*date_len+j*para_len+len)])
range_dict[key].append(
tree.result[
int(d * date_len + j * para_len) : int(
d * date_len + j * para_len + len(fields["times"])
)
]
)

@abstractmethod
def add_coverage(self, mars_metadata, coords, values):
pass
Expand Down
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.1.2"
__version__ = "0.1.3"

0 comments on commit 66c829e

Please sign in to comment.