diff --git a/pyproject.toml b/pyproject.toml index 41ce0b3..d7c9d8a 100644 --- a/pyproject.toml +++ b/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", diff --git a/src/mistsim/beam.py b/src/mistsim/beam.py index 6d335fc..4df488e 100644 --- a/src/mistsim/beam.py +++ b/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,6 +148,10 @@ 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, @@ -86,4 +159,5 @@ def __init__( horizon=horizon, beam_rot=beam_rot, beam_tilt=beam_tilt, + horizon_frame=horizon_frame, ) diff --git a/tests/test_beam.py b/tests/test_beam.py index adfc7ba..246e9de 100644 --- a/tests/test_beam.py +++ b/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] diff --git a/uv.lock b/uv.lock index 8e844b1..921f8d4 100644 --- a/uv.lock +++ b/uv.lock @@ -562,7 +562,7 @@ toml = [ [[package]] name = "croissant-sim" version = "5.2.1" -source = { git = "https://github.com/christianhbye/croissant.git?rev=v5.3.0.dev3#58bf56e8af77a5a07fd0669ec6ba38900549e50c" } +source = { git = "https://github.com/christianhbye/croissant.git?rev=v5.3.0.dev4#a83a9864162b64466a85654fc1a4a2072f92e743" } dependencies = [ { name = "astropy" }, { name = "equinox" }, @@ -2030,7 +2030,7 @@ dev = [ [package.metadata] requires-dist = [ - { name = "croissant-sim", git = "https://github.com/christianhbye/croissant.git?rev=v5.3.0.dev3" }, + { name = "croissant-sim", git = "https://github.com/christianhbye/croissant.git?rev=v5.3.0.dev4" }, { name = "equinox" }, { name = "jax" }, { name = "numpy" },