Tickets/sp 1619 (#114)

rubin_scheduler/scheduler/basis_functions/mask_basis_funcs.py

@@ -14,7 +14,7 @@
 
 from rubin_scheduler.scheduler.basis_functions import BaseBasisFunction
 from rubin_scheduler.scheduler.utils import HpInLsstFov, IntRounded
-from rubin_scheduler.utils import DEFAULT_NSIDE, _angular_separation
+from rubin_scheduler.utils import DEFAULT_NSIDE, _angular_separation, _hp_grow_mask
 
 
 class SolarElongMaskBasisFunction(BaseBasisFunction):
@@ -197,7 +197,7 @@ class AltAzShadowMaskBasisFunction(BaseBasisFunction):
         For sequences, try the length of the sequence + some buffer.
         Note that this just sets up the shadow *at* the shadow_minutes time
         (and not all times up to shadow_minutes).
-    altaz_limit_pad : `float`
+    pad : `float`
         The value by which to pad the telescope limits, to avoid
         healpix values mapping into pointings from the field tesselations
         which are actually out of bounds. This should typically be
@@ -212,18 +212,20 @@ def __init__(
         min_az=0,
         max_az=360,
         shadow_minutes=40.0,
-        altaz_limit_pad=2.0,
+        pad=3.0,
+        scale=1000,
     ):
         super().__init__(nside=nside)
         self.min_alt = np.radians(min_alt)
         self.max_alt = np.radians(max_alt)
         self.min_az = np.radians(min_az)
         self.max_az = np.radians(max_az)
         self.shadow_time = shadow_minutes / 60.0 / 24.0  # To days
-        self.altaz_limit_pad = np.radians(altaz_limit_pad)
+        self.pad = np.radians(pad)
 
         self.r_min_alt = IntRounded(self.min_alt)
         self.r_max_alt = IntRounded(self.max_alt)
+        self.scale = scale
 
     def _calc_value(self, conditions, indx=None):
         # Basis function value will be 0 except where masked (then np.nan)
@@ -307,15 +309,12 @@ def _calc_value(self, conditions, indx=None):
                 out_of_bounds = np.where((future_az - conditions.tel_az_limits[0]) % (two_pi) > az_range)[0]
                 result[out_of_bounds] = np.nan
 
-        # Grow the resulting mask by self.altaz_limit_pad (approximately),
-        # to avoid pushing field centers
-        # outside the boundaries of what is actually reachable.
-        if self.altaz_limit_pad > 0:
-            # Can't smooth a map with nan
-            map_to_smooth = np.where(np.isnan(result), hp.UNSEEN, 1)
-            map_back = hp.smoothing(map_to_smooth, fwhm=self.altaz_limit_pad * 2)
-            # Replace values and nans
-            result = np.where(map_back < 0.9, np.nan, result)
+        # Grow the resulting mask by self.pad, to avoid field centers
+        # falling outside the boundaries of what is actually reachable.
+        if self.pad > 0:
+            mask_indx = np.where(np.isnan(result))[0]
+            to_mask_indx = _hp_grow_mask(self.nside, tuple(mask_indx), scale=self.scale, grow_size=self.pad)
+            result[to_mask_indx] = np.nan
 
         return result
 

rubin_scheduler/utils/healpy_utils.py

@@ -7,13 +7,17 @@
     "_healbin",
     "moc2array",
     "hp_grow_argsort",
+    "_hp_grow_mask",
 )
 
 import warnings
+from functools import lru_cache
 
 import healpy as hp
 import numpy as np
 
