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)