Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Refactor walk_tree to no longer depend on order of polytope tree #38

Merged
merged 2 commits into from
Sep 16, 2024
Merged

Refactor walk_tree to no longer depend on order of polytope tree #38

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
cf79be9
Refactor walk_tree to no longer depend on order of polytope tree
awarde96 Sep 12, 2024
42ee808
Remove coords when no data available
awarde96 Sep 13, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
54 changes: 31 additions & 23 deletions covjsonkit/encoder/BoundingBox.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -83,14 +83,15 @@ def from_polytope(self, result):
coords = {}
mars_metadata = {}
range_dict = {}
lat = 0
param = 0
number = [0]
step = 0
dates = [0]
levels = [0]
fields = {}
fields["lat"] = 0
fields["param"] = 0
fields["number"] = [0]
fields["step"] = 0
fields["dates"] = []
fields["levels"] = [0]

self.walk_tree(result, lat, coords, mars_metadata, param, range_dict, number, step, dates, levels)
self.walk_tree(result, fields, coords, mars_metadata, range_dict)

self.add_reference(
{
Expand All @@ -104,28 +105,34 @@ def from_polytope(self, result):

combined_dict = {}

for date in range_dict:
for date in fields["dates"]:
if date not in combined_dict:
combined_dict[date] = {}
for level in range_dict[date]:
for num in range_dict[date][level]:
for level in fields["levels"]:
for num in fields["number"]:
if num not in combined_dict[date]:
combined_dict[date][num] = {}
for para in range_dict[date][level][num]:
for para in fields["param"]:
if para not in combined_dict[date][num]:
combined_dict[date][num][para] = {}
for s, value in range_dict[date][level][num][para].items():
if s not in combined_dict[date][num][para]:
combined_dict[date][num][para][s] = value
else:
# Cocatenate arrays
combined_dict[date][num][para][s] += value

levels = []
for date in range_dict.keys():
for level in range_dict[date].keys():
levels.append(level)
break
# for s, value in range_dict[date][level][num][para].items():
for s in fields["step"]:
key = (date, level, num, para, s)
for k, v in range_dict.items():
if k == key:
if s not in combined_dict[date][num][para]:
combined_dict[date][num][para][s] = v
else:
# Cocatenate arrays
combined_dict[date][num][para][s] += v

levels = fields["levels"]
for para in fields["param"]:
self.add_parameter(para)
# for date in range_dict.keys():
# for level in range_dict[date].keys():
# levels.append(level)
# break

for date in coords.keys():
coord = coords[date]["composite"]
Expand All @@ -134,6 +141,7 @@ def from_polytope(self, result):
for cor in coord:
coords[date]["composite"].append([cor[0], cor[1], level])

print(combined_dict)
for date in combined_dict.keys():
for num in combined_dict[date].keys():
val_dict = {}
Expand Down
55 changes: 31 additions & 24 deletions covjsonkit/encoder/TimeSeries.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,9 +44,9 @@ def add_range(self, coverage, values):
coverage["ranges"][param]["dataType"] = "float"
coverage["ranges"][param]["shape"] = [len(values[parameter])]
coverage["ranges"][param]["axisNames"] = [str(param)]
coverage["ranges"][param]["values"] = [
values[parameter][val][0] for val in values[parameter].keys()
] # values[parameter]
coverage["ranges"][param]["values"] = values[
parameter
] # [values[parameter][val][0] for val in values[parameter].keys()]

def add_mars_metadata(self, coverage, metadata):
coverage["mars:metadata"] = metadata
Expand Down Expand Up @@ -88,14 +88,15 @@ def from_polytope(self, result):
coords = {}
mars_metadata = {}
range_dict = {}
lat = 0
param = 0
number = [0]
step = 0
levels = [0]
dates = [0]
fields = {}
fields["lat"] = 0
fields["param"] = 0
fields["number"] = [0]
fields["step"] = 0
fields["dates"] = []
fields["levels"] = [0]

self.walk_tree(result, lat, coords, mars_metadata, param, range_dict, number, step, dates, levels)
self.walk_tree(result, fields, coords, mars_metadata, range_dict)

self.add_reference(
{
Expand All @@ -109,23 +110,21 @@ def from_polytope(self, result):

coordinates = {}

levels = []
for date in range_dict.keys():
for level in range_dict[date].keys():
levels.append(level)
break
levels = fields["levels"]
for para in fields["param"]:
self.add_parameter(para)

for date in range_dict.keys():
for date in fields["dates"]:
coordinates[date] = {
"x": [coords[date]["composite"][0][0]],
"y": [coords[date]["composite"][0][1]],
"z": [levels[0]],
}
coordinates[date]["t"] = []
for level in range_dict[date].keys():
for num in range_dict[date][level].keys():
for para in range_dict[date][level][num].keys():
for step in range_dict[date][level][num][para].keys():
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)
Expand All @@ -136,14 +135,22 @@ def from_polytope(self, result):
break
break

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

return self.covjson

Expand Down
167 changes: 84 additions & 83 deletions covjsonkit/encoder/encoder.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -93,96 +93,97 @@ def get_json(self):
# self.covjson = self.pydantic_coverage.model_dump_json(exclude_none=True, indent=4)
return orjson.dumps(self.covjson)

def walk_tree(self, tree, lat, coords, mars_metadata, param, range_dict, number, step, dates, levels):
def walk_tree(self, tree, fields, coords, mars_metadata, range_dict):
def create_composite_key(date, level, num, para, s):
return (date, level, num, para, s)

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

def handle_specific_axes(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", "time"]:
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
return None

def calculate_index_bounds(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(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:
# 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 != "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 == "levelist":
levels = c.values
for date in range_dict.keys():
for level in levels:
if level not in range_dict[date]:
range_dict[date][level] = {}
if c.axis.name == "param":
param = c.values
for date in range_dict.keys():
if range_dict[date] == {}:
range_dict[date] = {0: {}}
for level in levels:
if range_dict[date][level] == {}:
range_dict[date][level] = {0: {}}
for num in number:
for para in param:
if para not in range_dict[date][level][num]:
range_dict[date][level][num][para] = {}
self.add_parameter(para)
if c.axis.name == "date" or c.axis.name == "time":
dates = [str(date) + "Z" for date in c.values]
mars_metadata["Forecast date"] = str(c.values[0])
for date in dates:
coords[date] = {}
coords[date]["composite"] = []
coords[date]["t"] = [date]
if date not in range_dict:
range_dict[date] = {}
if c.axis.name == "number":
number = c.values
for date in dates:
for level in levels:
if level not in range_dict[date]:
range_dict[date][level] = {}
for num in number:
range_dict[date][level][num] = {}
if c.axis.name == "step":
step = c.values
for date in dates:
for level in levels:
for num in number:
for para in param:
for s in step:
range_dict[date][level][num][para][s] = []

self.walk_tree(c, lat, coords, mars_metadata, param, range_dict, number, step, dates, levels)
for child in tree.children:
handle_non_leaf_node(child)
result = handle_specific_axes(child)
if result is not None:
if child.axis.name == "latitude":
fields["lat"] = result
elif child.axis.name == "levelist":
fields["levels"] = result
elif child.axis.name == "param":
fields["param"] = result
elif child.axis.name in ["date", "time"]:
fields["dates"].extend(result)
elif child.axis.name == "number":
fields["number"] = result
elif child.axis.name == "step":
fields["step"] = result

self.walk_tree(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):
range_dict.pop(dates[0], None)
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(date, level, num, para, s)
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]
level_len = len(tree.result) / len(levels)
num_len = level_len / len(number)
para_len = num_len / len(param)
step_len = para_len / len(step)

for date in dates:
for val in tree.values:
coords[date]["composite"].append([lat, val])

for l, level in enumerate(levels): # noqa: E741
for i, num in enumerate(number):
for j, para in enumerate(param):
for k, s in enumerate(step):
range_dict[dates[0]][level][num][para][s].extend(
tree.result[
int(l * level_len)
+ int(i * num_len)
+ int(j * para_len)
+ int(k * step_len) : int(l * level_len)
+ int(i * num_len)
+ int(j * para_len)
+ int((k + 1) * step_len)
]
level_len = len(tree.result) / len(fields["levels"])
num_len = level_len / len(fields["number"])
para_len = num_len / len(fields["param"])
step_len = para_len / len(fields["step"])

append_composite_coords(fields["dates"][-1], tree.values, fields["lat"], coords)

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, s in enumerate(fields["step"]):
start_index, end_index = calculate_index_bounds(
level_len, num_len, para_len, step_len, l, i, j, k
)
key = create_composite_key(fields["dates"][-1], level, num, para, s)
if key not in range_dict:
range_dict[key] = []
range_dict[key].extend(tree.result[start_index:end_index])

@abstractmethod
def add_coverage(self, mars_metadata, coords, values):
Expand Down
Loading