Readded Prediction.at(parameter, level)

1 files changed, 13 insertions, 0 deletions

diff --git a/windy/point_forecast.py b/windy/point_forecast.py
index 844cef3..34d13b5 100644
--- a/windy/point_forecast.py
+++ b/windy/point_forecast.py
@@ -133,6 +133,10 @@ class Prediction:
 	def parameters(self) -> tuple:
 		return self._response.parameters
 
+	def at(self, parameter, level):
+		level = self._response.raw_predictions['level'][self._index].index(level)
+		return self._response.raw_predictions[parameter][self._index][level]
+
 	def __iter__(self):
 		return iter(self.parameters)
 
@@ -202,6 +206,15 @@ class Response:
 					profile[0], profile[new_index] = profile[new_index], profile[0]
 				profiles.append(np.array(profile) * units[parameter])
 			self.raw_predictions[parameter] = profiles
+		levels = []
+		for index in range(len(self.timestamps)):
+			profile = list(all_levels)
+			if pressure_sort:
+				new_index = pressure_sort[index]
+				profile[0], profile[new_index] = profile[new_index], profile[0]
+			levels.append(profile)
+		self.raw_predictions['level'] = levels
+
 
 	def __len__(self):
 		return len(self.timestamps)