Add slmsuite Hologram algorithm for BeamShaping

README.rst

@@ -54,6 +54,16 @@ Tested against the simulated SLM of slmsuite only (``pytest``). The Meadowlark a
 validated on real devices yet.
 
 
+BeamShaping algorithm
+=====================
+
+With ``pip install pymodaq_plugins_slmsuite[beamshaping]`` the ``slmsuite Hologram`` algorithm is added to the
+BeamShaping extension of ``pymodaq_plugins_beam_shaping`` (found through the ``pymodaq_plugins_beam_shaping.algorithms``
+entry point). It computes the phase mask with ``slmsuite.holography.algorithms.Hologram``, with the methods GS,
+WGS-Leonardo, WGS-Kim, WGS-Nogrette, WGS-Wu and WGS-tanh. The GPU is used if ``cupy`` is installed.
+The conjugate gradient method is not exposed yet.
+
+
 Installation instructions
 =========================
 

hatch_build.py

@@ -8,3 +8,6 @@
 class PluginInfoTomlHook(MetadataHookInterface):
     def update(self, metadata: dict) -> None:
         update_metadata_from_toml(metadata, here)
+        # algorithms discovered by the BeamShaping extension of pymodaq_plugins_beam_shaping
+        metadata.setdefault('entry-points', {})['pymodaq_plugins_beam_shaping.algorithms'] = {
+            'slmsuite': 'pymodaq_plugins_slmsuite.algorithms'}

pyproject.toml

@@ -26,6 +26,8 @@ maintainers = [
     {name = "Weber Sébastien", email = "sebastien.weber@cemes.fr"},
 ]
 
+optional-dependencies = {beamshaping = ["pymodaq_plugins_beam_shaping>=5.0.0"]}
+
 # nottodo: leave everything below as is!
 
 dynamic = ["version", "urls", "entry-points", "classifiers"]

src/pymodaq_plugins_slmsuite/algorithms/__init__.py

@@ -0,0 +1,7 @@
+"""Algorithms for the BeamShaping extension (pymodaq_plugins_beam_shaping)
+
+This package is declared in the entry point group *pymodaq_plugins_beam_shaping.algorithms*: BeamShaping imports it at
+startup, which registers the algorithms below in its algorithm factory. It needs pymodaq_plugins_beam_shaping to be
+installed (extra *beamshaping* of this plugin).
+"""
+from pymodaq_plugins_slmsuite.algorithms import slmsuite_hologram  # noqa: F401  (registers the algorithm)

src/pymodaq_plugins_slmsuite/algorithms/slmsuite_hologram.py

