diff options
| -rw-r--r-- | windy/point_forecast.py | 72 |
1 files changed, 53 insertions, 19 deletions
diff --git a/windy/point_forecast.py b/windy/point_forecast.py index 308d432..7e9304f 100644 --- a/windy/point_forecast.py +++ b/windy/point_forecast.py @@ -5,6 +5,7 @@ Module provides tools to deal with Windy's point forecast API. from dataclasses import dataclass from datetime import datetime from enum import Enum +from functools import reduce import numpy as np @@ -117,7 +118,8 @@ class Request: class Prediction: """ - Predicted values for each of the requested parameters along with their associated time point. + Predicted values for each of the requested parameters along with their associated time point. Effectively a time + slice of the entire Response. """ def __init__(self, response, index=0): self._response = response @@ -131,48 +133,80 @@ class Prediction: def parameters(self) -> tuple: return self._response.parameters - @property - def levels(self) -> tuple: - return self._response.levels + def levels(self, parameter) -> tuple: + return self._response.levels(parameter) def __iter__(self): return iter(self.parameters) - def __getitem__(self, key): - return self._response.values[key][self._index] + def __getitem__(self, parameter): + return self._response.raw_predictions[parameter][self._index] class Response: """ - Wraps raw JSON response from the Windy's API to allow for easier access, converts all values to pint's - Quantities, and converts all timestamps into datetime objects. + Parses raw JSON response from Windy's API to allow easier access to prepared pint-based vertical profiles of + each of the parameters from the requested forecast scope. Values are first keyed by the parameter name + (each of the self.parameters), then indexed by the predictions (respectively to self.timestamps), and then + indexed in numpy array by the levels (respectively to self.levels(parameter)). - Can be used in a for-loop to access all samples via Prediction: + Wrapper to access parameters of a certain prediction time point is available: >>> for prediction in response: >>> print(prediction.timestamp, prediction['temp']) - Otherwise, timestamps list and samples dictionary are available for direct access. + Otherwise, timestamps list and raw_predicitions structured as described above are available for direct access. """ _INTERNAL_FIELDS = ('ts', 'units', 'warning') def __init__(self, registry, raw): self.timestamps = [datetime.fromtimestamp(x // 1000) for x in raw['ts']] - count = len(self.timestamps) - split = [tuple(x.split("-")) for x in raw if x not in self._INTERNAL_FIELDS] - self.parameters = tuple({x for x, _ in split}) - self.levels = tuple(sorted({Level(x) for _, x in split})) - if 'pressure' in self.parameters: - for level in self.levels: + levels = {} + for key in raw: + if key in self._INTERNAL_FIELDS: + continue + parameter, level = key.split("-") + level = Level(level) + if parameter in levels: + levels[parameter].add(level) + else: + levels[parameter] = {level} + all_levels = tuple(sorted(reduce(lambda x, y: x | y, levels.values()))) + parameters = tuple(levels.keys()) + for parameter in levels: + levels[parameter] = tuple(sorted(levels[parameter])) + units = {x: registry(_convert_notation(raw['units'][f'{x}-{levels[x][0]}'])) for x in parameters} + if 'pressure' in parameters: + for level in all_levels: if level is Level.SURFACE: continue - raw[f'pressure-{level}'] = [level.pressure() for _ in range(len(self.timestamps))] - units = {x: registry(_convert_notation(raw['units'][f'{x}-surface'])) for x in self.parameters} # Don't guess surface - self.values = {p: [[raw[f'{p}-{l}'][i] for l in self.levels] * units[p] for i in range(count)] for p in self.parameters} + pressure = (level.pressure() * registry.hPa).m_as(units['pressure']) + raw[f'pressure-{level}'] = [pressure for _ in range(len(self.timestamps))] + levels['pressure'] = all_levels + self.raw_predictions = {} + for parameter in parameters: + profiles = [] + for index in range(len(self.timestamps)): + profile = [] + for level in levels[parameter]: + try: + profile.append(raw[f'{parameter}-{level}'][index]) + except KeyError: + pass + profiles.append(np.array(profile) * units[parameter]) + self.raw_predictions[parameter] = profiles + self._levels = levels def __len__(self): return len(self.timestamps) + @property + def parameters(self) -> tuple: + return tuple(self.raw_predictions.keys()) + + def levels(self, parameter) -> tuple: + return self._levels[parameter] + def predictions(self) -> Prediction: """ Yields Prediction for each time point available in this Response. |