diff --git a/README.rst b/README.rst index 5ff552f..0f458fd 100644 --- a/README.rst +++ b/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 ========================= diff --git a/hatch_build.py b/hatch_build.py index a0208ee..cf48e3e 100644 --- a/hatch_build.py +++ b/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'} diff --git a/pyproject.toml b/pyproject.toml index ddb9bee..57072c0 100644 --- a/pyproject.toml +++ b/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"] diff --git a/src/pymodaq_plugins_slmsuite/algorithms/__init__.py b/src/pymodaq_plugins_slmsuite/algorithms/__init__.py new file mode 100644 index 0000000..76804c3 --- /dev/null +++ b/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) diff --git a/src/pymodaq_plugins_slmsuite/algorithms/slmsuite_hologram.py b/src/pymodaq_plugins_slmsuite/algorithms/slmsuite_hologram.py new file mode 100644 index 0000000..0ab06df --- /dev/null +++ b/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 diff --git a/tests/test_slmsuite_algorithm.py b/tests/test_slmsuite_algorithm.py new file mode 100644 index 0000000..e33af89 --- /dev/null +++ b/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