@@ -0,0 +1,121 @@
+from typing import Optional
+
+import numpy as np
+
+from pymodaq_gui.parameter import Parameter
+from pymodaq_utils.logger import set_logger, get_module_name
+
+from pymodaq_plugins_beam_shaping.algorithms import AlgoBase
+from pymodaq_plugins_beam_shaping.algorithms.factory import AlgorithmFactory
+from pymodaq_plugins_beam_shaping.algorithms.utils import AlgoType, LensSetup
+
+from slmsuite.holography.algorithms import Hologram
+from slmsuite.misc.xp import as_numpy
+
+logger = set_logger(get_module_name(__file__))
+
+# Parameters used by each method of Hologram.optimize with the default value slmsuite gives them (ALGORITHM_DEFAULTS in
+# slmsuite/holography/algorithms/_header.py, checked by the tests). The GPU (cupy) is used by slmsuite if installed.
+METHOD_DEFAULTS: dict[str, dict[str, float]] = {
+    'GS': {},
+    'WGS-Leonardo': {'feedback_exponent': 0.8},
+    'WGS-Kim': {'feedback_exponent': 0.8, 'fix_phase_iteration': 10},
+    'WGS-Nogrette': {'feedback_factor': 0.1},
+    'WGS-Wu': {'feedback_exponent': 0.5},
+    'WGS-tanh': {'feedback_factor': 0.2, 'feedback_exponent': 0.5},
+}
+
+
+@AlgorithmFactory.register_algorithm()
+class SlmsuiteHologram(AlgoBase):
+    """ Gerchberg-Saxton and weighted Gerchberg-Saxton phase retrieval, computed by slmsuite
+
+    The phase is computed with slmsuite.holography.algorithms.Hologram (``optimize``), in the Fourier plane of a lens
+    as for the Gerchberg-Saxton algorithm of BeamShaping: the input beam amplitude is the nearfield amplitude and the
+    target amplitude is the farfield one. Each step runs ``Iterations per step`` iterations of the chosen method and
+    keeps the state of the optimization (weights, fixed farfield phase for WGS-Kim) from step to step. Changing the
+    input beam, the target, or the initial phase restarts the optimization.
+
+    * GS: Gerchberg-Saxton
+    * WGS-Leonardo: https://doi.org/10.1364/OE.15.001913
+    * WGS-Kim: https://doi.org/10.1364/OL.44.003178
+    * WGS-Nogrette: https://doi.org/10.1103/PhysRevX.4.021034
+    * WGS-Wu: https://doi.org/10.1364/OE.413723
+    * WGS-tanh: weights updated with a hyperbolic tangent
+
+    The weighted methods give a uniform intensity on a set of spots or regions better than GS. See the documentation of
+    slmsuite (``Hologram.optimize``) for the details of the parameters.
+    """
+    ALGO_NAME = 'slmsuite Hologram'
+    ALGOTYPE = AlgoType.AMPLITUDE
+    SETUP_TYPE = LensSetup.TwoF
+    ITERATIVE = True
+    MANUAL_LOOP = True
+
+    params = [
+        {'title': 'Method:', 'name': 'method', 'type': 'list', 'limits': list(METHOD_DEFAULTS), 'value': 'GS'},
+        {'title': 'Iterations per step:', 'name': 'iterations', 'type': 'int', 'value': 1, 'min': 1},
+        {'title': 'Feedback exponent:', 'name': 'feedback_exponent', 'type': 'float', 'value': 0.8, 'visible': False,
+         'tip': 'Power of the weighting (WGS-Leonardo, WGS-Kim, WGS-Wu, WGS-tanh)'},
+        {'title': 'Feedback factor:', 'name': 'feedback_factor', 'type': 'float', 'value': 0.1, 'visible': False,
+         'tip': 'Gain of the weighting (WGS-Nogrette), maximal relative change of the weights (WGS-tanh)'},
+        {'title': 'Fix phase at iteration:', 'name': 'fix_phase_iteration', 'type': 'int', 'value': 10, 'min': 0,
+         'visible': False,
+         'tip': 'WGS-Kim: the farfield phase is fixed after this number of iterations'},
+    ]
+
+    def __init__(self, parent=None, *args, **kwargs):
+        super().__init__(parent, *args, **kwargs)
+        self._hologram: Optional[Hologram] = None
+        self._last_phase: Optional[np.ndarray] = None  # modulator phase set by the last step
+
+    def value_changed(self, param: Parameter):
+        if param.name() == 'method':
+            self._update_settings_for_method(param.value())
+        if self.parent_app is not None:
+            self.parent_app.algo_settings_changed()
+
+    def _update_settings_for_method(self, method: str):
+        """Show the settings used by the method and reset them to its defaults"""
+        defaults = METHOD_DEFAULTS[method]
+        with self.settings.treeChangeBlocker():
+            for name in ('feedback_exponent', 'feedback_factor', 'fix_phase_iteration'):
+                child = self.settings.child(name)
+                child.setOpts(visible=name in defaults)
+                if name in defaults:
+                    child.setValue(defaults[name])
+
+    def do_things_after_set_input(self):
+        self._hologram = None
+
+    def do_things_after_set_target(self):
+        self._hologram = None
+
+    def _method_kwargs(self) -> dict[str, float]:
+        method = self.settings['method']
+        return {name: self.settings[name] for name in METHOD_DEFAULTS[method]}
+
+    def _build_hologram(self) -> Hologram:
+        """Start a new optimization from the current phase of the modulator field"""
+        shape = self._modulator_field.shape
+        return Hologram(target=np.asarray(self._target_field.amplitude, dtype=float),
+                        amp=np.asarray(self._input_field.amplitude, dtype=float),
+                        phase=np.asarray(self._modulator_field.phase, dtype=float),
+                        slm_shape=shape)
+
+    def compute_phase(self, do_step=True, **kwargs):
+        # the modulator phase differs from the last one we set when the user reset the initial phase: start again
+        if (self._hologram is None or self._last_phase is None or
+                not np.array_equal(self._modulator_field.phase, self._last_phase)):
+            self._hologram = self._build_hologram()
+
+        self._hologram.optimize(method=self.settings['method'], maxiter=self.settings['iterations'],
+                                verbose=False, **self._method_kwargs())
+
+        self.set_phase_in_modulator_plane(np.asarray(as_numpy(self._hologram.get_phase()), dtype=float))
+        self._last_phase = self._modulator_field.phase.copy()
+        self.compute_forward_fft(update_plots=False)  # the output field shown corresponds to the new phase
+
+
+if __name__ == '__main__':
+    pass

