diff --git a/covjsonkit/encoder/BoundingBox.py b/covjsonkit/encoder/BoundingBox.py index a33b82a..ea38e44 100644 --- a/covjsonkit/encoder/BoundingBox.py +++ b/covjsonkit/encoder/BoundingBox.py @@ -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( { @@ -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"] @@ -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 = {} diff --git a/covjsonkit/encoder/TimeSeries.py b/covjsonkit/encoder/TimeSeries.py index 3821cbf..0cbd643 100644 --- a/covjsonkit/encoder/TimeSeries.py +++ b/covjsonkit/encoder/TimeSeries.py @@ -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 @@ -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( { @@ -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) @@ -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 diff --git a/covjsonkit/encoder/encoder.py b/covjsonkit/encoder/encoder.py index 0ccdf91..05bf386 100644 --- a/covjsonkit/encoder/encoder.py +++ b/covjsonkit/encoder/encoder.py @@ -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):