Mosaicing

demo/mosaicing.py

+# SPDX-License-Identifier: GPL-3.0-or-later
+# Copyright(C) 2019-2021 Max-Planck-Society
+# Author: Philipp Arras
+
+import numpy as np
+from functools import reduce
+from operator import add
+import matplotlib.pyplot as plt
+
+import nifty7 as ift
+import resolve as rve
+
+
+@rve.onlymaster
+def imsave(fname, field):
+    plt.imsave(fname, field.val.T, origin="lower")
+    plt.close()
+
+
+def main():
+    rve.set_nthreads(2)
+    rve.set_wgridding(False)
+    diffusefluxlevel = 20
+
+    interobs_extended = {
+        f"extended{ii}": rve.Observation.load(f"/data/mosaicing/extended_field{ii}.npz")
+        for ii in range(5, 135)
+    }
+    interobs_compact = {
+        f"compact{ii}": rve.Observation.load(f"/data/mosaicing/compact_field{ii}.npz")
+        for ii in range(5, 153)
+    }
+    interobs = {**interobs_compact, **interobs_extended}
+    sdobs = {"sd0": rve.SingleDishObservation.load("/data/mosaicing/totalpower.npz")}
+    assert len(set(sdobs.keys()) & set(interobs.keys())) == 0
+
+    rve.set_epsilon(1 / 10 / max(o.max_snr() for o in interobs.values()))
+
+    # Compute phase center and define domain
+    # Convention: The phase center is the coordinate of the [nx/2, ny/2] pixel
+    xs, ys = [], []
+    for vv in interobs.values():
+        xs.append(vv.direction.phase_center[0])
+        ys.append(vv.direction.phase_center[1])
+    for vv in sdobs.values():
+        foo, bar = vv.pointings.phase_centers.T
+        xs.extend(foo)
+        ys.extend(bar)
+    xs, ys = np.array(xs), np.array(ys)
+    global_phase_center_casa = -2.7180339078174175, -0.5210930907116116
+    global_phase_center = (max(xs) + min(xs)) / 2, (max(ys) + min(ys)) / 2
+    margin = 1 * rve.ARCMIN2RAD
+    xfov = 2 * max(abs(xs - global_phase_center[0])) + 2 * margin
+    yfov = 2 * max(abs(ys - global_phase_center[1])) + 2 * margin
+    if rve.mpi.master:
+        print(f"Field of view: {xfov/rve.ARCMIN2RAD:.2f}' x {yfov/rve.ARCMIN2RAD:.2f}'")
+    npix = np.array([1800, 1800])
+    fov = np.array([xfov, yfov])
+    dom = ift.RGSpace(npix, fov / npix)
+    # End compute phase center and define domain
+
+    # Plot pointings
+    xs, ys = [], []
+    for vv in interobs.values():
+        xs.append(vv.direction.phase_center[0])
+        ys.append(vv.direction.phase_center[1])
+    xs, ys = np.array(xs), np.array(ys)
+    plt.scatter(xs, ys, label="Interferometer")
+    for kk, vv in sdobs.items():
+        plt.scatter(
+            *vv.pointings.phase_centers.T, label="Single dish " + kk, alpha=0.2, s=1
+        )
+    plt.scatter(*global_phase_center_casa, label="Phase center CASA")
+    plt.scatter(*global_phase_center, label="Phase center resolve")
+    plt.xlim([global_phase_center[0] - xfov / 2, global_phase_center[0] + xfov / 2])
+    plt.ylim([global_phase_center[1] - yfov / 2, global_phase_center[1] + yfov / 2])
+    plt.gca().invert_xaxis()
+    plt.gca().set_aspect("equal")
+    plt.legend()
+    if rve.mpi.master:
+        plt.savefig("pointings.png")
+    plt.close()
+    # End plot pointings
+
+    # Single dish likelihood
+    single_dish_response = []
+    for kk, oo in sdobs.items():
+        assert np.min(oo.freq) / np.max(oo.freq) > 0.99
+        R = rve.SingleDishResponse(
+            oo,
+            dom,
+            lambda x: rve.alma_beam_func(10.7, 0.75, np.mean(oo.freq), x),
+            global_phase_center,
+        )
+        single_dish_response.append(R.ducktape_left(kk))
+    single_dish_response = reduce(add, single_dish_response)
+    dsd = ift.MultiField.from_dict({kk: o.vis for kk, o in sdobs.items()})
+    invcovsd = ift.makeOp(
+        ift.MultiField.from_dict({kk: o.weight for kk, o in sdobs.items()})
+    )
+    lhsd = ift.GaussianEnergy(dsd, invcovsd) @ single_dish_response
+    imsave("sd_dirty.png", single_dish_response.adjoint(invcovsd(dsd)))
+    # End single dish likelihood
+
+    logsky = ift.SimpleCorrelatedField(
+        dom, diffusefluxlevel, (1, 0.1), (1.5, 1), (1.2, 0.4), (0.2, 0.2), (-1.5, 0.5)
+    )
+    sky = logsky.exp()
+
+    # Interferometer likelihood
+    beam_directions = {}
+    for kk, obs in interobs.items():
+        assert np.min(obs.freq) / np.max(obs.freq) > 0.99
+        beam_directions[kk] = rve.BeamDirection(
+            obs.direction.phase_center[0] - global_phase_center[0],
+            obs.direction.phase_center[1] - global_phase_center[1],
+            lambda x: rve.alma_beam_func(10.7, 0.75, np.mean(obs.freq), x),
+            0.1,
+        )
+    skyslicer = rve.SkySlicer(logsky.target, beam_directions)
+    ift.extra.check_linear_operator(skyslicer)
+    lhinter = rve.ImagingLikelihood(interobs, skyslicer)
+    # End interferometer likelihood
+
+    lh = (lhinter + lhsd) @ sky
+
+    # Plots
+    R = rve.StokesIResponse(interobs, skyslicer.target) @ skyslicer
+    vis = ift.MultiField.from_dict({kk: o.vis for kk, o in interobs.items()})
+    weight = ift.MultiField.from_dict({kk: o.weight for kk, o in interobs.items()})
+    imsave("total_dirty.png", R.adjoint(vis * weight))
+    imsave("skyslicer_adjoint.png", skyslicer.adjoint(ift.full(skyslicer.target, 1.0)))
+    # End plots
+
+    plotter = rve.Plotter("png", "plots")
+    plotter.add("logsky", logsky)
+    plotter.add("sky", logsky.exp(), cmap="inferno_r")
+    plotter.add("power spectrum logsky", logsky.power_spectrum)
+    # plotter.add_histogram(
+    #     "normalized residuals (original weights)", lh.normalized_residual
+    # )
+    ham = ift.StandardHamiltonian(
+        lh, ift.AbsDeltaEnergyController(0.5, 5, 300, name="Sampling")
+    )
+    fld = 0.1 * ift.from_random(lh.domain)
+    state = rve.MinimizationState(fld, [])
+    mini = ift.NewtonCG(ift.GradientNormController(name="newton", iteration_limit=5))
+    # TODO Add minisanity to resolve
+    # TODO Fix histogram plots
+    for ii in range(20):
+        state = rve.simple_minimize(ham, state.mean, 3, mini)
+        plotter.plot(f"iter{ii}", state)
+        state.save(f"iter{ii}")
+
+
+if __name__ == "__main__":
+    main()