def _main(dest=sys.stdout): from pfb.parser import create_parser args = create_parser().parse_args() if not args.nthreads: import multiprocessing args.nthreads = multiprocessing.cpu_count() if not args.mem_limit: import psutil args.mem_limit = int(psutil.virtual_memory()[0] / 1e9) # 100% of memory by default import numpy as np import numba import numexpr import dask import dask.array as da from daskms import xds_from_ms, xds_from_table from astropy.io import fits from pfb.utils.fits import (set_wcs, load_fits, save_fits, compare_headers, data_from_header) from pfb.utils.restoration import fitcleanbeam from pfb.utils.misc import Gaussian2D from pfb.operators.gridder import Gridder from pfb.operators.psf import PSF from pfb.deconv.sara import sara from pfb.deconv.clean import clean from pfb.deconv.spotless import spotless from pfb.deconv.nnls import nnls from pfb.opt.pcg import pcg if not isinstance(args.ms, list): args.ms = [args.ms] pyscilog.log_to_file(args.outfile + '.log') pyscilog.enable_memory_logging(level=3) GD = vars(args) print('Input Options:', file=log) for key in GD.keys(): print(' %25s = %s' % (key, GD[key]), file=log) # get max uv coords over all fields uvw = [] u_max = 0.0 v_max = 0.0 all_freqs = [] for ims in args.ms: xds = xds_from_ms(ims, group_cols=('FIELD_ID', 'DATA_DESC_ID'), columns=('UVW'), chunks={'row': args.row_chunks}) spws = xds_from_table(ims + "::SPECTRAL_WINDOW", group_cols="__row__") spws = dask.compute(spws)[0] for ds in xds: uvw = ds.UVW.data u_max = da.maximum(u_max, abs(uvw[:, 0]).max()) v_max = da.maximum(v_max, abs(uvw[:, 1]).max()) uv_max = da.maximum(u_max, v_max) spw = spws[ds.DATA_DESC_ID] tmp_freq = spw.CHAN_FREQ.data.squeeze() all_freqs.append(list([tmp_freq])) uv_max = u_max.compute() del uvw # get Nyquist cell size from africanus.constants import c as lightspeed all_freqs = dask.compute(all_freqs) freq = np.unique(all_freqs) cell_N = 1.0 / (2 * uv_max * freq.max() / lightspeed) if args.cell_size is not None: cell_rad = args.cell_size * np.pi / 60 / 60 / 180 if cell_N / cell_rad < 1: raise ValueError( "Requested cell size too small. " "Super resolution factor = ", cell_N / cell_rad) print("Super resolution factor = %f" % (cell_N / cell_rad), file=dest) else: cell_rad = cell_N / args.super_resolution_factor args.cell_size = cell_rad * 60 * 60 * 180 / np.pi print("Cell size set to %5.5e arcseconds" % args.cell_size, file=dest) if args.nx is None or args.ny is None: from ducc0.fft import good_size fov = args.fov * 3600 npix = int(fov / args.cell_size) if npix % 2: npix += 1 args.nx = good_size(npix) args.ny = good_size(npix) if args.nband is None: args.nband = freq.size print("Image size set to (%i, %i, %i)" % (args.nband, args.nx, args.ny), file=dest) # mask if args.mask is not None: mask_array = load_fits(args.mask, dtype=args.real_type).squeeze() if mask_array.shape != (args.nx, args.ny): raise ValueError("Mask has incorrect shape.") # add freq axis mask_array = mask_array[None] def mask(x): return mask_array * x else: mask_array = None def mask(x): return x # init gridder R = Gridder( args.ms, args.nx, args.ny, args.cell_size, nband=args.nband, nthreads=args.nthreads, do_wstacking=args.do_wstacking, row_chunks=args.row_chunks, psf_oversize=args.psf_oversize, data_column=args.data_column, epsilon=args.epsilon, weight_column=args.weight_column, imaging_weight_column=args.imaging_weight_column, model_column=args.model_column, flag_column=args.flag_column, weighting=args.weighting, robust=args.robust, mem_limit=int( 0.8 * args.mem_limit)) # assumes gridding accounts for 80% memory freq_out = R.freq_out radec = R.radec print("PSF size set to (%i, %i, %i)" % (args.nband, R.nx_psf, R.ny_psf), file=dest) # get headers hdr = set_wcs(args.cell_size / 3600, args.cell_size / 3600, args.nx, args.ny, radec, freq_out) hdr_mfs = set_wcs(args.cell_size / 3600, args.cell_size / 3600, args.nx, args.ny, radec, np.mean(freq_out)) hdr_psf = set_wcs(args.cell_size / 3600, args.cell_size / 3600, R.nx_psf, R.ny_psf, radec, freq_out) hdr_psf_mfs = set_wcs(args.cell_size / 3600, args.cell_size / 3600, R.nx_psf, R.ny_psf, radec, np.mean(freq_out)) # psf if args.psf is not None: try: compare_headers(hdr_psf, fits.getheader(args.psf)) psf = load_fits(args.psf, dtype=args.real_type).squeeze() except BaseException: raise psf = R.make_psf() save_fits(args.outfile + '_psf.fits', psf, hdr_psf) else: psf = R.make_psf() save_fits(args.outfile + '_psf.fits', psf, hdr_psf) # Normalising by wsum (so that the PSF always sums to 1) results in the # most intuitive sig_21 values and by far the least bookkeeping. # However, we won't save the cubes that way as it destroys information # about the noise in image space. Note only the MFS images will have the # usual units of Jy/beam. wsums = np.amax(psf.reshape(args.nband, R.nx_psf * R.ny_psf), axis=1) wsum = np.sum(wsums) psf /= wsum psf_mfs = np.sum(psf, axis=0) # fit restoring psf GaussPar = fitcleanbeam(psf_mfs[None], level=0.5, pixsize=1.0) GaussPars = fitcleanbeam(psf, level=0.5, pixsize=1.0) cpsf_mfs = np.zeros(psf_mfs.shape, dtype=args.real_type) cpsf = np.zeros(psf.shape, dtype=args.real_type) lpsf = np.arange(-R.nx_psf / 2, R.nx_psf / 2) mpsf = np.arange(-R.ny_psf / 2, R.ny_psf / 2) xx, yy = np.meshgrid(lpsf, mpsf, indexing='ij') cpsf_mfs = Gaussian2D(xx, yy, GaussPar[0], normalise=False) for v in range(args.nband): cpsf[v] = Gaussian2D(xx, yy, GaussPars[v], normalise=False) from pfb.utils.fits import add_beampars GaussPar = list(GaussPar[0]) GaussPar[0] *= args.cell_size / 3600 GaussPar[1] *= args.cell_size / 3600 GaussPar = tuple(GaussPar) hdr_psf_mfs = add_beampars(hdr_psf_mfs, GaussPar) save_fits(args.outfile + '_cpsf_mfs.fits', cpsf_mfs, hdr_psf_mfs) save_fits(args.outfile + '_psf_mfs.fits', psf_mfs, hdr_psf_mfs) GaussPars = list(GaussPars) for b in range(args.nband): GaussPars[b] = list(GaussPars[b]) GaussPars[b][0] *= args.cell_size / 3600 GaussPars[b][1] *= args.cell_size / 3600 GaussPars[b] = tuple(GaussPars[b]) GaussPars = tuple(GaussPars) hdr_psf = add_beampars(hdr_psf, GaussPar, GaussPars) save_fits(args.outfile + '_cpsf.fits', cpsf, hdr_psf) # dirty if args.dirty is not None: try: compare_headers(hdr, fits.getheader(args.dirty)) dirty = load_fits(args.dirty).squeeze() except BaseException: raise dirty = R.make_dirty() save_fits(args.outfile + '_dirty.fits', dirty, hdr) else: dirty = R.make_dirty() save_fits(args.outfile + '_dirty.fits', dirty, hdr) dirty /= wsum dirty_mfs = np.sum(dirty, axis=0) save_fits(args.outfile + '_dirty_mfs.fits', dirty_mfs, hdr_mfs) quit() # initial model and residual if args.x0 is not None: try: compare_headers(hdr, fits.getheader(args.x0)) model = load_fits(args.x0, dtype=args.real_type).squeeze() if args.first_residual is not None: try: compare_headers(hdr, fits.getheader(args.first_residual)) residual = load_fits(args.first_residual, dtype=args.real_type).squeeze() except BaseException: residual = R.make_residual(model) save_fits(args.outfile + '_first_residual.fits', residual, hdr) else: residual = R.make_residual(model) save_fits(args.outfile + '_first_residual.fits', residual, hdr) residual /= wsum except BaseException: model = np.zeros((args.nband, args.nx, args.ny)) residual = dirty.copy() else: model = np.zeros((args.nband, args.nx, args.ny)) residual = dirty.copy() residual_mfs = np.sum(residual, axis=0) save_fits(args.outfile + '_first_residual_mfs.fits', residual_mfs, hdr_mfs) # smooth beam if args.beam_model is not None: if args.beam_model[-5:] == '.fits': beam_image = load_fits(args.beam_model, dtype=args.real_type).squeeze() if beam_image.shape != (args.nband, args.nx, args.ny): raise ValueError("Beam has incorrect shape") elif args.beam_model == "JimBeam": from katbeam import JimBeam if args.band.lower() == 'l': beam = JimBeam('MKAT-AA-L-JIM-2020') else: beam = JimBeam('MKAT-AA-UHF-JIM-2020') beam_image = np.zeros((args.nband, args.nx, args.ny), dtype=args.real_type) 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 xx, yy = np.meshgrid(l_coord, m_coord, indexing='ij') for v in range(args.nband): beam_image[v] = beam.I(xx, yy, freq_out[v]) def beam(x): return beam_image * x else: beam_image = None def beam(x): return x if args.init_nnls: print("Initialising with NNLS", file=log) model = nnls(psf, model, residual, mask=mask_array, beam_image=beam_image, hdr=hdr, hdr_mfs=hdr_mfs, outfile=args.outfile, maxit=1, nthreads=args.nthreads) residual = R.make_residual(beam(mask(model))) / wsum residual_mfs = np.sum(residual, axis=0) # deconvolve rmax = np.abs(residual_mfs).max() rms = np.std(residual_mfs) redo_dirty = False print("Peak of initial residual is %f and rms is %f" % (rmax, rms), file=dest) for i in range(0, args.maxit): # run minor cycle of choice modelp = model.copy() if args.deconv_mode == 'sara': model = sara(psf, model, residual, mask=mask_array, beam_image=beam_image, hessian=R.convolve, wsum=wsum, adapt_sig21=args.adapt_sig21, hdr=hdr, hdr_mfs=hdr_mfs, outfile=args.outfile, cpsf=cpsf, nthreads=args.nthreads, sig_21=args.sig_21, sigma_frac=args.sigma_frac, maxit=args.minormaxit, tol=args.minortol, gamma=args.gamma, psi_levels=args.psi_levels, psi_basis=args.psi_basis, pdtol=args.pdtol, pdmaxit=args.pdmaxit, pdverbose=args.pdverbose, positivity=args.positivity, cgtol=args.cgtol, cgminit=args.cgminit, cgmaxit=args.cgmaxit, cgverbose=args.cgverbose, pmtol=args.pmtol, pmmaxit=args.pmmaxit, pmverbose=args.pmverbose) elif args.deconv_mode == 'clean': model = clean(psf, model, residual, mask=mask_array, beam=beam_image, nthreads=args.nthreads, maxit=args.minormaxit, gamma=args.gamma, peak_factor=args.peak_factor, threshold=args.threshold, hbgamma=args.hbgamma, hbpf=args.hbpf, hbmaxit=args.hbmaxit, hbverbose=args.hbverbose) elif args.deconv_mode == 'spotless': model = spotless(psf, model, residual, mask=mask_array, beam_image=beam_image, hessian=R.convolve, wsum=wsum, adapt_sig21=args.adapt_sig21, cpsf=cpsf_mfs, hdr=hdr, hdr_mfs=hdr_mfs, outfile=args.outfile, sig_21=args.sig_21, sigma_frac=args.sigma_frac, nthreads=args.nthreads, gamma=args.gamma, peak_factor=args.peak_factor, maxit=args.minormaxit, tol=args.minortol, threshold=args.threshold, positivity=args.positivity, hbgamma=args.hbgamma, hbpf=args.hbpf, hbmaxit=args.hbmaxit, hbverbose=args.hbverbose, pdtol=args.pdtol, pdmaxit=args.pdmaxit, pdverbose=args.pdverbose, cgtol=args.cgtol, cgminit=args.cgminit, cgmaxit=args.cgmaxit, cgverbose=args.cgverbose, pmtol=args.pmtol, pmmaxit=args.pmmaxit, pmverbose=args.pmverbose) else: raise ValueError("Unknown deconvolution mode ", args.deconv_mode) # get residual if redo_dirty: # Need to do this if weights or Jones has changed # (eg. if we change robustness factor, reweight or calibrate) psf = R.make_psf() wsums = np.amax(psf.reshape(args.nband, R.nx_psf * R.ny_psf), axis=1) wsum = np.sum(wsums) psf /= wsum dirty = R.make_dirty() / wsum # compute in image space # residual = dirty - R.convolve(beam(mask(model))) / wsum residual = R.make_residual(beam(mask(model))) / wsum residual_mfs = np.sum(residual, axis=0) # save current iteration model_mfs = np.mean(model, axis=0) save_fits(args.outfile + '_major' + str(i + 1) + '_model_mfs.fits', model_mfs, hdr_mfs) save_fits(args.outfile + '_major' + str(i + 1) + '_model.fits', model, hdr) save_fits(args.outfile + '_major' + str(i + 1) + '_residual_mfs.fits', residual_mfs, hdr_mfs) save_fits(args.outfile + '_major' + str(i + 1) + '_residual.fits', residual * wsum, hdr) # check stopping criteria rmax = np.abs(residual_mfs).max() rms = np.std(residual_mfs) eps = np.linalg.norm(model - modelp) / np.linalg.norm(model) print("At iteration %i peak of residual is %f, rms is %f, current " "eps is %f" % (i + 1, rmax, rms, eps), file=dest) if eps < args.tol: break if args.mop_flux: print("Mopping flux", file=dest) # vague Gaussian prior on x def hess(x): return mask(beam(R.convolve(mask(beam(x))))) / wsum + 1e-6 * x def M(x): return x / 1e-6 # preconditioner x = pcg(hess, mask(beam(residual)), np.zeros(residual.shape, dtype=residual.dtype), M=M, tol=0.1 * args.cgtol, maxit=args.cgmaxit, minit=args.cgminit, verbosity=args.cgverbose) model += x # residual = dirty - R.convolve(beam(mask(model))) / wsum residual = R.make_residual(beam(mask(model))) / wsum save_fits(args.outfile + '_mopped_model.fits', model, hdr) save_fits(args.outfile + '_mopped_residual.fits', residual, hdr) model_mfs = np.mean(model, axis=0) save_fits(args.outfile + '_mopped_model_mfs.fits', model_mfs, hdr_mfs) residual_mfs = np.sum(residual, axis=0) save_fits(args.outfile + '_mopped_residual_mfs.fits', residual_mfs, hdr_mfs) rmax = np.abs(residual_mfs).max() rms = np.std(residual_mfs) print("After mopping flux peak of residual is %f, rms is %f" % (rmax, rms), file=dest) # if args.interp_model: # nband = args.nband # order = args.spectral_poly_order # phi.trim_fat(model) # I = np.argwhere(phi.mask).squeeze() # Ix = I[:, 0] # Iy = I[:, 1] # npix = I.shape[0] # # get components # beta = model[:, Ix, Iy] # # fit integrated polynomial to model components # # we are given frequencies at bin centers, convert to bin edges # ref_freq = np.mean(freq_out) # delta_freq = freq_out[1] - freq_out[0] # wlow = (freq_out - delta_freq/2.0)/ref_freq # whigh = (freq_out + delta_freq/2.0)/ref_freq # wdiff = whigh - wlow # # set design matrix for each component # Xdesign = np.zeros([freq_out.size, args.spectral_poly_order]) # for i in range(1, args.spectral_poly_order+1): # Xdesign[:, i-1] = (whigh**i - wlow**i)/(i*wdiff) # weights = psf_max[:, None] # dirty_comps = Xdesign.T.dot(weights*beta) # hess_comps = Xdesign.T.dot(weights*Xdesign) # comps = np.linalg.solve(hess_comps, dirty_comps) # np.savez(args.outfile + "spectral_comps", comps=comps, ref_freq=ref_freq, mask=np.any(model, axis=0)) if args.write_model: print("Writing model", file=dest) R.write_model(model) if args.make_restored: print("Making restored", file=dest) cpsfo = PSF(cpsf, residual.shape, nthreads=args.nthreads) restored = cpsfo.convolve(model) # residual needs to be in Jy/beam before adding to convolved model wsums = np.amax(psf.reshape(-1, R.nx_psf * R.ny_psf), axis=1) restored += residual / wsums[:, None, None] save_fits(args.outfile + '_restored.fits', restored, hdr) restored_mfs = np.mean(restored, axis=0) save_fits(args.outfile + '_restored_mfs.fits', restored_mfs, hdr_mfs) residual_mfs = np.sum(residual, axis=0)
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 _binterp(**kw): args = OmegaConf.create(kw) OmegaConf.set_struct(args, True) from pfb.utils.fits import save_fits import dask import dask.array as da import numpy as np from numba import jit from astropy.io import fits import warnings from africanus.rime import parallactic_angles from pfb.utils.fits import load_fits, save_fits, data_from_header from daskms import xds_from_ms, xds_from_table if args.ms is None: if args.beam_model.lower() == 'jimbeam': for image in args.image: mhdr = fits.getheader(image) l_coord, ref_l = data_from_header(mhdr, axis=1) l_coord -= ref_l m_coord, ref_m = data_from_header(mhdr, axis=2) m_coord -= ref_m if mhdr["CTYPE4"].lower() == 'freq': freq_axis = 4 stokes_axis = 3 elif mhdr["CTYPE3"].lower() == 'freq': freq_axis = 3 stokes_axis = 4 else: raise ValueError("Freq axis must be 3rd or 4th") freq, ref_freq = data_from_header(mhdr, axis=freq_axis) 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((freq.size, l_coord.size, m_coord.size), dtype=args.out_dtype) for v in range(freq.size): # freq must be in MHz beam_image[v] = beam.I(xx, yy, freq[v] / 1e6).astype( args.out_dtype) if args.output_dir in image: idx = len(args.output_dir) iname = image[idx::] outname = iname + '.' + args.postfix else: outname = image + '.' + args.postfix beam_image = np.expand_dims(beam_image, axis=3 - stokes_axis + 1) save_fits(args.output_dir + outname, beam_image, mhdr, dtype=args.out_dtype) else: raise NotImplementedError("Not there yet, sorry") print("All done here.", file=log) # @jit(nopython=True, nogil=True, cache=True) # def _unflagged_counts(flags, time_idx, out): # for i in range(time_idx.size): # ilow = time_idx[i] # ihigh = time_idx[i+1] # out[i] = np.sum(~flags[ilow:ihigh]) # return out # def extract_dde_info(args, freqs): # """ # Computes paralactic angles, antenna scaling and pointing information # required for beam interpolation. # """ # # get ms info required to compute paralactic angles and weighted sum # nband = freqs.size # if args.ms is not None: # utimes = [] # unflag_counts = [] # ant_pos = None # phase_dir = None # for ms_name in args.ms: # # get antenna positions # ant = xds_from_table(ms_name + '::ANTENNA')[0].compute() # if ant_pos is None: # ant_pos = ant['POSITION'].data # else: # check all are the same # tmp = ant['POSITION'] # if not np.array_equal(ant_pos, tmp): # raise ValueError( # "Antenna positions not the same across measurement sets") # # get phase center for field # field = xds_from_table(ms_name + '::FIELD')[0].compute() # if phase_dir is None: # phase_dir = field['PHASE_DIR'][args.field].data.squeeze() # else: # tmp = field['PHASE_DIR'][args.field].data.squeeze() # if not np.array_equal(phase_dir, tmp): # raise ValueError( # 'Phase direction not the same across measurement sets') # # get unique times and count flags # xds = xds_from_ms(ms_name, columns=["TIME", "FLAG_ROW"], group_cols=[ # "FIELD_ID"])[args.field] # utime, time_idx = np.unique( # xds.TIME.data.compute(), return_index=True) # ntime = utime.size # # extract subset of times # if args.sparsify_time > 1: # I = np.arange(0, ntime, args.sparsify_time) # utime = utime[I] # time_idx = time_idx[I] # ntime = utime.size # utimes.append(utime) # flags = xds.FLAG_ROW.data.compute() # unflag_count = _unflagged_counts(flags.astype( # np.int32), time_idx, np.zeros(ntime, dtype=np.int32)) # unflag_counts.append(unflag_count) # utimes = np.concatenate(utimes) # unflag_counts = np.concatenate(unflag_counts) # ntimes = utimes.size # # compute paralactic angles # parangles = parallactic_angles(utimes, ant_pos, phase_dir) # # mean over antanna nant -> 1 # parangles = np.mean(parangles, axis=1, keepdims=True) # nant = 1 # # beam_cube_dde requirements # ant_scale = np.ones((nant, nband, 2), dtype=np.float64) # point_errs = np.zeros((ntimes, nant, nband, 2), dtype=np.float64) # return (parangles, # da.from_array(ant_scale, chunks=ant_scale.shape), # point_errs, # unflag_counts, # True) # else: # ntimes = 1 # nant = 1 # parangles = np.zeros((ntimes, nant,), dtype=np.float64) # ant_scale = np.ones((nant, nband, 2), dtype=np.float64) # point_errs = np.zeros((ntimes, nant, nband, 2), dtype=np.float64) # unflag_counts = np.array([1]) # return (parangles, ant_scale, point_errs, unflag_counts, False) # def make_power_beam(args, lm_source, freqs, use_dask): # print("Loading fits beam patterns from %s" % args.beam_model) # from glob import glob # paths = glob(args.beam_model + '**_**.fits') # beam_hdr = None # if args.corr_type == 'linear': # corr1 = 'XX' # corr2 = 'YY' # elif args.corr_type == 'circular': # corr1 = 'LL' # corr2 = 'RR' # else: # raise KeyError( # "Unknown corr_type supplied. Only 'linear' or 'circular' supported") # for path in paths: # if corr1.lower() in path[-10::]: # if 're' in path[-7::]: # corr1_re = load_fits(path) # if beam_hdr is None: # beam_hdr = fits.getheader(path) # elif 'im' in path[-7::]: # corr1_im = load_fits(path) # else: # raise NotImplementedError("Only re/im patterns supported") # elif corr2.lower() in path[-10::]: # if 're' in path[-7::]: # corr2_re = load_fits(path) # elif 'im' in path[-7::]: # corr2_im = load_fits(path) # else: # raise NotImplementedError("Only re/im patterns supported") # # get power beam # beam_amp = (corr1_re**2 + corr1_im**2 + corr2_re**2 + corr2_im**2)/2.0 # # get cube in correct shape for interpolation code # beam_amp = np.ascontiguousarray(np.transpose(beam_amp, (1, 2, 0)) # [:, :, :, None, None]) # # get cube info # if beam_hdr['CUNIT1'].lower() != "deg": # raise ValueError("Beam image units must be in degrees") # npix_l = beam_hdr['NAXIS1'] # refpix_l = beam_hdr['CRPIX1'] # delta_l = beam_hdr['CDELT1'] # l_min = (1 - refpix_l)*delta_l # l_max = (1 + npix_l - refpix_l)*delta_l # if beam_hdr['CUNIT2'].lower() != "deg": # raise ValueError("Beam image units must be in degrees") # npix_m = beam_hdr['NAXIS2'] # refpix_m = beam_hdr['CRPIX2'] # delta_m = beam_hdr['CDELT2'] # m_min = (1 - refpix_m)*delta_m # m_max = (1 + npix_m - refpix_m)*delta_m # if (l_min > lm_source[:, 0].min() or m_min > lm_source[:, 1].min() or # l_max < lm_source[:, 0].max() or m_max < lm_source[:, 1].max()): # raise ValueError("The supplied beam is not large enough") # beam_extents = np.array([[l_min, l_max], [m_min, m_max]]) # # get frequencies # if beam_hdr["CTYPE3"].lower() != 'freq': # raise ValueError( # "Cubes are assumed to be in format [nchan, nx, ny]") # nchan = beam_hdr['NAXIS3'] # refpix = beam_hdr['CRPIX3'] # delta = beam_hdr['CDELT3'] # assumes units are Hz # freq0 = beam_hdr['CRVAL3'] # bfreqs = freq0 + np.arange(1 - refpix, 1 + nchan - refpix) * delta # if bfreqs[0] > freqs[0] or bfreqs[-1] < freqs[-1]: # warnings.warn("The supplied beam does not have sufficient " # "bandwidth. Beam frequencies:") # with np.printoptions(precision=2): # print(bfreqs) # if use_dask: # return (da.from_array(beam_amp, chunks=beam_amp.shape), # da.from_array(beam_extents, chunks=beam_extents.shape), # da.from_array(bfreqs, bfreqs.shape)) # else: # return beam_amp, beam_extents, bfreqs # def interpolate_beam(ll, mm, freqs, args): # """ # Interpolate beam to image coordinates and optionally compute average # over time if MS is provoded # """ # nband = freqs.size # print("Interpolating beam") # parangles, ant_scale, point_errs, unflag_counts, use_dask = extract_dde_info( # args, freqs) # lm_source = np.vstack((ll.ravel(), mm.ravel())).T # beam_amp, beam_extents, bfreqs = make_power_beam( # args, lm_source, freqs, use_dask) # # interpolate beam # if use_dask: # from africanus.rime.dask import beam_cube_dde # lm_source = da.from_array(lm_source, chunks=lm_source.shape) # freqs = da.from_array(freqs, chunks=freqs.shape) # # compute ncpu images at a time to avoid memory errors # ntimes = parangles.shape[0] # I = np.arange(0, ntimes, args.ncpu) # nchunks = I.size # I = np.append(I, ntimes) # beam_image = np.zeros((ll.size, 1, nband), dtype=beam_amp.dtype) # for i in range(nchunks): # ilow = I[i] # ihigh = I[i+1] # part_parangles = da.from_array( # parangles[ilow:ihigh], chunks=(1, 1)) # part_point_errs = da.from_array( # point_errs[ilow:ihigh], chunks=(1, 1, freqs.size, 2)) # # interpolate and remove redundant axes # part_beam_image = beam_cube_dde(beam_amp, beam_extents, bfreqs, # lm_source, part_parangles, part_point_errs, # ant_scale, freqs).compute()[:, :, 0, :, 0, 0] # # weighted sum over time # beam_image += np.sum(part_beam_image * # unflag_counts[None, ilow:ihigh, None], axis=1, keepdims=True) # # normalise by sum of weights # beam_image /= np.sum(unflag_counts) # # remove time axis # beam_image = beam_image[:, 0, :] # else: # from africanus.rime.fast_beam_cubes import beam_cube_dde # beam_image = beam_cube_dde(beam_amp, beam_extents, bfreqs, # lm_source, parangles, point_errs, # ant_scale, freqs).squeeze() # # swap source and freq axes and reshape to image shape # beam_source = np.transpose(beam_image, axes=(1, 0)) # return beam_source.squeeze().reshape((freqs.size, *ll.shape)) # def main(args): # # get coord info # hdr = fits.getheader(args.image) # 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 # if hdr["CTYPE4"].lower() == 'freq': # freq_axis = 4 # elif hdr["CTYPE3"].lower() == 'freq': # freq_axis = 3 # else: # raise ValueError("Freq axis must be 3rd or 4th") # freqs, ref_freq = data_from_header(hdr, axis=freq_axis) # xx, yy = np.meshgrid(l_coord, m_coord, indexing='ij') # # interpolate primary beam to fits header and optionally average over time # beam_image = interpolate_beam(xx, yy, freqs, args) # # save power beam # save_fits(args.output_filename, beam_image, hdr) # print("Wrote interpolated beam cube to %s \n" % args.output_filename) # return
def main(): # ------------------------------------------------- # Options and some error checking parser = OptionParser(usage='%prog [options] input_fits') parser.add_option('--band', dest='band', help='Select [U]HF or [L]-band (default = L-band)', default='L') parser.add_option( '--freq', dest='freq', help= 'Frequency in MHz at which to evaluate beam model (default = get from input FITS header)', default='') parser.add_option( '--pbcut', dest='pbcut', help= 'Primary beam gain level beyond which to blank output images (default = 0.3)', default=0.3) parser.add_option( '--noavg', dest='azavg', help= 'Do not azimuthally-average the primary beam pattern (default = do this)', default=True, action='store_false') parser.add_option( '--nopbcorfits', dest='savepbcor', help= 'Do not save primary beam corrected image (default = save corrected image)', action='store_false', default=True) parser.add_option( '--nopbfits', dest='savepb', help='Do not save primary beam image (default = save PB image)', action='store_false', default=True) parser.add_option( '--nowtfits', dest='savewt', help='Do not save weight image (default = save weight image)', action='store_false', default=True) parser.add_option( '--pbcorname', dest='pbcor_fits', help= 'Filename for primary beam corrected image (default = based on input image)', default='') parser.add_option( '--pbname', dest='pb_fits', help='Filename for primary beam image (default = based on input image)', default='') parser.add_option( '--wtname', dest='wt_fits', help='Filename for weight image (default = based on input image)', default='') parser.add_option( '--overwrite', '-f', dest='overwrite', help='Overwrite any existing output files (default = do not overwrite)', action='store_true', default=False) (options, args) = parser.parse_args() # Input FITS file if len(args) != 1: msg('Please provide a FITS image') sys.exit() else: input_fits = args[0].rstrip('/') # MeerKAT band band = options.band[0].lower() if band not in ['l', 'u']: msg('Please check requested band') sys.exit() # Frequency freq = options.freq # Primary beam cut level pbcut = float(options.pbcut) # Azimuthal averaging azavg = options.azavg if azavg: try: from skued import azimuthal_average as aa except: msg('scikit-ued not found, azimuthal averaging is not available.') msg('Try: pip install scikit-ued') azavg = False # Output files savepbcor = options.savepbcor savepb = options.savepb savewt = options.savewt if [savepbcor, savepb, savewt] == [False, False, False]: msg('Nothing to do, please check your options') sys.exit() # Generate output names if not provided pbcor_fits = options.pbcor_fits pb_fits = options.pb_fits wt_fits = options.wt_fits if pbcor_fits == '': pbcor_fits = input_fits.replace('.fits', '.pbcor.fits') check_name(input_fits, pbcor_fits) if pb_fits == '': pb_fits = input_fits.replace('.fits', '.pb.fits') check_name(input_fits, pb_fits) if wt_fits == '': wt_fits = input_fits.replace('.fits', '.wt.fits') check_name(input_fits, wt_fits) # Bail out if some files will be overwritten overwrite = options.overwrite if not overwrite: file_check = [] if savepbcor: file_check.append(check_file(pbcor_fits)) if savepb: file_check.append(check_file(pb_fits)) if savewt: file_check.append(check_file(wt_fits)) if True in file_check: sys.exit() pol = 'I' # hardwired Stokes I beam for now # ------------------------------------------------- # Set up band if band == 'l': beam_model = 'MKAT-AA-L-JIM-2020' band = 'L-band' elif band == 'u': beam_model = 'MKAT-AA-UHF-JIM-2020' band = 'UHF' msg('Band is ' + band) msg('Beam model is ' + beam_model) beam = JimBeam(beam_model) # Get header info msg('Reading FITS image') msg(' <--- ' + input_fits) nx, ny, dx, dy, fitsfreq = get_header(input_fits) if nx != ny or abs(dx) != abs(dy): msg('Can only handle square images / pixels') sys.exit() extent = nx * dx # degrees if freq == '': freq = fitsfreq / 1e6 else: freq = float(freq) msg('Evaluating beam at ' + str(round(freq, 4)) + ' MHz') interval = numpy.linspace(-extent / 2.0, extent / 2.0, nx) xx, yy = numpy.meshgrid(interval, interval) beam_image = beam.I(xx, yy, freq) msg('Masking beam beyond the ' + str(pbcut) + ' level') mask = beam_image < pbcut beam_image[mask] = numpy.nan if azavg: msg('Azimuthally averaging the beam pattern') x0 = int(nx / 2) y0 = int(ny / 2) radius, average = aa(beam_image, center=(x0, y0)) # This can probably be sped up... for y in range(0, ny): for x in range(0, nx): val = (((float(y) - y0)**2.0) + ((float(x) - x0)**2.0))**0.5 beam_image[y][x] = average[int(val)] if savepbcor: msg('Correcting image') input_image = get_image(input_fits) pbcor_image = input_image / beam_image msg('Writing primary beam corrected image') msg(' ---> ' + pbcor_fits) copyfile(input_fits, pbcor_fits) flush_fits(pbcor_image, pbcor_fits) if savepb: msg('Writing primary beam image') msg(' ---> ' + pb_fits) copyfile(input_fits, pb_fits) flush_fits(beam_image, pb_fits) if savewt: msg('Writing weight (pb^2) image') msg(' ---> ' + wt_fits) copyfile(input_fits, wt_fits) flush_fits(beam_image**2.0, wt_fits) msg('Done')
def test_forwardmodel(do_beam, do_gains, tmp_path_factory): test_dir = tmp_path_factory.mktemp("test_pfb") packratt.get('/test/ms/2021-06-24/elwood/test_ascii_1h60.0s.MS.tar', str(test_dir)) import numpy as np np.random.seed(420) from numpy.testing import assert_allclose from pyrap.tables import table ms = table(str(test_dir / 'test_ascii_1h60.0s.MS'), readonly=False) spw = table(str(test_dir / 'test_ascii_1h60.0s.MS::SPECTRAL_WINDOW')) utime = np.unique(ms.getcol('TIME')) freq = spw.getcol('CHAN_FREQ').squeeze() freq0 = np.mean(freq) ntime = utime.size nchan = freq.size nant = np.maximum( ms.getcol('ANTENNA1').max(), ms.getcol('ANTENNA2').max()) + 1 ncorr = ms.getcol('FLAG').shape[-1] uvw = ms.getcol('UVW') nrow = uvw.shape[0] u_max = abs(uvw[:, 0]).max() v_max = abs(uvw[:, 1]).max() uv_max = np.maximum(u_max, v_max) # image size from africanus.constants import c as lightspeed cell_N = 1.0 / (2 * uv_max * freq.max() / lightspeed) srf = 2.0 cell_rad = cell_N / srf cell_size = cell_rad * 180 / np.pi print("Cell size set to %5.5e arcseconds" % cell_size) fov = 2 npix = int(fov / cell_size) if npix % 2: npix += 1 nx = npix ny = npix print("Image size set to (%i, %i, %i)" % (nchan, nx, ny)) # model model = np.zeros((nchan, nx, ny), dtype=np.float64) nsource = 10 Ix = np.random.randint(0, npix, nsource) Iy = np.random.randint(0, npix, nsource) alpha = -0.7 + 0.1 * np.random.randn(nsource) I0 = 1.0 + np.abs(np.random.randn(nsource)) for i in range(nsource): model[:, Ix[i], Iy[i]] = I0[i] * (freq / freq0)**alpha[i] if do_beam: # primary beam from katbeam import JimBeam beam = JimBeam('MKAT-AA-L-JIM-2020') l_coord = -np.arange(-(nx // 2), nx // 2) * cell_size m_coord = np.arange(-(ny // 2), ny // 2) * cell_size xx, yy = np.meshgrid(l_coord, m_coord, indexing='ij') pbeam = np.zeros((nchan, nx, ny), dtype=np.float64) for i in range(nchan): pbeam[i] = beam.I(xx, yy, freq[i] / 1e6) # freq in MHz model_att = pbeam * model bm = 'JimBeam' else: model_att = model bm = None # model vis from ducc0.wgridder import dirty2ms model_vis = np.zeros((nrow, nchan, ncorr), dtype=np.complex128) for c in range(nchan): model_vis[:, c:c + 1, 0] = dirty2ms(uvw, freq[c:c + 1], model_att[c], pixsize_x=cell_rad, pixsize_y=cell_rad, epsilon=1e-8, do_wstacking=True, nthreads=8) model_vis[:, c, -1] = model_vis[:, c, 0] ms.putcol('MODEL_DATA', model_vis.astype(np.complex64)) if do_gains: t = (utime - utime.min()) / (utime.max() - utime.min()) nu = 2.5 * (freq / freq0 - 1.0) from africanus.gps.utils import abs_diff tt = abs_diff(t, t) lt = 0.25 Kt = 0.1 * np.exp(-tt**2 / (2 * lt**2)) Lt = np.linalg.cholesky(Kt + 1e-10 * np.eye(ntime)) vv = abs_diff(nu, nu) lv = 0.1 Kv = 0.1 * np.exp(-vv**2 / (2 * lv**2)) Lv = np.linalg.cholesky(Kv + 1e-10 * np.eye(nchan)) L = (Lt, Lv) from pfb.utils.misc import kron_matvec jones = np.zeros((ntime, nant, nchan, 1, ncorr), dtype=np.complex128) for p in range(nant): for c in [0, -1]: # for now only diagonal xi_amp = np.random.randn(ntime, nchan) amp = np.exp(-nu[None, :]**2 + kron_matvec(L, xi_amp).reshape(ntime, nchan)) xi_phase = np.random.randn(ntime, nchan) phase = kron_matvec(L, xi_phase).reshape(ntime, nchan) jones[:, p, :, 0, c] = amp * np.exp(1.0j * phase) # corrupted vis model_vis = model_vis.reshape(nrow, nchan, 1, 2, 2) from africanus.calibration.utils import chunkify_rows time = ms.getcol('TIME') row_chunks, tbin_idx, tbin_counts = chunkify_rows(time, ntime) ant1 = ms.getcol('ANTENNA1') ant2 = ms.getcol('ANTENNA2') from africanus.calibration.utils import corrupt_vis vis = corrupt_vis(tbin_idx, tbin_counts, ant1, ant2, jones, model_vis).reshape(nrow, nchan, ncorr) model_vis[:, :, 0, 0, 0] = 1.0 + 0j model_vis[:, :, 0, -1, -1] = 1.0 + 0j muellercol = corrupt_vis(tbin_idx, tbin_counts, ant1, ant2, jones, model_vis).reshape(nrow, nchan, ncorr) ms.putcol('DATA', vis.astype(np.complex64)) ms.putcol('CORRECTED_DATA', muellercol.astype(np.complex64)) ms.close() mcol = 'CORRECTED_DATA' else: ms.putcol('DATA', model_vis.astype(np.complex64)) mcol = None from pfb.workers.grid.dirty import _dirty _dirty(ms=str(test_dir / 'test_ascii_1h60.0s.MS'), data_column="DATA", weight_column='WEIGHT', imaging_weight_column=None, flag_column='FLAG', mueller_column=mcol, row_chunks=None, epsilon=1e-5, wstack=True, mock=False, double_accum=True, output_filename=str(test_dir / 'test'), nband=nchan, field_of_view=fov, super_resolution_factor=srf, cell_size=None, nx=None, ny=None, output_type='f4', nworkers=1, nthreads_per_worker=1, nvthreads=8, mem_limit=8, nthreads=8, host_address=None) from pfb.workers.grid.psf import _psf _psf(ms=str(test_dir / 'test_ascii_1h60.0s.MS'), data_column="DATA", weight_column='WEIGHT', imaging_weight_column=None, flag_column='FLAG', mueller_column=mcol, row_out_chunk=-1, row_chunks=None, epsilon=1e-5, wstack=True, mock=False, psf_oversize=2, double_accum=True, output_filename=str(test_dir / 'test'), nband=nchan, field_of_view=fov, super_resolution_factor=srf, cell_size=None, nx=None, ny=None, output_type='f4', nworkers=1, nthreads_per_worker=1, nvthreads=8, mem_limit=8, nthreads=8, host_address=None) # solve for model using pcg and mask mask = np.any(model, axis=0) from astropy.io import fits from pfb.utils.fits import save_fits hdr = fits.getheader(str(test_dir / 'test_dirty.fits')) save_fits(str(test_dir / 'test_model.fits'), model, hdr) save_fits(str(test_dir / 'test_mask.fits'), mask, hdr) from pfb.workers.deconv.forward import _forward _forward(residual=str(test_dir / 'test_dirty.fits'), psf=str(test_dir / 'test_psf.fits'), mask=str(test_dir / 'test_mask.fits'), beam_model=bm, band='L', weight_table=str(test_dir / 'test.zarr'), output_filename=str(test_dir / 'test'), nband=nchan, output_type='f4', epsilon=1e-5, sigmainv=0.0, wstack=True, double_accum=True, cg_tol=1e-6, cg_minit=10, cg_maxit=100, cg_verbose=0, cg_report_freq=10, backtrack=False, nworkers=1, nthreads_per_worker=1, nvthreads=1, mem_limit=8, nthreads=1, host_address=None) # get inferred model from pfb.utils.fits import load_fits model_inferred = load_fits(str(test_dir / 'test_update.fits')).squeeze() for i in range(nsource): if do_beam: beam = pbeam[:, Ix[i], Iy[i]] assert_allclose( 0.0, beam * (model_inferred[:, Ix[i], Iy[i]] - model[:, Ix[i], Iy[i]]), atol=1e-4) else: assert_allclose(0.0, model_inferred[:, Ix[i], Iy[i]] - model[:, Ix[i], Iy[i]], atol=1e-4)