tests/test_slmsuite_algorithm.py

@@ -0,0 +1,95 @@
+import numpy as np
+import pytest
+
+pytest.importorskip('pymodaq_plugins_beam_shaping')
+
+from slmsuite.holography.algorithms import ALGORITHM_DEFAULTS
+
+from pymodaq_plugins_beam_shaping.algorithms import factory
+from pymodaq_plugins_beam_shaping.algorithms.factory import AlgorithmFactory
+from pymodaq_plugins_beam_shaping.field import Field
+
+from pymodaq_plugins_slmsuite.algorithms import slmsuite_hologram
+from pymodaq_plugins_slmsuite.algorithms.slmsuite_hologram import SlmsuiteHologram, METHOD_DEFAULTS
+
+
+@pytest.fixture
+def algo(qtbot):
+    algo = SlmsuiteHologram()
+    shape = algo.shape
+    y, x = np.indices(shape)
+    r2 = (x - shape[1] / 2) ** 2 + (y - shape[0] / 2) ** 2
+    algo.set_input_field(Field('input', amplitude=np.exp(-r2 / (2 * (shape[0] / 6) ** 2))))
+
+    target = np.zeros(shape)
+    cy, cx = shape[0] // 2, shape[1] // 2
+    for dy, dx in ((-12, -12), (-12, 12), (12, -12), (12, 12)):
+        target[cy + dy, cx + dx] = 1
+    algo.set_target_field(Field('target', amplitude=target))
+    algo.define_input_phase(np.zeros(shape))
+    return algo
+
+
+def efficiency_on_spots(algo):
+    target = algo._target_field.amplitude > 0
+    intensity = algo.intensity_output
+    return intensity[target].sum() / intensity.sum()
+
+
+def test_registered():
+    assert AlgorithmFactory.algorithms_registry[SlmsuiteHologram.ALGO_NAME] is SlmsuiteHologram
+
+
+@pytest.mark.parametrize('method, defaults', METHOD_DEFAULTS.items())
+def test_method_defaults_follow_slmsuite(method, defaults):
+    slm_defaults = ALGORITHM_DEFAULTS[method]
+    for name, value in defaults.items():
+        assert slm_defaults[name] == value
+
+
+@pytest.mark.parametrize('method', list(METHOD_DEFAULTS))
+def test_method_concentrates_light_on_spots(algo, method):
+    algo.settings.child('method').setValue(method)
+    algo.settings.child('iterations').setValue(5)
+    algo.compute_phase()
+    first = efficiency_on_spots(algo)
+    for _ in range(3):
+        algo.compute_phase()
+    assert efficiency_on_spots(algo) >= first > 0.05
+    assert np.all(np.isfinite(algo.modulator_field.phase))
+
+
+def test_optimization_state_is_kept_between_steps(algo):
+    algo.compute_phase()
+    hologram = algo._hologram
+    algo.compute_phase()
+    assert algo._hologram is hologram
+
+
+def test_restart_when_initial_phase_or_target_changes(algo):
+    algo.compute_phase()
+    hologram = algo._hologram
+
+    algo.define_input_phase(np.random.random(algo.shape))
+    algo.compute_phase()
+    assert algo._hologram is not hologram
+
+    hologram = algo._hologram
+    algo.do_things_after_set_target()
+    algo.compute_phase()
+    assert algo._hologram is not hologram
+
+
+def test_method_changes_visible_settings(algo):
+    algo.settings.child('method').setValue('WGS-Kim')
+    assert algo.settings.child('fix_phase_iteration').opts['visible']
+    assert algo.settings.child('feedback_exponent').opts['visible']
+    assert not algo.settings.child('feedback_factor').opts['visible']
+    algo.settings.child('method').setValue('GS')
+    assert not algo.settings.child('feedback_exponent').opts['visible']
+
+
+def test_discovered_through_entry_point():
+    modules = factory.register_external_algorithms()
+    assert slmsuite_hologram.__package__ in [m.__name__ for m in modules]
+    assert AlgorithmFactory.algorithms_registry[SlmsuiteHologram.ALGO_NAME] is SlmsuiteHologram