swapped order of transmission matrix axes, now t_ba

openwfs/algorithms/dual_reference.py

@@ -251,8 +251,8 @@ def execute(self, *, capture_intermediate_results: bool = False, progress_bar=No
             # two halves of the wavefront together. For that, we need the
             # relative phase between the two sides, which we extract from
             # the measurements of the flat wavefronts.
-            relative = t_side_0[self._zero_indices[0]] + np.conjugate(t_side_1[self._zero_indices[1]])
-            factor = (relative / np.abs(relative)).reshape((1, *self.feedback.data_shape))
+            relative = t_side_0[..., self._zero_indices[0]] + np.conjugate(t_side_1[..., self._zero_indices[1]])
+            factor = np.expand_dims(relative / np.abs(relative), -1)
 
         t_full = self.compute_t_set(t_side_0, 0) + self.compute_t_set(factor * t_side_1, 1)
 
@@ -304,11 +304,11 @@ def compute_t_set(self, t, side) -> nd:
         """
         Compute the transmission matrix in SLM space from transmission matrix in input mode space.
 
-        TODO: make sure t is in t_ba form, so that this is equivalent to np.tensordot(t, cobasis, 1)
+        Equivalent to ``np.tensordot(t, cobasis[side], 1)``
 
         Args:
             t: transmission matrix in mode-index space. The first axis corresponds to the input modes.
         Returns:
             nd: The transmission matrix in SLM space. The last two axes correspond to SLM coordinates
         """
-        return np.tensordot(self.cobasis[side], t, ((0,), (0,)))
+        return np.tensordot(t, self.cobasis[side], 1)

openwfs/algorithms/utilities.py

@@ -154,9 +154,9 @@ def analyze_phase_stepping(measurements: np.ndarray, axis: int):
     Args:
         measurements(ndarray): array of phase stepping measurements.
             The array holds measured intensities
-            with the first one or more dimensions corresponding to the segments(pixels) of the SLM,
+            with the first one or more dimensions (a1, a2, ...) corresponding to the segments of the SLM,
             one dimension corresponding to the phase steps,
-            and the last zero or more dimensions corresponding to the individual targets
+            and the last zero or more dimensions (b1, b2, ...) corresponding to the individual targets
             where the feedback was measured.
         axis(int): indicates which axis holds the phase steps.
 
@@ -172,24 +172,22 @@ def analyze_phase_stepping(measurements: np.ndarray, axis: int):
 
         \\frac{1}{phase_{steps}} \\sum I_p  \\exp(-i 2\\pi p / phase_{steps}) = A^* B
 
-    The value of A^* B for each set of measurements is stored in the `field` attribute of the return
-    value.
-    Other attributes hold an estimate of the signal-to-noise ratio,
-    and an estimate of the maximum enhancement that can be expected
-    if these measurements are used for wavefront shaping.
+    Returns:
+        WFSResult: The result of the analysis. The attribute `t` holds the complex transmission matrix.
+            Note that the dimensions of t are reversed with respect to the input, so t has shape b1×b2×...×a1×a2×...
+            Other attributes hold fidelity estimates (see WFSResult).
     """
     phase_steps = measurements.shape[axis]
-    n = int(np.prod(measurements.shape[:axis]))
-    # M = np.prod(measurements.shape[axis + 1:])
-    segments = tuple(range(axis))
+    a_count = int(np.prod(measurements.shape[:axis]))
+    a_axes = tuple(range(axis))
 
     # Fourier transform the phase stepping measurements
     t_f = np.fft.fft(measurements, axis=axis) / phase_steps
     t = np.take(t_f, 1, axis=axis)
 
     # compute the effect of amplitude variations.
     # for perfectly developed speckle, and homogeneous illumination, this factor will be pi/4
-    amplitude_factor = np.mean(np.abs(t), segments) ** 2 / np.mean(np.abs(t) ** 2, segments)
+    fidelity_amplitude = np.mean(np.abs(t), a_axes) ** 2 / np.mean(np.abs(t) ** 2, a_axes)
 
     # estimate the calibration error
     # we first construct a matrix that can be used to fit
@@ -203,6 +201,9 @@ def analyze_phase_stepping(measurements: np.ndarray, axis: int):
         cc = ff_inv @ np.take(t_f, k, axis=axis).ravel()
         signal_energy = signal_energy + np.sum(np.abs(ff @ cc) ** 2)
         c[k] = cc[0]
+    fidelity_calibration = np.abs(c[1]) ** 2 / np.sum(np.abs(c[1:]) ** 2)
+
+    # TODO: use the pinv fit to estimate the signal strength (requires special treatment of offset)
 
     # Estimate the error due to noise
     # The signal consists of the response with incorrect modulation,
@@ -212,26 +213,28 @@ def analyze_phase_stepping(measurements: np.ndarray, axis: int):
     # average over all targets to get the most accurate result (assuming all targets are similar)
     axes = tuple([i for i in range(t_f.ndim) if i != axis])
     energies = np.sum(np.abs(t_f) ** 2, axis=axes)
-    offset_energy = energies[0]
     total_energy = np.sum(energies)
+    offset_energy = energies[0]
     signal_energy = energies[1] + energies[-1]
     if phase_steps > 3:
         # estimate the noise energy as the energy that is not explained
         # by the signal or the offset.
         noise_energy = (total_energy - signal_energy - offset_energy) / (phase_steps - 3)
-        noise_factor = np.abs(np.maximum(signal_energy - 2 * noise_energy, 0.0) / signal_energy)
+        fidelity_noise = np.abs(np.maximum(signal_energy - 2 * noise_energy, 0.0) / signal_energy)
     else:
-        noise_factor = 1.0  # cannot estimate reliably
+        fidelity_noise = 1.0  # cannot estimate reliably
 
-    calibration_fidelity = np.abs(c[1]) ** 2 / np.sum(np.abs(c[1:]) ** 2)
+    # convert from t_ab to t_ba form
+    a_destination = np.array(a_axes) + t.ndim - len(a_axes)
+    t = np.moveaxis(t, a_axes, a_destination)
 
     return WFSResult(
         t,
         axis=axis,
-        fidelity_amplitude=amplitude_factor,
-        fidelity_noise=noise_factor,
-        fidelity_calibration=calibration_fidelity,
-        n=n,
+        fidelity_amplitude=fidelity_amplitude,
+        fidelity_noise=fidelity_noise,
+        fidelity_calibration=fidelity_calibration,
+        n=a_count,
     )
 
 

openwfs/simulation/transmission.py

@@ -36,7 +36,8 @@ def __init__(
         Gaussian beam.
 
         Args:
-            t: Transmission matrix.
+            t: Transmission matrix. Must have the form (*feedback_shape, height, width), where feedback_shape
+                is the shape of the feedback signal and may be 0 or more dimensional.
             aberrations: An array containing the aberrations in radians. Can be used instead of a transmission matrix,
                 equivalent to specifying ``t = np.exp(1j * aberrations) / (aberrations.shape[0] * aberrations.shape[1])``.
             slm:
@@ -51,7 +52,7 @@ def __init__(
 
         # transmission matrix (normalized so that the maximum transmission is 1)
         self.t = t if t is not None else np.exp(1.0j * aberrations) / (aberrations.shape[0] * aberrations.shape[1])
-        self.slm = slm if slm is not None else SLM(self.t.shape[0:2])
+        self.slm = slm if slm is not None else SLM(self.t.shape[-2:])
 
         super().__init__(self.slm.field, multi_threaded=multi_threaded)
         self.beam_amplitude = beam_amplitude
@@ -71,9 +72,9 @@ def _fetch(self, incident_field):  # noqa
             np.ndarray: A numpy array containing the calculated intensity in the focus.
 
         """
-        field = np.tensordot(incident_field * self.beam_amplitude, self.t, 2)
+        field = np.tensordot(self.t, incident_field * self.beam_amplitude, 2)
         return np.abs(field) ** 2
 
     @property
     def data_shape(self):
-        return self.t.shape[2:]
+        return self.t.shape[:-2]

tests/__init__.py

@@ -0,0 +1,8 @@
+import numpy as np
+
+
+def complex_random(shape):
+    """
+    Create an array with normally distributed complex elements with variance 1.
+    """
+    return np.sqrt(0.5) * (np.random.normal(size=shape) + 1j * np.random.normal(size=shape))

tests/test_algorithms_troubleshoot.py

@@ -2,6 +2,7 @@
 import numpy as np
 import pytest
 
+from . import complex_random
 from .test_simulation import phase_response_test_function, lookup_table_test_function
 from ..openwfs.algorithms import StepwiseSequential
 from ..openwfs.algorithms.troubleshoot import (
@@ -14,11 +15,33 @@
     measure_modulated_light,
     measure_modulated_light_dual_phase_stepping,
 )
+from ..openwfs.algorithms.utilities import analyze_phase_stepping
 from ..openwfs.processors import SingleRoi
 from ..openwfs.simulation import Camera
 from ..openwfs.simulation import SimulatedWFS, StaticSource, SLM, Microscope
 
 
+@pytest.mark.parametrize("phase_steps", [5, 6, 10, 20])
+@pytest.mark.parametrize("noise", [0.0, 0.5, 1.0, 2.0, 4.0])
+def test_analyze_phase_stepping(phase_steps, noise):
+    """Test the analyze_phase_stepping function"""
+    # TODO: find out why there is a (small) systematic error when the noise is high
+    # Construct a perfect phase stepping signal
+    np.random.seed(123)
+    t = complex_random((10000, 1))
+    ref = np.exp(np.arange(phase_steps) * 2j * np.pi / phase_steps)
+    signal = np.abs(ref + t) ** 2
+    signal += np.random.normal(scale=noise, size=signal.shape)
+    result = analyze_phase_stepping(signal, axis=1)
+    # the signal energy for a signal 2·cos(φ) is 2 P  (with P the number of phase steps), distributed over 2 bins,
+    # giving P per bin.
+    # the noise energy is σ²·P, distributed over P bins, giving σ² per bin.
+    tolerance = 4 / np.sqrt(t.size)
+    theoretical_fidelity = phase_steps / (phase_steps + noise**2)
+    print(f"noise fidelity. theoretical: {theoretical_fidelity}  measured: {result.fidelity_noise}")
+    assert np.isclose(result.fidelity_noise, theoretical_fidelity, rtol=tolerance)
+
+
 def test_signal_std():
     """
     Test signal std, corrected for (uncorrelated) noise in the signal.