+from .tree_utils import _build_tree, _xyz_from_ra_dec
+
 
 def _hpid2_ra_dec(nside, hpids, **kwargs):
     """
@@ -347,3 +351,62 @@ def hp_grow_argsort(in_map, ignore_nan=True):
         valid_neighbors_mask[(neighbors_of_current[sub_indx[0]], sub_indx[1])] = False
 
     return ordered_hp
+
+
+def _tree_for_mask(nside, scale=1000):
+    """Build a kdTree for HEALpixels"""
+    ra, dec = _hpid2_ra_dec(nside, np.arange(hp.nside2npix(nside)))
+    tree = _build_tree(ra, dec, scale=scale)
+    return tree
+
+
+@lru_cache(maxsize=10)
+def _hp_grow_mask(nside, masked_indx_tuple, grow_size=np.radians(2.0), scale=1000):
+    """Grow a HEALpix mask.
+
+    Would be nice if healpy.query_disc was vectorized, but here we are.
+
+    Parameters
+    ----------
+    nside : int
+        HEALpix nside to use
+    masked_indx_tuple : `tuple`
+        Indices that are currently masked. Needs to be tuple of ints
+        so cache can hash it.
+    scale : `int`
+        Scale passed to _build_tree to maintain cross-platform repeatability.
+        Default 1000.
+
+    Returns
+    -------
+    out_array : `np.array`
+        The input mask with any masked pixels grown the correct amount
+    """
+
+    # Set nside and kdTree as attributes so they will be cached for later
+    if not hasattr(_hp_grow_mask, "nside"):
+        _hp_grow_mask.nside = nside
+        _hp_grow_mask.tree = _tree_for_mask(nside, scale=scale)
+        _hp_grow_mask.scale = scale
+    if (_hp_grow_mask.nside != nside) | (_hp_grow_mask.scale != scale):
+        _hp_grow_mask.nside = nside
+        _hp_grow_mask.tree = _tree_for_mask(nside, scale=scale)
+        _hp_grow_mask.scale = scale
+
+    # Where are we masked
+    # Technically might be able to reach into the tree object and
+    # get the points that way and save an _hpid2_ra_dec call.
+    ra, dec = _hpid2_ra_dec(nside, np.arange(hp.nside2npix(nside))[list(masked_indx_tuple)])
+    x, y, z = _xyz_from_ra_dec(ra, dec)
+    if scale is not None:
+        x = np.round(x * scale).astype(int)
+        y = np.round(y * scale).astype(int)
+        z = np.round(z * scale).astype(int)
+        grow_size = np.round(grow_size * scale).astype(int)
+    # If there are lots of points, may want to use query_ball_tree
+    # instead for speed.
+    lists_of_neighbors = _hp_grow_mask.tree.query_ball_point(np.vstack([x, y, z]).T, grow_size)
+
+    u_indx = np.unique(np.concatenate(lists_of_neighbors))
+
+    return u_indx

tests/scheduler/test_basisfuncs.py

@@ -207,30 +207,30 @@ def test_AltAzShadowMask(self):
         conditions.tel_alt_limits = [np.radians(-100), np.radians(100)]
         # With no (real) limits, including no limits in conditions
         bf = basis_functions.AltAzShadowMaskBasisFunction(
-            nside=nside, min_alt=0, max_alt=90, min_az=0, max_az=360, shadow_minutes=0, altaz_limit_pad=0
+            nside=nside, min_alt=0, max_alt=90, min_az=0, max_az=360, shadow_minutes=0, pad=0
         )
         result = bf(conditions)
         self.assertTrue(np.all(np.isnan(result[np.where(conditions.alt < 0)])))
         self.assertTrue(np.all(result[np.where(conditions.alt >= 0)] == 0))
         # Set altitude limits but still no padding
         bf = basis_functions.AltAzShadowMaskBasisFunction(
-            nside=nside, min_alt=40, max_alt=60, min_az=0, max_az=360, shadow_minutes=0, altaz_limit_pad=0
+            nside=nside, min_alt=40, max_alt=60, min_az=0, max_az=360, shadow_minutes=0, pad=0
         )
         result = bf(conditions)
         good = np.where((conditions.alt > np.radians(40)) & (conditions.alt < np.radians(60)), True, False)
         self.assertTrue(np.all(np.isnan(result[~good])))
         self.assertTrue(np.all(result[good] == 0))
         # And set azimuth limits but no padding
         bf = basis_functions.AltAzShadowMaskBasisFunction(
-            nside=nside, min_alt=-90, max_alt=100, min_az=90, max_az=180, shadow_minutes=0, altaz_limit_pad=0
+            nside=nside, min_alt=-90, max_alt=100, min_az=90, max_az=180, shadow_minutes=0, pad=0
         )
         result = bf(conditions)
         good = np.where((conditions.az > np.radians(90)) & (conditions.az < np.radians(180)), True, False)
         self.assertTrue(np.all(np.isnan(result[~good])))
         self.assertTrue(np.all(result[good] == 0))
         # And set azimuth limits - order sensitive
         bf = basis_functions.AltAzShadowMaskBasisFunction(
-            nside=nside, min_alt=-90, max_alt=90, min_az=180, max_az=90, shadow_minutes=0, altaz_limit_pad=0
+            nside=nside, min_alt=-90, max_alt=90, min_az=180, max_az=90, shadow_minutes=0, pad=0
         )
         result = bf(conditions)
         good = np.where((conditions.az > np.radians(180)) | (conditions.az < np.radians(90)), True, False)
@@ -249,7 +249,7 @@ def test_AltAzShadowMask(self):
         # Set multiple altitude limits from the conditions
         conditions.sky_alt_limits = [[np.radians(40), np.radians(60)], [np.radians(80), np.radians(90)]]
         bf = basis_functions.AltAzShadowMaskBasisFunction(
-            nside=nside, min_alt=-90, max_alt=90, min_az=0, max_az=360, shadow_minutes=0, altaz_limit_pad=2
+            nside=nside, min_alt=-90, max_alt=90, min_az=0, max_az=360, shadow_minutes=0, pad=3
         )
         result = bf(conditions)
         # Conditions value get padded
@@ -282,7 +282,7 @@ def test_AltAzShadowMask(self):
         conditions.sky_alt_limits = None
         conditions.sky_az_limits = [[np.radians(90), np.radians(270)]]
         bf = basis_functions.AltAzShadowMaskBasisFunction(
-            nside=nside, min_alt=-90, max_alt=90, min_az=0, max_az=360, shadow_minutes=0, altaz_limit_pad=2
+            nside=nside, min_alt=-90, max_alt=90, min_az=0, max_az=360, shadow_minutes=0, pad=2
         )
         result = bf(conditions)
         # To more accurately count in azimuth, remove very high altitudes

tests/utils/test_healutils.py

@@ -77,6 +77,50 @@ def test_bin_deg(self):
         self.assertEqual(map3[hpid], 0.0)
         self.assertEqual(hp.maptype(map3), 0)
 
+    def test_mask_grow(self):
+        """Test we can grow a healpix mask map"""
+
+        nside = 32
+        nan_indx = tuple([0, 100])
+        scale = hp.nside2resol(nside)
+
+        to_mask = utils._hp_grow_mask(nside, nan_indx, grow_size=scale * 2)
+
+        # That should have made some things mask
+        assert 100 > np.size(to_mask) > 5
+
+        # Test another nside
+        nside = 128
+        nan_indx = tuple([0, 100])
+        scale = hp.nside2resol(nside)
+        to_mask = utils._hp_grow_mask(nside, nan_indx, grow_size=scale * 2)
+
+        # That should have made some things mask
+        assert 100 > np.size(to_mask) > 5
+
+        # Do a silly loop check to make sure things got masked properly
+        # Need to turn off the machine precision kwarg
+        nside = 128
+        nan_indx = tuple([0, 100])
+        scale = hp.nside2resol(nside) * 5
+        to_mask = utils._hp_grow_mask(nside, nan_indx, grow_size=scale, scale=None)
+
+        ra, dec = utils._hpid2_ra_dec(nside, np.arange(hp.nside2npix(nside)))
+        to_mask_set = set(to_mask)
+        all_close = []
+        for indx in nan_indx:
+            distances = utils._angular_separation(ra, dec, ra[indx], dec[indx])
+            close = np.where(distances <= scale)[0]
+            all_close.extend(close)
+            assert set(close).issubset(to_mask_set)
+        all_close = np.unique(all_close)
+        # Make sure we aren't including any extra pixels
+        assert np.size(all_close) == np.size(to_mask)
+
+        # Check that 0 distance doesn't mask anything new
+        to_mask = utils._hp_grow_mask(nside, nan_indx, grow_size=0)
+        assert np.size(to_mask) == np.size(nan_indx)
+
 
 if __name__ == "__main__":
     unittest.main()