def test_rechunking(ms, tmp_path_factory, prechunking, postchunking): store = tmp_path_factory.mktemp("zarr_store") ref_datasets = xds_from_ms(ms) for i, ds in enumerate(ref_datasets): chunks = ds.chunks row = sum(chunks["row"]) chan = sum(chunks["chan"]) corr = sum(chunks["corr"]) ref_datasets[i] = ds.assign_coords( row=np.arange(row), chan=np.arange(chan), corr=np.arange(corr), dummy=np.arange(10) # Orphan coordinate. ) chunked_datasets = [ds.chunk(prechunking) for ds in ref_datasets] dask.compute(xds_to_zarr(chunked_datasets, store)) rechunked_datasets = [ds.chunk(postchunking) for ds in xds_from_zarr(store)] dask.compute(xds_to_zarr(rechunked_datasets, store, rechunk=True)) rechunked_datasets = xds_from_zarr(store) assert all([ds.equals(rds) for ds, rds in zip(rechunked_datasets, ref_datasets)])
def test_add_datavars(ms, tmp_path_factory, prechunking, postchunking): store = tmp_path_factory.mktemp("zarr_store") ref_datasets = xds_from_ms(ms) for i, ds in enumerate(ref_datasets): chunks = ds.chunks row = sum(chunks["row"]) chan = sum(chunks["chan"]) corr = sum(chunks["corr"]) ref_datasets[i] = ds.assign_coords( row=np.arange(row), chan=np.arange(chan), corr=np.arange(corr), dummy=np.arange(10) # Orphan coordinate. ) chunked_datasets = [ds.chunk(prechunking) for ds in ref_datasets] dask.compute(xds_to_zarr(chunked_datasets, store)) rechunked_datasets = [ds.chunk(postchunking) for ds in xds_from_zarr(store)] augmented_datasets = [ds.assign({"DUMMY": (("row", "chan", "corr"), da.zeros_like(ds.DATA.data))}) for ds in rechunked_datasets] dask.compute(xds_to_zarr(augmented_datasets, store, rechunk=True)) augmented_datasets = xds_from_zarr(store) assert all([ds.DUMMY.chunks == cds.DATA.chunks for ds, cds in zip(augmented_datasets, chunked_datasets)])
def test_xds_to_zarr_coords(tmp_path_factory): zarr_store = tmp_path_factory.mktemp("zarr_coords") / "test.zarr" data = da.ones((100, 16, 4), chunks=(10, 4, 1), dtype=np.complex64) rowid = da.arange(100, chunks=10) data_vars = {"DATA": (("row", "chan", "corr"), data)} coords = { "ROWID": (("row",), rowid), "chan": (("chan",), np.arange(16)), "foo": (("foo",), np.arange(4)), } ds = [Dataset(data_vars, coords=coords)] writes = xds_to_zarr(ds, zarr_store) dask.compute(writes) rds = xds_from_zarr(zarr_store) assert len(ds) == len(rds) for ods, nds in zip(ds, rds): for c, v in ods.data_vars.items(): assert_array_equal(v.data, getattr(nds, c).data) for c, v in ods.coords.items(): assert_array_equal(v.data, getattr(nds, c).data)
def test_basic_roundtrip(tmp_path): path = tmp_path / "test.zarr" # We need >10 datasets to be sure roundtripping is consistent. xdsl = [Dataset({'x': (('y',), da.ones(i))}) for i in range(1, 12)] dask.compute(xds_to_zarr(xdsl, path)) xdsl = xds_from_zarr(path) dask.compute(xds_to_zarr(xdsl, path))
def test_xds_to_zarr(ms, spw_table, ant_table, tmp_path_factory): zarr_store = tmp_path_factory.mktemp("zarr_store") / "test.zarr" spw_store = zarr_store.parent / f"{zarr_store.name}::SPECTRAL_WINDOW" ant_store = zarr_store.parent / f"{zarr_store.name}::ANTENNA" ms_datasets = xds_from_ms(ms) spw_datasets = xds_from_table(spw_table, group_cols="__row__") ant_datasets = xds_from_table(ant_table) for i, ds in enumerate(ms_datasets): dims = ds.dims row, chan, corr = (dims[d] for d in ("row", "chan", "corr")) ms_datasets[i] = ds.assign_coords( **{ "chan": (("chan", ), np.arange(chan)), "corr": (("corr", ), np.arange(corr)), }) main_zarr_writes = xds_to_zarr(ms_datasets, zarr_store) assert len(ms_datasets) == len(main_zarr_writes) for ms_ds, zw_ds in zip(ms_datasets, main_zarr_writes): for k, _ in ms_ds.attrs[DASKMS_PARTITION_KEY]: assert getattr(ms_ds, k) == getattr(zw_ds, k) writes = [main_zarr_writes] writes.extend(xds_to_zarr(spw_datasets, spw_store)) writes.extend(xds_to_zarr(ant_datasets, ant_store)) dask.compute(writes) zarr_datasets = xds_from_zarr(zarr_store, chunks={"row": 1}) for ms_ds, zarr_ds in zip(ms_datasets, zarr_datasets): # Check data variables assert ms_ds.data_vars, "MS Dataset has no variables" for name, var in ms_ds.data_vars.items(): zdata = getattr(zarr_ds, name).data assert type(zdata) is type(var.data) # noqa assert_array_equal(var.data, zdata) # Check coordinates assert ms_ds.coords, "MS Datset has no coordinates" for name, var in ms_ds.coords.items(): zdata = getattr(zarr_ds, name).data assert type(zdata) is type(var.data) # noqa assert_array_equal(var.data, zdata) # Check dataset attributes for k, v in ms_ds.attrs.items(): zattr = getattr(zarr_ds, k) assert_array_equal(zattr, v)
def test_storage_zarr(ms, tmp_path_factory): zarr_store = tmp_path_factory.mktemp("zarr") / "test.zarr" oxdsl = xds_from_ms(ms) writes = xds_to_zarr(oxdsl, zarr_store) dask.compute(writes) oxdsl = xds_from_zarr(zarr_store) writes = xds_to_storage_table(oxdsl, zarr_store) oxdsl = dask.compute(oxdsl)[0] dask.compute(writes) xdsl = dask.compute(xds_from_zarr(zarr_store))[0] assert all([xds.equals(oxds) for xds, oxds in zip(xdsl, oxdsl)])
def xds_from_storage_table(store, **kwargs): from daskms.utils import dataset_type typ = dataset_type(store) if typ == "casa": return xds_from_table(store, **kwargs) elif typ == "zarr": from daskms.experimental.zarr import xds_from_zarr return xds_from_zarr(store, **kwargs) elif typ == "parquet": from daskms.experimental.arrow import xds_from_parquet return xds_from_parquet(store, **kwargs) else: raise TypeError(f"Unknown dataset {typ}")
def xds_from_storage_ms(store, **kwargs): if not isinstance(store, DaskMSStore): store = DaskMSStore(store, **kwargs.pop("storage_options", {})) typ = store.type() if typ == "casa": return xds_from_ms(store, **kwargs) elif typ == "zarr": from daskms.experimental.zarr import xds_from_zarr return xds_from_zarr(store, **kwargs) elif typ == "parquet": from daskms.experimental.arrow import xds_from_parquet return xds_from_parquet(store, **kwargs) else: raise TypeError(f"Unknown dataset {typ}")
def test_zarr_string_array(tmp_path_factory): zarr_store = tmp_path_factory.mktemp("string-arrays") / "test.zarr" data = ["hello", "this", "strange new world", "full of", "interesting", "stuff"] data = np.array(data, dtype=object).reshape(3, 2) data = da.from_array(data, chunks=((2, 1), (1, 1))) datasets = [Dataset({"DATA": (("x", "y"), data)})] writes = xds_to_zarr(datasets, zarr_store) dask.compute(writes) new_datasets = xds_from_zarr(zarr_store) assert len(new_datasets) == len(datasets) for nds, ds in zip(new_datasets, datasets): assert_array_equal(nds.DATA.data, ds.DATA.data)
def convolver(x): model = da.from_array(x, chunks=(1, nx, ny), name=False) xds = xds_from_zarr(args.weight_table, chunks={ 'row': row_chunk, 'chan': bin_counts })[0] convolvedim = hessian(xds.UVW.data, freqs, model, freq_bin_idx, freq_bin_counts, cell_rad, weights=xds.WEIGHT.data.astype( args.output_type), nthreads=args.nvthreads, epsilon=args.epsilon, do_wstacking=args.wstack, double_accum=args.double_accum) return convolvedim
def test_multiprocess_create(ms, tmp_path_factory): zarr_store = tmp_path_factory.mktemp("zarr_store") / "test.zarr" ms_datasets = xds_from_ms(ms) for i, ds in enumerate(ms_datasets): ms_datasets[i] = ds.chunk({"row": 1}) writes = xds_to_zarr(ms_datasets, zarr_store) dask.compute(writes, scheduler="processes") zds = xds_from_zarr(zarr_store) for zds, msds in zip(zds, ms_datasets): for k, v in msds.data_vars.items(): assert_array_equal(v, getattr(zds, k)) for k, v in msds.coords.items(): assert_array_equal(v, getattr(zds, k)) for k, v in msds.attrs.items(): assert_array_equal(v, getattr(zds, k))
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)
def _forward(**kw): args = OmegaConf.create(kw) OmegaConf.set_struct(args, True) import numpy as np import numexpr as ne import dask import dask.array as da from dask.distributed import performance_report from pfb.utils.fits import load_fits, set_wcs, save_fits, data_from_header from pfb.opt.hogbom import hogbom from astropy.io import fits print("Loading residual", file=log) residual = load_fits(args.residual, dtype=args.output_type).squeeze() nband, nx, ny = residual.shape hdr = fits.getheader(args.residual) print("Loading psf", file=log) psf = load_fits(args.psf, dtype=args.output_type).squeeze() _, nx_psf, ny_psf = psf.shape hdr_psf = fits.getheader(args.psf) wsums = np.amax(psf.reshape(-1, nx_psf*ny_psf), axis=1) wsum = np.sum(wsums) psf /= wsum psf_mfs = np.sum(psf, axis=0) assert (psf_mfs.max() - 1.0) < 1e-4 residual /= wsum residual_mfs = np.sum(residual, axis=0) # get info required to set WCS ra = np.deg2rad(hdr['CRVAL1']) dec = np.deg2rad(hdr['CRVAL2']) radec = [ra, dec] cell_deg = np.abs(hdr['CDELT1']) if cell_deg != np.abs(hdr['CDELT2']): raise NotImplementedError('cell sizes have to be equal') cell_rad = np.deg2rad(cell_deg) l_coord, ref_l = data_from_header(hdr, axis=1) l_coord -= ref_l m_coord, ref_m = data_from_header(hdr, axis=2) m_coord -= ref_m freq_out, ref_freq = data_from_header(hdr, axis=3) hdr_mfs = set_wcs(cell_deg, cell_deg, nx, ny, radec, ref_freq) save_fits(args.output_filename + '_residual_mfs.fits', residual_mfs, hdr_mfs, dtype=args.output_type) rms = np.std(residual_mfs) rmax = np.abs(residual_mfs).max() print("Initial peak residual = %f, rms = %f" % (rmax, rms), file=log) # load beam if args.beam_model is not None: if args.beam_model.endswith('.fits'): # beam already interpolated bhdr = fits.getheader(args.beam_model) l_coord_beam, ref_lb = data_from_header(bhdr, axis=1) l_coord_beam -= ref_lb if not np.array_equal(l_coord_beam, l_coord): raise ValueError("l coordinates of beam model do not match those of image. Use power_beam_maker to interpolate to fits header.") m_coord_beam, ref_mb = data_from_header(bhdr, axis=2) m_coord_beam -= ref_mb if not np.array_equal(m_coord_beam, m_coord): raise ValueError("m coordinates of beam model do not match those of image. Use power_beam_maker to interpolate to fits header.") freq_beam, _ = data_from_header(bhdr, axis=freq_axis) if not np.array_equal(freq_out, freq_beam): raise ValueError("Freqs of beam model do not match those of image. Use power_beam_maker to interpolate to fits header.") beam_image = load_fits(args.beam_model, dtype=args.output_type).squeeze() elif args.beam_model.lower() == "jimbeam": from katbeam import JimBeam if args.band.lower() == 'l': beam = JimBeam('MKAT-AA-L-JIM-2020') elif args.band.lower() == 'uhf': beam = JimBeam('MKAT-AA-UHF-JIM-2020') else: raise ValueError("Unkown band %s"%args.band[i]) xx, yy = np.meshgrid(l_coord, m_coord, indexing='ij') beam_image = np.zeros(residual.shape, dtype=args.output_type) for v in range(freq_out.size): # freq must be in MHz beam_image[v] = beam.I(xx, yy, freq_out[v]/1e6).astype(args.output_type) else: beam_image = np.ones((nband, nx, ny), dtype=args.output_type) if args.mask is not None: mask = load_fits(args.mask).squeeze() assert mask.shape == (nx, ny) beam_image *= mask[None, :, :] beam_image = da.from_array(beam_image, chunks=(1, -1, -1)) # if weight table is provided we use the vis space Hessian approximation if args.weight_table is not None: print("Solving for update using vis space approximation", file=log) normfact = wsum from pfb.utils.misc import plan_row_chunk from daskms.experimental.zarr import xds_from_zarr xds = xds_from_zarr(args.weight_table)[0] nrow = xds.row.size freq = xds.chan.data nchan = freq.size # bin edges fmin = freq.min() fmax = freq.max() fbins = np.linspace(fmin, fmax, nband + 1) # chan <-> band mapping band_mapping = {} chan_chunks = {} freq_bin_idx = {} freq_bin_counts = {} band_map = np.zeros(freq.size, dtype=np.int32) for band in range(nband): indl = freq >= fbins[band] indu = freq < fbins[band + 1] + 1e-6 band_map = np.where(indl & indu, band, band_map) # to dask arrays bands, bin_counts = np.unique(band_map, return_counts=True) band_mapping = tuple(bands) chan_chunks = {'chan': tuple(bin_counts)} freq = da.from_array(freq, chunks=tuple(bin_counts)) bin_idx = np.append(np.array([0]), np.cumsum(bin_counts))[0:-1] freq_bin_idx = da.from_array(bin_idx, chunks=1) freq_bin_counts = da.from_array(bin_counts, chunks=1) max_chan_chunk = bin_counts.max() bin_counts = tuple(bin_counts) # the first factor of 3 accounts for the intermediate visibilities # produced in Hessian (i.e. complex data + real weights) memory_per_row = (3 * max_chan_chunk * xds.WEIGHT.data.itemsize + 3 * xds.UVW.data.itemsize) # get approx image size pixel_bytes = np.dtype(args.output_type).itemsize band_size = nx * ny * pixel_bytes if args.host_address is None: # nworker bands on single node row_chunk = plan_row_chunk(args.mem_limit/args.nworkers, band_size, nrow, memory_per_row, args.nthreads_per_worker) else: # single band per node row_chunk = plan_row_chunk(args.mem_limit, band_size, nrow, memory_per_row, args.nthreads_per_worker) print("nrows = %i, row chunks set to %i for a total of %i chunks per node" % (nrow, row_chunk, int(np.ceil(nrow / row_chunk))), file=log) residual = da.from_array(residual, chunks=(1, -1, -1)) x0 = da.zeros((nband, nx, ny), chunks=(1, -1, -1), dtype=residual.dtype) xds = xds_from_zarr(args.weight_table, chunks={'row': -1, #row_chunk, 'chan': bin_counts})[0] from pfb.opt.pcg import pcg_wgt model = pcg_wgt(xds.UVW.data, xds.WEIGHT.data.astype(args.output_type), residual, x0, beam_image, freq, freq_bin_idx, freq_bin_counts, cell_rad, args.wstack, args.epsilon, args.double_accum, args.nvthreads, args.sigmainv, wsum, args.cg_tol, args.cg_maxit, args.cg_minit, args.cg_verbose, args.cg_report_freq, args.backtrack).compute() else: # we use the image space approximation print("Solving for update using image space approximation", file=log) normfact = 1.0 from pfb.operators.psf import hessian from ducc0.fft import r2c iFs = np.fft.ifftshift npad_xl = (nx_psf - nx)//2 npad_xr = nx_psf - nx - npad_xl npad_yl = (ny_psf - ny)//2 npad_yr = ny_psf - ny - npad_yl padding = ((0, 0), (npad_xl, npad_xr), (npad_yl, npad_yr)) unpad_x = slice(npad_xl, -npad_xr) unpad_y = slice(npad_yl, -npad_yr) lastsize = ny + np.sum(padding[-1]) psf_pad = iFs(psf, axes=(1, 2)) psfhat = r2c(psf_pad, axes=(1, 2), forward=True, nthreads=nthreads, inorm=0) psfhat = da.from_array(psfhat, chunks=(1, -1, -1)) residual = da.from_array(residual, chunks=(1, -1, -1)) x0 = da.zeros((nband, nx, ny), chunks=(1, -1, -1)) from pfb.opt.pcg import pcg_psf model = pcg_psf(psfhat, residual, x0, beam_image, args.sigmainv, args.nvthreads, padding, unpad_x, unpad_y, lastsize, args.cg_tol, args.cg_maxit, args.cg_minit, args.cg_verbose, args.cg_report_freq, args.backtrack).compute() print("Saving results", file=log) save_fits(args.output_filename + '_update.fits', model, hdr) model_mfs = np.mean(model, axis=0) save_fits(args.output_filename + '_update_mfs.fits', model_mfs, hdr_mfs) print("All done here.", file=log)
def _clean(**kw): args = OmegaConf.create(kw) OmegaConf.set_struct(args, True) import numpy as np import numexpr as ne import dask import dask.array as da from dask.distributed import performance_report from pfb.utils.fits import load_fits, set_wcs, save_fits, data_from_header from pfb.opt.hogbom import hogbom from astropy.io import fits print("Loading dirty", file=log) dirty = load_fits(args.dirty, dtype=args.output_type).squeeze() nband, nx, ny = dirty.shape hdr = fits.getheader(args.dirty) print("Loading psf", file=log) psf = load_fits(args.psf, dtype=args.output_type).squeeze() _, nx_psf, ny_psf = psf.shape hdr_psf = fits.getheader(args.psf) wsums = np.amax(psf.reshape(-1, nx_psf * ny_psf), axis=1) wsum = np.sum(wsums) psf /= wsum psf_mfs = np.sum(psf, axis=0) assert (psf_mfs.max() - 1.0) < 1e-4 dirty /= wsum dirty_mfs = np.sum(dirty, axis=0) # get info required to set WCS ra = np.deg2rad(hdr['CRVAL1']) dec = np.deg2rad(hdr['CRVAL2']) radec = [ra, dec] cell_deg = np.abs(hdr['CDELT1']) if cell_deg != np.abs(hdr['CDELT2']): raise NotImplementedError('cell sizes have to be equal') cell_rad = np.deg2rad(cell_deg) freq_out, ref_freq = data_from_header(hdr, axis=3) hdr_mfs = set_wcs(cell_deg, cell_deg, nx, ny, radec, ref_freq) save_fits(args.output_filename + '_dirty_mfs.fits', dirty_mfs, hdr_mfs, dtype=args.output_type) # set up Hessian approximation if args.weight_table is not None: normfact = wsum from africanus.gridding.wgridder.dask import hessian from pfb.utils.misc import plan_row_chunk from daskms.experimental.zarr import xds_from_zarr xds = xds_from_zarr(args.weight_table)[0] nrow = xds.row.size freqs = xds.chan.data nchan = freqs.size # bin edges fmin = freqs.min() fmax = freqs.max() fbins = np.linspace(fmin, fmax, nband + 1) # chan <-> band mapping band_mapping = {} chan_chunks = {} freq_bin_idx = {} freq_bin_counts = {} band_map = np.zeros(freqs.size, dtype=np.int32) for band in range(nband): indl = freqs >= fbins[band] indu = freqs < fbins[band + 1] + 1e-6 band_map = np.where(indl & indu, band, band_map) # to dask arrays bands, bin_counts = np.unique(band_map, return_counts=True) band_mapping = tuple(bands) chan_chunks = {'chan': tuple(bin_counts)} freqs = da.from_array(freqs, chunks=tuple(bin_counts)) bin_idx = np.append(np.array([0]), np.cumsum(bin_counts))[0:-1] freq_bin_idx = da.from_array(bin_idx, chunks=1) freq_bin_counts = da.from_array(bin_counts, chunks=1) max_chan_chunk = bin_counts.max() bin_counts = tuple(bin_counts) # the first factor of 3 accounts for the intermediate visibilities # produced in Hessian (i.e. complex data + real weights) memory_per_row = (3 * max_chan_chunk * xds.WEIGHT.data.itemsize + 3 * xds.UVW.data.itemsize) # get approx image size pixel_bytes = np.dtype(args.output_type).itemsize band_size = nx * ny * pixel_bytes if args.host_address is None: # nworker bands on single node row_chunk = plan_row_chunk(args.mem_limit / args.nworkers, band_size, nrow, memory_per_row, args.nthreads_per_worker) else: # single band per node row_chunk = plan_row_chunk(args.mem_limit, band_size, nrow, memory_per_row, args.nthreads_per_worker) print( "nrows = %i, row chunks set to %i for a total of %i chunks per node" % (nrow, row_chunk, int(np.ceil(nrow / row_chunk))), file=log) def convolver(x): model = da.from_array(x, chunks=(1, nx, ny), name=False) xds = xds_from_zarr(args.weight_table, chunks={ 'row': row_chunk, 'chan': bin_counts })[0] convolvedim = hessian(xds.UVW.data, freqs, model, freq_bin_idx, freq_bin_counts, cell_rad, weights=xds.WEIGHT.data.astype( args.output_type), nthreads=args.nvthreads, epsilon=args.epsilon, do_wstacking=args.wstack, double_accum=args.double_accum) return convolvedim else: normfact = 1.0 from pfb.operators.psf import hessian from ducc0.fft import r2c iFs = np.fft.ifftshift npad_xl = (nx_psf - nx) // 2 npad_xr = nx_psf - nx - npad_xl npad_yl = (ny_psf - ny) // 2 npad_yr = ny_psf - ny - npad_yl padding = ((0, 0), (npad_xl, npad_xr), (npad_yl, npad_yr)) unpad_x = slice(npad_xl, -npad_xr) unpad_y = slice(npad_yl, -npad_yr) lastsize = ny + np.sum(padding[-1]) psf_pad = iFs(psf, axes=(1, 2)) psfhat = r2c(psf_pad, axes=(1, 2), forward=True, nthreads=nthreads, inorm=0) psfhat = da.from_array(psfhat, chunks=(1, -1, -1)) def convolver(x): model = da.from_array(x, chunks=(1, nx, ny), name=False) convolvedim = hessian(model, psfhat, padding, nvthreads, unpad_x, unpad_y, lastsize) return convolvedim # psfo = PSF(psf, dirty.shape, nthreads=args.nthreads) # def convolver(x): return psfo.convolve(x) rms = np.std(dirty_mfs) rmax = np.abs(dirty_mfs).max() print("Iter %i: peak residual = %f, rms = %f" % (0, rmax, rms), file=log) residual = dirty.copy() residual_mfs = dirty_mfs.copy() model = np.zeros_like(residual) for k in range(args.nmiter): print("Running Hogbom", file=log) x = hogbom(residual, psf, gamma=args.hb_gamma, pf=args.hb_peak_factor, maxit=args.hb_maxit, verbosity=args.hb_verbose, report_freq=args.hb_report_freq) model += x print("Getting residual", file=log) convimage = convolver(model) dask.visualize(convimage, filename=args.output_filename + '_hessian' + str(k) + '_graph.pdf', optimize_graph=False) with performance_report(filename=args.output_filename + '_hessian' + str(k) + '_per.html'): convimage = dask.compute(convimage, optimize_graph=False)[0] ne.evaluate('dirty - convimage/normfact', out=residual, casting='same_kind') ne.evaluate('sum(residual, axis=0)', out=residual_mfs, casting='same_kind') rms = np.std(residual_mfs) rmax = np.abs(residual_mfs).max() print("Iter %i: peak residual = %f, rms = %f" % (k + 1, rmax, rms), file=log) print("Saving results", file=log) save_fits(args.output_filename + '_model.fits', model, hdr) model_mfs = np.mean(model, axis=0) save_fits(args.output_filename + '_model_mfs.fits', model_mfs, hdr_mfs) save_fits(args.output_filename + '_residual.fits', residual * wsums[:, None, None], hdr) save_fits(args.output_filename + '_residual.fits', residual_mfs, hdr_mfs) print("All done here.", file=log)