Slim down dependencies and tidy up install docs

.github/workflows/docs.yml

@@ -24,10 +24,10 @@ jobs:
         uses: astral-sh/setup-uv@v5
 
       - name: Install documentation dependencies
-        run: uv sync --extra docs
+        run: uv sync --group docs
 
       - name: Build documentation
-        run: uv run mkdocs build --strict
+        run: uv run --group docs mkdocs build --strict
 
       - name: Upload documentation
         if: github.event_name != 'pull_request'

.github/workflows/publish.yml

@@ -82,8 +82,8 @@ jobs:
         with:
           python-version: ${{ matrix.python-version }}
           cache-dependency-glob: "pyproject.toml"
-      - name: Install the project with `dev`
-        run: uv sync --extra dev
+      - name: Install the project with the `dev` group
+        run: uv sync
       - name: Run tests
         run: uv run pytest tests
 

.github/workflows/test.yml

@@ -22,7 +22,7 @@ jobs:
         with:
           python-version: ${{ matrix.python-version }}
           cache-dependency-glob: "pyproject.toml"
-      - name: Install the project with `dev`
-        run: uv sync --extra dev
+      - name: Install the project with the `dev` group
+        run: uv sync
       - name: Run tests
         run: uv run pytest tests

PtyLab/Engines/BaseEngine.py

@@ -1919,8 +1919,15 @@ def position_update_to_change_in_z(self, loop):
         z_{\mathrm{target}} = \frac{z}{s}
         $$
 
-        The distance update is filtered through an Adam optimizer before being
-        applied to `reconstruction.zo`. After updating the propagation distance,
+        Rather than jumping to $z_{\mathrm{target}}$, the distance moves a fixed
+        fraction $\alpha = 0.2$ of the way there on each call, so that noisy position
+        estimates do not make it jump:
+
+        $$
+        z \leftarrow z + \alpha \, (z_{\mathrm{target}} - z)
+        $$
+
+        After updating the propagation distance,
         the corrected scan positions are rescaled around their center so that the
         global scale change is transferred from the position correction to the
         propagation distance.
@@ -1932,25 +1939,7 @@ def position_update_to_change_in_z(self, loop):
         Notes:
             This method is used together with position correction and is called
             periodically when `map_position_to_z_change` is enabled.
-
-            The method requires JAX, which is imported when the function is
-            called.
         """
-        import jax
-        from jax.experimental import optimizers
-
-        if not hasattr(self, "optlib"):
-            self.i_z_optimizer = 0
-            # from itertools import count
-            # count
-            op_init, op_update, op_get = optimizers.adam(3e-3)
-            state = op_init(self.reconstruction.zo)
-            self.optlib = {"op_update": op_update, "op_get": op_get, "state": state}
-        else:
-            state = self.optlib["state"]
-            op_get = self.optlib["op_get"]
-            op_update = self.optlib["op_update"]
-
         X0 = self.reconstruction.encoder_corrected
         Y0 = self.experimentalData.encoder
         msqdisplacement = np.linalg.norm(1e6 * X0 - 1e6 * Y0)
@@ -1962,27 +1951,16 @@ def position_update_to_change_in_z(self, loop):
         # now, find the scaling with respect to the original one
         factor = np.std(X0) / np.std(Y0)
 
-        # update z
-        new_z = self.reconstruction.zo / factor
-        step = new_z - self.reconstruction.zo
-        self.logger.info(f"Naive estimate of new z: {new_z:.3f}, stepsize {step:.3f}")
-        step = 5 * step
-        # check if the thing should be updated.
-        if abs(step) < 1e-4:  # if it's too small, just truncate it,
-            # it may be that the distance changed due to some other update.
-            # Take that into account as if we don't the steps will be super large.
-            self.i_z_optimizer += 1
-            step = self.reconstruction.zo - op_get(state)
-            self.optlib["state"] = op_update(self.i_z_optimizer, -step, state)
-
+        # target distance implied by the scale change of the positions
+        target_z = self.reconstruction.zo / factor
+        step = target_z - self.reconstruction.zo
+        self.logger.info(f"Naive estimate of new z: {target_z:.3f}, stepsize {step:.3f}")
+        if abs(step) < 2e-5:  # same skip threshold as before (it used 5 * step < 1e-4)
             self.logger.info("Skipping update as step is too small")
-            # as we're only updating it for sake of good measure, we don't have to update anything else.
             return
-        # now, as we're actually updating, we can increase the step
-        self.i_z_optimizer += 1
-        self.optlib["state"] = op_update(self.i_z_optimizer, -step, state)
-        # get the new value
-        z_new = float(jax.device_get(op_get(self.optlib["state"])))
+        # damped update: move a fraction of the way toward the target each call
+        damping = 0.2
+        z_new = self.reconstruction.zo + damping * step
 
         self.logger.info(f"Loop: {loop} step: {step}")
         self.logger.info(

PtyLab/Engines/zPIE.py

@@ -343,22 +343,3 @@ def probeUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray):
             frac * DELTA, axis=(0, 1, 3), keepdims=True
         )
         return r
-
-    # def display_focus_bokeh(self):
-    #     # first, make a document
-    #     from bokeh.plotting import figure, output_file, save
-    #     from bokeh.io import hplot
-    #     from pathlib import Path
-    #
-    #     folder = Path('plots/zPIE.html')
-    #     output_file(folder)
-    #     s1, s2 = self.reconstruction.make_alignment_plot(False)
-    #     s3 = figure(width=250, height=250, title='TV per focus')
-    #
-    #
-    #     s3.xaxix.axis_label = 'Distance [um]'
-    #     s3.yaxix.axis_label = 'TV'
-    #     s3.line(self.dz*1e6, self.merit, )
-    #     pass
-    #
-    #

PtyLab/ExperimentalData/ExperimentalData.py

@@ -1,22 +1,20 @@
-import numpy as np
-from PtyLab.utils.gpuUtils import (transfer_fields_to_cpu,
-                                   transfer_fields_to_gpu)
-
-try:
-    import pyqtgraph as pg
-except ImportError:
-    print("Cannot use pyqtgraph")
-# from pathlib import Path
 import logging
 
 import matplotlib.pyplot as plt
-# import tables
+import numpy as np
+
 from PtyLab.io import readHdf5
-from PtyLab.utils.gpuUtils import (getArrayModule, transfer_fields_to_cpu,
-                                   transfer_fields_to_gpu)
+from PtyLab.utils.gpuUtils import (
+    getArrayModule,
+    transfer_fields_to_cpu,
+    transfer_fields_to_gpu,
+)
+
 # from PtyLab.io import readExample
 from PtyLab.utils.visualisation import setColorMap, show3Dslider
 
+# from pathlib import Path
+
 
 class ExperimentalData:
     """
