def test_swapaxes():
x = np.random.normal(0, 10, size=(10, 12, 7))
d = da.from_array(x, chunks=(4, 5, 2))
assert_eq(np.swapaxes(x, 0, 1), da.swapaxes(d, 0, 1))
assert_eq(np.swapaxes(x, 2, 1), da.swapaxes(d, 2, 1))
assert_eq(x.swapaxes(2, 1), d.swapaxes(2, 1))
assert_eq(x.swapaxes(0, 0), d.swapaxes(0, 0))
assert_eq(x.swapaxes(1, 2), d.swapaxes(1, 2))
assert_eq(x.swapaxes(0, -1), d.swapaxes(0, -1))
assert_eq(x.swapaxes(-1, 1), d.swapaxes(-1, 1))
assert d.swapaxes(0, 1).name == d.swapaxes(0, 1).name
assert d.swapaxes(0, 1).name != d.swapaxes(1, 0).name
def test_swapaxes():
x = np.random.normal(0, 10, size=(10, 12, 7))
d = da.from_array(x, chunks=(4, 5, 2))
assert_eq(np.swapaxes(x, 0, 1), da.swapaxes(d, 0, 1))
assert_eq(np.swapaxes(x, 2, 1), da.swapaxes(d, 2, 1))
assert_eq(x.swapaxes(2, 1), d.swapaxes(2, 1))
assert_eq(x.swapaxes(0, 0), d.swapaxes(0, 0))
assert_eq(x.swapaxes(1, 2), d.swapaxes(1, 2))
assert_eq(x.swapaxes(0, -1), d.swapaxes(0, -1))
assert_eq(x.swapaxes(-1, 1), d.swapaxes(-1, 1))
assert d.swapaxes(0, 1).name == d.swapaxes(0, 1).name
assert d.swapaxes(0, 1).name != d.swapaxes(1, 0).name
def _calc_unitcell_vectors(self):
"""The vectors that define the shape of the unit cell in each frame
Returns
-------
vectors : da.ndarray, shape(n_frames, 3, 3)
Vectors defining the shape of the unit cell in each frame.
The semantics of this array are that the shape of the unit cell
in frame ``i`` are given by the three vectors, ``value[i, 0, :]``,
``value[i, 1, :]``, and ``value[i, 2, :]``.
"""
if self.unitcell_lengths is None or self.unitcell_angles is None:
return None
v1, v2, v3 = lengths_and_angles_to_box_vectors(
self._unitcell_lengths[:, 0], # a
self._unitcell_lengths[:, 1], # b
self._unitcell_lengths[:, 2], # c
self._unitcell_angles[:, 0], # alpha
self._unitcell_angles[:, 1], # beta
self._unitcell_angles[:, 2], # gamma
)
return da.swapaxes(da.dstack((v1, v2, v3)), 1, 2)
def _jones2col(**kw):
args = OmegaConf.create(kw)
from omegaconf import ListConfig
if not isinstance(args.ms, list) and not isinstance(args.ms, ListConfig):
args.ms = [args.ms]
OmegaConf.set_struct(args, True)
import numpy as np
from daskms.experimental.zarr import xds_from_zarr
from daskms import xds_from_ms, xds_to_table
import dask.array as da
import dask
from africanus.calibration.utils import chunkify_rows
from africanus.calibration.utils.dask import corrupt_vis
# get net gains
G = xds_from_zarr(args.gain_table + '::G')
# chunking info
t_chunks = G[0].t_chunk.data
if len(t_chunks) > 1:
t_chunks = G[0].t_chunk.data[1:-1] - G[0].t_chunk.data[0:-2]
assert (t_chunks == t_chunks[0]).all()
utpc = t_chunks[0]
else:
utpc = t_chunks[0]
times = xds_from_ms(args.ms[0], columns=['TIME'])[0].get('TIME').data.compute()
row_chunks, tbin_idx, tbin_counts = chunkify_rows(times, utimes_per_chunk=utpc, daskify_idx=True)
f_chunks = G[0].f_chunk.data
if len(f_chunks) > 1:
f_chunks = G[0].f_chunk.data[1:-1] - G[0].f_chunk.data[0:-2]
assert (f_chunks == f_chunks[0]).all()
chan_chunks = f_chunks[0]
else:
if f_chunks[0]:
chan_chunks = f_chunks[0]
else:
chan_chunks = -1
columns = ('DATA', 'FLAG', 'FLAG_ROW', 'ANTENNA1', 'ANTENNA2')
if args.acol is not None:
columns += (args.acol,)
# open MS
xds = xds_from_ms(args.ms[0], chunks={'row': row_chunks, 'chan': chan_chunks},
columns=columns,
group_cols=('FIELD_ID', 'DATA_DESC_ID', 'SCAN_NUMBER'))
# Current hack probably only works for single field and DDID
try:
assert len(xds) == len(G)
except Exception as e:
raise ValueError("Number of datasets in gains do not "
"match those in MS")
# assuming scans are aligned
out_data = []
for g, ds in zip(G, xds):
try:
assert g.SCAN_NUMBER == ds.SCAN_NUMBER
except Exception as e:
raise ValueError("Scans not aligned")
nrow = ds.dims['row']
nchan = ds.dims['chan']
ncorr = ds.dims['corr']
# need to swap axes for africanus
jones = da.swapaxes(g.gains.data, 1, 2)
flag = ds.FLAG.data
frow = ds.FLAG_ROW.data
ant1 = ds.ANTENNA1.data
ant2 = ds.ANTENNA2.data
frow = (frow | (ant1 == ant2))
flag = (flag[:, :, 0] | flag[:, :, -1])
flag = da.logical_or(flag, frow[:, None])
if args.acol is not None:
acol = ds.get(args.acol).data.reshape(nrow, nchan, 1, ncorr)
else:
acol = da.ones((nrow, nchan, 1, ncorr),
chunks=(row_chunks, chan_chunks, 1, -1),
dtype=jones.dtype)
cvis = corrupt_vis(tbin_idx, tbin_counts, ant1, ant2, jones, acol)
# compare where unflagged
if args.compareto is not None:
flag = flag.compute()
vis = ds.get(args.compareto).values[~flag]
print("Max abs difference = ", np.abs(cvis.compute()[~flag] - vis).max())
quit()
out_ds = ds.assign(**{args.mueller_column: (("row", "chan", "corr"), cvis)})
out_data.append(out_ds)
writes = xds_to_table(out_data, args.ms[0], columns=[args.mueller_column])
dask.compute(writes)
