Skip to content

Commit 97d0097

Browse files
scpmwscpmw
committed
ARL adjustments for crocodile
We can get rid of a lot of strange coordinate transformations now. The results aren't exactly the same as before (~1% difference), but that might be coincidence? Also a bunch of stylistic changes in here - nothing wrong with having accessors, especially if it means we can save a bunch of variables.
1 parent fc62851 commit 97d0097

8 files changed

Lines changed: 141 additions & 378 deletions

File tree

‎.gitignore‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
/.cproject
55
/.project
66
/.eggs
7+
/crocodile.egg-info
78

89
*.pyc
910
/build
@@ -22,11 +23,14 @@
2223
/docs/notebooks
2324
/docs/anna
2425
/docs/ben
26+
/docs/arl
2527
.ipynb_checkpoints
28+
2629
out.fits
2730
dirty.fits
31+
psf.fits
32+
2833
/doc
2934
/.ipynb_checkpoints
30-
psf.fits
3135
.DS_Store
3236
/data

‎arl/data_models.py‎

Lines changed: 25 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,13 @@ def __init__(self):
7373
self.data = None
7474
self.wcs = None
7575

76+
@property
77+
def nchan(self): return self.data.shape[0]
78+
@property
79+
def npol(self): return self.data.shape[1]
80+
@property
81+
def npixel(self): return self.data.shape[3]
82+
7683

7784
class SkyComponent:
7885
""" A single SkyComponent with direction, flux, shape, and params for the shape
@@ -88,6 +95,11 @@ def __init__(self):
8895
self.shape = None # str e.g. 'Point' 'Gaussian'
8996
self.params = None # numpy.array shape dependent parameters
9097

98+
@property
99+
def nchan(self): return self.flux.shape[0]
100+
@property
101+
def npol(self): return self.flux.shape[1]
102+
91103

92104
class SkyModel:
93105
""" A skymodel consisting of a list of images and a list of skycomponents
@@ -160,7 +172,7 @@ def weight(self): return self.data['weight']
160172

161173
def uvw_lambda(self, channel=0):
162174
""" Calculates baseline coordinates in wavelengths. """
163-
return self.data['uvw'] * (self.frequency[channel] / const.c).value
175+
return self.data['uvw'] * self.frequency[channel] / const.c.value
164176

165177

166178
class QA:
@@ -174,3 +186,15 @@ def __init__(self, origin=None, data=None, context=None):
174186
self.origin = origin # Name of function originating QA assessment
175187
self.data = data # Dictionary containing standard fields
176188
self.context = context # Context string (TBD)
189+
190+
def assert_same_chan_pol(o1, o2):
191+
"""
192+
Assert that two entities indexed over channels and polarisations
193+
have the same number of them.
194+
"""
195+
assert o1.npol == o2.npol, \
196+
"%s and %s have different number of polarisations: %d != %d" % \
197+
(type(o1).__name__, type(o2).__name__, o1.npol, o2.npol)
198+
assert o1.nchan == o2.nchan, \
199+
"%s and %s have different number of channels: %d != %d" % \
200+
(type(o1).__name__, type(o2).__name__, o1.nchan, o2.nchan)

‎arl/fourier_transforms.py‎

Lines changed: 58 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -3,15 +3,17 @@
33
# Synthesis imaging functions
44
#
55

6+
import functools
67
import numpy
78
import pylru
9+
import copy
810
from astropy import constants as const
911
from astropy import units as units
1012
from astropy import wcs
1113
from astropy.coordinates import CartesianRepresentation
1214

1315
from crocodile.simulate import simulate_point, skycoord_to_lmn
14-
from crocodile.synthesis import wcacheimg, wcachefwd, wkernaf, doimg, dopredict
16+
from crocodile.synthesis import w_cache_imaging, w_cache_predict, w_kernel, w_conj_kernel_fn, do_imaging, do_predict
1517