@@ -81,10 +79,10 @@ class ExperimentalData:
 
         magnification (float):
             Microscope magnification. Available for FPM datasets, Used to convert the detector pixel size to the sample-plane pixel size.
-        
+
         NA (float or None):
             Microscope numerical aperture optional for FPM. If not provided, it is estimated from the Fourier-space pupil diameter during reconstruction.
-        
+
         energyAtPos (np.ndarray):
             Integrated diffraction intensity for each measurement frame, obtained by summing the ptychogram over the detector pixels. Used to normalize the reconstruction error for each scan position.
 
@@ -107,9 +105,7 @@ def __init__(self, filename=None, operationMode="CPM"):
         self.logger = logging.getLogger("ExperimentalData")
         self.logger.debug("Initializing ExperimentalData object")
 
-        self.operationMode = (
-            operationMode  # Select the data schema for CPM or FPM.
-        )
+        self.operationMode = operationMode  # Select the data schema for CPM or FPM.
         self._setFields()
         if filename is not None:
             self.loadData(filename)
@@ -222,7 +218,7 @@ def loadData(self, filename=None):
         )
         # Expose the loaded dataset fields as ExperimentalData attributes.
         attributesToSet = measurementDict.keys()
-        
+
         # self.logger.setLevel(logging.DEBUG)
         for a in attributesToSet:
             # make sure that property is not an attribtue
@@ -251,15 +247,15 @@ def reduce_positions(self, start, end):
                 Index at which to stop the selection. The measurement at this
                 index is not included.
 
-        Notes: 
+        Notes:
             This method modifies ``self.ptychogram`` and ``self.encoder`` in place.
         """
-        self.ptychogram = self.ptychogram[start: end]
-        self.encoder = self.encoder[start: end]
+        self.ptychogram = self.ptychogram[start:end]
+        self.encoder = self.encoder[start:end]
         self._setData()
 
     def cropCenter(self, size):
-        '''
+        """
         Crop each diffraction pattern to a centered square region.
 
         The ptychogram is cropped along its two detector dimensions while the
@@ -279,20 +275,22 @@ def cropCenter(self, size):
 
             Derived detector quantities such as ``Nd`` and ``Ld`` are currently
             not recomputed by this method.
-        '''
+        """
         if not isinstance(size, int):
-            raise TypeError('Crop value is not valid. Int expected')
+            raise TypeError("Crop value is not valid. Int expected")
 
         x = self.ptychogram.shape[-1]
         startx = x // 2 - (size // 2)
 
         startx += 1
 
-        self.ptychogram = self.ptychogram[..., startx: startx + size, startx: startx + size]
+        self.ptychogram = self.ptychogram[
+            ..., startx : startx + size, startx : startx + size
+        ]
         self._setData()
 
     def binData(self, binning):
-        '''
+        """
         Spatially bin each diffraction pattern by averaging neighboring pixels.
 
         Each ``binning × binning`` detector region is replaced by its mean value,
@@ -308,7 +306,7 @@ def binData(self, binning):
 
             The current implementation performs mean binning, so each output pixel
             contains the average intensity of the corresponding input pixel block.
-        '''
+        """
         Ndp = self.ptychogram.shape[0]
         Ny = self.ptychogram.shape[1]
         Nx = self.ptychogram.shape[2]
@@ -452,7 +450,6 @@ def _move_data_to_gpu(self):
         """Move all required fata to the GPU"""
         transfer_fields_to_gpu(self, self.fields_to_transfer, self.logger)
 
-
     def relative_intensity(self, index):
         """
         Return the normalized mean intensity of a selected ptychogram frame.
@@ -475,7 +472,9 @@ def relative_intensity(self, index):
 
             The normalized intensities are cached after the first call.
         """
-        if not hasattr(self, '_relative_intensity'):
-            self._relative_intensity = self.ptychogram.mean((-2,-1))
-            self._relative_intensity /= (self._relative_intensity.mean() + 2*self._relative_intensity.std())
+        if not hasattr(self, "_relative_intensity"):
+            self._relative_intensity = self.ptychogram.mean((-2, -1))
+            self._relative_intensity /= (
+                self._relative_intensity.mean() + 2 * self._relative_intensity.std()
+            )
         return self._relative_intensity[index]