deps: bump croissant-sim to v5.3.0.dev4

pyproject.toml

@@ -17,7 +17,7 @@ classifiers = [
   "Topic :: Scientific/Engineering :: Astronomy",
 ]
 dependencies = [
-  "croissant-sim @ git+https://github.com/christianhbye/croissant.git@v5.3.0.dev3",
+  "croissant-sim @ git+https://github.com/christianhbye/croissant.git@v5.3.0.dev4",
   "equinox",
   "jax",
   "numpy",

src/mistsim/beam.py

@@ -1,6 +1,35 @@
 import warnings
 
 import croissant as cro
+import jax.numpy as jnp
+from croissant.horizon import _horizon_in_beam_frame
+
+# mistsim's ground grid is the beam grid as it is at beam_az_rot = 0
+# (phi = 0 North, phi = 90 deg West, compass azimuth A = -phi). Croissant's
+# ground grid has phi = 0 East (A = 90 - phi). The same direction therefore
+# sits 90 deg further round on croissant's grid.
+_GROUND_GRID_OFFSET_DEG = 90.0
+
+
+def _to_croissant_ground_grid(horizon, sampling, spatial_shape):
+    """Move a mask from mistsim's ground grid onto croissant's.
+
+    Croissant's own periodic shift does the work: it is exact when 90 deg
+    is a whole number of grid columns (the 1-deg MWSS grid, and every
+    HEALPix ring) and linear interpolation otherwise. Scalars and
+    theta-only masks come back unchanged.
+    """
+    nside = None
+    if sampling == "healpix":
+        nside = int(round((spatial_shape[0] / 12) ** 0.5))
+    return _horizon_in_beam_frame(
+        jnp.asarray(horizon),
+        "topocentric",
+        _GROUND_GRID_OFFSET_DEG,
+        sampling,
+        spatial_shape,
+        nside,
+    )
 
 
 class Beam(cro.Beam):
@@ -15,12 +44,19 @@ def __init__(
         beam_az_rot=0.0,
         beam_tilt=0.0,
         lmax=None,
+        horizon_frame="topocentric",
     ):
         """
         Beam pattern object. Holds the beam pattern in local antenna
         coordinates and associated metadata. The beam must be defined
         on the grid specified by the `sampling` scheme.
 
+        Theta is colatitude from zenith and phi is right-handed about
+        the zenith, from the beam's x axis towards its y axis. The x
+        axis points to compass azimuth `beam_az_rot`, so a beam-grid
+        direction has compass azimuth ``A = beam_az_rot - degrees(phi)``
+        (mod 360): with ``beam_az_rot = 0``, phi = 90 deg is West.
+
         Note that the `lmax` parameter is no longer used. The `lmax` is
         automatically determined from the shape of the input data and
         the sampling scheme. To change the `lmax` the simulation runs
@@ -42,13 +78,19 @@ def __init__(
             "mwss", which is a 1 deg equiangular sampling in theta and
             phi and includes the poles.
         horizon : array_like or None
-            The horizon mask: a boolean array specified for each
-            (theta, phi) direction (or pixel), with the same shape as
-            the last two (one for healpix) axes of data. It is an array
-            with True values for directions that are above the horizon
-            and False for directions that are below the horizon.
-            If None, it is assumed that the horizon is at
-            theta = 90 degrees.
+            Visible fractions in [0, 1] for each (theta, phi) direction
+            (or pixel), broadcastable to the spatial axes of data. Zero
+            blocks a sample, one keeps it, and fractional values weight
+            partially visible cells; boolean masks are accepted. With the
+            default ``horizon_frame="topocentric"`` the mask is fixed to
+            the ground. It is given on the beam grid as it is at
+            ``beam_az_rot = 0`` (phi = 0 North, phi = 90 deg West), so a
+            direction at compass azimuth ``A`` sits at ``phi = -A``
+            (mod 360 deg), and it stays there whatever `beam_az_rot` is.
+            If None, the horizon is at theta = 90 degrees with fractional
+            boundary cells (croissant's default). For a horizon given as
+            a function of azimuth, ``croissant.horizon_weights`` builds
+            fractional boundary weights on regular grids.
         beam_az_rot : float
             Azimuthal rotation of the beam in degrees. The rotation is
             defined in the astronomy convention, i.e., the angle
@@ -62,11 +104,38 @@ def __init__(
         lmax : int or None
             Removed. Will be ignored if provided and raise a
             FutureWarning.
+        horizon_frame : {"topocentric", "beam"}
+            Which frame `horizon` is fixed to. **Use the default,
+            ``"topocentric"``.** A horizon mask describes what blocks the
+            sky from where the antenna stands: terrain, buildings, the
+            ground itself. All of these are fixed to the ground, so the
+            mask must not turn when the beam does. Structures attached
+            to the antenna are not a horizon mask: they belong in the
+            beam pattern itself, from the EM simulation.
+
+            ``"beam"`` applies the mask on the rotated beam grid, so it
+            turns with `beam_az_rot`. It exists only for masks a caller
+            has already counter-rotated into the beam frame by hand (the
+            only correct way to handle terrain before this option
+            existed). It is the same as ``"topocentric"`` when
+            ``beam_az_rot = 0`` and for theta-only masks.
+
+            mistsim moves a topocentric mask onto croissant's ground grid
+            (phi = 0 East) before passing it on; croissant then rotates
+            it into the beam frame by periodic linear interpolation in
+            phi. Both steps are exact when the shifts are whole grid
+            columns (e.g. the 1-deg MWSS grid with whole-degree
+            `beam_az_rot`) and soften sharp edges otherwise. The `horizon`
+            attribute holds croissant's ground-grid weights;
+            ``horizon_in_beam_frame`` gives the weights actually applied.
 
         Raises
         ------
         FutureWarning
             If `lmax` is not None.
+        ValueError
+            If `horizon_frame` is not "beam" or "topocentric" (raised
+            by croissant).
 
         """
         if lmax is not None:
@@ -79,11 +148,16 @@ def __init__(
             )
         # croissant expects X-axis along East
         beam_rot = beam_az_rot - 90
+        if horizon is not None and horizon_frame == "topocentric":
+            horizon = _to_croissant_ground_grid(
+                horizon, sampling, jnp.shape(data)[1:]
+            )
         super().__init__(
             data,
             freqs,
             sampling=sampling,
             horizon=horizon,
             beam_rot=beam_rot,
             beam_tilt=beam_tilt,
+            horizon_frame=horizon_frame,
         )

tests/test_beam.py

@@ -57,3 +57,173 @@ def test_beam_az_rot_matches_croissant():
         np.array(cro_beam.compute_alm()),
         atol=1e-12,
     )
