summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--windy/point_forecast.py52
1 files changed, 36 insertions, 16 deletions
diff --git a/windy/point_forecast.py b/windy/point_forecast.py
index 5a590e2..73f0786 100644
--- a/windy/point_forecast.py
+++ b/windy/point_forecast.py
@@ -161,29 +161,42 @@ class Response:
def __init__(self, registry, raw):
self.timestamps = [datetime.fromtimestamp(x // 1000) for x in raw['ts']]
- 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}
+
+ def _levels_per_parameter(): # (raw, self._INTERNAL_FIELDS)
+ 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}
+ return levels
+
+ levels = _levels_per_parameter()
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}
- pressure_sort = []
- if 'pressure' in parameters:
+ parameters = tuple(sorted(levels.keys()))
+
+ def _parse_unit_map(): # (registry, raw, levels, parameters)
+ return {x: registry(_convert_notation(raw['units'][f'{x}-{levels[x][0]}'])) for x in parameters}
+
+ units = _parse_unit_map()
+
+ def _generate_raw_pressure(): # (all_levels, registry, units, self.timestamps)
+ generated = {}
for level in all_levels:
if level is Level.SURFACE:
continue
pressure = (level.pressure() * registry.hPa).m_as(units['pressure'])
- raw[f'pressure-{level}'] = [pressure for _ in range(len(self.timestamps))]
- levels['pressure'] = all_levels
+ generated[f'pressure-{level}'] = [pressure for _ in range(len(self.timestamps))]
+ return generated
+
+ def _prepare_pressure_sort(): # (self.timestamps, raw, all_levels)
+ pressure_sort = []
for index in range(len(self.timestamps)):
surface = raw[f'pressure-{Level.SURFACE}'][index]
new_index = 0
@@ -191,6 +204,13 @@ class Response:
if surface > raw[f'pressure-{level}'][index]:
break
pressure_sort.append(new_index)
+ return pressure_sort
+
+ pressure_sort = []
+ if 'pressure' in parameters:
+ raw.update(_generate_raw_pressure())
+ levels['pressure'] = all_levels
+ pressure_sort = _prepare_pressure_sort()
self.raw_predictions = {}
for parameter in parameters:
profiles = []