From 7587e42259405cd2cca758639879f05b9f456f04 Mon Sep 17 00:00:00 2001
From: Aki <please@ignore.pl>
Date: Sat, 8 Oct 2022 11:07:54 +0200
Subject: Extended error handling and edge cases in Response initialization

---
 windy/point_forecast.py | 72 ++++++++++++++++++++++++++++++++++++-------------
 1 file 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.
-- 
cgit v1.1