Add visibility forward fitting

.editorconfig

@@ -12,6 +12,6 @@ trim_trailing_whitespace=true
 indent_style=space
 indent_size=4
 
-[*.yml]
+[{*.yml,*.toml}]
 indent_style=space
 indent_size=2

.github/workflows/ci.yml

@@ -2,15 +2,15 @@ name: CI
 
 on:
   push:
-#    branches:
-#      - 'main'
-#      - '*.*'
-#      - '!*backport*'
-#    tags:
-#      - 'v*'
-#      - '!*dev*'
-#      - '!*pre*'
-#      - '!*post*'
+    branches:
+      - 'main'
+      - '*.*'
+      - '!*backport*'
+    tags:
+      - 'v*'
+      - '!*dev*'
+      - '!*pre*'
+      - '!*post*'
   pull_request:
   # Allow manual runs through the web UI
   workflow_dispatch:

changelog/83.feature.rst

@@ -0,0 +1 @@
+Add a basic visibility forward fitting method (`xrayvision.vis_forward_fit.forward_fit.vis_forward_fit`).

docs/reference/forward_fit.rst

@@ -0,0 +1,14 @@
+.. vis_forward_fit:
+
+Vis Forward Fit ('xrayvision.vis_forward_fit')
+**********************************************
+
+The ``vis_forward_fit`` submodule contains the visibility forward fitting methods
+
+.. automodapi:: xrayvision.vis_forward_fit
+
+.. automodapi:: xrayvision.vis_forward_fit.forward_fit
+   :include-all-objects:
+
+.. automodapi:: xrayvision.vis_forward_fit.sources
+   :include-all-objects:

docs/reference/index.rst

@@ -11,5 +11,6 @@ Reference
    transform
    utils
    visibility
+   forward_fit
 
    ../whatsnew/index

examples/rhessi.py

@@ -4,6 +4,7 @@
 ======================================
 
 Create images from RHESSI visibility data
