Playing with the screens package

make_irf.py

import numpy as np
from astropy import units as u
from screens.fields import dynamic_field

R = np.random.default_rng()

sig = 5 * u.mas

n_points = 256
th = np.append(0, R.normal(0, 10, n_points - 1)) * u.mas
th_perp = np.append(0, R.normal(0, 10, n_points - 1)) * u.mas
a = 0.5 * np.exp(-((th / sig)**2 + (th_perp / sig)**2) / 2)

a = a.to_value(u.one)
realization = a * R.standard_normal(th.shape + (2,)).view('c16').squeeze(-1)
realization[np.where(th == 0)] = 1
realization /= np.sqrt((np.abs(realization)**2).sum())

nchan = 4096
n_subint = 512

sample_rate = 3.125 * u.MHz
fobs = 350 * u.MHz
f = fobs + np.fft.fftshift(np.fft.fftfreq(nchan, 1/sample_rate))

dt = 8
t = np.arange(-n_subint*dt//2, n_subint*dt//2, dt) * u.s

d_eff = 2.5 * u.kpc
mu_eff = 13 * u.mas / u.yr

dyn_wave = dynamic_field(th, th_perp, realization,
                         d_eff, mu_eff, f, t[:, None])

wave_field = dyn_wave[...].sum(tuple(range(0, dyn_wave.ndim - 2)))
dynspec = wave_field.real**2 + wave_field.imag**2
dynspec += 0.1 * R.normal(size=dynspec.shape)

sec = np.fft.fft2(dynspec)
sec /= sec[0, 0]
tau = np.fft.fftfreq(dynspec.shape[1], f[1] - f[0]).to(u.us)
fd = np.fft.fftfreq(dynspec.shape[0], t[1] - t[0]).to(u.mHz)

sec = np.fft.fftshift(sec)  # Secondary spectrum
tau = np.fft.fftshift(tau) << tau.unit
fd = np.fft.fftshift(fd) << fd.unit