From eac1eccbc971d6db729e035a6623c4a3101d6092 Mon Sep 17 00:00:00 2001 From: Eric Neilsen Date: Wed, 25 Oct 2023 11:17:30 -0700 Subject: [PATCH 1/4] Add ROI to base survey --- rubin_sim/scheduler/surveys/base_survey.py | 40 +++++++++++++++++----- 1 file changed, 32 insertions(+), 8 deletions(-) diff --git a/rubin_sim/scheduler/surveys/base_survey.py b/rubin_sim/scheduler/surveys/base_survey.py index 346b0c01c..c031a2919 100644 --- a/rubin_sim/scheduler/surveys/base_survey.py +++ b/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) From 901a4d15c7ef4426c404d58afd41f0eec4d36dc9 Mon Sep 17 00:00:00 2001 From: Eric Neilsen Date: Wed, 25 Oct 2023 11:19:02 -0700 Subject: [PATCH 2/4] Define ROI in DD surveys --- rubin_sim/scheduler/surveys/dd_surveys.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/rubin_sim/scheduler/surveys/dd_surveys.py b/rubin_sim/scheduler/surveys/dd_surveys.py index 280ea3744..2a106505d 100644 --- a/rubin_sim/scheduler/surveys/dd_surveys.py +++ b/rubin_sim/scheduler/surveys/dd_surveys.py @@ -3,6 +3,7 @@ import copy import logging import random +from functools import cached_property import numpy as np @@ -10,7 +11,7 @@ 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? """ From c9809ec4364aec1fe79715fbe81f47210e613ea0 Mon Sep 17 00:00:00 2001 From: Eric Neilsen Date: Wed, 25 Oct 2023 11:37:42 -0700 Subject: [PATCH 3/4] Add mixin to turn a map basis function into a scalar one --- .../basis_functions/basis_functions.py | 60 +++++++++++++++++++ tests/scheduler/test_scalarbf.py | 57 ++++++++++++++++++ 2 files changed, 117 insertions(+) create mode 100644 tests/scheduler/test_scalarbf.py diff --git a/rubin_sim/scheduler/basis_functions/basis_functions.py b/rubin_sim/scheduler/basis_functions/basis_functions.py index b6fd3c6a8..f1c5c09a0 100644 --- a/rubin_sim/scheduler/basis_functions/basis_functions.py +++ b/rubin_sim/scheduler/basis_functions/basis_functions.py @@ -1,5 +1,6 @@ __all__ = ( "BaseBasisFunction", + "HealpixLimitedBasisFunctionMixin", "ConstantBasisFunction", "DelayStartBasisFunction", "TargetMapBasisFunction", @@ -7,11 +8,13 @@ "AvoidFastRevists", "VisitRepeatBasisFunction", "M5DiffBasisFunction", + "M5DiffAtHpixBasisFunction", "StrictFilterBasisFunction", "GoalStrictFilterBasisFunction", "FilterChangeBasisFunction", "SlewtimeBasisFunction", "AggressiveSlewtimeBasisFunction", + "SkybrightnessLimitAtHpixBasisFunction", "SkybrightnessLimitBasisFunction", "CablewrapUnwrapBasisFunction", "CadenceEnhanceBasisFunction", @@ -174,6 +177,53 @@ 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""" @@ -899,6 +949,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 +1320,12 @@ def _calc_value(self, conditions, indx=None): return result +class SkybrightnessLimitAtHpixBasisFunction( + HealpixLimitedBasisFunctionMixin, SkybrightnessLimitBasisFunction +): + pass + + class CablewrapUnwrapBasisFunction(BaseBasisFunction): """ Parameters diff --git a/tests/scheduler/test_scalarbf.py b/tests/scheduler/test_scalarbf.py new file mode 100644 index 000000000..5c548268b --- /dev/null +++ b/tests/scheduler/test_scalarbf.py @@ -0,0 +1,57 @@ +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 BaseBasisFunction + + +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 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() From 9124d6603c113393dac78059e7cc38766f89ea15 Mon Sep 17 00:00:00 2001 From: Eric Neilsen Date: Wed, 25 Oct 2023 14:44:14 -0700 Subject: [PATCH 4/4] ROI tests --- .../basis_functions/basis_functions.py | 11 +++++ tests/scheduler/test_scalarbf.py | 12 +---- tests/scheduler/test_surveys.py | 46 +++++++++++++++++++ 3 files changed, 58 insertions(+), 11 deletions(-) diff --git a/rubin_sim/scheduler/basis_functions/basis_functions.py b/rubin_sim/scheduler/basis_functions/basis_functions.py index f1c5c09a0..de92df621 100644 --- a/rubin_sim/scheduler/basis_functions/basis_functions.py +++ b/rubin_sim/scheduler/basis_functions/basis_functions.py @@ -2,6 +2,7 @@ "BaseBasisFunction", "HealpixLimitedBasisFunctionMixin", "ConstantBasisFunction", + "SimpleArrayBasisFunction", "DelayStartBasisFunction", "TargetMapBasisFunction", "AvoidLongGapsBasisFunction", @@ -231,6 +232,16 @@ 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""" diff --git a/tests/scheduler/test_scalarbf.py b/tests/scheduler/test_scalarbf.py index 5c548268b..9585e34a7 100644 --- a/tests/scheduler/test_scalarbf.py +++ b/tests/scheduler/test_scalarbf.py @@ -4,17 +4,7 @@ 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 BaseBasisFunction - - -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 +from rubin_sim.scheduler.basis_functions import SimpleArrayBasisFunction class SimpleArrayAtHpixBasisFunction(HealpixLimitedBasisFunctionMixin, SimpleArrayBasisFunction): diff --git a/tests/scheduler/test_surveys.py b/tests/scheduler/test_surveys.py index 51ea8dcab..8eceeabbd 100644 --- a/tests/scheduler/test_surveys.py +++ b/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()