Scale the FPM object plot by its real-space pixel size, not dxo

PtyLab/Monitor/Monitor.py

@@ -205,6 +205,34 @@ def initializeMonitors(self):
         if self.verboseLevel == "high":
             self.diffractionDataMonitor = DiffractionDataPlot()
 
+    @property
+    def objectPixelSize(self):
+        """Pixel size of the object estimate that is plotted."""
+        if self.reconstruction.data.operationMode == "FPM":
+            return self.reconstruction.dxo_fpm
+        return self.reconstruction.dxo
+
+    @property
+    def probePixelSize(self):
+        """Axis step of the probe panel."""
+        if self.reconstruction.data.operationMode == "FPM":
+            return self.reconstruction.dfp
+        return self.reconstruction.dxp
+
+    @property
+    def probeAxisUnit(self):
+        """Unit of the probe panel axes: a length for CPM, a spatial frequency for FPM."""
+        if self.reconstruction.data.operationMode == "FPM":
+            return "1/um"
+        return "mm"
+
+    @property
+    def probeLabel(self):
+        """Title of the probe panel. FPM estimates a pupil instead of a probe."""
+        if self.reconstruction.data.operationMode == "FPM":
+            return "Pupil estimate"
+        return "Probe estimate"
+
     def updateObjectProbeErrorMonitor(
         self,
         error,
@@ -226,15 +254,16 @@ def updateObjectProbeErrorMonitor(
             object_estimate,
             self.reconstruction,
             objectPlot=self.objectPlot,
-            pixelSize=self.reconstruction.dxo,
+            pixelSize=self.objectPixelSize,
             axisUnit="mm",
             amplitudeScalingFactor=self.objectPlotContrast,
         )
         self.defaultMonitor.updateProbe(
             probe_estimate,
             self.reconstruction,
-            pixelSize=self.reconstruction.dxp,
-            axisUnit="mm",
+            pixelSize=self.probePixelSize,
+            axisUnit=self.probeAxisUnit,
+            label=self.probeLabel,
             amplitudeScalingFactor=self.probePlotContrast,
         )
         self.defaultMonitor.update_z(zo)

PtyLab/Monitor/Plots.py

@@ -122,7 +122,12 @@ def updateObject(
         self.im_object.autoscale()
 
     def updateProbe(
-        self, probe_estimate, optimizable, amplitudeScalingFactor=1, **kwargs
+        self,
+        probe_estimate,
+        optimizable,
+        amplitudeScalingFactor=1,
+        label="Probe estimate",
+        **kwargs,
     ):
 
         # from PtyLab.Operators.Operators import fft2c
@@ -137,6 +142,7 @@ def updateProbe(
 
         if self.firstrun:
             self.im_probe = complexPlot(PE, ax=self.ax_probe, **kwargs)
+            self.txt_purityProbe = self.ax_probe.set_title(label)
             # self.im_probe_ff = complexPlot(PE_ff, self.ax_probe_ff, **kwargs)
         else:
             self.im_probe.set_data(PE)
@@ -146,7 +152,7 @@ def updateProbe(
                 and optimizable.purityProbe == optimizable.purityProbe
             ):
                 self.txt_purityProbe.set_text(
-                    "Probe estimate\nPurity: %.2f" % (100 * optimizable.purityProbe) + "%"
+                    "%s\nPurity: %.2f" % (label, 100 * optimizable.purityProbe) + "%"
                 )
         self.im_probe.autoscale()
 

PtyLab/Reconstruction/Reconstruction.py

@@ -649,6 +649,22 @@ def Lo(self):
         """Field of view (entrance pupil plane)"""
         return self.No * self.dxo
 
+    @property
+    def dxo_fpm(self):
+        """Real-space object pixel size for FPM.
+        """
+        return self.dxp * self.Np / self.No
+
+    @property
+    def Lo_fpm(self):
+        """Real-space field of view of the FPM object, equal to that of the raw images."""
+        return self.No * self.dxo_fpm
+
+    @property
+    def dfp(self):
+        """Spatial-frequency pixel size of the probe grid, 1 / Lp."""
+        return 1 / self.Lp
+
     @property
     def xo(self):
         """object coordinates 1D"""

PtyLab/utils/visualisation.py

@@ -94,22 +94,54 @@ def complex2rgb_vectorized(probe, **kwargs):
     return probe_rgb
 
 
+# Axis units for the image plots. The reciprocal ones are for Fourier-space
+# quantities such as the FPM pupil, whose axes are spatial frequencies.
+unitRatio = {
+    "pixel": 1,
+    "m": 1,
+    "cm": 1e2,
+    "mm": 1e3,
+    "um": 1e6,
+    "1/m": 1,
+    "1/mm": 1e-3,
+    "1/um": 1e-6,
+}
+
+
+def plotExtent(pixelSize, axisUnit, shape):
+    """
+    Extent for imshow, expressed in axisUnit.
+
+    Real-space axes run from zero, as they always have. Reciprocal axes are
+    centred on zero frequency instead, which is where the pupil sits.
+
+    :param pixelSize: sample spacing of the array, in SI units
+    :param str axisUnit: any key of unitRatio
+    :param shape: shape of the array that is plotted
+    :return: [left, right, bottom, top] for imshow
+    """
+    step = pixelSize * unitRatio[axisUnit]
+    width, height = step * shape[1], step * shape[0]
+    if axisUnit.startswith("1/"):
+        return [-width / 2, width / 2, height / 2, -height / 2]
+    return [0, width, height, 0]
+
+
 def complexPlot(rgb, ax=None, pixelSize=1, axisUnit="pixel"):
     """
     Plot a 2D complex plot (hue for phase, brightness for amplitude). Input array need to be prepared by using
     the complex2rgb function.
     :param rgb: a rgb array that is converted from a 2D complex np.ndarray by using complex2rgb
     :param ax: Optional axis to plot in
     :param pixelSize: pixelSize in x and y, to display the physical dimension of the plot
-    :param str axisUnit: Options: default 'pixel', 'm', 'cm', 'mm', 'um'
+    :param str axisUnit: Options: default 'pixel', 'm', 'cm', 'mm', 'um', and the
+        reciprocal '1/m', '1/mm', '1/um' for Fourier-space quantities
     :return: An hsv plot
     """
 
     if not ax:
         fig, ax = plt.subplots()
-    unitRatio = {"pixel": 1, "m": 1, "cm": 1e2, "mm": 1e3, "um": 1e6}
-    pixelSize = pixelSize * unitRatio[axisUnit]
-    extent = [0, pixelSize * rgb.shape[1], pixelSize * rgb.shape[0], 0]
+    extent = plotExtent(pixelSize, axisUnit, rgb.shape)
 
     im = ax.imshow(rgb, extent=extent, interpolation=None)
     ax.set_ylabel(axisUnit)
@@ -199,9 +231,7 @@ def absplot(
     U = np.abs(asNumpyArray(u))
     if not ax:
         fig, ax = plt.subplots()
-    unitRatio = {"pixel": 1, "m": 1, "cm": 1e2, "mm": 1e3, "um": 1e6}
-    pixelSize = pixelSize * unitRatio[axisUnit]
-    extent = [0, pixelSize * U.shape[1], pixelSize * U.shape[0], 0]
+    extent = plotExtent(pixelSize, axisUnit, U.shape)
 
     if amplitudeScalingFactor != 1:
         U[U > amplitudeScalingFactor * np.max(U)] = amplitudeScalingFactor * np.max(U)

tests/Reconstruction/test_fpm_sampling.py

@@ -0,0 +1,122 @@
+import h5py
+import numpy as np
+import pytest
+from numpy.testing import assert_allclose
+
+from PtyLab.ExperimentalData.ExperimentalData import ExperimentalData
+from PtyLab.Monitor.Monitor import Monitor
+from PtyLab.Params.Params import Params
+from PtyLab.Reconstruction.Reconstruction import Reconstruction
+from PtyLab.utils.visualisation import plotExtent
+
+
+@pytest.fixture
+def fpm_reconstruction(tmp_path):
+    """A minimal FPM dataset, sampled like the LED-array microscope examples."""
+    rng = np.random.default_rng(42)
+    Nd, N_frames = 64, 25
+
+    # LEDs on a 5x5 grid, 60 mm below the sample
+    led = np.linspace(-4e-3, 4e-3, 5)
+    encoder = np.stack(np.meshgrid(led, led), axis=-1).reshape(-1, 2)
+
+    hdf5_path = tmp_path / "fpm.hdf5"
+    with h5py.File(hdf5_path, "w") as hf:
+        hf.create_dataset(
+            "ptychogram", data=rng.random((N_frames, Nd, Nd)).astype(np.float32)
+        )
+        hf.create_dataset("encoder", data=encoder)
+        hf.create_dataset("dxd", data=np.array(5.5e-6))
+        hf.create_dataset("magnification", data=np.array(4.0))
+        hf.create_dataset("wavelength", data=np.array(625e-9))
+        hf.create_dataset("zled", data=np.array(60e-3))
+        hf.create_dataset("NA", data=np.array(0.1))
+
+    data = ExperimentalData(hdf5_path, operationMode="FPM")
+    return Reconstruction(data, Params())
+
+
+def test_fpm_object_sampling_preserves_field_of_view(fpm_reconstruction):
+    """The enlarged FPM object adds bandwidth, not field of view."""
+    reconstruction = fpm_reconstruction
+    assert reconstruction.No > reconstruction.Np
+
+    assert_allclose(
+        reconstruction.dxo_fpm,
+        reconstruction.dxp * reconstruction.Np / reconstruction.No,
+    )
+    assert reconstruction.dxo_fpm < reconstruction.dxp
+    assert_allclose(reconstruction.Lo_fpm, reconstruction.Np * reconstruction.dxp)
+
+
+def test_monitor_plots_fpm_object_with_fpm_pixel_size(fpm_reconstruction):
+    monitor = Monitor()
+    monitor.reconstruction = fpm_reconstruction
+
+    assert_allclose(monitor.objectPixelSize, fpm_reconstruction.dxo_fpm)
+    # the plotted extent is the field of view of the raw images, not No * dxo
+    assert_allclose(
+        fpm_reconstruction.No * monitor.objectPixelSize,
+        fpm_reconstruction.Np * fpm_reconstruction.dxp,
+    )
+
+
+def test_monitor_plots_cpm_object_with_dxo(generate_simu_hdf5):
+    data = ExperimentalData("example:simulation_cpm")
+    reconstruction = Reconstruction(data, Params())
+
+    monitor = Monitor()
+    monitor.reconstruction = reconstruction
+
+    assert_allclose(monitor.objectPixelSize, reconstruction.dxo)
+
+
+def test_fpm_pupil_sampling_matches_the_numerical_aperture(fpm_reconstruction):
+    """The pupil grid step is what puts the NA cut-off where the code puts it."""
+    reconstruction = fpm_reconstruction
+    assert_allclose(reconstruction.dfp, 1 / (reconstruction.Np * reconstruction.dxp))
+
+    # radius of the aperture, in pixels, derived two independent ways
+    assert_allclose(
+        (reconstruction.NA / reconstruction.wavelength) / reconstruction.dfp,
+        (reconstruction.data.entrancePupilDiameter / 2) / reconstruction.dxp,
+    )
+
+
+def test_monitor_plots_fpm_pupil_in_reciprocal_units(fpm_reconstruction):
+    """For FPM the probe panel shows the pupil, which lives in Fourier space."""
+    monitor = Monitor()
+    monitor.reconstruction = fpm_reconstruction
+
+    assert monitor.probeLabel == "Pupil estimate"
+    assert monitor.probeAxisUnit == "1/um"
+    assert_allclose(monitor.probePixelSize, fpm_reconstruction.dfp)
+
+    # the axis spans the bandwidth the low-resolution grid can carry, 1 / dxp
+    assert_allclose(
+        fpm_reconstruction.Np * monitor.probePixelSize, 1 / fpm_reconstruction.dxp
+    )
+
+
+def test_monitor_probe_panel_unchanged_for_cpm(generate_simu_hdf5):
+    """CPM reconstructs a real-space probe, so that panel must not move."""
+    data = ExperimentalData("example:simulation_cpm")
+    reconstruction = Reconstruction(data, Params())
+
+    monitor = Monitor()
+    monitor.reconstruction = reconstruction
+
+    assert monitor.probeLabel == "Probe estimate"
+    assert monitor.probeAxisUnit == "mm"
+    assert_allclose(monitor.probePixelSize, reconstruction.dxp)
+
+
+def test_only_reciprocal_axes_are_centred_on_zero(generate_simu_hdf5):
+    """Real-space extents keep the historical corner origin."""
+    shape = (8, 8)
+
+    assert plotExtent(2e-6, "mm", shape) == [0, 2e-6 * 8 * 1e3, 2e-6 * 8 * 1e3, 0]
+
+    left, right, bottom, top = plotExtent(3e3, "1/um", shape)
+    assert left == -right and top == -bottom
+    assert_allclose(right, 3e3 * 8 * 1e-6 / 2)