Merge pull request #380 from lsst/tickets/PREOPS-4519

Use the same code for basis function & reward area
lsst · Nov 7, 2023 · 5a74b42 · 5a74b42
2 parents 9bb4052 + 02f2e18
commit 5a74b42
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 26 deletions.
diff --git a/rubin_sim/scheduler/surveys/base_survey.py b/rubin_sim/scheduler/surveys/base_survey.py
@@ -234,28 +234,25 @@ def __repr__(self):
         return repr
 
     def _reward_to_scalars(self, reward):
-        try:
-            pix_area = hp.nside2pixarea(self.nside, degrees=True)
-        except AttributeError:
-            pix_area = None
-
         if np.isscalar(reward):
+            scalar_reward = reward
             if np.isnan(reward) or reward == -np.inf:
                 unmasked_area = 0
+                scalar_reward = -np.inf
             else:
-                unmasked_area = pix_area * hp.nside2npix(self.nside)
+                try:
+                    pix_area = hp.nside2pixarea(self.nside, degrees=True)
+                    unmasked_area = pix_area * np.count_nonzero(self.roi_mask)
+                except AttributeError:
+                    unmasked_area = 4 * np.pi * (180 / np.pi) ** 2
         else:
-            unmasked_area = pix_area * np.count_nonzero(reward > -np.inf)
-
-        if np.isscalar(reward):
-            scalar_reward = reward
-        elif unmasked_area == 0:
-            scalar_reward = -np.inf
-        else:
-            scalar_reward = np.nanmax(reward)
-
-        if np.isnan(scalar_reward):
-            scalar_reward = -np.inf
+            reward_in_roi = np.where(self.roi_mask, reward, -np.inf)
+            pix_area = hp.nside2pixarea(self.nside, degrees=True)
+            unmasked_area = pix_area * np.count_nonzero(reward_in_roi > -np.inf)
+            if unmasked_area == 0:
+                scalar_reward = -np.inf
+            else:
+                scalar_reward = np.nanmax(reward_in_roi)
 
         return scalar_reward, unmasked_area
 
@@ -293,18 +290,13 @@ def make_reward_df(self, conditions, accum=True):
 
         short_labels = self.bf_short_labels
 
-        _, scalar_area = self._reward_to_scalars(np.where(self.roi_mask, 1, -np.inf))
+        _, scalar_area = self._reward_to_scalars(1)
 
         for weight, basis_function in zip(full_basis_weights, self.basis_functions):
             bf_label.append(short_labels[basis_function.label()])
             bf_class.append(basis_function.__class__.__name__)
             bf_reward = basis_function(conditions)
-            if np.isscalar(bf_reward):
-                max_reward = bf_reward
-                basis_area = scalar_area
-            else:
-                bf_reward[~self.roi_mask] = -np.inf
-                max_reward, basis_area = self._reward_to_scalars(bf_reward)
+            max_reward, basis_area = self._reward_to_scalars(bf_reward)
 
             if basis_area == 0:
                 this_feasibility = False
@@ -322,8 +314,6 @@ def make_reward_df(self, conditions, accum=True):
                 test_survey.basis_functions = basis_functions
                 test_survey.basis_weights = basis_weights
                 this_accum_reward = test_survey.calc_reward_function(conditions)
-                if not np.isscalar(this_accum_reward):
-                    this_accum_reward[~self.roi_mask] = -np.inf
                 accum_reward, accum_area = self._reward_to_scalars(this_accum_reward)
                 accum_rewards.append(accum_reward)
                 accum_areas.append(accum_area)

diff --git a/tests/scheduler/test_surveys.py b/tests/scheduler/test_surveys.py
@@ -77,6 +77,7 @@ def test_roi(self):
         # Make sure it still works as expected if no ROI is set
         weights = [1] * num_bfs
         survey = surveys.BaseMarkovSurvey(bfs, weights)
+        reward_df = survey.make_reward_df(conditions)
         for value, max_basis_reward in zip(bf_values.max(axis=1), reward_df["max_basis_reward"]):
             self.assertEqual(max_basis_reward, value)