Visibility Projectors

from serval.containers import SkyModel, TIRSPowerBeam
from serval.core import PowerBeamVisProjector, SingleVisibilityGenerator, SkyVisProjector
from serval import sht
from serval.utils import (
    plane_wave_mag_from_bandlimit,
    harmonic_point_source,
    power_law_sky,
)
import numpy as np
import matplotlib.pyplot as plt
from astropy import units
import matplotlib as mpl

plt.style.use("./example_style.mplstyle")
mpl.rcParams["text.usetex"] = False
_ = np.seterr(divide="ignore")

This notebook shows a more typical example of using Serval to produce projectors for batched m-modes and visibilities for a single baseline over multiple frequency channels. We'll look at two methods, a beam projector and a sky projector. The former fixes the sky but allows quick batch simulations with arbitrary beams. The latter fixes the beams but allows quick batch simulations with arbitrary skies.

Set up simulation settings

We set up the simulation in a similar way to that shown in the "Simple Usage" notebook.

beam_lmax = 40  # Maximum harmonic degree for power primary beam
baseline_lmax = 50  # Maximum harmonic degree for baseline term
sky_lmax = beam_lmax + baseline_lmax  # Resulting maximum harmonic degree for sky term

# Telescope position in TIRS lat, lon in radians.
telescope_latitude = np.radians(-30.0)
telescope_longitude = np.radians(21.0)
sky_epoch = "J2026.0"

But now we work on limited ranges of sky \(m\) values over multiple frequency channels.

sky_absm_limits = (0, 10)

nchannels = 64
frequencies_MHz = np.linspace(400, 800, nchannels)

Simulated Sky Containers

Sky models for a power-law and point-source sky are generated over multiple frequency channels, here just assuming a power-law scaling with frequency for both. The containers keep the shared frequency axis and CIRS metadata attached to the harmonic coefficients.

frequency_scaling = (frequencies_MHz / 600) ** -1

pl_sky_grid = power_law_sky(sky_lmax, power_law_index=2, seed=123)
single_pl_sky_alm = sht.analysis(pl_sky_grid, return_sparse=False)
pl_sky = SkyModel(
    lmax=sky_lmax,
    frequencies_MHz=frequencies_MHz,
    alm=frequency_scaling[:, None, None] * single_pl_sky_alm[None, ...],
    map_unit="K",
    polarisation="I",
    coordinate_basis="CIRS",
    epoch=sky_epoch,
)

single_ps_sky_alm = harmonic_point_source(180, 0, sky_lmax, return_sparse=False)
ps_sky = SkyModel(
    lmax=sky_lmax,
    frequencies_MHz=frequencies_MHz,
    alm=frequency_scaling[:, None, None] * single_ps_sky_alm[None, ...],
    map_unit="K",
    polarisation="I",
    coordinate_basis="CIRS",
    epoch=sky_epoch,
)

Power Primary Beam Container

Gaussian primary beams whose FWHM goes as \(1/ u\) are now generated and wrapped in a TIRSPowerBeam container. The projectors still receive raw coefficients later, but the beam stays container-backed during construction and plotting.

# D_eff gives FWHM = lambda/D_eff ≈ 20° at 600 MHz, scaling correctly with frequency
D_eff = 0.5 / np.radians(20)
power_beam = TIRSPowerBeam.from_gaussian(
    D_eff=D_eff,
    lmax=beam_lmax,
    frequencies_MHz=frequencies_MHz,
    pol_product="XX",
    sky_pol="I",
    latitude=telescope_latitude,
    longitude=telescope_longitude,
    altitude=np.pi / 2,
    azimuth=0,
)

Baseline Term

A baseline with equal E and N projection and fits within the imposed band-limits is generated.

bs_pw_mag = plane_wave_mag_from_bandlimit(baseline_lmax) / 2 / np.pi
baseline_length = (frequencies_MHz.max() * units.MHz).to(
    "m", equivalencies=units.spectral()
) * bs_pw_mag  # Compute band-limited baseline length
baseline_unit_enu = np.array(
    [1.0, 1.0, 0]
)  # Baseline direction in East-North-Up coordinates.
baseline_unit_enu /= np.linalg.norm(baseline_unit_enu)
baseline_enu = tuple(baseline_unit_enu * baseline_length.value)  # Rescale to length
print(baseline_enu)

(np.float64(0.9586757156054164), np.float64(0.9586757156054164), np.float64(0.0))

Power Beam Visibility Projector

We now generate a visibility projector that fixes a sky but allows quick \(m\)-mode and visibility generation as we vary power primary beams. The projector API is still array-based, so we pass SkyModel.as_coeff() and TIRSPowerBeam.as_coeff() at the boundary.

