From e47eecc846c1580a25d6671ba7b309ff7fbee79d Mon Sep 17 00:00:00 2001
From: Joshua Greaves <joshgreaves@google.com>
Date: Fri, 16 Dec 2022 04:35:58 +0000
Subject: [PATCH] Fix bug introduced by bug fix in sklearn's GP.

PiperOrigin-RevId: 495767145
---
 balloon_learning_environment/env/balloon/acs.py |  4 ++--
 balloon_learning_environment/env/features.py    |  3 ++-
 balloon_learning_environment/env/wind_gp.py     | 16 +++++++++++-----
 3 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/balloon_learning_environment/env/balloon/acs.py b/balloon_learning_environment/env/balloon/acs.py
index a052c46..ff9f513 100644
--- a/balloon_learning_environment/env/balloon/acs.py
+++ b/balloon_learning_environment/env/balloon/acs.py
@@ -21,14 +21,14 @@
 from scipy import interpolate
 
 
-_PRESSURE_RATIO_TO_POWER: interpolate.interpolate.interp1d = (
+_PRESSURE_RATIO_TO_POWER: interpolate.interp1d = (
     interpolate.interp1d(
         np.array([1.0, 1.05, 1.2, 1.25, 1.35]),  # pressure_ratio
         np.array([100.0, 100.0, 300.0, 400.0, 400.0]),  # power
         fill_value='extrapolate'))
 
 
-_PRESSURE_RATIO_POWER_TO_EFFICIENCY: interpolate.interpolate.interp2d = (
+_PRESSURE_RATIO_POWER_TO_EFFICIENCY: interpolate.interp2d = (
     interpolate.interp2d(
         np.linspace(1.05, 1.35, 13),  # pressure_ratio
         np.linspace(100.0, 400.0, 4),  # power
diff --git a/balloon_learning_environment/env/features.py b/balloon_learning_environment/env/features.py
index b5d4cc5..aabe1aa 100644
--- a/balloon_learning_environment/env/features.py
+++ b/balloon_learning_environment/env/features.py
@@ -138,7 +138,8 @@ def observe(self, observation: simulator_data.SimulatorObservation) -> None:
   def get_features(self) -> np.ndarray:
     """Gets the current feature vector given all observations."""
 
-  @abc.abstractproperty
+  @property
+  @abc.abstractmethod
   def observation_space(self) -> gym.Space:
     """Gets the observation space specification for the feature vector."""
 
diff --git a/balloon_learning_environment/env/wind_gp.py b/balloon_learning_environment/env/wind_gp.py
index 0f1447c..a19fbfc 100644
--- a/balloon_learning_environment/env/wind_gp.py
+++ b/balloon_learning_environment/env/wind_gp.py
@@ -19,9 +19,6 @@
 This lets us query any point (x, y, p, t) in the wind field for its value,
 as well as the model confidence's in this value.
 
----- Open issues
-* The forecast is not used
-
 """
 
 import datetime as dt
@@ -190,8 +187,6 @@ def query_batch(self, locations: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
 
       # Output should be a N x 2 set of predictions about local measurements,
       # and a N-sized vector of standard deviations.
-      # TODO(bellemare): Determine why deviations is a single number per sample,
-      # instead of two (since we have two values being predicted).
       means, deviations = self.model.predict(locations, return_std=True)
 
       # Deviations are std.dev., convert to variance and normalize.
@@ -200,6 +195,17 @@ def query_batch(self, locations: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
       # like 0.07 from the GP, but that doesn't seem to match the Loon code.
       deviations = deviations**2 / _SIGMA_EXP_SQUARED
 
+      # Previously, there was a bug in sklearn which meant that the GP
+      # only returned 1 deviation per location when it should have
+      # returned 2. This has since been fixed, but people running with
+      # older code won't need to apply the next block because the shapes
+      # will work as expected.
+      assert deviations.ndim in (1, 2)
+      if deviations.ndim == 2:
+        # We use the standard deviation as a measure of uncertainty,
+        # and take the average of the two components as a proxy for this.
+        deviations = np.mean(deviations, axis=1)
+
       # TODO(bellemare): Sal says this needs normalizing so that the lower bound
       # is really zero.