OPSIM-1111: feature - Add altitude and azimuth limits to Conditions, along with new basis function + shadow mask

rubin_scheduler/scheduler/basis_functions/mask_basis_funcs.py

@@ -9,8 +9,11 @@
     "MaskAzimuthBasisFunction",
     "SolarElongationMaskBasisFunction",
     "AreaCheckMaskBasisFunction",
+    "AltAzShadowMaskBasisFunction",
 )
 
+import warnings
+
 import healpy as hp
 import matplotlib.pylab as plt
 import numpy as np
@@ -194,6 +197,90 @@
         return result
 
 
+class AltAzShadowMaskBasisFunction(BaseBasisFunction):
+    """Mask out range altitudes and azimuths, then extend the
+    mask so if observations are taken in pairs, the second in the pair will
+    not have moved into a masked region.
+
+    Masks any alt/az regions as specified by the conditions object, then
+    applies any additional altitude masking as suppied by the kwargs.
+    This mask is then extended using `shadow minutes`.
+
+    Parameters
+    ----------
+    nside : `int`
+        HEALpix nside. Default None will look up the package-wide default.
+    min_alt : `float`
+        Minimum altitude to apply to the mask. Default 20 (degrees).
+    max_alt : `float`
+        Maximum altitude to allow. Default 82 (degrees).
+    shadow_minutes : `float`
+        How long to extend masked area in longitude. Default 40 (minutes).
+    """
+
+    def __init__(
+        self,
+        nside=None,
+        min_alt=20.0,
+        max_alt=82.0,
+        shadow_minutes=40.0,
+    ):
+        super(AltAzShadowMaskBasisFunction, self).__init__(nside=nside)
+        self.min_alt = np.radians(min_alt)
+        self.max_alt = np.radians(max_alt)
+        self.shadow_time = shadow_minutes / 60.0 / 24.0  # To days
+
+    def _calc_value(self, conditions, indx=None):
+        # Mask everything to start
+        result = np.zeros(hp.nside2npix(self.nside), dtype=float) + np.nan
+
+        in_range_alt = np.zeros(hp.nside2npix(self.nside), dtype=int)
+        in_range_az = np.zeros(hp.nside2npix(self.nside), dtype=int)
+
+        # Compute the alt,az values in the future. Use the conditions object
+        # so the results are cached and can be used by other surveys is needed.
+        # Technically this could fail if the masked region is very narrow or shadow time
+        # is very large.
+        future_alt, future_az = conditions.future_alt_az(np.max(conditions.mjd + self.shadow_time))
+
+        # apply limits from the conditions object
+        for limits in conditions.tel_alt_limits:
+            good = np.where(
+                (IntRounded(conditions.alt) >= IntRounded(np.min(limits)))
+                & (IntRounded(conditions.alt) <= IntRounded(np.max(limits)))
+            )[0]
+            in_range_alt[good] += 1
+            good = np.where(
+                (IntRounded(future_alt) >= IntRounded(np.min(limits)))
+                & (IntRounded(future_alt) <= IntRounded(np.max(limits)))
+            )[0]
+            in_range_alt[good] += 1
+
+        for limits in conditions.tel_az_limits:
+            good = np.where(
+                (IntRounded(conditions.az) >= IntRounded(np.min(limits)))
+                & (IntRounded(conditions.az) <= IntRounded(np.max(limits)))
+            )[0]
+            in_range_az[good] += 1
+            good = np.where(
+                (IntRounded(future_az) >= IntRounded(np.min(limits)))
+                & (IntRounded(future_az) <= IntRounded(np.max(limits)))
+            )[0]
+            in_range_az[good] += 1
+
+        passed_all = np.where((in_range_alt > 1) & (in_range_az > 1))[0]
+        result[passed_all] = 0
+
+        # Apply additional alt constraint in case we want to be more conservative than the limit
+        result[np.where(IntRounded(conditions.alt) < IntRounded(self.min_alt))] = np.nan
+        result[np.where(IntRounded(conditions.alt) > IntRounded(self.max_alt))] = np.nan
+
+        result[np.where(IntRounded(future_alt) < IntRounded(self.min_alt))] = np.nan
+        result[np.where(IntRounded(future_alt) > IntRounded(self.max_alt))] = np.nan
+
+        return result
+
+
 class ZenithShadowMaskBasisFunction(BaseBasisFunction):
     """Mask the zenith, and things that will soon pass near zenith. Useful for making sure
     observations will not be too close to zenith when they need to be observed again (e.g. for a pair).
@@ -217,6 +304,11 @@
         penalty=np.nan,
         site="LSST",
     ):
+        warnings.warn(
+            "Deprecating ZenithShadowMaskBasisFunction in favor of AltAzShadowMaskBasisFunction.",
+            DeprecationWarning,
+        )
+
         super(ZenithShadowMaskBasisFunction, self).__init__(nside=nside)
         self.update_on_newobs = False
 

rubin_scheduler/scheduler/example/example_scheduler.py

@@ -402,12 +402,10 @@ def blob_for_long(
         # Masks, give these 0 weight
         bfs.append(
             (
-                bf.ZenithShadowMaskBasisFunction(
+                bf.AltAzShadowMaskBasisFunction(
                     nside=nside,
                     shadow_minutes=shadow_minutes,
                     max_alt=max_alt,
-                    penalty=np.nan,
-                    site="LSST",
                 ),
                 0.0,
             )
@@ -606,7 +604,7 @@ def gen_greedy_surveys(
         # Masks, give these 0 weight
         bfs.append(
             (
-                bf.ZenithShadowMaskBasisFunction(nside=nside, shadow_minutes=shadow_minutes, max_alt=max_alt),
+                bf.AltAzShadowMaskBasisFunction(nside=nside, shadow_minutes=shadow_minutes, max_alt=max_alt),
                 0,
             )
         )
@@ -824,12 +822,10 @@ def generate_blobs(
         # Masks, give these 0 weight
         bfs.append(
             (
-                bf.ZenithShadowMaskBasisFunction(
+                bf.AltAzShadowMaskBasisFunction(
                     nside=nside,
                     shadow_minutes=shadow_minutes,
                     max_alt=max_alt,
-                    penalty=np.nan,
-                    site="LSST",
                 ),
                 0.0,
             )
@@ -1033,12 +1029,10 @@ def generate_twi_blobs(
         # Masks, give these 0 weight
         bfs.append(
             (
-                bf.ZenithShadowMaskBasisFunction(
+                bf.AltAzShadowMaskBasisFunction(
                     nside=nside,
                     shadow_minutes=shadow_minutes,
                     max_alt=max_alt,
-                    penalty=np.nan,
-                    site="LSST",
                 ),
                 0.0,
             )
@@ -1232,7 +1226,7 @@ def generate_twilight_near_sun(
         bfs.append((bf.NearSunTwilightBasisFunction(nside=nside, max_airmass=max_airmass), 0))
         bfs.append(
             (
-                bf.ZenithShadowMaskBasisFunction(nside=nside, shadow_minutes=shadow_minutes, max_alt=max_alt),
+                bf.AltAzShadowMaskBasisFunction(nside=nside, shadow_minutes=shadow_minutes, max_alt=max_alt),
                 0,
             )
         )

rubin_scheduler/scheduler/features/conditions.py

@@ -1,5 +1,6 @@
 __all__ = ("Conditions",)
 
+import functools
 import warnings
 from io import StringIO
 
@@ -169,6 +170,12 @@
             Dictionary of planet name and coordinate e.g., 'venus_RA', 'mars_dec'
         scheduled_observations : np.array
             A list of MJD times when there are scheduled observations. Defaults to empty array.
+        tel_az_limits : list of float pairs
+            A list of lists giving valid azimuth ranges. e.g., [0, 2*np.pi] would
+            mean all azimuth values are valid, while [[0, np.pi/2], [3*np.pi/2, 2*np.pi]]
+            would mean anywhere in the south is invalid.  Radians.
+        tel_alt_limits : list of float pairs
+            A list of lists giving valid altitude ranges. Radians.
 
         Attributes (calculated on demand and cached)
         ------------------------------------------
@@ -283,6 +290,10 @@
         self.tel_az = None
         self.cumulative_azimuth_rad = None
 
+        # Telescope limits
+        self.tel_az_limits = None
+        self.tel_alt_limits = None
+
         # Full sky cloud map
         self._cloud_map = None
         self._HA = None
@@ -387,6 +398,26 @@
             self._mjd,
         )
 
+    @functools.lru_cache(maxsize=10)
+    def future_alt_az(self, mjd):
+        """Compute the altitude and azimuth for a future time.
+
+        Returns
+        -------
+        altitude : `np.array`
+            The altutude of each healpix at MJD (radians)
+        azimuth : : `np.array`
+            The azimuth of each healpix at MJD (radians)
+        """
+        alt, az = _approx_ra_dec2_alt_az(
+            self.ra,
+            self.dec,
+            self.site.latitude_rad,
+            self.site.longitude_rad,
+            mjd,
+        )
+        return alt, az
+
     @property
     def mjd(self):
         return self._mjd

rubin_scheduler/scheduler/model_observatory/kinem_model.py

@@ -201,6 +201,8 @@
         azimuth_accel=7.0,
         azimuth_decel=7.0,
         settle_time=3.0,
+        az_limits=None,
+        alt_limits=None,
     ):
         """Parameters to define the TELESCOPE movement and position.
 