+
 """
 
 import astropy.units as apu

examples/stix.py

@@ -16,6 +16,8 @@
 from xrayvision.clean import vis_clean
 from xrayvision.imaging import vis_psf_map, vis_to_map
 from xrayvision.mem import mem, resistant_mean
+from xrayvision.vis_forward_fit.forward_fit import vis_forward_fit
+from xrayvision.vis_forward_fit.sources import Source, SourceList
 
 ###############################################################################
 # Create images from STIX visibility data.
@@ -33,13 +35,13 @@
 ###############################################################################
 # Lets have a look at the point spread function (PSF) or dirty beam
 
-psf_map = vis_psf_map(stix_vis, shape=(129, 129) * apu.pixel, pixel_size=2 * apu.arcsec / apu.pix, scheme="uniform")
+psf_map = vis_psf_map(stix_vis, shape=(129, 129) * apu.pixel, pixel_size=1 * apu.arcsec / apu.pix, scheme="uniform")
 psf_map.plot()
 
 ###############################################################################
 # Back projection
 
-backproj_map = vis_to_map(stix_vis, shape=(129, 129) * apu.pixel, pixel_size=2 * apu.arcsec / apu.pix, scheme="uniform")
+backproj_map = vis_to_map(stix_vis, shape=(129, 129) * apu.pixel, pixel_size=1 * apu.arcsec / apu.pix, scheme="uniform")
 backproj_map.plot()
 
 ###############################################################################
@@ -48,7 +50,7 @@
 clean_map, model_map, resid_map = vis_clean(
     stix_vis,
     shape=[129, 129] * apu.pixel,
-    pixel_size=[2, 2] * apu.arcsec / apu.pix,
+    pixel_size=[1, 1] * apu.arcsec / apu.pix,
     clean_beam_width=20 * apu.arcsec,
     niter=100,
 )
@@ -62,17 +64,42 @@
 percent_lambda = 2 / (snr_value**2 + 90)
 
 mem_map = mem(
-    stix_vis, shape=[129, 129] * apu.pixel, pixel_size=[2, 2] * apu.arcsec / apu.pix, percent_lambda=percent_lambda
+    stix_vis, shape=[129, 129] * apu.pixel, pixel_size=[1, 1] * apu.arcsec / apu.pix, percent_lambda=percent_lambda
+)
+
+###############################################################################
+# VIS_FWD_FIT
+
+sources = SourceList(
+    [
+        Source(
+            "elliptical",
+            15 * stix_vis.visibilities.unit,
+            1 * apu.arcsec,
+            2 * apu.arcsec,
+            5 * apu.arcsec,
+            2 * apu.arcsec,
+            1,
+        )
+    ]
+)
+
+vis_fwd_map = vis_forward_fit(stix_vis, sources, shape=[129, 129] * apu.pixel, pixel_size=[1, 1] * apu.arcsec / apu.pix)
+
+vis_fwd_pso_map = vis_forward_fit(
+    stix_vis, sources, method="PSO", shape=[129, 129] * apu.pixel, pixel_size=[1, 1] * apu.arcsec / apu.pix
 )
-mem_map.plot()
 
 ###############################################################################
 # Comparison
 fig = plt.figure(figsize=(10, 10))
-fig.add_subplot(221, projection=psf_map)
-fig.add_subplot(222, projection=backproj_map)
-fig.add_subplot(223, projection=clean_map)
-fig.add_subplot(224, projection=mem_map)
+fig.add_subplot(231, projection=psf_map)
+fig.add_subplot(232, projection=backproj_map)
+fig.add_subplot(233, projection=clean_map)
+fig.add_subplot(234, projection=mem_map)
+fig.add_subplot(235, projection=mem_map)
+fig.add_subplot(236, projection=mem_map)
+
 axs = fig.get_axes()
 psf_map.plot(axes=axs[0])
 axs[0].set_title("PSF")
@@ -82,4 +109,9 @@
 axs[2].set_title("Clean")
 mem_map.plot(axes=axs[3])
 axs[3].set_title("MEM")
+vis_fwd_map.plot(axes=axs[4])
+axs[4].set_title("VIS_FWRDFIT")
+vis_fwd_pso_map.plot(axes=axs[5])
+axs[5].set_title("VIS_FWRDFIT_PSO")
+
 plt.show()

pyproject.toml

@@ -20,7 +20,7 @@ dependencies = [
     "numpy>=1.24.0",
     "packaging>=23.0",
     "scipy>=1.13",
-    "xarray>=2023.5.0"
+    "xarray>=2023.5.0",
 ]
 dynamic = ["version"]
 keywords = ["solar", "physics", "solar", "sun", "x-rays"]
@@ -42,7 +42,10 @@ classifiers = [
 map = [
     "sunpy[map]>=5.1.0"
 ]
-all = ["xrayvisim[map]"]
+pso = [
+        "pymoo>=0.6.1.3"
+]
+all = ["xrayvisim[map,pso]"]
 tests = [
     "matplotlib>=3.8.0",
     "pytest-astropy>=0.11.0",

pytest.ini

@@ -40,3 +40,6 @@ filterwarnings =
     # Until update code need to ignore missing WCS
     ignore:.*:sunpy.util.exceptions.SunpyMetadataWarning
     ignore:.*divide by zero.*:RuntimeWarning
+    ignore:The.*feasible.*:DeprecationWarning
+    # Can't use ipython in PyCharm debugger without this
+    ignore::DeprecationWarning

ruff.toml

@@ -8,7 +8,7 @@ exclude = [
 ]
 
 [lint]
-select = ["E", "F", "W", "UP", "PT"]
+select = ["E", "F", "W", "UP", "PT", "C"]
 extend-ignore = [
   # pycodestyle (E, W)
   "E501", # LineTooLong                             # TODO! fix
@@ -27,7 +27,12 @@ extend-ignore = [
 "docs/conf.py" = ["E402"]
 "docs/*.py" = [
   "INP001",  # Implicit-namespace-package. The examples are not a package.
+  "D100"
 ]
+"examples/*.py" = [
+  "D"
+]
+
 "__init__.py" = ["E402", "F401", "F403"]
 "test_*.py" = ["B011", "D", "E402", "PGH001", "S101"]
 # Need to import clients to register them, but don't use them in file

xrayvision/clean.py

@@ -198,7 +198,6 @@ def vis_clean(
     map :
         Return a `sunpy.map.Map` by default or array only if `False`
     """
