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))