Merge pull request #376 from lsst/tickets/PREOPS-4124

tickets/PREOPS-4124 Support restricting the scheduler reward summary stats to a region of interest

rubin_sim/scheduler/basis_functions/basis_functions.py

@@ -1,17 +1,21 @@
 __all__ = (
     "BaseBasisFunction",
+    "HealpixLimitedBasisFunctionMixin",
     "ConstantBasisFunction",
+    "SimpleArrayBasisFunction",
     "DelayStartBasisFunction",
     "TargetMapBasisFunction",
     "AvoidLongGapsBasisFunction",
     "AvoidFastRevists",
     "VisitRepeatBasisFunction",
     "M5DiffBasisFunction",
+    "M5DiffAtHpixBasisFunction",
     "StrictFilterBasisFunction",
     "GoalStrictFilterBasisFunction",
     "FilterChangeBasisFunction",
     "SlewtimeBasisFunction",
     "AggressiveSlewtimeBasisFunction",
+    "SkybrightnessLimitAtHpixBasisFunction",
     "SkybrightnessLimitBasisFunction",
     "CablewrapUnwrapBasisFunction",
     "CadenceEnhanceBasisFunction",
@@ -174,13 +178,70 @@ def label(self):
         return label
 
 
+class HealpixLimitedBasisFunctionMixin:
+    """A mixin to limit a basis function to a set of Healpix pixels."""
+
+    def __init__(self, hpid, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.hpid = hpid
+
+    def check_feasibility(self, conditions):
+        """Check the feasibility of the current set of conditions.
+
+        Parameters
+        ----------
+        conditions : `rubin_sim.scheduler.features.Conditions`
+            The conditions for which to test feasibility.
+
+        Returns
+        -------
+        feasibility : `bool`
+            True if the current set of conditions is feasible, False otherwise.
+        """
+
+        if super().check_feasibility(conditions):
+            if self.recalc or (self.update_on_mjd and conditions.mjd != self.mjd_last):
+                value = self._calc_value(conditions)
+            else:
+                value = super().value
+
+            feasibility = np.nanmax(value) > -np.inf
+        else:
+            feasibility = False
+        return feasibility
+
+    def _calc_value(self, conditions, all_sky=False, **kwargs):
+        all_sky_value = super()._calc_value(conditions, **kwargs)
+
+        if all_sky:
+            return all_sky_value
+
+        if np.isscalar(all_sky_value):
+            value = all_sky_value
+        else:
+            assert len(all_sky_value) == hp.nside2npix(self.nside)
+            value = all_sky_value[self.hpid]
+
+        return value
+
+
 class ConstantBasisFunction(BaseBasisFunction):
     """Just add a constant"""
 
     def __call__(self, conditions, **kwargs):
         return 1
 
 
+class SimpleArrayBasisFunction(BaseBasisFunction):
+    def __init__(self, value, *args, **kwargs):
+        self.assigned_value = value
+        super().__init__(*args, **kwargs)
+
+    def _calc_value(self, conditions, **kwargs):
+        self.value = self.assigned_value
+        return self.value
+
+
 class DelayStartBasisFunction(BaseBasisFunction):
     """Force things to not run before a given night"""
 
@@ -899,6 +960,10 @@ def _calc_value(self, conditions, indx=None):
         return result
 
 
+class M5DiffAtHpixBasisFunction(HealpixLimitedBasisFunctionMixin, M5DiffBasisFunction):
+    pass
+
+
 class StrictFilterBasisFunction(BaseBasisFunction):
     """Remove the bonus for staying in the same filter if certain conditions are met.
 
@@ -1266,6 +1331,12 @@ def _calc_value(self, conditions, indx=None):
         return result
 
 
+class SkybrightnessLimitAtHpixBasisFunction(
+    HealpixLimitedBasisFunctionMixin, SkybrightnessLimitBasisFunction
+):
+    pass
+
+
 class CablewrapUnwrapBasisFunction(BaseBasisFunction):
     """
     Parameters

rubin_sim/scheduler/surveys/base_survey.py

@@ -1,6 +1,7 @@
 __all__ = ("BaseSurvey", "BaseMarkovSurvey")
 
 from copy import copy, deepcopy
+from functools import cached_property
 
 import healpy as hp
 import numpy as np
@@ -96,6 +97,20 @@ def __init__(
         # Scheduled observations
         self.scheduled_obs = scheduled_obs
 
+    @cached_property
+    def roi_hpid(self):
+        return None
+
+    @cached_property
+    def roi_mask(self):
+        if self.roi_hpid is None:
+            mask = np.ones(hp.nside2npix(self.nside), dtype=bool)
+        else:
+            mask = np.zeros(hp.nside2npix(self.nside), dtype=bool)
+            mask[self.roi_hpid] = True
+
+        return mask
+
     def _generate_survey_name(self):
         self.survey_name = ""
 
@@ -225,7 +240,10 @@ def _reward_to_scalars(self, reward):
             pix_area = None
 
         if np.isscalar(reward):
-            unmasked_area = pix_area * hp.nside2npix(self.nside)
+            if np.isnan(reward) or reward == -np.inf:
+                unmasked_area = 0
+            else:
+                unmasked_area = pix_area * hp.nside2npix(self.nside)
         else:
             unmasked_area = pix_area * np.count_nonzero(reward > -np.inf)
 
@@ -236,6 +254,9 @@ def _reward_to_scalars(self, reward):
         else:
             scalar_reward = np.nanmax(reward)
 
+        if np.isnan(scalar_reward):
+            scalar_reward = -np.inf
+
         return scalar_reward, unmasked_area
 
     def make_reward_df(self, conditions, accum=True):
@@ -272,9 +293,7 @@ def make_reward_df(self, conditions, accum=True):
 
         short_labels = self.bf_short_labels
 
-        # Only count the part of the sky high enough to observe.
-        horizon_mask = np.where(conditions.alt > np.radians(20), 1, np.nan)
-        _, scalar_area = self._reward_to_scalars(horizon_mask)
+        _, scalar_area = self._reward_to_scalars(np.where(self.roi_mask, 1, -np.inf))
 
         for weight, basis_function in zip(full_basis_weights, self.basis_functions):
             bf_label.append(short_labels[basis_function.label()])
@@ -284,14 +303,17 @@ def make_reward_df(self, conditions, accum=True):
                 max_reward = bf_reward
                 basis_area = scalar_area
             else:
-                bf_reward = bf_reward * horizon_mask
+                bf_reward[~self.roi_mask] = -np.inf
                 max_reward, basis_area = self._reward_to_scalars(bf_reward)
 
-            max_rewards.append(max_reward)
-            basis_areas.append(basis_area)
+            if basis_area == 0:
+                this_feasibility = False
+            else:
+                this_feasibility = np.array(basis_function.check_feasibility(conditions)).any()
 
-            this_feasibility = np.array(basis_function.check_feasibility(conditions)).any()
             feasibility.append(this_feasibility)
+            max_rewards.append(max_reward)
+            basis_areas.append(basis_area)
 
             if accum:
                 basis_functions.append(basis_function)
@@ -300,6 +322,8 @@ 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)

rubin_sim/scheduler/surveys/dd_surveys.py

@@ -3,14 +3,15 @@
 import copy
 import logging
 import random
+from functools import cached_property
 
 import numpy as np
 
 import rubin_sim.scheduler.basis_functions as basis_functions
 from rubin_sim.scheduler import features
 from rubin_sim.scheduler.surveys import BaseSurvey
 from rubin_sim.scheduler.utils import empty_observation
-from rubin_sim.utils import ddf_locations
+from rubin_sim.utils import ddf_locations, ra_dec2_hpid
 
 log = logging.getLogger(__name__)
 
@@ -111,6 +112,11 @@ def __init__(
             self.extra_features["Ntot"] = features.N_obs_survey()
             self.extra_features["N_survey"] = features.N_obs_survey(note=self.survey_name)
 
+    @cached_property
+    def roi_hpid(self):
+        hpid = ra_dec2_hpid(self.nside, np.degrees(self.ra), np.degrees(self.dec))
+        return hpid
+
     def check_continue(self, observation, conditions):
         # feasibility basis functions?
         """

tests/scheduler/test_scalarbf.py

@@ -0,0 +1,47 @@
+import unittest
+import healpy as hp
+import numpy as np
+from rubin_sim.scheduler.features import Conditions
+from rubin_sim.scheduler.utils import set_default_nside
+from rubin_sim.scheduler.basis_functions import HealpixLimitedBasisFunctionMixin
+from rubin_sim.scheduler.basis_functions import SimpleArrayBasisFunction
+
+
+class SimpleArrayAtHpixBasisFunction(HealpixLimitedBasisFunctionMixin, SimpleArrayBasisFunction):
+    pass
+
+
+class TestHealpixLimitedBasisFunctionMixin(unittest.TestCase):
+    def setUp(self):
+        self.hpid = 2111
+        random_seed = 6563
+        self.nside = set_default_nside()
+        self.npix = hp.nside2npix(self.nside)
+        self.rng = np.random.default_rng(seed=random_seed)
+        self.all_values = self.rng.random(self.npix)
+        self.value = self.all_values[self.hpid]
+        self.conditions = Conditions(nside=self.nside, mjd=60200.2)
+
+    def test_array_data(self):
+        basis_function = SimpleArrayBasisFunction(self.all_values)
+        test_values = basis_function(self.conditions)
+        self.assertTrue(np.array_equal(test_values, self.all_values))
+
+    def test_scalar_data(self):
+        basis_function = SimpleArrayAtHpixBasisFunction(self.hpid, self.all_values)
+        test_value = basis_function(self.conditions)
+        self.assertEqual(test_value, self.value)
+
+        feasible = basis_function(self.conditions)
+        self.assertTrue(feasible)
+
+    def test_infeasible_at_hpix(self):
+        values_invalid_at_hpix = self.all_values.copy()
+        values_invalid_at_hpix[self.hpid] = -np.inf
+        basis_function = SimpleArrayAtHpixBasisFunction(self.hpid, values_invalid_at_hpix)
+        feasible = basis_function.check_feasibility(self.conditions)
+        self.assertFalse(feasible)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/scheduler/test_surveys.py

@@ -2,11 +2,15 @@
 import unittest
 
 import pandas as pd
+import numpy as np
+import healpy as hp
 
 import rubin_sim.scheduler.basis_functions as basis_functions
 import rubin_sim.scheduler.surveys as surveys
 from rubin_sim.data import get_data_dir
+from rubin_sim.scheduler.utils import set_default_nside
 from rubin_sim.scheduler.model_observatory import ModelObservatory
+from rubin_sim.scheduler.basis_functions import SimpleArrayBasisFunction
 
 
 class TestSurveys(unittest.TestCase):
@@ -34,6 +38,48 @@ def test_field_survey(self):
         self.assertIsInstance(reward_df, pd.DataFrame)
         reward_df = survey.make_reward_df(conditions, accum=False)
 
+    def test_roi(self):
+        random_seed = 6563
+        infeasible_hpix = 123
+        nside = set_default_nside()
+        npix = hp.nside2npix(nside)
+        rng = np.random.default_rng(seed=random_seed)
+        num_bfs = 3
+        bf_values = rng.random((num_bfs, npix))
+        bf_values[:, infeasible_hpix] = -np.inf
+        bfs = [SimpleArrayBasisFunction(values) for values in bf_values]
+
+        observatory = ModelObservatory()
+        conditions = observatory.return_conditions()
+
+        # A few cases with an ROI with one valid healpix
+        for i in range(3):
+            hpix = rng.integers(npix)
+            ra, decl = hp.pix2ang(nside, hpix, lonlat=True)
+            survey = surveys.FieldSurvey(bfs, RA=ra, dec=decl, reward_value=1)
+            reward_df = survey.make_reward_df(conditions)
+            for value, max_basis_reward in zip(bf_values[:, hpix], reward_df["max_basis_reward"]):
+                self.assertEqual(max_basis_reward, value)
+
+        # One case with an ROI with only an infeasible healpix
+        ra, decl = hp.pix2ang(nside, infeasible_hpix, lonlat=True)
+        survey = surveys.FieldSurvey(bfs, RA=ra, dec=decl, reward_value=1)
+        reward_df = survey.make_reward_df(conditions)
+        for max_basis_reward in reward_df["max_basis_reward"]:
+            self.assertEqual(max_basis_reward, -np.inf)
+
+        for area in reward_df["basis_area"]:
+            self.assertEqual(area, 0.0)
+
+        for feasible in reward_df["feasible"]:
+            self.assertFalse(feasible)
+
+        # Make sure it still works as expected if no ROI is set
+        weights = [1] * num_bfs
+        survey = surveys.BaseMarkovSurvey(bfs, weights)
+        for value, max_basis_reward in zip(bf_values.max(axis=1), reward_df["max_basis_reward"]):
+            self.assertEqual(max_basis_reward, value)
+
 
 if __name__ == "__main__":
     unittest.main()