1618
from arl.data_models import *
1719
from arl.image_operations import create_image_from_array
@@ -62,7 +64,7 @@ def create_wcs_from_visibility(vis: Visibility, params={}):
6264
# Beware of python indexing order! wcs and the array have opposite ordering
6365
shape = [len(vis.frequency), npol, npixel, npixel]
6466
w = wcs.WCS(naxis=4)
65-
w.wcs.cdelt = [-cellsize * 180.0 / numpy.pi, cellsize * 180.0 / numpy.pi, 1.0, channelwidth.value]
67+
w.wcs.cdelt = [cellsize * 180.0 / numpy.pi, cellsize * 180.0 / numpy.pi, 1.0, channelwidth.value]
6668
w.wcs.crpix = [npixel // 2 + 1, npixel // 2 + 1, 1.0, 1.0]
6769
w.wcs.ctype = ["RA---SIN", "DEC--SIN", 'STOKES', 'FREQ']
6870
w.wcs.crval = [phasecentre.ra.value, phasecentre.dec.value, 1.0, reffrequency.value]
@@ -74,6 +76,20 @@ def create_wcs_from_visibility(vis: Visibility, params={}):
7476
return shape, reffrequency, cellsize, w, imagecentre
7577

7678

79+
def w_cache_size(vis:Visibility, wstep, frequency=None):
80+
"""
81+
Determine optimal w-kernel cache size for de/gridding the
82+
given visibilities.
83+
"""
84+
85+
if frequency is None:
86+
frequency = max(vis.frequency)
87+
88+
wmax = numpy.max(numpy.abs(vis.w))
89+
wmax *= float(frequency) / const.c.value
90+
return int(numpy.ceil(wmax / wstep))
91+
92+
7793
def invert_visibility(vis: Visibility, params={}):
7894
"""Invert to make dirty Image and PSF
7995
@@ -98,11 +114,13 @@ def invert_visibility(vis: Visibility, params={}):
98114
log.debug("invert_visibility: Gridding by w projection")
99115

100116
wstep = get_parameter(params, "wstep", 10000.0)
101-
102-
wcachesize = int(numpy.ceil(numpy.abs(vis.data['uvw'][:, 2]).max() * reffrequency.value / (const.c.value * wstep)))
117+
wcachesize = w_cache_size(vis, wstep)
103118
log.debug("invert_visibility: Making w-kernel cache of %d kernels" % wcachesize)
104-
wcache = pylru.FunctionCacheManager(lambda iw: wkernaf(N=256, theta=theta, w=iw * wstep, s=15, Qpx=4), 10000)
105-
imgfn = lambda *x: wcacheimg(*x, wstep=wstep, wcache=wcache)
119+
120+
cache_fn = w_conj_kernel_fn(pylru.FunctionCacheManager(w_kernel, wcachesize))
121+
imgfn = functools.partial(w_cache_imaging,
122+
wstep=wstep, kernel_cache=cache_fn,
123+
NpixFF=256, NpixKern=15, Qpx=4)
106124
else:
107125
raise NotImplementedError("gridding algorithm %s not supported" % gridding_algorithm)
108126

@@ -123,9 +141,8 @@ def invert_visibility(vis: Visibility, params={}):
123141
for pol in range(npol):
124142
log.debug('invert_visibility: Inverting channel %d, polarisation %d' % (channel, pol))
125143
d[channel, pol, :, :], p[channel, 0, :, :], pmax = \
126-
doimg(theta, 1.0 / cellsize, vis.data['uvw'] *
127-
(vis.frequency[channel] / const.c).value,
128-
vis.data['vis'][:, channel, pol], imgfn=imgfn)
144+
do_imaging(theta, 1.0 / cellsize, vis.uvw_lambda(channel),
145+
vis.vis[:, channel, pol], imgfn=imgfn)
129146
assert pmax > 0.0, ("No data gridded for channel %d" % channel)
130147
else:
131148
raise NotImplementedError("mode %s not supported" % spectral_mode)
@@ -138,9 +155,8 @@ def invert_visibility(vis: Visibility, params={}):
138155

139156
return dirty, psf, pmax
140157

141-
142158
def predict_visibility(vis: Visibility, sm: SkyModel, params={}) -> Visibility:
143-
"""Predict the visibility (in place) from a SkyModel including both components and images
159+
"""Predict the visibility from a SkyModel including both components and images
144160
145161
:param vis:
146162
:type Visibility: Visibility to be processed
@@ -150,8 +166,10 @@ def predict_visibility(vis: Visibility, sm: SkyModel, params={}) -> Visibility:
150166
"""
151167
shape, reffrequency, cellsize, w, imagecentre = create_wcs_from_visibility(vis, params=params)
152168

153-
vshape = vis.data['vis'].shape
154-
vis.data['vis'] = numpy.zeros(vshape)
169+
# Create copy of visibilities
170+
vis = copy.copy(vis)
171+
vis.data = copy.copy(vis.data)
172+
vis.data['vis'] = numpy.zeros(vis.vis.shape)
155173

156174
spectral_mode = get_parameter(params, 'spectral_mode', 'channel')
157175
log.debug('predict_visibility: spectral mode is %s' % spectral_mode)
@@ -160,75 +178,59 @@ def predict_visibility(vis: Visibility, sm: SkyModel, params={}) -> Visibility:
160178
log.debug("predict_visibility: Predicting Visibility from sky model images")
161179

162180
for im in sm.images:
163-
wimage = im.wcs
164-
165-
ishape = sm.images[0].data.shape
166-
nchan = ishape[0]
167-
npol = ishape[1]
168-
npixel = ishape[3]
169-
vshape = vis.data['vis'].shape
170-
171-
assert ishape[0] == vshape[1], "Image and visibility have different number of polarisations: %d %d" % (
172-
ishape[1], vshape[2])
173-
assert ishape[0] == len(vis.frequency), "Image and visibility have different number of channels %d %d" % \
174-
(ishape[0], len(vis.frequency))
175-
cellsize = abs(wimage.wcs.cdelt[0]) * numpy.pi / 180.0
176-
theta = npixel * cellsize
181+
assert_same_chan_pol(vis, im)
182+
183+
# Determine image size
184+
cellsize = abs(im.wcs.wcs.cdelt[0]) * numpy.pi / 180.0
185+
theta = im.npixel * cellsize
177186
log.debug("predict_visibility: Image cellsize %f radians" % cellsize)
178187
log.debug("predict_visibility: Field of view %f radians" % theta)
179188
assert (theta / numpy.sqrt(2) < 1.0), "Field of view larger than celestial sphere"
180189

190+
# Parameterise imaging
181191
wstep = get_parameter(params, "wstep", 10000.0)
182-
wcachesize = int(numpy.ceil(numpy.abs(vis.data['uvw'][:, 2]).max() * reffrequency.value / const.c.value /
183-
wstep))
192+
wcachesize = w_cache_size(vis, wstep)
184193
log.debug("predict_visibility: Making w-kernel cache of %d kernels" % wcachesize)
185-
wcache = pylru.FunctionCacheManager(lambda iw: wkernaf(N=256, theta=theta, w=iw * wstep, s=15, Qpx=4),
186-
10000)
187-
predfn = lambda *x: wcachefwd(*x, wstep=wstep, wcache=wcache)
194+
195+
cache_fn = w_conj_kernel_fn(pylru.FunctionCacheManager(w_kernel, wcachesize))
196+
predfn = functools.partial(w_cache_predict,
197+
wstep=wstep, kernel_cache=cache_fn,
198+
NpixFF=256, NpixKern=15, Qpx=4)
188199

189200
spectral_mode = get_parameter(params, 'spectral_mode', 'channel')
190201
log.debug('predict_visibility: spectral mode is %s' % spectral_mode)
191202

192203
if spectral_mode == 'channel':
193-
for channel in range(nchan):
194-
uvw = vis.data['uvw'] * (vis.frequency[channel] / const.c).value
195-
for pol in range(npol):
204+
for channel in range(im.nchan):
205+
uvw = vis.uvw_lambda(channel)
206+
for pol in range(im.npol):
196207
log.debug('predict_visibility: Predicting from image channel %d, polarisation %d' % (
197208
channel, pol))
198-
puvw, dv = dopredict(theta, 1.0 / cellsize, uvw, sm.images[0].data[channel, pol, :, :],
199-
predfn=predfn)
200-
vis.data['vis'][:, channel, pol] = vis.data['vis'][:, channel, pol] + dv
209+
img = sm.images[0].data[channel, pol, :, :]
210+
dv = do_predict(theta, 1.0 / cellsize, uvw, img, predfn)
211+
vis.vis[:, channel, pol] += dv
201212
else:
202213
raise NotImplementedError("mode %s not supported" % spectral_mode)
203214

204-
log.debug("predict_visibility: Finished predicting Visibility from sky model images")
215+
log.debug("predict_visibility: Finished predicting Visibility from sky model images")
205216

206217
if len(sm.components):
207218
log.debug("predict_visibility: Predicting Visibility from sky model components")
208219

209-
for icomp in range(len(sm.components)):
210-
comp = sm.components[icomp]
211-
cshape = comp.flux.shape
212-
assert len(cshape) == 2, "Flux should be two dimensional (pol, freq)"
213-
nchan = cshape[0]
214-
npol = cshape[1]
215-
216-
vshape = vis.data['vis'].shape
217-
log.debug("predict_visibility: visibility shape = %s" % str(vis.data['vis'].shape))
218-
219-
assert vshape[1] == nchan, "Component %d and visibility %d have different number of channels" % \
220-
(cshape[0], len(vis.frequency))
221-
assert vshape[2] == npol, "Component %d and visibility %d have different number of polarisations" % (
222-
npol, vshape[2])
223-
# dc = comp.direction.represent_as(CartesianRepresentation)
220+
for icomp, comp in enumerate(sm.components):
221+
222+
log.debug("predict_visibility: visibility shape = %s" % str(vis.vis.shape))
223+
assert_same_chan_pol(vis, comp)
224+
224225
l,m,n = skycoord_to_lmn(comp.direction, vis.phasecentre)
225226
log.debug('fourier_transforms.predict_visibility: Cartesian representation of component %d = (%f, %f, %f)'
226227
% (icomp, l,m,n))
228+
227229
if spectral_mode =='channel':
228-
for channel in range(nchan):
230+
for channel in range(comp.nchan):
229231
uvw = vis.uvw_lambda(channel)
230232
phasor = simulate_point(uvw, l, m)
231-
for pol in range(npol):
233+
for pol in range(comp.npol):
232234
log.debug(
233235
'predict_visibility: Predicting from component %d channel %d, polarisation %d' % (
234236
icomp, channel,

‎arl/test_support.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -151,6 +151,12 @@ def create_named_configuration(name: str = 'LOWBD2', params={}):
151151
location=location,
152152
mount='altaz',
153153
names='VLA_%d')
154+
elif name == 'VLAA_north':
155+
location = EarthLocation(lon="-107.6184", lat="90.000", height=2124.0)
156+
fc = create_configuration_from_file(antfile=crocodile_path("data/configurations/VLA_A_hor_xyz.csv"),
157+
location=location,
158+
mount='altaz',
159+
names='VLA_%d')
154160
else:
155161
fc = Configuration()
156162
raise UserWarning("No such Configuration %s" % name)

‎arl/visibility_operations.py‎

Lines changed: 12 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -281,24 +281,21 @@ def sum_visibility(vis: Visibility, direction: SkyCoord, params={}) -> numpy.arr
281281
:returns: flux[nch,npol], weight[nch,pol]
282282
"""
283283
log_parameters(params)
284-
dc = direction.represent_as(CartesianRepresentation)
284+
l,m,n = skycoord_to_lmn(direction, vis.phasecentre)
285285
log.debug('sum_visibility: Cartesian representation of direction = (%f, %f, %f)' % (
286-
dc.x, dc.y, dc.z))
287-
nchan = vis.data['vis'].shape[1]
288-
npol = vis.data['vis'].shape[2]
289-
flux = numpy.zeros([nchan, npol])
290-
weight = numpy.zeros([nchan, npol])
291-
for channel in range(nchan):
292-
uvw = vis.data['uvw'] * (vis.frequency[channel] / const.c).value
293-
uvw[:, 2] *= -1.0
294-
phasor = numpy.conj(simulate_point(uvw, dc.z, dc.y))
295-
for pol in range(npol):
286+
l,m,n))
287+
flux = numpy.zeros([vis.nchan, vis.npol])
288+
weight = numpy.zeros([vis.nchan, vis.npol])
289+
for channel in range(vis.nchan):
290+
uvw = vis.uvw_lambda(channel)
291+
phasor = numpy.conj(simulate_point(uvw, l,m))
292+
for pol in range(vis.npol):
296293
log.debug('sum_visibility: Summing visibility for channel %d, polarisation %d' % (
297294
channel, pol))
298-
flux[channel, pol] = flux[channel, pol] + \
299-
numpy.sum(numpy.real(vis.data['vis'][:, channel, pol] *
300-
vis.data['weight'][:, channel, pol] * phasor))
301-
weight[channel, pol] = weight[channel, pol] + numpy.sum(vis.data['weight'][:, channel, pol])
295+
ws = vis.weight[:, channel, pol]
296+
wvis = ws * vis.vis[:, channel, pol]
297+
flux[channel, pol] += numpy.sum(numpy.real(wvis * phasor))
298+
weight[channel, pol] += numpy.sum(ws)
302299
flux[weight > 0.0] = flux[weight > 0.0] / weight[weight > 0.0]
303300
flux[weight <= 0.0] = 0.0
304301
return flux, weight

‎examples/arl/imaging.ipynb‎

Lines changed: 28 additions & 303 deletions
Large diffs are not rendered by default.

‎tests/test_pipelines.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,4 +85,9 @@ def test_spectral_line_imaging(self):
8585

8686

8787
if __name__ == '__main__':
88+
import sys
89+
import logging
90+
log = logging.getLogger()
91+
log.setLevel(logging.DEBUG)
92+
log.addHandler(logging.StreamHandler(sys.stdout))
8893
unittest.main()

‎tests/test_visibility_operations.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -67,8 +67,8 @@ def test_phase_rotation_identity(self):
6767
rotatedvis = phaserotate_visibility(
6868
phaserotate_visibility(self.vismodel, newphasecentre),
6969
self.phasecentre)
70-
assert_allclose(rotatedvis.uvw, original_uvw, rtol=1e-11)
71-
assert_allclose(rotatedvis.vis, original_vis, rtol=1e-11)
70+
assert_allclose(rotatedvis.uvw, original_uvw, rtol=1e-10)
71+
assert_allclose(rotatedvis.vis, original_vis, rtol=1e-10)
7272

7373
def test_phase_rotation(self):
7474

0 commit comments

Comments
 (0)