Converted to numpy array and transponed the values

The current format will allow easy access to vertical profiles which in turn
will ease up dewpoint calculations.

1 files changed, 20 insertions, 8 deletions

diff --git a/windy/point_forecast.py b/windy/point_forecast.py
index 0132751..308d432 100644
--- a/windy/point_forecast.py
+++ b/windy/point_forecast.py
@@ -6,6 +6,8 @@ from dataclasses import dataclass
 from datetime import datetime
 from enum import Enum
 
+import numpy as np
+
 
 def _json(value):
 	try:
@@ -82,6 +84,11 @@ class Level(_StrEnum):
 	H200 = "200h"
 	H150 = "150h"
 
+	def pressure(self):
+		if self is Level.SURFACE:
+			raise ValueError
+		return float(self.value[:-1])
+
 
 @dataclass
 class Request:
@@ -132,7 +139,7 @@ class Prediction:
 		return iter(self.parameters)
 
 	def __getitem__(self, key):
-		return self._response.samples[key][self._index]
+		return self._response.values[key][self._index]
 
 
 class Response:
@@ -143,7 +150,7 @@ class Response:
 	Can be used in a for-loop to access all samples via Prediction:
 
 		>>> for prediction in response:
-		>>>     print(prediction.timestamp, prediction['temp-surface'])
+		>>>     print(prediction.timestamp, prediction['temp'])
 
 	Otherwise, timestamps list and samples dictionary are available for direct access.
 	"""
@@ -151,12 +158,17 @@ class Response:
 
 	def __init__(self, registry, raw):
 		self.timestamps = [datetime.fromtimestamp(x // 1000) for x in raw['ts']]
-		self.parameters = tuple(x for x in raw if x not in self._INTERNAL_FIELDS)
-		self.levels = tuple(sorted({Level(x.split("-")[1]) for x in self.parameters}))
-		self.samples = {}
-		for parameter in self.parameters:
-			unit = registry(_convert_notation(raw['units'][parameter]))
-			self.samples[parameter] = [x * unit for x in raw[parameter]]
+		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:
+				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}
 
 	def __len__(self):
 		return len(self.timestamps)