def test_jones_multiply(self):
""" Verify jones matrix multiplication code against NumPy """
# Generate random matrices
def rmat(shape):
return np.random.random(size=shape) + \
np.random.random(size=shape)*1j
N = 100
shape = (100, 2,2)
A, B = rmat(shape), rmat(shape)
AM = [np.matrix(A[i,:,:]) for i in range(N)]
BM = [np.matrix(B[i,:,:]) for i in range(N)]
C = CPUSolver.jones_multiply(A, B, jones_shape='2x2')
for Am, Bm, Cm in zip(AM, BM, C):
assert np.allclose(Am*Bm, Cm)
C = CPUSolver.jones_multiply(A, B, hermitian=True, jones_shape='2x2')
for Am, Bm, Cm in zip(AM, BM, C):
assert np.allclose(Am*Bm.H, Cm)
def test_jones_multiply(self):
""" Verify jones matrix multiplication code against NumPy """
# Generate random matrices
def rmat(shape):
return np.random.random(size=shape) + \
np.random.random(size=shape)*1j
N = 100
shape = (100, 2, 2)
A, B = rmat(shape), rmat(shape)
AM = [np.matrix(A[i, :, :]) for i in range(N)]
BM = [np.matrix(B[i, :, :]) for i in range(N)]
C = CPUSolver.jones_multiply(A, B, jones_shape='2x2')
for Am, Bm, Cm in zip(AM, BM, C):
assert np.allclose(Am * Bm, Cm)
C = CPUSolver.jones_multiply(A, B, hermitian=True, jones_shape='2x2')
for Am, Bm, Cm in zip(AM, BM, C):
assert np.allclose(Am * Bm.H, Cm)
def solvers(slvr_cfg, **kwargs):
""" Returns CPU and GPU solvers for computing the RIME """
cpu_slvr_cfg = slvr_cfg.copy()
cpu_slvr_cfg[Options.DATA_SOURCE] = Options.DATA_SOURCE_TEST
cpu_slvr_cfg[Options.VERSION] = Options.VERSION_FOUR
cpu_slvr_cfg.update(kwargs)
gpu_slvr_cfg = slvr_cfg.copy()
gpu_slvr_cfg[Options.DATA_SOURCE] = Options.DATA_SOURCE_EMPTY
gpu_slvr_cfg[Options.VERSION] = Options.VERSION_FOUR
gpu_slvr_cfg.update(kwargs)
return montblanc.factory.rime_solver(gpu_slvr_cfg), CPUSolver(cpu_slvr_cfg)
def get_v4_output(self, slvr_cfg, **kwargs):
# Get visibilities from the v4 solver
gpu_slvr_cfg = slvr_cfg.copy()
gpu_slvr_cfg.update(**kwargs)
gpu_slvr_cfg[Options.VERSION] = Options.VERSION_FOUR
with FitsBeam(base_beam_file) as fb:
beam_shape = fb.shape
gpu_slvr_cfg[Options.E_BEAM_WIDTH] = beam_shape[0]
gpu_slvr_cfg[Options.E_BEAM_HEIGHT] = beam_shape[1]
gpu_slvr_cfg[Options.E_BEAM_DEPTH] = beam_shape[2]
from montblanc.impl.rime.v4.cpu.CPUSolver import CPUSolver
with montblanc.rime_solver(gpu_slvr_cfg) as slvr:
cpu_slvr_cfg = slvr.config().copy()
cpu_slvr_cfg[Options.DATA_SOURCE] = Options.DATA_SOURCE_EMPTY
cpu_slvr = CPUSolver(cpu_slvr_cfg)
# Create and transfer lm to the solver
lm = np.empty(shape=slvr.lm.shape, dtype=slvr.lm.dtype)
l, m = lm[:,0], lm[:,1]
l[:] = _L
m[:] = _M
slvr.transfer_lm(lm)
# Create and transfer stoke and alpha to the solver
stokes = np.empty(shape=slvr.stokes.shape, dtype=slvr.stokes.dtype)
alpha = np.empty(shape=slvr.alpha.shape, dtype=slvr.alpha.dtype)
I, Q, U, V = stokes[:,:,0], stokes[:,:,1], stokes[:,:,2], stokes[:,:,3]
I[:] = _I
Q[:] = _Q
U[:] = _U
V[:] = _V
alpha[:] = _ALPHA
slvr.transfer_stokes(stokes)
slvr.transfer_alpha(alpha)
fb.reconfigure_frequency_axes(slvr.retrieve_frequency())
# Create the beam from FITS file
ebeam = np.zeros(shape=slvr.E_beam.shape, dtype=slvr.E_beam.dtype)
ebeam.real[:] = fb.real()
ebeam.imag[:] = fb.imag()
slvr.transfer_E_beam(ebeam)
# Configure the beam extents
ll, lm, lf, ul, um, uf = fb.beam_extents
slvr.set_beam_ll(ll)
slvr.set_beam_lm(lm)
slvr.set_beam_lfreq(lf)
slvr.set_beam_ul(ul)
slvr.set_beam_um(um)
slvr.set_beam_ufreq(uf)
# Set the reference frequency
ref_freq = np.full(slvr.ref_frequency.shape, _REF_FREQ, dtype=slvr.ref_frequency.dtype)
slvr.transfer_ref_frequency(ref_freq)
from montblanc.solvers import copy_solver
copy_solver(slvr, cpu_slvr)
cpu_slvr.compute_E_beam()
slvr.solve()
return slvr.X2, slvr.retrieve_model_vis()
def test_sqrt_multiply(self):
"""
Confirm that multiplying the square root
of the brightness matrix into the
per antenna jones terms results in the same
jones matrices as multiplying the
brightness matrix into the per baseline
jones matrices.
"""
slvr_cfg = montblanc.rime_solver_cfg(na=14,
ntime=10,
nchan=16,
sources=montblanc.sources(
point=10, gaussian=10),
dtype=Options.DTYPE_DOUBLE,
pipeline=Pipeline([]))
with CPUSolver(slvr_cfg) as cpu_slvr:
nsrc, ntime, na, nbl, nchan = cpu_slvr.dim_global_size(
'nsrc', 'ntime', 'na', 'nbl', 'nchan')
# Calculate per baseline antenna pair indexes
ant0, ant1 = cpu_slvr.ap_idx(src=True, chan=True)
# Get the brightness matrix
B = cpu_slvr.compute_b_jones()
# Fill in the jones matrix with random values
cpu_slvr.jones[:] = np.random.random(
size=cpu_slvr.jones.shape).astype(cpu_slvr.jones.dtype) + \
np.random.random(
size=cpu_slvr.jones.shape).astype(cpu_slvr.jones.dtype)
# Superfluous really, but makes below readable
assert cpu_slvr.jones.shape == (nsrc, ntime, na, nchan, 4)
# Get per baseline jones matrices from
# the per antenna jones matrices
J2, J1 = cpu_slvr.jones[ant0], cpu_slvr.jones[ant1]
assert J1.shape == (nsrc, ntime, nbl, nchan, 4)
assert J2.shape == (nsrc, ntime, nbl, nchan, 4)
# Tile the brightness term over the baseline dimension
# and transpose so that polarisations are last
JB = np.tile(B[:, :, np.newaxis, :, :], (1, 1, nbl, 1, 1))
assert JB.shape == (nsrc, ntime, nbl, nchan, 4)
# Calculate the first result using the classic equation
# J2.B.J1^H
res_one = cpu_slvr.jones_multiply(J2, JB)
res_one = cpu_slvr.jones_multiply(res_one, J1, hermitian=True)
# Compute the square root of the
# brightness matrix
B_sqrt = cpu_slvr.compute_b_sqrt_jones(B)
# Tile the brightness square root term over
# the antenna dimension and transpose so that
# polarisations are last
JBsqrt = np.tile(B_sqrt[:, :, np.newaxis, :, :], (1, 1, na, 1, 1))
assert JBsqrt.shape == (nsrc, ntime, na, nchan, 4)
# Multiply the square root of the brightness matrix
# into the per antenna jones terms
J = (cpu_slvr.jones_multiply(cpu_slvr.jones, JBsqrt).reshape(
nsrc, ntime, na, nchan, 4))
# Get per baseline jones matrices from
# the per antenna jones matrices
J2, J1 = J[ant0], J[ant1]
assert J2.shape == (nsrc, ntime, nbl, nchan, 4)
assert J1.shape == (nsrc, ntime, nbl, nchan, 4)
# Calculate the first result using the optimised version
# (J2.sqrt(B)).(J1.sqrt(B))^H == J2.sqrt(B).sqrt(B)^H.J1^H
# == J2.sqrt(B).sqrt(B).J1^H
# == J2.B.J1^H
res_two = cpu_slvr.jones_multiply(J2, J1, hermitian=True)
# Results from two different methods should be the same
self.assertTrue(np.allclose(res_one, res_two))
