diff --git a/covjsonkit/decoder/VerticalProfile.py b/covjsonkit/decoder/VerticalProfile.py index d45246a..2bc6f3a 100644 --- a/covjsonkit/decoder/VerticalProfile.py +++ b/covjsonkit/decoder/VerticalProfile.py @@ -1,3 +1,4 @@ +import pandas as pd import xarray as xr from .decoder import Decoder @@ -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] diff --git a/covjsonkit/encoder/VerticalProfile.py b/covjsonkit/encoder/VerticalProfile.py index 2cf8658..63df9ad 100644 --- a/covjsonkit/encoder/VerticalProfile.py +++ b/covjsonkit/encoder/VerticalProfile.py @@ -1,4 +1,8 @@ import logging +import time +from datetime import datetime, timedelta + +import pandas as pd from .encoder import Encoder @@ -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( { @@ -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 diff --git a/covjsonkit/version.py b/covjsonkit/version.py index 573b11c..4d04613 100644 --- a/covjsonkit/version.py +++ b/covjsonkit/version.py @@ -1 +1 @@ -__version__ = "0.0.27" +__version__ = "0.0.28"