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_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_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_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_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 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 src_perms(slvr_cfg, permute_weights=False):
"""
Permute the source types and return a SolverConfiguration suitable
for use as input to the solver function/factory.
Parameters:
slvr_cfg : RimeSolverConfiguration
Configuration containing other sensible defaults to include
in the returned permutation.
e.g. {'na': 14, 'ntime': 20, 'nchan': 48}
{ 'point': 0, 'gaussian': 00, 'sersic': 20 }
{ 'point': 0, 'gaussian': 20, 'sersic': 0 }
{ 'point': 20, 'gaussian': 0, 'sersic': 0 }
{ 'point': 0, 'gaussian': 20, 'sersic': 20 }
{ 'point': 20, 'gaussian': 20, 'sersic': 0 }
{ 'point': 20, 'gaussian': 0, 'sersic': 20 }
{ 'point': 20, 'gaussian': 20, 'sersic': 20 }
>>> slvr_cfg = RimeSolverConfiguration(na=14, ntime=20, nchan=48)
>>> for p_slvr_cfg in src_perms(slvr_cfg, True)
>>> with solver(p_slvr_cfg) as slvr:
>>> slvr.solve()
"""
if slvr_cfg is None:
slvr_cfg = montblanc.rime_solver_cfg()
from montblanc.src_types import SOURCE_VAR_TYPES
src_types = SOURCE_VAR_TYPES.keys()
count = 0
weight_vector = [True, False] if permute_weights is True else [False]
for wv in weight_vector:
count = 0
for p in itertools.product([0, 20], repeat=len(src_types)):
# Nasty, but works to avoid the (0,0,0) case
count += 1
if count == 1:
continue
params = montblanc.rime_solver_cfg(**slvr_cfg)
params[Options.WEIGHT_VECTOR] = wv
src_dict = {s: p[i] for i,s in enumerate(src_types)}
params[Options.SOURCES] = montblanc.sources(**src_dict)
yield params
def src_perms(slvr_cfg, permute_weights=False):
"""
Permute the source types and return a SolverConfiguration suitable
for use as input to the solver function/factory.
Parameters:
slvr_cfg : RimeSolverConfiguration
Configuration containing other sensible defaults to include
in the returned permutation.
e.g. {'na': 14, 'ntime': 20, 'nchan': 48}
{ 'point': 0, 'gaussian': 00, 'sersic': 20 }
{ 'point': 0, 'gaussian': 20, 'sersic': 0 }
{ 'point': 20, 'gaussian': 0, 'sersic': 0 }
{ 'point': 0, 'gaussian': 20, 'sersic': 20 }
{ 'point': 20, 'gaussian': 20, 'sersic': 0 }
{ 'point': 20, 'gaussian': 0, 'sersic': 20 }
{ 'point': 20, 'gaussian': 20, 'sersic': 20 }
>>> slvr_cfg = montblanc.rime_solver_cfg(na=14, ntime=20, nchan=48)
>>> for p_slvr_cfg in src_perms(slvr_cfg, True)
>>> with solver(p_slvr_cfg) as slvr:
>>> slvr.solve()
"""
if slvr_cfg is None:
slvr_cfg = montblanc.rime_solver_cfg()
from montblanc.src_types import SOURCE_VAR_TYPES
src_types = SOURCE_VAR_TYPES.keys()
count = 0
weight_vector = [True, False] if permute_weights is True else [False]
for wv in weight_vector:
count = 0
for p in itertools.product([0, 20], repeat=len(src_types)):
# Nasty, but works to avoid the (0,0,0) case
count += 1
if count == 1:
continue
params = montblanc.rime_solver_cfg(**slvr_cfg)
params[Options.WEIGHT_VECTOR] = wv
src_dict = {s: p[i] for i, s in enumerate(src_types)}
params[Options.SOURCES] = montblanc.sources(**src_dict)
yield params
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_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_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 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_cmp_visibilities(self):
""" Test visibilities produced by montblanc and meqtrees """
if not os.path.exists(msfile):
raise unittest.SkipTest("MeasurementSet '{ms}' required "
"for this test is not present".format(ms=msfile))
slvr_cfg = montblanc.rime_solver_cfg(
msfile=msfile,
sources=montblanc.sources(point=1, gaussian=0, sersic=0),
dtype='double', version=Options.VERSION_FOUR,
mem_budget=2*1024*1024*1024)
# Test the v4 and v5 residuals agree
v4_chi, v4_vis = self.get_v4_output(slvr_cfg,
vis_output=Options.VISIBILITY_OUTPUT_RESIDUALS)
v5_chi, v5_vis = self.get_v5_output(slvr_cfg,
vis_output=Options.VISIBILITY_OUTPUT_RESIDUALS)
self.assertTrue(np.allclose(v4_vis, v5_vis))
self.assertTrue(np.allclose(v4_chi, v5_chi))
v4_chi, v4_vis = self.get_v4_output(slvr_cfg)
v5_chi, v5_vis = self.get_v5_output(slvr_cfg)
# Test that v4 and v5 model visibilities agree
self.assertTrue(np.allclose(v4_vis, v5_vis))
self.assertTrue(np.allclose(v4_chi, v5_chi))
# Test that meqtrees agrees with v5
try:
meq_vis = self.get_meq_vis(slvr_cfg[Options.MS_FILE])
self.assertTrue(np.allclose(v5_vis, meq_vis))
except Exception as e:
montblanc.log.exception("Unable to run MeqTrees for "
"purposes of comparing model visibilities. "
"This will not be treated as a test failure")
if __name__ == '__main__':
import sys
import argparse
parser = argparse.ArgumentParser(description='RIME MS test script')
parser.add_argument('msfile', help='Measurement Set File')
parser.add_argument('-v','--version',dest='version', type=str,
default=Options.VERSION_FOUR, choices=Options.VALID_VERSIONS,
help='RIME Pipeline Version.')
args = parser.parse_args(sys.argv[1:])
# Get the solver.
slvr_cfg = montblanc.rime_solver_cfg(msfile=args.msfile,
sources=montblanc.sources(point=1, gaussian=0, sersic=0),
dtype='double', version=args.version)
with montblanc.rime_solver(slvr_cfg) as slvr:
if args.version in [Options.VERSION_TWO]:
lm = np.empty(shape=slvr.lm.shape, dtype=slvr.lm.dtype)
l, m = lm[0,:], lm[1,:]
l[:] = 0.1
m[:] = 0.25
slvr.transfer_lm(lm)
B = np.empty(shape=slvr.brightness.shape, dtype=slvr.brightness.dtype)
I, Q, U, V, alpha = B[0,:,:], B[1,:,:], B[2,:,:], B[3,:,:], B[4,:,:]
I[:] = 2
Q[:] = 1
'--version',
dest='version',
type=str,
default=Options.VERSION_FIVE,
choices=[Options.VERSION_FIVE],
help='RIME Pipeline Version.')
args = parser.parse_args(sys.argv[1:])
# Set the logging level
montblanc.log.setLevel(logging.INFO)
slvr_cfg = montblanc.rime_solver_cfg(
msfile=args.msfile,
sources=montblanc.sources(point=args.npsrc,
gaussian=args.ngsrc,
sersic=args.nssrc),
init_weights='weight',
weight_vector=False,
dtype='float',
auto_correlations=args.auto_correlations,
version=args.version)
with montblanc.rime_solver(slvr_cfg) as slvr:
# Random point source coordinates in the l,m,n (brightness image) domain
slvr.lm[:] = mbu.random_like(slvr.lm) * 0.1
# Need a positive semi-definite brightness matrix
I, Q, U, V = slvr.stokes[:, :,
0], slvr.stokes[:, :,
1], slvr.stokes[:, :,
# nssrc : number of sersic sources
# init_weights : (1) None (2) 'sigma' or (3) 'weight'. Either
# (1) do not initialise the weight vector, or
# (2) initialise from the MS 'SIGMA' tables, or
# (3) initialise from the MS 'WEIGHT' tables.
# weight_vector : indicates whether a weight vector should be used to
# compute the chi squared or a single sigma squared value
# nparams : number of sersic parameters with respect which to
# compute the chi_squared gradient
# store_cpu : indicates whether copies of the data passed into the
# solver transfer_* methods should be stored on the solver object
global slvr
slvr_cfg = montblanc.rime_solver_cfg(msfile=args.msfile,
sources=montblanc.sources(point=0, gaussian=0, sersic=args.nssrc),
init_weights=None, weight_vector=False,# mem_budget=1024*1024*1024,
sersic_gradient=True, dtype='double', version='v5')
with montblanc.rime_solver(slvr_cfg) as slvr:
nsrc, nssrc, ntime, nchan = slvr.dim_local_size('nsrc', 'nssrc', 'ntime', 'nchan')
true_values = np.empty((nssrc,6)) # position, flux, scalelength, ellipticity
# Read from catalog file, sources position, flux, scalelength and ellipticity
catalog = np.loadtxt('catalog1000.txt')[:nsrc,:]
# Source coordinates in the l,m (brightness image) domain
l = catalog[:,0]
m = catalog[:,1]
lm = mbu.shape_list([l,m], shape=slvr.lm.shape, dtype=slvr.lm.dtype)
parser.add_argument('msfile', help='Measurement Set File')
parser.add_argument('-s', '--sky-file', dest='sky_file', type=str, required=True, help='Sky Model File')
parser.add_argument('-c','--count',dest='count', type=int, default=10, help='Number of Iterations')
parser.add_argument('-v','--version',dest='version', type=str, default='v4', choices=[Options.VERSION_FOUR],
help='RIME Pipeline Version.')
args = parser.parse_args(sys.argv[1:])
# Set the logging level
montblanc.log.setLevel(logging.WARN)
# Parse the sky model file
sky_parse = mbu.parse_sky_model(args.sky_file)
sources = montblanc.sources(point=sky_parse.src_counts.get('npsrc', 0),
gaussian=sky_parse.src_counts.get('ngsrc', 0),
sersic=sky_parse.src_counts.get('nssrc', 0))
slvr_cfg = montblanc.rime_solver_cfg(msfile=args.msfile,
sources=sources, init_weights=None, weight_vector=False,
version=args.version)
with montblanc.rime_solver(slvr_cfg) as slvr:
# Get the lm coordinates
lm = sky_parse.shape_arrays(['l','m'], slvr.lm.shape, slvr.lm.dtype)
# Get the stokes and alpha parameters
stokes, alpha = repeat_brightness_over_time(slvr, sky_parse)
# If there are gaussian sources, create their
# shape matrix and transfer it.
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))
args = parser.parse_args(sys.argv[1:])
threshold = float(args.threshold)
print "Threshold: ", threshold
#Set visibility noise variance (muJy)
time_acc = 60
efficiency = 0.9
channel_bandwidth_hz = 240e6
SEFD = 400e6
sigma = (SEFD * SEFD) / (2. * time_acc * channel_bandwidth_hz * efficiency *
efficiency)
# Get the RIME solver
slvr_cfg = montblanc.rime_solver_cfg(msfile=args.msfile,
sources=montblanc.sources(point=0,
gaussian=0,
sersic=1),
init_weights=None,
weight_vector=False,
dtype='double',
version='v4')
with montblanc.rime_solver(slvr_cfg) as slvr:
nsrc, nssrc, ntime, nbl, nchan = slvr.dim_local_size(
'nsrc', 'nssrc', 'ntime', 'nbl', 'nchan')
# Read from modes file, sources flux, scalelength and ellipticity
data = np.loadtxt('best_modes.txt')
modes = []
modes_SNR = []
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))
import sys
import argparse
parser = argparse.ArgumentParser(description='RIME MS test script')
parser.add_argument('msfile', help='Measurement Set File')
parser.add_argument('-np','--npsrc',dest='npsrc', type=int, default=1, help='Number of Point Sources')
parser.add_argument('-ng','--ngsrc',dest='ngsrc', type=int, default=0, help='Number of Gaussian Sources')
parser.add_argument('-ns','--nssrc',dest='nssrc', type=int, default=0, help='Number of Gaussian Sources')
parser.add_argument('-c','--count',dest='count', type=int, default=10, help='Number of Iterations')
parser.add_argument('-v','--version',dest='version', type=str,
default=Options.VERSION_FOUR, choices=Options.VALID_VERSIONS, help='version')
args = parser.parse_args(sys.argv[1:])
slvr_cfg = montblanc.rime_solver_cfg(msfile=args.msfile,
sources=montblanc.sources(point=args.npsrc, gaussian=args.ngsrc, sersic=args.nssrc),
dtype='double', version=args.version)
# Get the solver.
with montblanc.rime_solver(slvr_cfg) as slvr:
nsrc = slvr.dim_global_size('nsrc')
# Create 1Jy point sources
if args.version == Options.VERSION_TWO:
lm = np.empty(shape=slvr.lm.shape, dtype=slvr.lm.dtype)
lm[0,:] = 0 # Set all l = 0
lm[1,:] = 0 # Set all m = 0
slvr.transfer_lm(lm)
brightness = np.empty(shape=slvr.brightness.shape, dtype=slvr.brightness.dtype)
brightness[0,:,:] = 1