beam_vis_proj = PowerBeamVisProjector(
    latitude=telescope_latitude,
    longitude=telescope_longitude,
    baseline_enu = baseline_enu,
    power_beam_lmax = beam_lmax,
    frequencies_MHz = frequencies_MHz,
    generate_baseline_cache_on_init=True,
    generate_gaunt_cache_on_init=True,
    sky_absm_limits = sky_absm_limits,
)

Computing 0.01G Gaunts (0.33 GB) took: 0.03s.


/home/devin/Dropbox/arcturus/Projects/radio/hirax/systematics/dev/serval/src/serval/core.py:980: UserWarning: The upper limit for the mmodes of the sky does not correspond to the optimal value given the baseline and beam which is 80. Consider changing it.
  self.set_bandlimits_from_baseline()

Now that the object is setup, we proceed by * Setting up a zarr store for the cached integration projector * Generating the projector with a chosen sky container via .as_coeff() * Writing this to the zarr store.

Projectors working on different sky \(m\) values can write to the same store, and therefore in parallel build up a full projector. Here we just do one set of \(m\) values.

# Power-law sky
beam_vis_proj.setup_zarr_store(store_location="data/pl_sky_beam_proj")
beam_vis_proj.beam_linear_mmode_generator(pl_sky.as_coeff())
beam_vis_proj.write_integrator_cache_to_store(store_location="data/pl_sky_beam_proj")

# Point Source sky
beam_vis_proj.setup_zarr_store(store_location="data/ps_sky_beam_proj")
beam_vis_proj.beam_linear_mmode_generator(ps_sky.as_coeff())
beam_vis_proj.write_integrator_cache_to_store(store_location="data/ps_sky_beam_proj")

Computing 0.01G Gaunts (0.33 GB) took: 0.03s.

Now we read these files back in and use them to generate visibilities and m-modes. Note that this is now very fast since the heavy lifting work is in generating the cached projectors; the only array conversion left is the explicit .as_coeff() call when we feed the beam into the projector.

# Power-law sky
pl_sky_beam_proj = PowerBeamVisProjector.from_zarr_store(
    store_location="data/pl_sky_beam_proj", sky_absm_limits=sky_absm_limits
)
pl_mm_beam_proj = pl_sky_beam_proj.mmodes(power_beam.as_coeff())
pl_vis_beam_proj = pl_sky_beam_proj.visibilities(power_beam.as_coeff())

# Point source sky
ps_sky_beam_proj = PowerBeamVisProjector.from_zarr_store(
    store_location="data/ps_sky_beam_proj", sky_absm_limits=sky_absm_limits
)
ps_mm_beam_proj = ps_sky_beam_proj.mmodes(power_beam.as_coeff())
ps_vis_beam_proj = ps_sky_beam_proj.visibilities(power_beam.as_coeff())

def plot_results(proj, vis, mmodes, label):
    eras = np.linspace(0, 360, 2*proj.sky_lmax+1, endpoint=False)
    ms = proj.sky_m_values

    extent_vis = (frequencies_MHz[0], frequencies_MHz[-1], eras[-1], eras[0])
    aspect_vis = np.abs((extent_vis[0] - extent_vis[1])/(extent_vis[2]-extent_vis[3]))
    extent_ms = (frequencies_MHz[0], frequencies_MHz[-1], ms[-1], ms[0])
    aspect_ms = np.abs((extent_ms[0] - extent_ms[1])/(extent_ms[2]-extent_ms[3]))

    fig, axs = plt.subplots(2, 2, figsize=(7.5, 6))

    im = axs[0, 0].imshow(vis.T.real, extent=extent_vis, aspect=aspect_vis)
    plt.colorbar(im, ax=axs[0, 0], label=r"Re[$\mathcal{V}(\theta_{\mathrm{ERA}})$]")
    axs[0, 0].set_ylabel(r"$\theta_{\mathrm{ERA}}$ [deg]")
    axs[0, 0].set_xlabel(r"Frequency [MHz]")

    im = axs[0, 1].imshow(vis.T.imag, extent=extent_vis, aspect=aspect_vis)
    plt.colorbar(im, ax=axs[0, 1], label=r"Im[$\mathcal{V}(\theta_{\mathrm{ERA}})$]")
    axs[0, 1].set_ylabel(r"$\theta_{\mathrm{ERA}}$ [deg]")
    axs[0, 1].set_xlabel(r"Frequency [MHz]")

    im = axs[1, 0].imshow(mmodes.T.real, extent=extent_ms, aspect=aspect_ms)
    plt.colorbar(im, ax=axs[1, 0], label=r"Re[$\mathcal{V}(m)$]")
    axs[1, 0].set_ylabel(r"$m$")
    axs[1, 0].set_xlabel(r"Frequency [MHz]")

    im = axs[1, 1].imshow(mmodes.T.imag, extent=extent_ms, aspect=aspect_ms)
    plt.colorbar(im, ax=axs[1, 1], label=r"Im[$\mathcal{V}(m)$]")
    axs[1, 1].set_ylabel(r"$m$")
    axs[1, 1].set_xlabel(r"Frequency [MHz]")

    fig.suptitle(label)
    fig.tight_layout()
