def E_beam_test_helper(self, beam_lw, beam_mh, beam_nud, dtype, cmp=None):
if cmp is None:
cmp = {}
slvr_cfg = montblanc.rime_solver_cfg(na=32,
ntime=50,
nchan=64,
sources=montblanc.sources(
point=10, gaussian=10),
beam_lw=beam_lw,
beam_mh=beam_mh,
beam_nud=beam_nud,
dtype=dtype,
pipeline=Pipeline([RimeEBeam()]))
gpu_slvr, cpu_slvr = solvers(slvr_cfg)
with gpu_slvr, cpu_slvr:
# Check that the beam cube dimensions are
# correctly configured
self.assertTrue(cpu_slvr.E_beam.shape == (beam_lw, beam_mh,
beam_nud, 4))
self.E_beam_test_impl(gpu_slvr, cpu_slvr, cmp=cmp)
def test_transpose(self):
slvr_cfg = montblanc.rime_solver_cfg(
na=14,
ntime=10,
nchan=16,
sources=montblanc.sources(point=10, gaussian=10),
weight_vector=True,
pipeline=Pipeline([MatrixTranspose()]),
data_source=Options.DATA_SOURCE_TEST,
version=Options.VERSION_FOUR)
with montblanc.factory.rime_solver(slvr_cfg) as gpu_slvr:
nsrc, nchan = gpu_slvr.dim_global_size('nsrc', 'nchan')
gpu_slvr.register_array(name='matrix_in',
shape=('nsrc', 'nchan'),
dtype='ft')
gpu_slvr.register_array(name='matrix_out',
shape=('nchan', 'nsrc'),
dtype='ft')
# Recreates existing arrays, but OK for testing purposes!
gpu_slvr.create_arrays()
matrix = np.random.random(size=(nsrc, nchan)).astype(gpu_slvr.ft)
gpu_slvr.transfer_matrix_in(matrix)
gpu_slvr.solve()
transposed_matrix = gpu_slvr.retrieve_matrix_out()
assert np.all(matrix == transposed_matrix.T)
def test_B_sqrt_double(self):
""" Test the B sqrt double kernel """
slvr_cfg = montblanc.rime_solver_cfg(na=7,
ntime=200,
nchan=320,
sources=montblanc.sources(
point=10, gaussian=10),
dtype=Options.DTYPE_DOUBLE,
pipeline=Pipeline([RimeBSqrt()]))
gpu_slvr, cpu_slvr = solvers(slvr_cfg)
with gpu_slvr, cpu_slvr:
self.B_sqrt_test_impl(gpu_slvr, cpu_slvr)
def test_sum_coherencies_residuals_double(self):
""" Test computation of double residuals """
slvr_cfg = montblanc.rime_solver_cfg(
na=14,
ntime=20,
nchan=48,
sources=montblanc.sources(point=10, gaussian=10),
dtype=Options.DTYPE_DOUBLE,
pipeline=Pipeline([RimeSumCoherencies()]),
vis_output=Options.VISIBILITY_OUTPUT_RESIDUALS)
gpu_slvr, cpu_slvr = solvers(slvr_cfg)
with gpu_slvr, cpu_slvr:
self.sum_coherencies_test_impl(gpu_slvr, cpu_slvr)
def test_sum_coherencies_double(self):
""" Test the coherency sum double kernel """
slvr_cfg = montblanc.rime_solver_cfg(na=14,
ntime=20,
nchan=48,
dtype=Options.DTYPE_DOUBLE)
for p_slvr_cfg in src_perms(slvr_cfg, permute_weights=True):
wv = p_slvr_cfg[Options.WEIGHT_VECTOR]
p_slvr_cfg['pipeline'] = Pipeline([RimeSumCoherencies()])
gpu_slvr, cpu_slvr = solvers(p_slvr_cfg)
with gpu_slvr, cpu_slvr:
self.sum_coherencies_test_impl(gpu_slvr, cpu_slvr)
def test_B_sqrt_float(self):
""" Test the B sqrt float kernel """
slvr_cfg = montblanc.rime_solver_cfg(na=7,
ntime=200,
nchan=320,
sources=montblanc.sources(
point=10, gaussian=10),
dtype=Options.DTYPE_FLOAT,
pipeline=Pipeline([RimeBSqrt()]))
gpu_slvr, cpu_slvr = solvers(slvr_cfg)
with gpu_slvr, cpu_slvr:
# This fails more often with an rtol of 1e-4
self.B_sqrt_test_impl(gpu_slvr, cpu_slvr, cmp={'rtol': 1e-3})
def test_sum_coherencies_float(self):
""" Test the coherency sum float kernel """
slvr_cfg = montblanc.rime_solver_cfg(na=14,
ntime=20,
nchan=48,
dtype=Options.DTYPE_FLOAT)
for p_slvr_cfg in src_perms(slvr_cfg, permute_weights=True):
p_slvr_cfg['pipeline'] = Pipeline([RimeSumCoherencies()])
gpu_slvr, cpu_slvr = solvers(p_slvr_cfg)
with gpu_slvr, cpu_slvr:
self.sum_coherencies_test_impl(gpu_slvr,
cpu_slvr,
cmp={'rtol': 1e-3})
def test_EK_double(self):
""" Double precision EK test """
slvr_cfg = montblanc.rime_solver_cfg(na=64,
ntime=10,
nchan=64,
sources=montblanc.sources(
point=10,
gaussian=10,
sersic=10),
dtype=Options.DTYPE_DOUBLE,
pipeline=Pipeline([RimeEK()]))
gpu_slvr, cpu_slvr = solvers(slvr_cfg)
with gpu_slvr, cpu_slvr:
self.EK_test_impl(gpu_slvr, cpu_slvr)
def test_EKBSqrt_double(self):
""" Double precision EKBSqrt test """
slvr_cfg = montblanc.rime_solver_cfg(
na=14,
ntime=20,
nchan=64,
sources=montblanc.sources(point=10, gaussian=10),
dtype=Options.DTYPE_DOUBLE,
pipeline=Pipeline([RimeEBeam(),
RimeBSqrt(),
RimeEKBSqrt()]))
gpu_slvr, cpu_slvr = solvers(slvr_cfg)
with gpu_slvr, cpu_slvr:
self.EKBSqrt_test_impl(gpu_slvr, cpu_slvr, cmp={'rtol': 1e-4})
def get_pipeline(slvr_cfg):
wv = slvr_cfg.get(Options.WEIGHT_VECTOR, False)
return Pipeline([RimeEK(), RimeGaussBSum(weight_vector=wv)])
def get_empty_pipeline(slvr_cfg):
""" Get an empty pipeline object """
return Pipeline([])
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))
def get_pipeline(slvr_cfg):
return Pipeline(
[RimeBSqrt(),
RimeEBeam(),
RimeEKBSqrt(),
RimeSumCoherencies()])