-
     dirty_map = vis_to_map(vis, shape=shape, pixel_size=pixel_size, **kwargs)
     dirty_beam_shape = [x.value * 3 + 1 if x.value * 3 % 2 == 0 else x.value * 3 for x in shape] * shape.unit
     dirty_beam = vis_psf_image(vis, shape=dirty_beam_shape, pixel_size=pixel_size, **kwargs)

xrayvision/conftest.py

@@ -1,9 +1,8 @@
 # Force MPL to use non-gui backends for testing.
-import matplotlib
 
 try:
     pass
 except ImportError:
     pass
-else:
-    matplotlib.use("Agg")
+# else:
+# matplotlib.use("Agg")

xrayvision/coordinates/frames.py

@@ -71,7 +71,7 @@ def projective_wcs_to_frame(wcs):
     observer = None
     for frame, attr_names in required_attrs.items():
         attrs = [getattr(wcs.wcs.aux, attr_name) for attr_name in attr_names]
-        if all([attr is not None for attr in attrs]):
+        if all(attr is not None for attr in attrs):
             kwargs = {"obstime": dateavg}
             if rsun is not None:
                 kwargs["rsun"] = rsun

xrayvision/imaging.py

@@ -245,7 +245,7 @@ def vis_psf_image(
     # Make sure psf is always odd so power is in exactly one pixel
     shape = [s // 2 * 2 + 1 for s in shape.to_value(apu.pix)] * shape.unit
     psf_arr = idft_map(
-        np.ones(vis.visibilities.shape) * vis.visibilities.unit,
+        np.ones(vis.visibilities.shape) * np.prod(pixel_size.value) * vis.visibilities.unit,
         u=vis.u,
         v=vis.v,
         shape=shape,

xrayvision/mem.py

@@ -133,7 +133,6 @@ def _estimate_flux(vis, shape, pixel, maxiter=1000, tol=1e-3):
     Estimated total flux
 
     """
-
     Hv, Lip, Visib = _prepare_for_optimise(pixel, shape, vis)
 
     # PROJECTED LANDWEBER
@@ -237,7 +236,6 @@ def _get_mean_visibilities(vis, shape, pixel):
     -------
     Mean Visibilities
     """
-
     if vis.amplitude_uncertainty is None:
         amplitude_uncertainty = np.ones_like(vis.visibilities)
     else:
@@ -374,7 +372,6 @@ def _proximal_operator(z, f, m, lamb, Lip, niter=250):
     -------
 
     """
-
     # INITIALIZATION OF THE DYKSTRA - LIKE SPLITTING
     x = z[:]
     p = np.zeros_like(x)
@@ -434,11 +431,11 @@ def _optimise_fb(Hv, Visib, Lip, flux, lambd, shape, pixel, maxiter, tol):
         Maximum number of iterations
     tol :
         Tolerance value used in the stopping rule ( || x - x_old || <= tol || x_old ||)
+
     Returns
     -------
     MEM Image
     """
-
     # 'f': value of the total flux of the image (taking into account the area of the pixel)
     f = flux / (pixel[0] * pixel[1])
     # 'm': total flux divided by the number of pixels of the image

xrayvision/tests/test_transform.py

@@ -1,7 +1,7 @@
 import astropy.units as apu
 import numpy as np
 import pytest
-from numpy.fft import fft2, fftshift, ifft2, ifftshift
+from numpy.fft import fft2, ifft2
 from numpy.testing import assert_allclose
 from scipy import signal
 
@@ -377,34 +377,36 @@ def test_phase_center_equivalence():
     assert np.allclose(data, img2)
 
 
-def test_fft_equivalence():
-    # Odd (3, 3) so symmetric and chose shape and pixel size so xy values will run from 0 to 2 the same as in fft
-    # TODO: add same kind of test for even for fft2 then A[n/2] has both pos and negative nyquist frequencies
-    #  e.g shape (2, 2), (3, 2), (2, 3)
-    shape = (3, 3)
-    pixel = (1, 1)
-    center = (1, 1)
+@pytest.mark.parametrize("dim", ((2, 3, 4, 5, 6)))
+def test_fft_equivalence(dim):
+    shape = np.array((dim, dim))
+    pixel = np.array((1, 1))
 
-    data = np.arange(np.prod(shape)).reshape(shape)
+    # In order to replicate the FFT need to make the product of x*u + y*v match mk/M + nl/N where M, N and the array
+    # dimensions so neex x, y to go from 0 to M-1, N-1 and u, v 0 to m-1/M, and n-1/N so need to chose parameters so
+    # this happens namely the center for u, v and the dft
+    uv_center = -1 / ((0 - shape / 2 + 0.5) / (shape * pixel))
+    xy_center = -1 * (0 - shape / 2 + 0.5)
     vv = generate_uv(
-        shape[0] * apu.pix, phase_center=center[0] * apu.arcsec, pixel_size=pixel[0] * apu.arcsec / apu.pix
+        shape[0] * apu.pix, phase_center=uv_center[0] * apu.arcsec, pixel_size=pixel[0] * apu.arcsec / apu.pix
     )
     uu = generate_uv(
-        shape[1] * apu.pix, phase_center=center[1] * apu.arcsec, pixel_size=pixel[1] * apu.arcsec / apu.pix
+        shape[1] * apu.pix, phase_center=uv_center[1] * apu.arcsec, pixel_size=pixel[1] * apu.arcsec / apu.pix
     )
     u, v = np.meshgrid(uu, vv)
     u = u.flatten()
     v = v.flatten()
 
-    vis = dft_map(data, u=u, v=v, pixel_size=pixel * apu.arcsec / apu.pix, phase_center=center * apu.arcsec)
+    data = np.arange(np.prod(shape)).reshape(shape)
+    vis = dft_map(data, u=u, v=v, pixel_size=pixel * apu.arcsec / apu.pix, phase_center=xy_center * apu.arcsec)
 
     ft = fft2(data)
-    fts = fftshift(ft)
     vis = vis.reshape(shape)
-    # Convention in xrayvison has the minus sign on the forward transform but numpy has it on reverse
+    # Convention in xrayvison has the minus sign on the forward transform but numpy has it on reverse and the first
+    #
+    # The rtol seems to vary as the size increase I'm assuming poor round off error handling in the naive DFT/IDFT
     vis_conj = np.conjugate(vis)
-    assert_allclose(fts, vis_conj, atol=1e-13)
+    assert_allclose(ft, vis_conj, rtol=1e-10, atol=1e-10)
 
-    vis_ft = ifftshift(vis_conj)
-    img = ifft2(vis_ft)
-    assert_allclose(np.real(img), data, atol=1e-14)
+    img = ifft2(vis_conj)
+    assert_allclose(np.real(img), data, rtol=1e-10, atol=1e-10)

xrayvision/tests/test_visibility.py

@@ -109,7 +109,7 @@ def test_vis_eq(visibilities):
 def test_meta_eq(vis_meta):
     meta = vis_meta
     assert meta == meta
-    meta = vm.VisMeta(dict())
+    meta = vm.VisMeta({})
     assert meta == meta
 
 

xrayvision/transform.py

@@ -161,10 +161,8 @@ def dft_map(
     x, y = np.meshgrid(x, y)
     uv = np.vstack([u, v])
     # Check units are correct for exp need to be dimensionless and then remove units for speed
-    if (uv[0, :] * x[0, 0]).unit == apu.dimensionless_unscaled and (
-        uv[1, :] * y[0, 0]
-    ).unit == apu.dimensionless_unscaled:
-        uv = uv.value  # type: ignore
+    if (uv.unit * x.unit) == apu.dimensionless_unscaled and (uv.unit * y.unit) == apu.dimensionless_unscaled:
+        uv = uv.value  # src_type: ignore
         x = x.value
         y = y.value
 
@@ -174,7 +172,7 @@ def dft_map(
                 2j * np.pi * (x[..., np.newaxis] * uv[np.newaxis, 0, :] + y[..., np.newaxis] * uv[np.newaxis, 1, :])
             ),
             axis=(0, 1),
-        )
+        ) * np.prod(pixel_size.value)
 
         return vis
     else:
@@ -245,7 +243,7 @@ def idft_map(
             axis=2,
         )
 
-        return np.real(image)
+        return np.real(image) / np.prod(pixel_size.value)
     else:
         raise UnitsError("Incompatible units on uv {uv.unit} should cancel with xy to leave a dimensionless quantity")