+
+
+def _mwss_grid():
+    """Degrees of colatitude and longitude on the 1-deg MWSS grid."""
+    theta = np.linspace(0.0, 180.0, 181)
+    phi = np.arange(360.0)
+    return theta, phi
+
+
+def _ground_sector_mask(lo=80.0, hi=100.0, theta_h=60.0):
+    """Ground mask blocking compass azimuth [lo, hi] below theta_h.
+
+    mistsim's ground grid is the beam grid at beam_az_rot = 0, so
+    compass azimuth A sits at phi = -A.
+    """
+    theta, phi = _mwss_grid()
+    azimuth = np.mod(-phi, 360.0)
+    in_sector = (azimuth >= lo) & (azimuth <= hi)
+    return np.where((theta[:, None] > theta_h) & in_sector[None, :], 0.0, 1.0)
+
+
+def _asymmetric_lobe():
+    """A beam with a lobe along its x axis (phi = 0)."""
+    theta, phi = _mwss_grid()
+    cos_phi = np.cos(np.deg2rad(phi))[None, :]
+    sin_theta = np.sin(np.deg2rad(theta))[:, None]
+    return jnp.asarray((1.0 + 0.8 * cos_phi * sin_theta)[None])
+
+
+def test_horizon_frame_defaults_to_topocentric():
+    data = jnp.ones((1, 181, 360))
+    beam = Beam(data, jnp.array([50.0]))
+    assert beam.horizon_frame == "topocentric"
+
+
+def test_horizon_frame_rejects_unknown_value():
+    data = jnp.ones((1, 181, 360))
+    with pytest.raises(ValueError, match="horizon_frame"):
+        Beam(data, jnp.array([50.0]), horizon_frame="ground")
+
+
+@pytest.mark.parametrize("beam_az_rot", [0.0, 40.0, 233.0])
+def test_default_mask_stays_on_the_ground(beam_az_rot):
+    """A ground-fixed sector lands at the same compass azimuth.
+
+    A beam-grid column phi_b points to compass azimuth
+    A = beam_az_rot - phi_b, which is ground column -A, so the weights
+    applied in the beam frame must equal the ground mask read there.
+    """
+    _, phi = _mwss_grid()
+    mask = _ground_sector_mask()
+    beam = Beam(
+        jnp.ones((1, 181, 360)),
+        jnp.array([50.0]),
+        horizon=jnp.asarray(mask),
+        beam_az_rot=beam_az_rot,
+    )
+    applied = np.asarray(beam.horizon_in_beam_frame)
+    azimuth = np.mod(beam_az_rot - phi, 360.0)
+    ground_col = np.mod(-azimuth, 360.0).astype(int)
+    np.testing.assert_array_equal(applied, mask[:, ground_col])
+    # the blocked columns are the East sector, whatever the rotation
+    blocked = azimuth[(applied == 0).any(axis=0)]
+    assert blocked.min() >= 80.0 and blocked.max() <= 100.0
+    assert blocked.size == 21
+
+
+def test_beam_frame_mask_rotates_with_the_beam():
+    """horizon_frame="beam" applies the mask as given, at any rotation."""
+    mask = _ground_sector_mask()
+    for beam_az_rot in (0.0, 40.0):
+        beam = Beam(
+            jnp.ones((1, 181, 360)),
+            jnp.array([50.0]),
+            horizon=jnp.asarray(mask),
+            beam_az_rot=beam_az_rot,
+            horizon_frame="beam",
+        )
+        np.testing.assert_array_equal(
+            np.asarray(beam.horizon_in_beam_frame), mask
+        )
+
+
+def test_frames_agree_at_zero_rotation_mwss():
+    """At beam_az_rot = 0 the new default changes nothing."""
+    data = _asymmetric_lobe()
+    mask = jnp.asarray(_ground_sector_mask())
+    topo = Beam(data, jnp.array([50.0]), horizon=mask)
+    old = Beam(data, jnp.array([50.0]), horizon=mask, horizon_frame="beam")
+    np.testing.assert_array_equal(
+        np.asarray(topo.horizon_in_beam_frame),
+        np.asarray(old.horizon_in_beam_frame),
+    )
+    np.testing.assert_array_equal(
+        np.asarray(topo.compute_fgnd()), np.asarray(old.compute_fgnd())
+    )
+
+
+def test_frames_agree_at_zero_rotation_healpix():
+    """The 90-deg grid shift is exact on every HEALPix ring."""
+    nside = 8
+    npix = 12 * nside**2
+    rng = np.random.default_rng(0)
+    mask = jnp.asarray(rng.uniform(size=npix))
+    data = jnp.ones((1, npix))
+    topo = Beam(data, jnp.array([50.0]), sampling="healpix", horizon=mask)
+    old = Beam(
+        data,
+        jnp.array([50.0]),
+        sampling="healpix",
+        horizon=mask,
+        horizon_frame="beam",
+    )
+    np.testing.assert_allclose(
+        np.asarray(topo.horizon_in_beam_frame),
+        np.asarray(old.horizon_in_beam_frame),
+        rtol=0,
+        atol=1e-15,
+    )
+
+
+@pytest.mark.parametrize("beam_az_rot", [0.0, 40.0])
+def test_theta_only_mask_is_frame_independent(beam_az_rot):
+    """The pipeline's theta-only masks behave the same in both frames."""
+    theta, _ = _mwss_grid()
+    mask = jnp.asarray((theta <= 80.0)[:, None].astype(float))
+    kw = dict(horizon=mask, beam_az_rot=beam_az_rot)
+    topo = Beam(_asymmetric_lobe(), jnp.array([50.0]), **kw)
+    old = Beam(
+        _asymmetric_lobe(), jnp.array([50.0]), horizon_frame="beam", **kw
+    )
+    np.testing.assert_array_equal(
+        np.asarray(topo.compute_fgnd()), np.asarray(old.compute_fgnd())
+    )
+
+
+def test_topocentric_fgnd_matches_manual_counter_rotation():
+    """An asymmetric beam sees the ground mask a caller would build.
+
+    With a lobe on the beam's x axis, the blocked fraction depends on
+    where the lobe points relative to the ground-fixed sector. The
+    default (topocentric) result must equal the beam-frame result with
+    the mask counter-rotated by hand.
+    """
+    _, phi = _mwss_grid()
+    data = _asymmetric_lobe()
+    mask = _ground_sector_mask()
+    fgnd = []
+    for beam_az_rot in (0.0, 90.0):
+        azimuth = np.mod(beam_az_rot - phi, 360.0)
+        manual = mask[:, np.mod(-azimuth, 360.0).astype(int)]
+        topo = Beam(
+            data,
+            jnp.array([50.0]),
+            horizon=jnp.asarray(mask),
+            beam_az_rot=beam_az_rot,
+        )
+        by_hand = Beam(
+            data,
+            jnp.array([50.0]),
+            horizon=jnp.asarray(manual),
+            beam_az_rot=beam_az_rot,
+            horizon_frame="beam",
+        )
+        f_topo = float(np.asarray(topo.compute_fgnd()).ravel()[0])
+        f_hand = float(np.asarray(by_hand.compute_fgnd()).ravel()[0])
+        np.testing.assert_allclose(f_topo, f_hand, rtol=1e-12)
+        fgnd.append(f_topo)
+    # the lobe points East at beam_az_rot = 90, into the sector
+    assert fgnd[1] > fgnd[0]