def test_radec_to_lm(ra, dec, expected_l, expected_m):
ra, dec = np.radians(ra), np.radians(dec)
ra0, dec0 = np.radians(0), np.radians(90)
l, m = radec_to_lm(ra, dec, ra0, dec0)
assert_allclose([l, m], [expected_l, expected_m], atol=1e-12)
def test_secondorder_multiple_sources(mockms, mockcomp1, Ax, Ay, ra, dec, x, y, xx, xy, yy):
u, v, w = mockms.u_lambda, mockms.v_lambda, mockms.w_lambda
l, m = radec_to_lm(ra, dec, mockms.ra0, mockms.dec0)
mockms.data = np.array([Ax, 0, 0, Ay]) * np.exp(
2j * np.pi * (u*l + v*m + w*(np.sqrt(1 - l**2 - m**2) - 1))
)[:, :, None]
# Additional sources
l, m = 0, 0.05
mockms.data += np.array([Ax, 0, 0, Ay]) * np.exp(
2j * np.pi * (u*l + v*m + w*(np.sqrt(1 - l**2 - m**2) - 1))
)[:, :, None]
l, m = 0.6, -0.5
mockms.data += np.array([Ax, 0, 0, Ay]) * np.exp(
2j * np.pi * (u*l + v*m + w*(np.sqrt(1 - l**2 - m**2) - 1))
)[:, :, None]
U, V = mockms.U, mockms.V
antphases = x * U + y * V + xx * U**2 + xy * U * V + yy * V**2
phases = antphases[mockms.ant1] - antphases[mockms.ant2]
mockms.data *= np.exp(-1j * phases)[:, None, None]
mockcomp1.ra = ra
mockcomp1.dec = dec
source = Model('mymodel', [mockcomp1])
solution = Solution(ncomp=1)
calibrate.solve(source, solution, mockms, 1)
calibrate.solve(source, solution, mockms, 2)
solphases = solution.phases(U, V)
assert_allclose(solution.get_params(0), [Ax, Ay], rtol=5e-2)
assert_allclose(solphases, antphases, atol=5e-2)
def test_firstorder__multiplesources_amplitudefit(mockms, mockcomp1, Ax, Ay, ra, dec):
u, v, w = mockms.u_lambda, mockms.v_lambda, mockms.w_lambda
l, m = radec_to_lm(ra, dec, mockms.ra0, mockms.dec0)
mockms.data = np.array([Ax, 0, 0, Ay]) * np.exp(
2j * np.pi * (u*l + v*m + w*(np.sqrt(1 - l**2 - m**2) - 1))
)[:, :, None]
# Add a source at l=0, m=0
# In part, this tests that we are filtering out autocorrelations
mockms.data[:, :, 0] += Ax
mockms.data[:, :, 3] += Ay
solution = Solution(ncomp=1)
mockcomp1.ra = ra
mockcomp1.dec = dec
source = Model('mymodel', [mockcomp1])
calibrate.solve(source, solution, mockms, 1)
params = solution.get_params(2)
assert_allclose(params[:2], [Ax, Ay], rtol=1e-2)
assert_allclose(params[2:4], [0, 0], atol=5e-6)
assert_allclose(params[4:], [0, 0, 0])
def test_secondorder(mockms, mockcomp1, Ax, Ay, ra, dec, x, y, xx, xy, yy):
u, v, w = mockms.u_lambda, mockms.v_lambda, mockms.w_lambda
l, m = radec_to_lm(ra, dec, mockms.ra0, mockms.dec0)
mockms.data = np.array([Ax, 0, 0, Ay]) * np.exp(
2j * np.pi * (u*l + v*m + w*(np.sqrt(1 - l**2 - m**2) - 1))
)[:, :, None]
U, V = mockms.U, mockms.V
phases = x * U + y * V + xx * U**2 + xy * U * V + yy * V**2
phases = phases[mockms.ant1] - phases[mockms.ant2]
mockms.data *= np.exp(-1j * phases)[:, None, None]
mockcomp1.ra = ra
mockcomp1.dec = dec
source = Model('mymodel', [mockcomp1])
solution = Solution(ncomp=1)
calibrate.solve(source, solution, mockms, 1)
calibrate.solve(source, solution, mockms, 2)
params = solution.get_params(2)
assert_allclose(params[:2], [Ax, Ay], rtol=5e-2)
assert_allclose(params[2:4], [x, y], atol=5e-6)
assert_allclose(params[4:], [xx, xy, yy], atol=5e-9)
def test_firstorder_multiple_sources_with_noise(mockms, mockcomp1, Ax, Ay, ra, dec, x, y):
u, v, w = mockms.u_lambda, mockms.v_lambda, mockms.w_lambda
l, m = radec_to_lm(ra, dec, mockms.ra0, mockms.dec0)
mockms.data = np.array([Ax, 0, 0, Ay]) * np.exp(
2j * np.pi * (u*l + v*m + w*(np.sqrt(1 - l**2 - m**2) - 1))
)[:, :, None]
# Additional sources
l, m = 0, 0.05
mockms.data += np.array([Ax, 0, 0, Ay]) * np.exp(
2j * np.pi * (u*l + v*m + w*(np.sqrt(1 - l**2 - m**2) - 1))
)[:, :, None]
l, m = 0.6, -0.5
mockms.data += np.array([Ax, 0, 0, Ay]) * np.exp(
2j * np.pi * (u*l + v*m + w*(np.sqrt(1 - l**2 - m**2) - 1))
)[:, :, None]
antphases = x * mockms.U + y * mockms.V
phases = antphases[mockms.ant1] - antphases[mockms.ant2]
mockms.data *= np.exp(-1j * phases)[:, None, None]
mockms.data += np.random.normal(0, 20, mockms.data.shape) + 1j *np.random.normal(0, 20, mockms.data.shape)
mockcomp1.ra = ra
mockcomp1.dec = dec
source = Model('mymodel', [mockcomp1])
solution = Solution(ncomp=1)
calibrate.solve(source, solution, mockms, 1)
solphases = solution.phases(mockms.U, mockms.V)
assert_allclose(solution.get_params(0), [Ax, Ay], rtol=5e-2)
assert_allclose(solphases, antphases, atol=5e-2)
def test_there_and_back_again(ra, dec):
ra, dec = np.radians(ra), np.radians(dec)
ra0, dec0 = np.radians(90), np.radians(-45)
l, m = radec_to_lm(ra, dec, ra0, dec0)
print(l, m)
ra_back, dec_back = lm_to_radec(l, m, ra0, dec0)
assert_allclose([ra_back, dec_back], [ra, dec])
def solve(src, solution, mset, order):
# Phase rotate onto source and average in frequency
uvw, rotated = phase_rotate(mset.uvw,
mset.data[:, :,
[True, False, False, True]], src.ra,
src.dec, mset.ra0, mset.dec0, mset.lambdas)
start = tm.time()
rotated = freq_average(rotated)[:, None, :]
elapsed = tm.time() - start
logger.debug("Frequency averaging elapsed: %g", elapsed)
# Create array of unscaled (flux = 1) point sources for each component
start = tm.time()
u_lambda, v_lambda, w_lambda = uvw.T[:, :, None] / mset.midlambda
models = np.empty(
(len(src.components), rotated.shape[0], rotated.shape[1]),
dtype=np.complex128)
for i, comp in enumerate(src.components):
l, m = radec_to_lm(comp.ra, comp.dec, src.ra, src.dec)
models[i] = np.exp(2j * np.pi *
(u_lambda * l + v_lambda * m + w_lambda *
(np.sqrt(1 - l**2 - m**2) - 1)))
elapsed = tm.time() - start
logger.debug("Model creation elapsed: %g", elapsed)
# Fit
logger.debug("Fitting source '%s'...", src.name)
if order == 1:
f = residuals.full_firstorder
elif order == 2:
f = residuals.full_secondorder
start = tm.time()
res = least_squares(
f,
x0=solution.get_params(order=order),
args=(mset.U, mset.V, mset.ant1, mset.ant2, rotated, models),
verbose=1,
x_scale=solution.x_scale(order=order),
)
logger.debug("Fit (order=%d) elapsed: %g", order, tm.time() - start)
logger.debug(res.message)
logger.debug("Fit params:" + " %g" * len(res.x), *res.x)
# For some reason with scipy.optimize.least_squares, res.cost is only half of chisquared
solution.chisquared = 2 * res.cost
solution.set_params(res.x)
# If fit failed to converge, mark it as failed
if not res.success:
logger.warning("Fit failed; marking solution as failed")
solution.failed = True
def interpolate(tec, sources, solutions, mset, oversample=1, smoothing_kernel=0):
# Extract dimensions and world coordinates
data = tec.data
height, width = data.shape[3:]
center = SkyCoord(tec.header['CRVAL1'], tec.header['CRVAL2'], unit=(units.degree, units.degree))
wcs = WCS(tec.header)
# Create lists of lm coordinates in the FITS projection for calibration directions
ras = np.array([src.ra for src in sources])
decs = np.array([src.dec for src in sources])
directions_lm = radec_to_lm(ras, decs, center.ra.rad, center.dec.rad)
# Solve phases for each antenna for each calibration direction
phases = np.empty((len(mset.antids), len(sources)))
for i, solution in enumerate(solutions):
phases[:, i] = solution.phases(mset.U, mset.V)
# Get oversampled l,m values for TEC file
xx, yy = np.meshgrid(range(0, oversample * width), range(0, oversample * height))
pixels = np.array([xx.flatten(), yy.flatten()]).T
ret = wcs.all_pix2world([[x / oversample - 1/oversample, y / oversample - 1/oversample, 0, 0, 0] for x, y in pixels], 0)
grid_lm = radec_to_lm(np.radians(ret.T[0]), np.radians(ret.T[1]), center.ra.rad, center.dec.rad)
from scipy.interpolate import Rbf
for i in mset.antids:
# Compute interpolated phases
phases_grid = griddata(directions_lm.T, phases[i], grid_lm.T, method='nearest', fill_value=0)
# phases_grid = Rbf(directions_lm[0], directions_lm[1], phases[i], smooth=0.1)(grid_lm[0], grid_lm[1])
# phases_grid = nearestneighbour(directions_lm[0], directions_lm[1], phases[i], grid_lm.T, maxradius=3.0)
phases_grid = np.reshape(phases_grid, (oversample * height, oversample * width)) # [ dec, ra ]
# Gaussian smooth
phases_grid = gaussian_filter(phases_grid, oversample * smoothing_kernel, mode='constant', cval=0)
# Downsample
phases_grid = phases_grid[oversample//2::oversample, oversample//2::oversample]
data[0, 0, i] = phases_grid / 8.44797245E9 * mset.midfreq
def test_phaserotate_to_source(mockms, ra, dec):
u, v, w = mockms.u_lambda, mockms.v_lambda, mockms.w_lambda
l, m = radec_to_lm(ra, dec, mockms.ra0, mockms.dec0)
mockms.data = np.array([1, 0, 0, 1]) * np.exp(
2j * np.pi * (u*l + v*m + w*(np.sqrt(1 - l**2 - m**2) - 1))
)[:, :, None]
uvw, data = phaserotate.phase_rotate(mockms.uvw, mockms.data, ra, dec, mockms.ra0, mockms.dec0, mockms.lambdas)
# Assert baseline lengths are unchanged
assert_allclose((uvw**2).sum(axis=1), (mockms.uvw**2).sum(axis=1))
# Assert rotated data now all has phase 0
desired = np.ones_like(data[:, :, [True, False, False, True]])
assert_allclose(data[:, :, [True, False, False, True]], desired)
def test_firstorder_amplitudefit_centered(mockms, mockcomp1, Ax, Ay):
u, v, w = mockms.u_lambda, mockms.v_lambda, mockms.w_lambda
l, m = radec_to_lm(mockms.ra0, mockms.dec0, mockms.ra0, mockms.dec0)
mockms.data = np.array([Ax, 0, 0, Ay]) * np.exp(
2j * np.pi * (u*l + v*m + w*(np.sqrt(1 - l**2 - m**2) - 1))
)[:, :, None]
solution = Solution(ncomp=1)
mockcomp1.ra = mockms.ra0
mockcomp1.dec = mockms.dec0
source = Model('mymodel', [mockcomp1])
calibrate.solve(source, solution, mockms, 1)
params = solution.get_params(2)
assert_allclose(params, [Ax, Ay, 0, 0, 0, 0, 0], atol=1e-15)