plot_results(beam_vis_proj, pl_vis_beam_proj, pl_mm_beam_proj, "Power-law Sky")
plot_results(beam_vis_proj, ps_vis_beam_proj, ps_mm_beam_proj, "Point Source Sky")

Sky Visibility Projector

Now we do the same for the sky projector, this time keeping the beam fixed and varying the sky container.

gb_sky_vis_proj = SkyVisProjector(
    latitude=telescope_latitude,
    longitude=telescope_longitude,
    baseline_enu = baseline_enu,
    sky_lmax = sky_lmax,
    frequencies_MHz = frequencies_MHz,
    generate_baseline_cache_on_init=True,
    generate_gaunt_cache_on_init=True,
    sky_absm_limits = sky_absm_limits,
)

Computing 0.01G Gaunts (0.33 GB) took: 0.03s.


/home/devin/Dropbox/arcturus/Projects/radio/hirax/systematics/dev/serval/src/serval/core.py:1204: UserWarning: The value of power_beam_lmax is changed to the difference of baseline_lmax from sky_lmax.
  self.generate_baseline_cache()

gb_sky_vis_proj.setup_zarr_store(store_location="data/gaussian_beam_sky_proj")
gb_sky_vis_proj.sky_linear_mmode_generator(power_beam.as_coeff())
gb_sky_vis_proj.write_integrator_cache_to_store(store_location="data/gaussian_beam_sky_proj")

Here we only need to generate the projector once since we only have one beam model, but we can still evaluate it quickly for both sky containers.

sky_vis_proj = SkyVisProjector.from_zarr_store(
    store_location="data/gaussian_beam_sky_proj", sky_absm_limits=sky_absm_limits
)

pl_mm_sky_proj = sky_vis_proj.mmodes(pl_sky.as_coeff())
ps_mm_sky_proj = sky_vis_proj.mmodes(ps_sky.as_coeff())

pl_vis_sky_proj = sky_vis_proj.visibilities(pl_sky.as_coeff())
ps_vis_sky_proj = sky_vis_proj.visibilities(ps_sky.as_coeff())

plot_results(sky_vis_proj, pl_vis_sky_proj, pl_mm_sky_proj, "Power-law Sky")
plot_results(sky_vis_proj, ps_vis_sky_proj, ps_mm_sky_proj, "Point Source Sky")

For a sanity check, we can compare the generated visibilities to those from a grid-based method for a few individual frequency channels.

for test_ind in [0, 1, 8, 13, 27, -1]:
    # Power-law sky
    single_visgen = SingleVisibilityGenerator(
        latitude=telescope_latitude,
        longitude=telescope_longitude,
        baseline_enu=baseline_enu,
        frequency_MHz=frequencies_MHz[test_ind],
        power_beam_alm=power_beam.as_coeff(freq_inds=test_ind),
        sky_alm=pl_sky.as_coeff(freq_inds=test_ind),
        method="grid",
        sky_absm_limits=sky_absm_limits,
    )
    assert np.allclose(single_visgen.visibilities(), pl_vis_beam_proj[test_ind])
    assert np.allclose(single_visgen.visibilities(), pl_vis_sky_proj[test_ind])
    single_visgen.sky_alm = ps_sky.as_coeff(freq_inds=test_ind)
    assert np.allclose(single_visgen.visibilities(), ps_vis_beam_proj[test_ind])
    assert np.allclose(single_visgen.visibilities(), ps_vis_sky_proj[test_ind])
print("All match!")

All match!

Compact Beam Visibility Projector

The beam visibility projector can be made much more compact by expressing the primary beam in a coordinate frame where it is close to azimuthally symmetric before decomposing it in spherical harmonics. For a perfectly azimuthally symmetric beam like the one here, you would only need the \(m = 0\) mode if you used a frame aligned with the pointing centre.

In general the beams will not be perfectly azimuthally symmetric but can still be quite symmetric, with only low-order \(m\) modes needed. The new TIRSPowerBeam.to_pointed_basis() helper does that basis rotation directly on the container and optionally truncates the pointed-basis \(m\) range for us.

pnt_beam_mmax = 3  # Max |m| = 3
pointing_altitude = np.pi / 2
pointing_azimuth = np.pi
pointing_boresight = 0.0

