Visibility Projectors
from serval.core import PowerBeamVisProjector, SingleVisibilityGenerator, SkyVisProjector
from serval import sht
from serval.utils import (
pointed_gaussian,
gaussian,
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
mpl.rcParams["text.usetex"] = True
This notebook shows a more typical example of using Serval to produce projectors for producing batched of mmodes and visibilities for a single baseline 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 maxmimum harmonic degree for sky term
# Telescope position in TIRS lat, lon in radians.
telescope_latitude = 0
telescope_longitude = 0
But now we work on limited ranges of sky \(m\)'s but over multiple frequecy channels
Simiulated Skies
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.
# Now we need sky models pwer channel, will assume power-law in frequency
frequency_scaling = (frequencies_MHz/600)**-1
pl_sky_grid = power_law_sky(sky_lmax, power_law_index=2)
single_pl_sky_alm = sht.analysis(pl_sky_grid, return_sparse=False)
pl_sky_alms = frequency_scaling[:, None, None]*single_pl_sky_alm[None, ...]
single_ps_sky_alm = harmonic_point_source(180, 0, sky_lmax).todense()
ps_sky_alms = frequency_scaling[:, None, None]*single_ps_sky_alm[None, ...]
/home/devin/Dropbox/arcturus/Projects/radio/hirax/systematics/dev/serval/src/serval/utils.py:82: RuntimeWarning: divide by zero encountered in power
power = ells ** (-power_law_index)
Power Primary Beams
Gaussian primary beams whose FWHM goes as \(1/\nu\) are now generated.
fwhms = np.radians(20) * (600 / frequencies_MHz)
beam_alms = np.stack(
[
sht.analysis(
pointed_gaussian(
altitude=np.pi / 2, # Pointing at zenith
azimuth=0,
latitude=telescope_latitude,
longitude=telescope_longitude,
fwhm=fwhm, # Frequency dependent FWHM
template=sht.grid_template(beam_lmax),
),
return_sparse=False,
)
for fwhm in fwhms
],
axis=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
baseline_enu
(0.9586757156054164, 0.9586757156054164, 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.
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 13790658 Gaunts ( 0.33097579200000005 GB) took: 171.123074 ms
Now that the object is setup, we proceed by * Setting up a zarr store for the cached integration projector * Generate the projector with a chosen sky. * Write this to the zarr store.
Projectors working on different sky \(m\)'s can write to the same store, and therefore in parallel build up a full projector. Here we just do one set of \(m\)'s.
# 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_alms)
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_alms)
beam_vis_proj.write_integrator_cache_to_store(store_location="data/ps_sky_beam_proj")
Now we read these files back in and use them to generate visibilities and mmodes. Note that this is now very fast since the heavy lifting work is in generating the cached projectors.
# 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(beam_alms)
pl_vis_beam_proj = pl_sky_beam_proj.visibilities(beam_alms)
# 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(beam_alms)
ps_vis_beam_proj = ps_sky_beam_proj.visibilities(beam_alms)
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_{\rm ERA})$]")
axs[0, 0].set_ylabel(r"$\theta_{\rm 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_{\rm ERA})$]")
axs[0, 1].set_ylabel(r"$\theta_{\rm 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
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 13790658 Gaunts ( 0.33097579200000005 GB) took: 149.96217099999998 ms
gb_sky_vis_proj.setup_zarr_store(store_location="data/gaussian_beam_sky_proj")
gb_sky_vis_proj.sky_linear_mmode_generator(beam_alms)
gb_sky_vis_proj.write_integrator_cache_to_store(store_location="data/gaussian_beam_sky_proj")
Here we only needed to generate the projector once since we only have one beam mdel, but we'll quickly run it for both sky models.
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_alms)
ps_mm_sky_proj = sky_vis_proj.mmodes(ps_sky_alms)
pl_vis_sky_proj = sky_vis_proj.visibilities(pl_sky_alms)
ps_vis_sky_proj = sky_vis_proj.visibilities(ps_sky_alms)
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 gird-baced 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=beam_alms[test_ind],
sky_alm=pl_sky_alms[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_alms[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 azimuthally symmetric before decomposing it in sherical harmonics. For a perfectly azimuthally symetric 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.
We can generate a projector that assumes your beam is compact in \(m\) in a frame aligned with its pointing, and then contract the integration projector with a wigner-d matrix that rotates the beam to be TIRS frame in the cache.
This creates a much more compact projector with the cost of having to fix a pointing centre for the beams to be used by the projector.
pnt_beam_mmax = 3 # Max |m| = 3
pnt_beam_alms = np.stack(
[
sht.set_bandlimits(sht.analysis(
gaussian(
(0, 0, 1),
fwhm=fwhm, # Frequency dependent FWHM
template=sht.grid_template(beam_lmax),
),
return_sparse=False,
), beam_lmax, pnt_beam_mmax)
for fwhm in fwhms
],
axis=0,
)
print(beam_alms.shape)
print(pnt_beam_alms.shape)
(64, 41, 81)
(64, 41, 7)
We see that this is much more compact. To show that these are pointing oreinted beams, we can plot one and see that it is pointing along the z-axis.
pnt_beam_example = sht.array_synthesis(sht.set_bandlimits(pnt_beam_alms[0], beam_lmax, beam_lmax))
pnt_beam_example.plot()
/home/devin/miniforge3/envs/serval_dev/lib/python3.11/site-packages/pyshtools/shclasses/shgrid.py:1464: UserWarning: FigureCanvasAgg is non-interactive, and thus cannot be shown
fig.show()
(<Figure size 640x768 with 2 Axes>,
array([<Axes: title={'center': 'Real component'}, xlabel='Longitude', ylabel='Latitude'>,
<Axes: title={'center': 'Imaginary component'}, xlabel='Longitude', ylabel='Latitude'>],
dtype=object))
Now we give the beam projector the 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=np.pi/2,
pointing_azimuth=np.pi,
pointing_boresight=np.pi,
)
Computing 13790658 Gaunts ( 0.33097579200000005 GB) took: 139.57617199999999 ms
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_alms)
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_alms)
pnt_beam_vis_proj.write_integrator_cache_to_store(store_location="data/ps_sky_pnt_beam_proj")
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(pnt_beam_alms)
pl_vis_pnt_beam_proj = pl_sky_pnt_beam_proj.visibilities(pnt_beam_alms)
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(pnt_beam_alms)
ps_vis_pnt_beam_proj = ps_sky_pnt_beam_proj.visibilities(pnt_beam_alms)
And we can check them again against the grid method which doesn't assume anything about the pointings:
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=beam_alms[test_ind],
sky_alm=pl_sky_alms[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_alms[test_ind]
assert np.allclose(single_visgen.visibilities(), ps_vis_pnt_beam_proj[test_ind])
print("All match!")
All match!