@@ -240,6 +242,14 @@
         self.telaz_accel_rad = np.radians(azimuth_accel)
         self.telaz_decel_rad = np.radians(azimuth_decel)
         self.mount_settletime = settle_time
+        if alt_limits is None:
+            self.alt_limits = [[self.telalt_minpos_rad, self.telalt_maxpos_rad]]
+        else:
+            self.alt_limits = np.radians(alt_limits)
+        if az_limits is None:
+            self.az_limits = [[0, 2.0 * np.pi]]
+        else:
+            self.az_limits = np.radians(az_limits)
 
     def setup_optics(self, ol_slope=1.0 / 3.5, cl_delay=[0.0, 36.0], cl_altlimit=[0.0, 9.0, 90.0]):
         """

rubin_scheduler/scheduler/model_observatory/model_observatory.py

@@ -366,6 +366,10 @@
 
         self.conditions.mjd_start = self.mjd_start
 
+        # Telescope limits
+        self.conditions.tel_az_limits = self.observatory.az_limits
+        self.conditions.tel_alt_limits = self.observatory.alt_limits
+
         # Planet positions from almanac
         self.conditions.planet_positions = self.almanac.get_planet_positions(self.mjd)
 

rubin_scheduler/scheduler/surveys/scripted_surveys.py

@@ -193,6 +193,22 @@
             & ((HA > observation["HA_max"]) | (HA < observation["HA_min"]))
             & (conditions.sun_alt < observation["sun_alt_max"])
         )[0]
+
+        # Also check the alt,az limits given by the conditions object
+        count = in_range * 0
+        for limits in conditions.tel_alt_limits:
+            ir = np.where((alt[in_range] >= np.min(limits)) & (alt[in_range] <= np.max(limits)))[0]
+            count[ir] += 1
+        good = np.where(count > 0)[0]
+        in_range = in_range[good]
+
+        count = in_range * 0
+        for limits in conditions.tel_az_limits:
+            ir = np.where((az[in_range] >= np.min(limits)) & (az[in_range] <= np.max(limits)))[0]
+            count[ir] += 1
+        good = np.where(count > 0)[0]
+        in_range = in_range[good]
+
         return in_range
 
     def _check_list(self, conditions):

tests/scheduler/test_baseline.py

@@ -81,7 +81,7 @@ def gen_greedy_surveys(nside):
         bfs.append(bf.StrictFilterBasisFunction(filtername=filtername))
         # Masks, give these 0 weight
         bfs.append(bf.AvoidDirectWind(nside=nside))
-        bfs.append(bf.ZenithShadowMaskBasisFunction(nside=nside, shadow_minutes=60.0, max_alt=76.0))
+        bfs.append(bf.AltAzShadowMaskBasisFunction(nside=nside, shadow_minutes=60.0, max_alt=76.0))
         bfs.append(bf.MoonAvoidanceBasisFunction(nside=nside, moon_distance=30.0))
         bfs.append(bf.CloudedOutBasisFunction())
 
@@ -144,7 +144,7 @@ def gen_blob_surveys(nside):
         bfs.append(bf.StrictFilterBasisFunction(filtername=filtername))
         # Masks, give these 0 weight
         bfs.append(bf.AvoidDirectWind(nside=nside))
-        bfs.append(bf.ZenithShadowMaskBasisFunction(nside=nside, shadow_minutes=60.0, max_alt=76.0))
+        bfs.append(bf.AltAzShadowMaskBasisFunction(nside=nside, shadow_minutes=60.0, max_alt=76.0))
         bfs.append(bf.MoonAvoidanceBasisFunction(nside=nside, moon_distance=30.0))
         bfs.append(bf.CloudedOutBasisFunction())
         # feasibility basis fucntions. Also give zero weight.

tests/scheduler/test_utils.py

@@ -6,7 +6,7 @@
 from rubin_scheduler.data import get_data_dir
 from rubin_scheduler.scheduler import sim_runner
 from rubin_scheduler.scheduler.example import example_scheduler, run_sched
-from rubin_scheduler.scheduler.model_observatory import ModelObservatory
+from rubin_scheduler.scheduler.model_observatory import KinemModel, ModelObservatory
 from rubin_scheduler.scheduler.utils import restore_scheduler, run_info_table, season_calc
 from rubin_scheduler.utils import survey_start_mjd
 
@@ -30,7 +30,7 @@ def test_example(self):
         """Test the example scheduler executes all the expected surveys"""
         mjd_start = survey_start_mjd()
         scheduler = example_scheduler(mjd_start=mjd_start)
-        observatory, scheduler, observations = run_sched(scheduler, mjd_start=mjd_start, survey_length=5)
+        observatory, scheduler, observations = run_sched(scheduler, mjd_start=mjd_start, survey_length=7)
         u_notes = np.unique(observations["note"])
 
         # Note that some of these may change and need to be updated if survey
@@ -84,6 +84,47 @@ def test_start_of_night_example(self):
         os.path.isfile(os.path.join(get_data_dir(), "scheduler/dust_maps/dust_nside_32.npz")),
         "Test data not available.",
     )
+    def test_altaz_limit(self):
+        """Test that setting some azimuth limits via the kinematic model works"""
+        mjd_start = survey_start_mjd()
+        scheduler = example_scheduler(mjd_start=mjd_start)
+        km = KinemModel(mjd0=mjd_start)
+        km.setup_telescope(az_limits=[[0.0, 90.0], [270.0, 360.0]])
+        mo = ModelObservatory(mjd_start=mjd_start, kinem_model=km)
+
+        mo, scheduler, observations = sim_runner(
+            mo,
+            scheduler,
+            survey_length=3.0,
+            verbose=False,
+            filename=None,
+        )
+
+        az = np.degrees(observations["az"])
+        forbidden = np.where((az > 90) & (az < 270))[0]
+        # Let a few pairs try to complete since by default we don't use an agressive shadow_minutes
+        n_forbidden = np.size(
+            [obs for obs in observations[forbidden]["note"] if (("pair_33" not in obs) | (", b" not in obs))]
+        )
+
+        assert n_forbidden == 0
+
+        km = KinemModel(mjd0=mjd_start)
+        km.setup_telescope(alt_limits=[[40.0, 70.0]])
+        mo = ModelObservatory(mjd_start=mjd_start, kinem_model=km)
+
+        mo, scheduler, observations = sim_runner(
+            mo,
+            scheduler,
+            survey_length=3.0,
+            verbose=False,
+            filename=None,
+        )
+        alt = np.degrees(observations["alt"])
+        n_forbidden = np.size(np.where((alt > 70) & (alt < 40))[0])
+
+        assert n_forbidden == 0
+
     def test_restore(self):
         """Test we can restore a scheduler properly"""
         # MJD set so it's in test data range