pointed_power_beam = power_beam.to_pointed_basis(
    latitude=telescope_latitude,
    longitude=telescope_longitude,
    altitude=pointing_altitude,
    azimuth=pointing_azimuth,
    boresight=pointing_boresight,
    mmax=pnt_beam_mmax,
)

print(power_beam.alm.shape)
print(pointed_power_beam.alm.shape)

(64, 41, 81)
(64, 41, 7)

We can see that this pointed-basis representation is much more compact. To show that these are pointing-oriented beams, we can plot one frequency slice directly from the pointed beam container.

_ = pointed_power_beam.bandlimited_to(
    lmax=pointed_power_beam.lmax,
    mmax=pointed_power_beam.lmax,
).as_grid(freq_inds=0, return_shgrid=True).plot(show=False)

Now we give the beam projector the same pointing and compactness information and generate visibilities and \(m\)-modes as before.

pnt_beam_vis_proj = PowerBeamVisProjector(
    latitude=telescope_latitude,
    longitude=telescope_longitude,
    baseline_enu=baseline_enu,
    power_beam_lmax=beam_lmax,
    frequencies_MHz=frequencies_MHz,
    generate_baseline_cache_on_init=True,
    generate_gaunt_cache_on_init=True,
    sky_absm_limits=sky_absm_limits,
    pointing_contract=True,
    pointed_beam_mmax=pnt_beam_mmax,
    pointing_altitude=pointing_altitude,
    pointing_azimuth=pointing_azimuth,
    pointing_boresight=pointing_boresight,
)

Computing 0.01G Gaunts (0.33 GB) took: 0.03s.


/home/devin/Dropbox/arcturus/Projects/radio/hirax/systematics/dev/serval/src/serval/core.py:980: UserWarning: The upper limit for the mmodes of the sky does not correspond to the optimal value given the baseline and beam which is 80. Consider changing it.
  self.set_bandlimits_from_baseline()

pnt_beam_vis_proj.setup_zarr_store(store_location="data/pl_sky_pnt_beam_proj")
pnt_beam_vis_proj.beam_linear_mmode_generator(pl_sky.as_coeff())
pnt_beam_vis_proj.write_integrator_cache_to_store(store_location="data/pl_sky_pnt_beam_proj")

pnt_beam_vis_proj.setup_zarr_store(store_location="data/ps_sky_pnt_beam_proj")
pnt_beam_vis_proj.beam_linear_mmode_generator(ps_sky.as_coeff())
pnt_beam_vis_proj.write_integrator_cache_to_store(store_location="data/ps_sky_pnt_beam_proj")

Computing 0.01G Gaunts (0.33 GB) took: 0.03s.

pl_sky_pnt_beam_proj = PowerBeamVisProjector.from_zarr_store(
    store_location="data/pl_sky_pnt_beam_proj", sky_absm_limits=sky_absm_limits
)
pl_mm_pnt_beam_proj = pl_sky_pnt_beam_proj.mmodes(pointed_power_beam.as_coeff())
pl_vis_pnt_beam_proj = pl_sky_pnt_beam_proj.visibilities(pointed_power_beam.as_coeff())

ps_sky_pnt_beam_proj = PowerBeamVisProjector.from_zarr_store(
    store_location="data/ps_sky_pnt_beam_proj", sky_absm_limits=sky_absm_limits
)
ps_mm_pnt_beam_proj = ps_sky_pnt_beam_proj.mmodes(pointed_power_beam.as_coeff())
ps_vis_pnt_beam_proj = ps_sky_pnt_beam_proj.visibilities(pointed_power_beam.as_coeff())

And we can check them again against the grid method, which does not assume anything about the pointing basis. The reference path below intentionally uses the original TIRS-basis beam container, while the compact projector consumes the pointed-basis coefficients derived from that same container lineage.

for test_ind in [0, 1, 8, 13, 27, -1]:
    # Power-law sky
    single_visgen = SingleVisibilityGenerator(
        latitude=telescope_latitude,
        longitude=telescope_longitude,
        baseline_enu=baseline_enu,
        frequency_MHz=frequencies_MHz[test_ind],
        power_beam_alm=power_beam.as_coeff(freq_inds=test_ind),
        sky_alm=pl_sky.as_coeff(freq_inds=test_ind),
        method="grid",
        sky_absm_limits=sky_absm_limits,
    )
    assert np.allclose(single_visgen.visibilities(), pl_vis_pnt_beam_proj[test_ind])
    single_visgen.sky_alm = ps_sky.as_coeff(freq_inds=test_ind)
    assert np.allclose(single_visgen.visibilities(), ps_vis_pnt_beam_proj[test_ind])
print("All match!")

All match!