diff --git a/openwfs/algorithms/dual_reference.py b/openwfs/algorithms/dual_reference.py index f4fb1c5..0eb882d 100644 --- a/openwfs/algorithms/dual_reference.py +++ b/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) diff --git a/openwfs/algorithms/utilities.py b/openwfs/algorithms/utilities.py index 235cd35..86e46c1 100644 --- a/openwfs/algorithms/utilities.py +++ b/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,16 +172,14 @@ 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 @@ -189,7 +187,7 @@ def analyze_phase_stepping(measurements: np.ndarray, axis: int): # 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, ) diff --git a/openwfs/simulation/transmission.py b/openwfs/simulation/transmission.py index 99eeaa1..bb52d3f 100644 --- a/openwfs/simulation/transmission.py +++ b/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] diff --git a/tests/__init__.py b/tests/__init__.py index e69de29..ac52d7f 100644 --- a/tests/__init__.py +++ b/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)) diff --git a/tests/test_algorithms_troubleshoot.py b/tests/test_algorithms_troubleshoot.py index 130a20f..fd83257 100644 --- a/tests/test_algorithms_troubleshoot.py +++ b/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. diff --git a/tests/test_wfs.py b/tests/test_wfs.py index 3e2afdc..1057090 100644 --- a/tests/test_wfs.py +++ b/tests/test_wfs.py @@ -7,6 +7,7 @@ from scipy.linalg import hadamard from scipy.ndimage import zoom +from . import complex_random from ..openwfs.algorithms import ( StepwiseSequential, FourierDualReference, @@ -20,10 +21,15 @@ from ..openwfs.simulation.mockdevices import GaussianNoise, Camera +@pytest.fixture(autouse=True) +def reset(): + np.random.seed(42) # for reproducibility + + @pytest.mark.parametrize("algorithm", ["ssa", "fourier"]) -@pytest.mark.parametrize("shape", [(4, 7), (10, 7), (20, 31)]) +@pytest.mark.parametrize("shape, feedback_shape", [((4, 7), (210,)), ((10, 7), (21, 10)), ((20, 31), (3, 7, 10))]) @pytest.mark.parametrize("noise", [0.0, 0.1]) -def test_multi_target_algorithms(shape: tuple[int, int], noise: float, algorithm: str): +def test_multi_target_algorithms(shape: tuple[int, int], feedback_shape: tuple[int, ...], noise: float, algorithm: str): """ Test the multi-target capable algorithms (SSA and Fourier dual ref). @@ -31,11 +37,11 @@ def test_multi_target_algorithms(shape: tuple[int, int], noise: float, algorithm and it also verifies that the enhancement and noise fidelity are estimated correctly by the algorithm. """ - M = 100 # number of targets + M = np.prod(feedback_shape) # number of targets phase_steps = 6 # create feedback object, with noise if needed - sim = SimulatedWFS(t=random_transmission_matrix((*shape, M))) + sim = SimulatedWFS(t=complex_random((*feedback_shape, *shape))) sim.slm.set_phases(0.0) I_0 = np.mean(sim.read()) feedback = GaussianNoise(sim, std=I_0 * noise) @@ -49,7 +55,10 @@ def test_multi_target_algorithms(shape: tuple[int, int], noise: float, algorithm phase_steps=phase_steps, ) N = np.prod(shape) # number of input modes - alg_fidelity = (N - 1) / N # SSA is inaccurate if N is low + # SSA is inaccurate if N is slightly lower because the reference beam varies with each segment. + # The variation is of order 1/N in the signal, so the fidelity is (N-1)/N. + # todo: check! + alg_fidelity = (N - 1) / N signal = (N - 1) / N**2 # for estimating SNR masks = [True] # use all segments for determining fidelity else: # 'fourier': @@ -78,23 +87,26 @@ def test_multi_target_algorithms(shape: tuple[int, int], noise: float, algorithm # For the dual reference algorithm, the left and right half will have a different overall amplitude factor. # This is corrected for by computing the correlations for the left and right half separately # - I_opt = np.zeros((M,)) - for b in range(M): - sim.slm.set_phases(-np.angle(result.t[:, :, b])) + I_opt = np.zeros(feedback_shape) + for b in np.ndindex(feedback_shape): + sim.slm.set_phases(-np.angle(result.t[b])) I_opt[b] = feedback.read()[b] - t_correlation = t_fidelity(result.t, sim.t, masks) + t_correlation = t_fidelity(result.t, sim.t, masks=masks) # Check the enhancement, noise fidelity and # the fidelity of the transmission matrix reconstruction coverage = N / np.prod(shape) - theoretical_noise_fidelity = signal / (signal + noise**2 / phase_steps) + print(signal) + print(noise) + theoretical_noise_fidelity = signal * phase_steps / (signal * phase_steps + noise**2) + enhancement = I_opt.mean() / I_0 theoretical_enhancement = np.pi / 4 * theoretical_noise_fidelity * alg_fidelity * (N - 1) + 1 estimated_enhancement = result.estimated_enhancement.mean() * alg_fidelity theoretical_t_correlation = theoretical_noise_fidelity * alg_fidelity * coverage estimated_t_correlation = result.fidelity_noise * result.fidelity_calibration * alg_fidelity * coverage - tolerance = 2.0 / np.sqrt(M) + tolerance = 4.0 / np.sqrt(M) # TODO: find out if this should be stricter print( f"\nenhancement: \ttheoretical= {theoretical_enhancement},\testimated={estimated_enhancement},\tactual: {enhancement}" ) @@ -135,14 +147,6 @@ def test_multi_target_algorithms(shape: tuple[int, int], noise: float, algorithm Expected {theoretical_noise_fidelity}, got {result.fidelity_noise}""" -def random_transmission_matrix(shape): - """ - Create a random transmission matrix with the given shape. - """ - np.random.seed(42) # for reproducibility - return np.random.normal(size=shape) + 1j * np.random.normal(size=shape) - - def half_mask(shape): """ Args: @@ -157,7 +161,9 @@ def half_mask(shape): return mask -def t_fidelity(t: np.ndarray, t_correct: np.ndarray, masks: Optional[tuple[np.ndarray, ...]] = (True,)) -> float: +def t_fidelity( + t: np.ndarray, t_correct: np.ndarray, *, columns: int = 2, masks: Optional[tuple[np.ndarray, ...]] = (True,) +) -> float: """ Compute the fidelity of the measured transmission matrix. @@ -171,14 +177,15 @@ def t_fidelity(t: np.ndarray, t_correct: np.ndarray, masks: Optional[tuple[np.nd Args: t: The measured transmission matrix. 'rows' of t are the *last* index t_correct: The correct transmission matrix + columns: The number of columns in the transmission matrix masks: Masks for the left and right half of the wavefront, or None to use the full wavefront """ fidelity = 0.0 norm = 0.0 - for r in range(t.shape[-1]): # each row + for r in np.ndindex(t.shape[:-columns]): # each row for m in masks: - rm = t[..., r][m] - rm_c = t_correct[..., r][m] + rm = t[r][m] + rm_c = t_correct[r][m] fidelity += abs(np.vdot(rm, rm_c) ** 2 / np.vdot(rm, rm)) norm += np.vdot(rm_c, rm_c) @@ -320,9 +327,7 @@ def test_flat_wf_response_fourier(optimized_reference, type): assert ( abs(field_correlation(result.t / np.abs(result.t), t / np.abs(t))) > 0.99 ), "The phases were not calculated correctly" - assert ( - t_fidelity(np.expand_dims(result.t, -1), np.expand_dims(t, -1), alg.masks) > 0.99 - ), "The amplitudes were not calculated correctly" + assert t_fidelity(result.t, t, masks=alg.masks) > 0.99, "The amplitudes were not calculated correctly" def test_flat_wf_response_ssa(): @@ -344,7 +349,7 @@ def test_flat_wf_response_ssa(): def test_multidimensional_feedback_ssa(): - aberrations = np.random.uniform(0.0, 2 * np.pi, (256, 256, 5, 2)) + aberrations = np.random.uniform(0.0, 2 * np.pi, (5, 2, 256, 256)) sim = SimulatedWFS(aberrations=aberrations) alg = StepwiseSequential(feedback=sim, slm=sim.slm) @@ -352,7 +357,7 @@ def test_multidimensional_feedback_ssa(): # compute the phase pattern to optimize the intensity in target 2,1 target = (2, 1) - optimised_wf = -np.angle(t[(..., *target)]) + optimised_wf = -np.angle(t[*target, ...]) # Calculate the enhancement factor # Note: technically this is not the enhancement, just the ratio after/before @@ -369,7 +374,7 @@ def test_multidimensional_feedback_ssa(): def test_multidimensional_feedback_fourier(): - aberrations = np.random.uniform(0.0, 2 * np.pi, (256, 256, 5, 2)) + aberrations = np.random.uniform(0.0, 2 * np.pi, (5, 2, 256, 256)) sim = SimulatedWFS(aberrations=aberrations) # input the camera as a feedback object, such that it is multidimensional @@ -377,7 +382,7 @@ def test_multidimensional_feedback_fourier(): t = alg.execute().t # compute the phase pattern to optimize the intensity in target 0 - optimised_wf = -np.angle(t[:, :, 2, 1]) + optimised_wf = -np.angle(t[2, 1, :, :]) # Calculate the enhancement factor # Note: technically this is not the enhancement, just the ratio after/before @@ -400,7 +405,7 @@ def test_simple_genetic(population_size: int, elite_size: int): Note: this is not very rigid test, as we currently don't have theoretical expectations for the performance. """ shape = (100, 71) - sim = SimulatedWFS(t=random_transmission_matrix(shape), multi_threaded=False) + sim = SimulatedWFS(t=complex_random(shape), multi_threaded=False) alg = SimpleGenetic( feedback=sim, slm=sim.slm, @@ -437,7 +442,7 @@ def test_dual_reference_ortho_split(basis_str: str, shape: tuple[int, int]): mask = half_mask(shape) phases_set = np.pad(phases, ((0, 0), (0, 0), (0, shape[1] // 2))) - sim = SimulatedWFS(t=random_transmission_matrix(shape)) + sim = SimulatedWFS(t=complex_random(shape)) alg = DualReference( feedback=sim, @@ -459,7 +464,7 @@ def test_dual_reference_ortho_split(basis_str: str, shape: tuple[int, int]): result_t_phase_only = np.exp(1j * np.angle(result.t)) assert np.abs(field_correlation(sim_t_phase_only, result_t_phase_only)) > 0.999 - assert t_fidelity(np.expand_dims(result.t, -1), np.expand_dims(sim.t, -1), alg.masks) > 0.9 + assert t_fidelity(result.t, sim.t, masks=alg.masks) > 0.9 def test_dual_reference_non_ortho_split():