def __init__(self, sigma0, nband, nx, ny, nthreads=8):
self.nthreads = nthreads
self.nx = nx
self.ny = ny
nx_psf = 2 * self.nx
npad_x = (nx_psf - nx) // 2
ny_psf = 2 * self.ny
npad_y = (ny_psf - ny) // 2
self.padding = ((0, 0), (npad_x, npad_x), (npad_y, npad_y))
self.ax = (1, 2)
self.unpad_x = slice(npad_x, -npad_x)
self.unpad_y = slice(npad_y, -npad_y)
self.lastsize = ny + np.sum(self.padding[-1])
# set length scales
length_scale = 0.5
K = make_kernel(nx_psf, ny_psf, sigma0, length_scale)
self.K = K
K_pad = iFs(self.K, axes=self.ax)
self.Khat = r2c(K_pad,
axes=self.ax,
forward=True,
nthreads=nthreads,
inorm=0)
self.Khatinv = np.where(self.Khat.real > 1e-14, 1.0 / self.Khat, 1e-14)
# get covariance in each dimension
# pixel coordinates
self.Kv = mock_array(nband) # np.eye(nband) * sigma0**2
self.Kvinv = mock_array(nband) # np.eye(nband) / sigma0**2
if nx == ny:
l_coord = m_coord = np.arange(-(nx // 2), nx // 2)
self.Kl = self.Km = expsq(l_coord, l_coord, 1.0, length_scale)
self.Klinv = self.Kminv = np.linalg.pinv(self.Kl,
hermitian=True,
rcond=1e-12)
self.Kl *= sigma0**2
self.Klinv /= sigma0**2
else:
l_coord = np.arange(-(nx // 2), nx // 2)
m_coord = np.arange(-(ny // 2), ny // 2)
self.Kl = expsq(l_coord, l_coord, sigma0, length_scale)
self.Km = expsq(m_coord, m_coord, 1.0, length_scale)
self.Klinv = np.linalg.pinv(self.Kl, hermitian=True, rcond=1e-12)
self.Kminv = np.linalg.pinv(self.Km, hermitian=True, rcond=1e-12)
# Kronecker matrices for "fast" matrix vector products
self.Kkron = (self.Kv, self.Kl, self.Km)
self.Kinvkron = (self.Kvinv, self.Klinv, self.Kminv)
def normal_gains(t, nu, s, n_time, n_ant, n_chan, n_dir, n_corr, sigma_f, lt,
lnu, ls):
"""Produce normal complex-gains based on the dimensions given."""
# Scale down domain
t = t / t.max() if t.max() != 0 else t
nu = nu / nu.max() if nu.max() != 0 else nu
s = s / s.max() if s.max() != 0 else s
# Make prior covariace matrices
Kt = expsq(t, t, sigma_f, lt)
Knu = expsq(nu, nu, 1.0, lnu)
Ks = expsq(s, s, 1.0, ls)
# Stack and get cholesky factors
K = np.array((Kt, Knu, Ks), dtype=object)
L = kt.kron_cholesky(K)
# Simulate independent gain per antenna and direction
gains = np.zeros((n_time, n_ant, n_chan, n_dir, n_corr),
dtype=np.complex128)
iterations = n_ant * n_corr
head = f"==> Generating (iterations={iterations}): "
for p in range(n_ant):
for c in range(n_corr):
# Progress Bar
progress_bar(head, iterations + 1, p * n_corr + c + 1)
# Generate random complex vector
xi = np.random.randn(n_time, n_chan, n_dir)/np.sqrt(2)\
+ 1.0j * np.random.randn(n_time, n_chan, n_dir)/np.sqrt(2)
# Apply to field
gains[:, p, :, :, c] = \
kt.kron_matvec(L, xi).reshape(n_time, n_chan, n_dir) + 1.0
print()
# Return normal-form gains
return gains
def normal_gains(t, nu, s, n_ant, n_corr, sigmaf, lt, lnu, ls):
"""Produce normal complex-gains based on the dimensions given."""
# Scale down domain
t = t / t.max() if t.max() != 0 else t
nu = nu / nu.max() if nu.max() != 0 else nu
s = s / s.max() if s.max() != 0 else s
# Get dimensions
n_time = t.size
n_chan = nu.size
n_dir = s.shape[0]
# Make prior covariace matrices
Kt = expsq(t, t, sigmaf, lt)
Knu = expsq(nu, nu, 1.0, lnu)
Ks = expsq(s, s, 1.0, ls)
# Stack and get cholesky factors
K = np.array((Kt, Knu, Ks), dtype=object)
L = kt.kron_cholesky(K)
# Simulate independent gain per antenna and direction
gains = np.zeros((n_time, n_ant, n_chan, n_dir, n_corr),
dtype=np.complex128)
for p in tqdm(range(n_ant)):
for c in range(n_corr):
# Generate random complex vector
xi = np.random.randn(n_time, n_chan, n_dir)/np.sqrt(2)\
+ 1.0j * np.random.randn(n_time, n_chan, n_dir)/np.sqrt(2)
# Apply to field
gains[:, p, :, :, c] = \
kt.kron_matvec(L, xi).reshape(n_time, n_chan, n_dir) + 1.0
# Return normal-form gains
return gains
def test_matvec(nx, ny, nband, sigma0, length_scale):
v = np.arange(-(nband // 2), nband // 2)
x = np.arange(-(nx // 2), nx // 2)
y = np.arange(-(ny // 2), ny // 2)
A1 = expsq(v, v, sigma0, length_scale)
A2 = expsq(x, x, 1.0, length_scale)
A3 = expsq(y, y, 1.0, length_scale)
sigma = 1e-13
C1 = np.linalg.pinv(A1, hermitian=True, rcond=sigma)
C2 = np.linalg.pinv(A2, hermitian=True, rcond=sigma)
C3 = np.linalg.pinv(A3, hermitian=True, rcond=sigma)
A = (A1, A2, A3)
C = (C1, C2, C3)
xi = np.random.randn(nband * nx * ny)
res = kron_matvec(A, xi)
rec = kron_matvec(C, res)
assert_array_almost_equal(rec, xi, decimal=5)