Work on multires CLEAN

playground/multires_clean.py

@@ -11,14 +11,53 @@ oversampling_factor = 4  # FIXME Why does this need to be so large?
 minor_conv = {"rel first peak": 0.2, "max components": 3000, "gain": 0.8}
 nmajor = 10
 weighting_scheme = "uniform"
-# weighting_scheme = "natural"
-
+oversampling_highres = 2
 
 do_wgridding = True
 epsilon = 1e-8
 nthreads = 6
 
 
+def highres_domain_and_lowrew_mask(dom, oversampling, xmin, xmax, ymin, ymax):
+    rve.assert_sky_domain(dom)
+    _, _, _, sdom = dom
+    if isinstance(xmin, str):
+        xmin = rve.str2rad(xmin)
+    if isinstance(xmax, str):
+        xmax = rve.str2rad(xmax)
+    if isinstance(ymin, str):
+        ymin = rve.str2rad(ymin)
+    if isinstance(ymax, str):
+        ymax = rve.str2rad(ymax)
+    if not isinstance(oversampling, int):
+        s = f"`oversampling` needs to be int. Got: {oversampling}"
+        raise TypeError(s)
+
+    Dstx, Dsty = sdom.distances
+    Nx, Ny = sdom.shape
+    xmin = np.round(Nx//2 + xmin/Dstx).astype(int)
+    xmax = np.round(Nx//2 + xmax/Dstx).astype(int)
+    ymin = np.round(Ny//2 + ymin/Dsty).astype(int)
+    ymax = np.round(Ny//2 + ymax/Dsty).astype(int)
+
+    mask_lowres = np.ones(sdom.shape)
+    mask_lowres[xmin:xmax, ymin:ymax] = 0.
+    mask_lowres = ift.makeField(sdom, mask_lowres)
+
+    center = ((xmax-xmin)//2 - Nx//2)*Dstx, ((ymax-ymin)//2 - Ny//2)*Dsty
+
+    npix = np.array([(xmax-xmin)*oversampling, (ymax-ymin)*oversampling])
+    assert (npix % 2 == 0).all()
+
+    dst = np.array(sdom.distances) / oversampling
+    dom_highres = ift.makeDomain(dom[:3] + (ift.RGSpace(npix, dst),))
+    rve.assert_sky_domain(dom_highres)
+    np.testing.assert_allclose(sdom.total_volume - mask_lowres.s_integrate(),
+                               dom_highres[-1].total_volume)
+    mask_lowres = ift.ContractionOperator(dom, (0, 1, 2)).adjoint(mask_lowres)
+    return dom_highres, mask_lowres, center_highres
+
+
 def main():
     _, ms = sys.argv
 
@@ -35,12 +74,20 @@ def main():
     assert all(dd == 1 for dd in dom.shape[:3])
     squeeze = ift.ContractionOperator(dom, (0, 1, 2))
 
+    dom_highres, mask_lowres, center_highres = highres_domain_and_lowrew_mask(dom, oversampling_highres, "-0.94arcmin", "-0.78arcmin", "-0.40arcmin", "-0.28arcmin")
+
     R = rve.InterferometryResponse(ms, dom, do_wgridding, epsilon, nthreads=nthreads)
+    R_highres = rve.InterferometryResponse(ms, dom_highres, do_wgridding, epsilon,
+                                           center=center_highres, nthreads=nthreads)
 
     dirty = get_dirty(ms, dom, ms.vis)
+    dirty_highres = get_dirty(ms, dom_highres, ms.vis, center=center_highres)
     psf = get_psf(ms, dom)
+    psf_highres = get_psf(ms, dom, center=center_highres)
 
     rve.ubik_tools.field2fits(dirty, "dirty.fits")
+    rve.ubik_tools.field2fits(dirty_highres, "dirty_highres.fits")
+    exit()
     rve.ubik_tools.field2fits(psf, "psf.fits")
 
     imajor = 0
@@ -144,7 +191,7 @@ def major_cycle(ms, dom, residuals, model, psf, convergence_criteria):
     return ift.makeField(dom, model.reshape(dom.shape)), 0
 
 
-def get_psf(ms, dom):
+def get_psf(ms, dom, center=(0., 0.)):
     psf = rve.dirty_image(
         ms,
         weighting_scheme,
@@ -153,13 +200,14 @@ def get_psf(ms, dom):
         epsilon,
         nthreads=nthreads,
         vis=ms.vis * 0 + 1,
+        center=(0., 0.)
     )
     return psf / psf.val.max()
 
 
-def get_dirty(ms, dom, vis, weighting_scheme=weighting_scheme):
+def get_dirty(ms, dom, vis, weighting_scheme=weighting_scheme, center=(0., 0.)):
     dirty = rve.dirty_image(
-        ms, weighting_scheme, dom, do_wgridding, epsilon, nthreads=nthreads, vis=vis
+        ms, weighting_scheme, dom, do_wgridding, epsilon, nthreads=nthreads, vis=vis, center=center
     )
     return dirty