def EKBSqrt_test_impl(self, gpu_slvr, cpu_slvr, cmp=None):
""" Type independent implementation of the EKBSqrt test """
if cmp is None:
cmp = {}
# Make the beam cube sufficiently large to contain the
# test values for the lm and pointing error coordinates
# specified in RimeSolver.py
S = 1
cpu_slvr.set_beam_ll(-S)
cpu_slvr.set_beam_lm(-S)
cpu_slvr.set_beam_ul(S)
cpu_slvr.set_beam_um(S)
copy_solver(cpu_slvr, gpu_slvr)
# Call the GPU solver
gpu_slvr.solve()
ekb_cpu = cpu_slvr.compute_ekb_sqrt_jones_per_ant()
ekb_gpu = gpu_slvr.retrieve_jones()
self.assertTrue(np.allclose(ekb_cpu, ekb_gpu, **cmp))
# Test that at a decent proportion of
# the calculated EKB terms are non-zero
non_zero = np.count_nonzero(ekb_cpu)
non_zero_ratio = non_zero / float(ekb_cpu.size)
self.assertTrue(non_zero_ratio > 0.85,
'Non-zero EKB ratio is %f.' % non_zero)
def B_sum_test_impl(self, gpu_slvr, cpu_slvr,
weight_vector=False, cmp=None):
""" Type independent implementation of the B Sum test """
if cmp is None:
cmp = {}
# This beam width produces reasonable values
# for testing the E term
cpu_slvr.set_beam_width(65*1e5)
cpu_slvr.set_sigma_sqrd(np.random.random(1)[0])
# Copy CPU solver data into the gpu solver
copy_solver(cpu_slvr, gpu_slvr)
# Call the GPU solver
gpu_slvr.solve()
cpu_slvr.solve()
ebk_vis_cpu = cpu_slvr.model_vis
ebk_vis_gpu = gpu_slvr.retrieve_model_vis()
self.assertTrue(np.allclose(ebk_vis_cpu, ebk_vis_gpu, **cmp))
# So technically the chi squared should be living
# on the GPU array, but RimeGaussBSum places it
# in the X2 property
# chi_sqrd_result = gpu_slvr.retrieve_X2()
chi_sqrd_result_cpu = cpu_slvr.X2
chi_sqrd_result_gpu = gpu_slvr.X2
self.assertTrue(np.allclose(chi_sqrd_result_cpu,
chi_sqrd_result_gpu, **cmp))
def E_beam_test_impl(self, gpu_slvr, cpu_slvr, cmp=None):
if cmp is None:
cmp = {}
# Make the beam cube sufficiently large to contain the
# test values for the lm and pointing error coordinates
# specified in RimeSolver.py
S = 1
cpu_slvr.set_beam_ll(-S)
cpu_slvr.set_beam_lm(-S)
cpu_slvr.set_beam_ul(S)
cpu_slvr.set_beam_um(S)
E_term_cpu = cpu_slvr.compute_E_beam()
copy_solver(cpu_slvr, gpu_slvr)
gpu_slvr.solve()
E_term_gpu = gpu_slvr.retrieve_jones()
self.assertTrue(np.allclose(E_term_cpu, E_term_gpu, **cmp))
# Test that at a decent proportion of
# the calculated E terms are non-zero
non_zero_E = np.count_nonzero(E_term_cpu)
non_zero_E_ratio = non_zero_E / float(E_term_cpu.size)
self.assertTrue(non_zero_E_ratio > 0.85,
'Non-zero E-term ratio is {r}.'.format(r=non_zero_E_ratio))
def E_beam_test_impl(self, gpu_slvr, cpu_slvr, cmp=None):
if cmp is None:
cmp = {}
# Make the beam cube sufficiently large to contain the
# test values for the lm and pointing error coordinates
# specified in RimeSolver.py
S = 1
cpu_slvr.set_beam_ll(-S)
cpu_slvr.set_beam_lm(-S)
cpu_slvr.set_beam_ul(S)
cpu_slvr.set_beam_um(S)
E_term_cpu = cpu_slvr.compute_E_beam()
copy_solver(cpu_slvr, gpu_slvr)
gpu_slvr.solve()
E_term_gpu = gpu_slvr.retrieve_jones()
self.assertTrue(np.allclose(E_term_cpu, E_term_gpu, **cmp))
# Test that at a decent proportion of
# the calculated E terms are non-zero
non_zero_E = np.count_nonzero(E_term_cpu)
non_zero_E_ratio = non_zero_E / float(E_term_cpu.size)
self.assertTrue(
non_zero_E_ratio > 0.85,
'Non-zero E-term ratio is {r}.'.format(r=non_zero_E_ratio))
def EKBSqrt_test_impl(self, gpu_slvr, cpu_slvr, cmp=None):
""" Type independent implementation of the EKBSqrt test """
if cmp is None:
cmp = {}
# Make the beam cube sufficiently large to contain the
# test values for the lm and pointing error coordinates
# specified in RimeSolver.py
S = 1
cpu_slvr.set_beam_ll(-S)
cpu_slvr.set_beam_lm(-S)
cpu_slvr.set_beam_ul(S)
cpu_slvr.set_beam_um(S)
copy_solver(cpu_slvr, gpu_slvr)
# Call the GPU solver
gpu_slvr.solve()
ekb_cpu = cpu_slvr.compute_ekb_sqrt_jones_per_ant()
ekb_gpu = gpu_slvr.retrieve_jones()
self.assertTrue(np.allclose(ekb_cpu, ekb_gpu, **cmp))
# Test that at a decent proportion of
# the calculated EKB terms are non-zero
non_zero = np.count_nonzero(ekb_cpu)
non_zero_ratio = non_zero / float(ekb_cpu.size)
self.assertTrue(non_zero_ratio > 0.85,
'Non-zero EKB ratio is %f.' % non_zero)
def B_sum_test_impl(self,
gpu_slvr,
cpu_slvr,
weight_vector=False,
cmp=None):
""" Type independent implementation of the B Sum test """
if cmp is None:
cmp = {}
# This beam width produces reasonable values
# for testing the E term
cpu_slvr.set_beam_width(65 * 1e5)
cpu_slvr.set_sigma_sqrd(np.random.random(1)[0])
# Copy CPU solver data into the gpu solver
copy_solver(cpu_slvr, gpu_slvr)
# Call the GPU solver
gpu_slvr.solve()
cpu_slvr.solve()
ebk_vis_cpu = cpu_slvr.model_vis
ebk_vis_gpu = gpu_slvr.retrieve_model_vis()
self.assertTrue(np.allclose(ebk_vis_cpu, ebk_vis_gpu, **cmp))
# So technically the chi squared should be living
# on the GPU array, but RimeGaussBSum places it
# in the X2 property
# chi_sqrd_result = gpu_slvr.retrieve_X2()
chi_sqrd_result_cpu = cpu_slvr.X2
chi_sqrd_result_gpu = gpu_slvr.X2
self.assertTrue(
np.allclose(chi_sqrd_result_cpu, chi_sqrd_result_gpu, **cmp))
def sum_coherencies_test_impl(self, gpu_slvr,
cpu_slvr, cmp=None):
""" Type independent implementation of the coherency sum test """
if cmp is None:
cmp = {}
# Randomise the sigma squared
cpu_slvr.set_sigma_sqrd(np.random.random(1)[0])
# The pipeline for this test case doesn't
# create the jones terms. Create some
# random terms and transfer them to the GPU
sh, dt = cpu_slvr.jones.shape, cpu_slvr.jones.dtype
cpu_slvr.jones[:] = (
np.random.random(size=sh).astype(dt) +
1j*np.random.random(size=sh).astype(dt))
copy_solver(cpu_slvr, gpu_slvr)
# Call the GPU solver
gpu_slvr.solve()
# Check that the CPU and GPU visibilities
# match each other
ekb_per_bl = cpu_slvr.compute_ekb_jones_per_bl(cpu_slvr.jones)
ekb_vis_cpu = cpu_slvr.compute_ekb_vis(ekb_per_bl)
gekb_vis_cpu = cpu_slvr.compute_gekb_vis(ekb_vis_cpu)
gekb_vis_gpu = gpu_slvr.retrieve_model_vis()
self.assertTrue(np.allclose(gekb_vis_cpu, gekb_vis_gpu, **cmp))
# Check that the chi squared sum terms
# match each other
chi_sqrd_sum_terms_cpu = cpu_slvr.compute_chi_sqrd_sum_terms(
vis=gekb_vis_cpu)
chi_sqrd_sum_terms_gpu = gpu_slvr.retrieve_chi_sqrd_result()
self.assertTrue(np.allclose(chi_sqrd_sum_terms_cpu,
chi_sqrd_sum_terms_gpu, **cmp))
chi_sqrd_result = cpu_slvr.compute_chi_sqrd(
chi_sqrd_terms=chi_sqrd_sum_terms_cpu)
self.assertTrue(np.allclose(chi_sqrd_result, gpu_slvr.X2, **cmp))
def EK_test_impl(self, gpu_slvr, cpu_slvr, cmp=None):
""" Type independent implementation of the EK test """
if cmp is None:
cmp = {}
# This beam width produces reasonable values
# for testing the E term
cpu_slvr.set_beam_width(65*1e5)
# Copy CPU solver data into the gpu solver
copy_solver(cpu_slvr, gpu_slvr)
# Call the GPU solver
gpu_slvr.solve()
ek_cpu = cpu_slvr.compute_ek_jones_scalar_per_ant()
ek_gpu = gpu_slvr.retrieve_jones_scalar()
# Test that the jones CPU calculation matches
# that of the GPU calculation
self.assertTrue(np.allclose(ek_cpu, ek_gpu, **cmp))
def EK_test_impl(self, gpu_slvr, cpu_slvr, cmp=None):
""" Type independent implementation of the EK test """
if cmp is None:
cmp = {}
# This beam width produces reasonable values
# for testing the E term
cpu_slvr.set_beam_width(65 * 1e5)
# Copy CPU solver data into the gpu solver
copy_solver(cpu_slvr, gpu_slvr)
# Call the GPU solver
gpu_slvr.solve()
ek_cpu = cpu_slvr.compute_ek_jones_scalar_per_ant()
ek_gpu = gpu_slvr.retrieve_jones_scalar()
# Test that the jones CPU calculation matches
# that of the GPU calculation
self.assertTrue(np.allclose(ek_cpu, ek_gpu, **cmp))
def sum_coherencies_test_impl(self, gpu_slvr, cpu_slvr, cmp=None):
""" Type independent implementation of the coherency sum test """
if cmp is None:
cmp = {}
# Randomise the sigma squared
cpu_slvr.set_sigma_sqrd(np.random.random(1)[0])
# The pipeline for this test case doesn't
# create the jones terms. Create some
# random terms and transfer them to the GPU
sh, dt = cpu_slvr.jones.shape, cpu_slvr.jones.dtype
cpu_slvr.jones[:] = (np.random.random(size=sh).astype(dt) +
1j * np.random.random(size=sh).astype(dt))
copy_solver(cpu_slvr, gpu_slvr)
# Call the GPU solver
gpu_slvr.solve()
# Check that the CPU and GPU visibilities
# match each other
ekb_per_bl = cpu_slvr.compute_ekb_jones_per_bl(cpu_slvr.jones)
ekb_vis_cpu = cpu_slvr.compute_ekb_vis(ekb_per_bl)
gekb_vis_cpu = cpu_slvr.compute_gekb_vis(ekb_vis_cpu)
gekb_vis_gpu = gpu_slvr.retrieve_model_vis()
self.assertTrue(np.allclose(gekb_vis_cpu, gekb_vis_gpu, **cmp))
# Check that the chi squared sum terms
# match each other
chi_sqrd_sum_terms_cpu = cpu_slvr.compute_chi_sqrd_sum_terms(
vis=gekb_vis_cpu)
chi_sqrd_sum_terms_gpu = gpu_slvr.retrieve_chi_sqrd_result()
self.assertTrue(
np.allclose(chi_sqrd_sum_terms_cpu, chi_sqrd_sum_terms_gpu, **cmp))
chi_sqrd_result = cpu_slvr.compute_chi_sqrd(
chi_sqrd_terms=chi_sqrd_sum_terms_cpu)
self.assertTrue(np.allclose(chi_sqrd_result, gpu_slvr.X2, **cmp))
def B_sqrt_test_impl(self, gpu_slvr, cpu_slvr, cmp=None):
""" Type independent implementation of the B square root test """
if cmp is None:
cmp = {}
copy_solver(cpu_slvr, gpu_slvr)
# Calculate CPU version of the B sqrt matrix
b_sqrt_cpu = cpu_slvr.compute_b_sqrt_jones()
# Call the GPU solver
gpu_slvr.solve()
# Get the GPU version of the B sqrt matrix
b_sqrt_gpu = gpu_slvr.retrieve_B_sqrt()
self.assertTrue(np.allclose(b_sqrt_cpu, b_sqrt_gpu, **cmp))
# TODO: Replace with np.einsum
# Pick 16 random points in the same and
# check our square roots are OK for that
nsrc, ntime, nchan = cpu_slvr.dim_global_size('nsrc', 'ntime', 'nchan')
N = 16
rand_srcs = [random.randrange(0, nsrc) for i in range(N)]
rand_t = [random.randrange(0, ntime) for i in range(N)]
rand_ch = [random.randrange(0, nchan) for i in range(N)]
b_cpu = cpu_slvr.compute_b_jones()
# Test that the square root of B
# multiplied by itself yields B.
# Also tests that the square root of B
# is the Hermitian of the square root of B
for src, t, ch in zip(rand_srcs, rand_t, rand_ch):
B_sqrt = b_sqrt_cpu[src,t,ch].reshape(2,2)
B = b_cpu[src,t,ch].reshape(2,2)
self.assertTrue(np.allclose(B, np.dot(B_sqrt, B_sqrt)))
self.assertTrue(np.all(B_sqrt == B_sqrt.conj().T))
def B_sqrt_test_impl(self, gpu_slvr, cpu_slvr, cmp=None):
""" Type independent implementation of the B square root test """
if cmp is None:
cmp = {}
copy_solver(cpu_slvr, gpu_slvr)
# Calculate CPU version of the B sqrt matrix
b_sqrt_cpu = cpu_slvr.compute_b_sqrt_jones()
# Call the GPU solver
gpu_slvr.solve()
# Get the GPU version of the B sqrt matrix
b_sqrt_gpu = gpu_slvr.retrieve_B_sqrt()
self.assertTrue(np.allclose(b_sqrt_cpu, b_sqrt_gpu, **cmp))
# TODO: Replace with np.einsum
# Pick 16 random points in the same and
# check our square roots are OK for that
nsrc, ntime, nchan = cpu_slvr.dim_global_size('nsrc', 'ntime', 'nchan')
N = 16
rand_srcs = [random.randrange(0, nsrc) for i in range(N)]
rand_t = [random.randrange(0, ntime) for i in range(N)]
rand_ch = [random.randrange(0, nchan) for i in range(N)]
b_cpu = cpu_slvr.compute_b_jones()
# Test that the square root of B
# multiplied by itself yields B.
# Also tests that the square root of B
# is the Hermitian of the square root of B
for src, t, ch in zip(rand_srcs, rand_t, rand_ch):
B_sqrt = b_sqrt_cpu[src, t, ch].reshape(2, 2)
B = b_cpu[src, t, ch].reshape(2, 2)
self.assertTrue(np.allclose(B, np.dot(B_sqrt, B_sqrt)))
self.assertTrue(np.all(B_sqrt == B_sqrt.conj().T))
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()
