From e5511a0662ca2a48fd1552abfa04cac1ec929a05 Mon Sep 17 00:00:00 2001 From: Chang Liu Date: Wed, 23 Sep 2026 11:53:31 +0200 Subject: [PATCH 01/14] minor change in quick start.md --- docs/getting-started/quickstart.md | 51 +++++++++++++++++++----------- 1 file changed, 33 insertions(+), 18 deletions(-) diff --git a/docs/getting-started/quickstart.md b/docs/getting-started/quickstart.md index 21a485b..6b4ec18 100644 --- a/docs/getting-started/quickstart.md +++ b/docs/getting-started/quickstart.md @@ -1,4 +1,17 @@ # Quick Start +PtyLab supports both **conventional ptychography (CPM)** and **Fourier ptychography (FPM)**. This quick start uses CPM as the main example. + +If you are new to PtyLab, a good place to start a ptychography reconstruction is the example script: +```text +example_scripts/exampleReconstructionCPM.py +``` +The script shows a complete CPM reconstruction workflow and can be used as a reference while reading the sections below. + +To get started without your own dataset, use the bundled simulation data `simu.hdf5`. If it has not yet been generated， run +```text +example_scripts/simulationData.py +``` +for example data generation. ## Using `easyInitialize` @@ -19,36 +32,38 @@ for loop, posLoop in engine.reconstruct(): pass reconstruction.saveResults("result.hdf5") + ``` +The alias `"example:simulation_cpm"` resolves to the synthetic CPM dataset `simu.hdf5`. + +Available example datasets: + +| Name | Description | +|------|-------------| +| `"example:simulation_cpm"` | Synthetic CPM dataset | +| `"example:simulation_fpm"` | Synthetic FPM dataset | + `easyInitialize` returns a 5-tuple for CPM: | Object | Type | Description | |--------|------|-------------| -| `experimentalData` | `ExperimentalData` | Diffraction data and geometry from the HDF5 file | -| `reconstruction` | `Reconstruction` | Mutable state: object array, probe array, scan positions | +| `experimentalData` | `ExperimentalData` | Diffraction data and experimental geometry from the HDF5 file | +| `reconstruction` | `Reconstruction` | Mutable reconstruction state, including object, probe, and scan positions | | `params` | `Params` | Shared configuration (propagator type, constraints, switches) | | `monitor` | `Monitor` | Real-time visualization during reconstruction | | `engine` | `BaseEngine` | The reconstruction algorithm instance (e.g. `mPIE`) | -## Using built-in example data - -To get started without your own dataset, use the bundled simulation data: +A minimal reconstruction can then be run as: ```python -experimentalData, reconstruction, params, monitor, engine = PtyLab.easyInitialize( - "example:simulation_cpm", - engine=Engines.ePIE, - operationMode="CPM", -) -``` +engine.numIterations = 50 -Available example datasets: +for loop, posLoop in engine.reconstruct(): + pass -| Name | Description | -|------|-------------| -| `"example:simulation_cpm"` | Synthetic CPM dataset | -| `"example:simulation_fpm"` | Synthetic FPM dataset | +reconstruction.saveResults("result.hdf5") +``` ## Headless mode @@ -116,7 +131,7 @@ reconstruction.load_object("previous_result.hdf5") ## Next steps - [CPM Workflow Overview](../cpm/overview.md) — understand the full reconstruction pipeline -- [Engines](../cpm/engines.md) — choose the right reconstruction algorithm - [Configuration Reference](../cpm/configuration.md) — all available `Params` options -- [FPM Workflow](../fpm/overview.md) — Fourier ptychography with LED arrays +- [Engines](../cpm/engines.md) — choose the right reconstruction algorithm - [Tutorial Notebooks](../tutorials/tutorial_CPM_sim.ipynb) — worked examples end to end +- [FPM Workflow](../fpm/overview.md) — Fourier ptychography with LED arrays From f7dda89c2d2bed9599f385522f3c23371628b69e Mon Sep 17 00:00:00 2001 From: Chang Liu Date: Wed, 23 Sep 2026 16:48:29 +0200 Subject: [PATCH 02/14] Enhance documentations: API reference - ExperimentalData --- PtyLab/ExperimentalData/ExperimentalData.py | 101 ++++++++++++++++---- PtyLab/__init__.py | 63 +++++++++++- docs/api/experimental-data.md | 5 +- 3 files changed, 150 insertions(+), 19 deletions(-) diff --git a/PtyLab/ExperimentalData/ExperimentalData.py b/PtyLab/ExperimentalData/ExperimentalData.py index c441ebb..3698107 100644 --- a/PtyLab/ExperimentalData/ExperimentalData.py +++ b/PtyLab/ExperimentalData/ExperimentalData.py @@ -218,7 +218,7 @@ def loadData(self, filename=None): def reduce_positions(self, start, end): """ - Restrict the dataset to a contiguous subset of measurement positions. + Restrict the dataset to a contiguous subset of measurement positions. The ptychogram and the corresponding encoder positions are sliced along their first dimension using standard Python slicing semantics. @@ -256,9 +256,6 @@ def cropCenter(self, size): Notes: This method modifies ``self.ptychogram`` in place. - - 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') @@ -273,8 +270,21 @@ def cropCenter(self, size): def binData(self, binning): ''' - :param binning: Binning parameter (int, e.g. 2) - :return: + Spatially bin each diffraction pattern by averaging neighboring pixels. + + Each ``binning × binning`` detector region is replaced by its mean value, + reducing both detector dimensions by the specified binning factor. + + Args: + binning (int): + Integer binning factor applied along both detector dimensions. + The detector dimensions must be divisible by this value. + + Notes: + This method modifies ``self.ptychogram`` in place. + + 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] @@ -283,7 +293,7 @@ def binData(self, binning): ptychogram_temp = np.copy(self.ptychogram) self.ptychogram = np.zeros((Ndp, Ny // binning, Nx // binning)) - # Loop through all dp + # Bin each diffraction pattern independently. for i in range(Ndp): temp = ptychogram_temp[i] reshaped_temp = temp.reshape(Ny // binning, binning, Nx // binning, binning) @@ -292,9 +302,40 @@ def binData(self, binning): def setOrientation(self, orientation, force_contiguous=True): """ - Sets the correct orientation. This function follows the ptypy convention. + Apply the detector orientation specified by the ptypy convention. + + The orientation is applied to the last two dimensions of + ``self.ptychogram`` using combinations of axis flips and transposition. + + Args: + orientation (int or None): + Orientation code following the ptypy convention: + + - ``0``: no transformation. + - ``1``: flip detector columns. + - ``2``: flip detector rows. + - ``3``: flip detector rows and columns. + - ``4``: transpose the detector dimensions. + - ``5``: transpose, then flip columns. + - ``6``: transpose, then flip rows. + - ``7``: transpose, then flip rows and columns. + + If None, no transformation is applied. + + force_contiguous (bool, optional): + If True, convert the transformed ptychogram to a contiguous + NumPy array. Defaults to True. + + Raises: + TypeError: + If ``orientation`` is not an integer or None. + + ValueError: + If ``orientation`` is not one of the supported values from + 0 to 7. - If orientation is None, it won't change the current orientation. + Notes: + This method modifies ``self.ptychogram`` in place. """ if orientation is None: # do not update. return @@ -333,7 +374,17 @@ def setOrientation(self, orientation, force_contiguous=True): self.ptychogram = np.ascontiguousarray(self.ptychogram) def _setData(self): + """ + Update detector geometry and dataset-derived quantities. + + This method derives detector coordinates, detector size, frame count, + per-frame integrated intensity, and the maximum probe-amplitude scale + from the current ``ptychogram`` and detector pixel size. + Notes: + This method should be called whenever the ptychogram shape or detector + sampling changes. + """ # Set the detector coordinates self.Nd = self.ptychogram.shape[-1] # Detector coordinates 1D @@ -352,7 +403,15 @@ def _setData(self): def showPtychogram(self): """ - show ptychogram. + Display the measured ptychogram stack on a logarithmic intensity scale. + + The diffraction patterns are clipped to non-negative values, converted to + ``log10(I + 1)`` for visualization, and displayed with an interactive + slider over the measurement frames. + + Notes: + This method is intended for data inspection only and does not modify + ``self.ptychogram``. """ xp = getArrayModule(self.ptychogram) print(f"Min max ptychogram: {np.min(self.ptychogram)}, {self.ptychogram.max()}") @@ -373,15 +432,25 @@ def _move_data_to_gpu(self): def relative_intensity(self, index): """ - Return the relative intensity of the ptychogram at index compared to the brightest one + Return the normalized mean intensity of a selected ptychogram frame. + + The mean detector intensity is computed for each frame and + normalized. - Parameters - ---------- - index + Args: + index (int): + Index of the measurement frame. + + Returns: + float: + Normalized mean intensity of the selected frame. + + Notes: + The normalization is defined as - Returns - ------- + ``I_rel = I_mean / (mean(I_mean) + 2 * std(I_mean))``. + The normalized intensities are cached after the first call. """ if not hasattr(self, '_relative_intensity'): self._relative_intensity = self.ptychogram.mean((-2,-1)) diff --git a/PtyLab/__init__.py b/PtyLab/__init__.py index 32d4446..55367cb 100644 --- a/PtyLab/__init__.py +++ b/PtyLab/__init__.py @@ -14,7 +14,34 @@ def easyInitialize( operationMode="CPM", dummyMonitor=False, ) -> Tuple[ExperimentalData, Reconstruction, Params, Monitor, Engines.BaseEngine]: - """Do a 'standard' initialization, and return the items you need with some sensible defaults.""" + ''' + Initialize the main PtyLab components for CPM or FPM reconstruction. + + Args: + filename (Path): + Path to the experimental data file. + + engine (Engines.BaseEngine, optional): + EReconstruction engine class to instantiate. + Defaults to ``Engines.ePIE``. + + operationMode (str, optional): + Operation mode, either ``"CPM"`` or ``"FPM"``. + Defaults to ``"CPM"``. + + dummyMonitor (bool, optional): + If True, use a dummy monitor without graphical output. + Defaults to False. + + Returns: + tuple: + Initialized PtyLab objects required for the reconstruction. + For FPM, additionally returns an ``IlluminationCalibration`` object. + + Raises: + NotImplementedError: + If ``operationMode`` is neither ``"CPM"`` nor ``"FPM"``. + ''' if operationMode == "CPM": return _easyInitializeCPM(filename, engine, operationMode, dummyMonitor) if operationMode == "FPM": @@ -24,6 +51,40 @@ def easyInitialize( def _easyInitializeCPM(filename, engine_function, operationMode, dummy_monitor=False): + ''' + Initialize the main PtyLab components for a conventional ptychography reconstruction. + + Args: + filename (str or Path): + Path to the experimental data file. + + engine_function (type[Engines.BaseEngine]): + Reconstruction engine class to instantiate, for example + ``Engines.ePIE`` or ``Engines.mPIE``. + + operationMode (str): + Ptychographic operation mode passed to ``ExperimentalData``. + For this helper, this is expected to be ``"CPM"``. + + dummy_monitor (bool, optional): + If True, use a ``DummyMonitor`` without graphical output. + Otherwise, initialize the standard graphical ``Monitor``. + Defaults to False. + + Returns: + tuple: + A tuple containing: + + - ``ExperimentalData``: loaded diffraction data and acquisition geometry. + - ``Reconstruction``: initialized reconstruction state. + - ``Params``: reconstruction parameters. + - ``Monitor`` or ``DummyMonitor``: reconstruction monitor. + - ``BaseEngine``: initialized reconstruction engine. + + Notes: + The object and probe are initialized by calling + ``reconstruction.initializeObjectProbe()`` before the engine is created. + ''' experimentalData = ExperimentalData(filename, operationMode) params = Params() if dummy_monitor: diff --git a/docs/api/experimental-data.md b/docs/api/experimental-data.md index a74e1af..e0318d9 100644 --- a/docs/api/experimental-data.md +++ b/docs/api/experimental-data.md @@ -1,5 +1,6 @@ # Experimental Data -::: PtyLab.ExperimentalData +::: PtyLab.ExperimentalData.ExperimentalData.ExperimentalData options: - show_submodules: true + show_root_heading: true + show_root_full_path: false \ No newline at end of file From 9827070cd8e5c86d0675649aac7bda434bc1fee9 Mon Sep 17 00:00:00 2001 From: Chang Liu Date: Thu, 24 Sep 2026 11:50:43 +0200 Subject: [PATCH 03/14] Improve API documentation for Reconstruction.py. --- PtyLab/ExperimentalData/ExperimentalData.py | 28 +- PtyLab/Reconstruction/Reconstruction.py | 590 +++++++++++++++++--- docs/api/reconstruction.md | 19 +- 3 files changed, 556 insertions(+), 81 deletions(-) diff --git a/PtyLab/ExperimentalData/ExperimentalData.py b/PtyLab/ExperimentalData/ExperimentalData.py index b44537b..ee356a0 100644 --- a/PtyLab/ExperimentalData/ExperimentalData.py +++ b/PtyLab/ExperimentalData/ExperimentalData.py @@ -22,10 +22,9 @@ class ExperimentalData: """ Store experimental data and geometry for a PtyLab reconstruction. - The class defines the experimental fields required for conventional - ptychography (CPM) or Fourier ptychography (FPM), loads them from an HDF5 - dataset, and derives basic detector and acquisition quantities used by the - reconstruction. + The class defines the dataset fields required for conventional ptychography + (CPM) and Fourier ptychography (FPM), loads them from an HDF5 dataset, and + derives detector and dataset quantities used during reconstruction. Args: filename (str or Path, optional): @@ -65,11 +64,32 @@ class ExperimentalData: zo (float): Sample-to-detector distance in meters. Available for CPM datasets. + entrancePupilDiameter (float or None): + Effective probe diameter used for probe initialization. Optional for CPM datasets, not used in FPM. + + spectralDensity (np.ndarray or None): + Spectral information used for polychromatic reconstruction. Optional for CPM datasets, not used in FPM. + + theta (float or None): + Sample tilt or incidence-angle parameter used for reflection-mode. Optional for CPM datasets, not used in FPM. + + emptyBeam (np.ndarray or None): + Reference illumination or probe image. Optional for CPM datasets, not used in FPM. + zled (float): LED-to-sample distance in meters. Available for FPM datasets. magnification (float): Microscope magnification. Available for FPM datasets. + + 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 intensity of each measurement frame. + + maxProbePower (float): + Maximum integrated-amplitude scale derived from the ptychogram. Raises: ValueError: diff --git a/PtyLab/Reconstruction/Reconstruction.py b/PtyLab/Reconstruction/Reconstruction.py index 28f9695..f4a384c 100644 --- a/PtyLab/Reconstruction/Reconstruction.py +++ b/PtyLab/Reconstruction/Reconstruction.py @@ -20,7 +20,31 @@ def calculate_pixel_positions(encoder_corrected, dxo, No, Np, asint): """ - Calculate the pixel positions. + Convert real-space scan positions to object-array pixel indices. + + The returned positions correspond to the upper-left corner of the + object patch illuminated by the probe. + + Args: + encoder_corrected (np.ndarray): + Corrected scan positions in meters, typically with shape + ``(numFrames, 2)``. + + dxo (float): + Object-plane pixel size in meters. + + No (int): + Number of pixels along one dimension of the object array. + + Np (int): + Number of pixels along one dimension of the probe array. + + asint (bool): + If True, return positions as integer pixel indices. + + Returns: + np.ndarray: + Pixel coordinates of the upper-left corner of each object patch. """ positions = np.round( encoder_corrected / dxo @@ -33,13 +57,176 @@ def calculate_pixel_positions(encoder_corrected, dxo, No, Np, asint): class Reconstruction(object): """ - This object will contain all the things that can be modified by a reconstruction. + Store and manage the mutable state of a PtyLab reconstruction. + + The reconstruction state is initialized from an ``ExperimentalData`` object + and a ``Params`` instance. Experimental quantities that may change during + reconstruction are copied from ``ExperimentalData``, while derived sampling, + coordinate grids, scan positions, object/probe settings, and reconstruction + state are maintained by this class. + + Args: + data (ExperimentalData): + Experimental data and acquisition geometry used to initialize the + reconstruction. + + params (Params): + Reconstruction parameters and algorithm settings. + + Attributes: + wavelength (float or np.ndarray): + Illumination wavelength in meters. + + dxd (float): + Detector pixel size in meters. + + zo (float): + Propagation distance used by the reconstruction. For CPM, this is the + sample-to-detector distance. Updating ``zo`` also updates ``dxp``. + + dxp (float): + Probe-plane pixel size in meters. For CPM, it is derived from the + wavelength, propagation distance, and detector field of view. For FPM, + it is derived from the detector pixel size and microscope magnification. + + theta (float or None): + Angular geometry parameter used for CPM, when provided by the + experimental dataset. + + spectralDensity (np.ndarray or None): + Spectral weight used for polychromatic CPM reconstruction. + + entrancePupilDiameter (float or None): + Effective probe or pupil diameter associated with the reconstruction. + + zled (float): + LED-to-sample distance in meters for FPM. + + NA (float or None): + Numerical aperture used for FPM. + + encoder_corrected (np.ndarray): + Measurement coordinates currently used by the reconstruction, + including any applied position corrections. + + positions0 (np.ndarray): + Reconstruction positions at initialization, in pixel coordinates. + + nlambda (int): + Number of reconstructed wavelength modes. + + nosm (int): + Number of object modes. + + npsm (int): + Number of probe modes. + + nslice (int): + Number of object slices. + + No (int): + Number of pixels along one dimension of the reconstructed object. + + initialObject (str): + Object initialization method. + + initialProbe (str): + Probe initialization method. + + object (np.ndarray): + Current complex-valued reconstructed object. + + probe (np.ndarray): + Current complex-valued reconstructed probe. + + error (np.ndarray): + Reconstruction error history. Available after reconstruction has + produced an error metric. - In itself, it's little more than a data holder. It is initialized with an ExperimentalData object. + purityProbe (float): + Current probe-purity value used by reconstruction diagnostics. - Some parameters which are "immutable" within the ExperimentalData can be modified - (e.g. zo modification by zPIE during the reconstruction routine). All of them - are defined in the listOfReconstructionProperties + purityObject (float): + Current object-purity value used by reconstruction diagnostics. + + Nd (int): + Number of detector pixels along one dimension. + + xd (np.ndarray): + One-dimensional detector-plane coordinates in meters. + + Xd (np.ndarray): + Two-dimensional detector-plane x-coordinate grid in meters. + + Yd (np.ndarray): + Two-dimensional detector-plane y-coordinate grid in meters. + + Ld (float): + Physical width of the detector grid in meters. + + Np (int): + Number of pixels along one dimension of the probe grid. + + xp (np.ndarray): + One-dimensional probe-plane coordinates in meters. + + Xp (np.ndarray): + Two-dimensional probe-plane x-coordinate grid in meters. + + Yp (np.ndarray): + Two-dimensional probe-plane y-coordinate grid in meters. + + Lp (float): + Physical field of view of the probe grid in meters. + + dxo (float): + Object-grid pixel size in meters. In the current implementation this + is equal to ``dxp``. For the real-space FPM object sampling, use + ``dxo_fpm``. + + xo (np.ndarray): + One-dimensional object-grid coordinates in meters. + + Xo (np.ndarray): + Two-dimensional object-grid x-coordinate array in meters. + + Yo (np.ndarray): + Two-dimensional object-grid y-coordinate array in meters. + + Lo (float): + Physical field of view associated with the object grid. + + dxo_fpm (float): + Real-space object pixel size for FPM in meters. + + Lo_fpm (float): + Real-space field of view of the FPM object in meters. + + dfp (float): + Spatial-frequency sampling of the FPM probe grid in inverse meters. + + positions (np.ndarray): + Reconstruction positions in pixel coordinates. Positions are given in + row-column order and refer to the upper-left corner of each + reconstructed patch. For CPM they refer to real-space object patches; + for FPM they refer to patches in the high-resolution Fourier-space + representation. + + NAd (float): + Effective detection numerical aperture. + + DoF (float): + Estimated depth of field in meters. + + TV (float): + Total-variation metric of the current reconstructed object. + + Notes: + ``Reconstruction`` contains quantities that may evolve during an + iterative reconstruction. + + Several geometric and sampling quantities are exposed as properties and + are derived from the experimental data and current reconstruction state. """ _Nd = None @@ -62,7 +249,16 @@ class Reconstruction(object): ] def __init__(self, data: ExperimentalData, params: Params): + """ + Initialize the reconstruction state from experimental data and parameters. + + Args: + data (ExperimentalData): + Experimental data and acquisition geometry. + params (Params): + Reconstruction parameters and algorithm settings. + """ self.zMomentum = 0 self.wavelength = None self._zo = None @@ -110,7 +306,22 @@ def __init__(self, data: ExperimentalData, params: Params): def copyAttributesFromExperiment(self, data: ExperimentalData): """ - Copy all the attributes from the experiment that are in listOfReconstructionProperties (CPM or FPM) + Copy reconstruction-relevant attributes from the experimental data. + + The attributes copied depend on the selected operation mode. The propagation distance + and corrected measurement positions are handled separately. + + Args: + data (ExperimentalData): + Experimental data object from which reconstruction parameters and + measurement positions are copied. + + Notes: + For CPM, ``zo`` is assigned after the other geometry parameters because + setting ``zo`` also updates the probe-plane sampling ``dxp``. + + ``encoder_corrected`` is initialized from ``data.encoder`` only if it + has not already been set, preserving any existing position corrections. """ self.logger.debug("Copying attributes from Experimental Data") if self.data.operationMode == "CPM": @@ -134,12 +345,23 @@ def copyAttributesFromExperiment(self, data: ExperimentalData): self.encoder_corrected = data.encoder.copy() def reset_positioncorrection(self): - """Reset the position corrections.""" + """ + Reset corrected measurement positions to the original encoder positions. + + The current ``encoder_corrected`` values are replaced by a copy of + ``ExperimentalData.encoder``, removing any position corrections applied + during reconstruction. + """ self.encoder_corrected = self.data.encoder.copy() @property def zo(self): - """Distance from sample to detector. Also updates all derived qualities.""" + """ + Propagation distance used by the reconstruction. + + For CPM, this represents the sample-to-detector distance. Updating + ``zo`` also updates the probe-plane pixel size ``dxp``. + """ return self._zo @zo.setter @@ -154,7 +376,22 @@ def zo(self, new_value): def computeParameters(self): """ - compute parameters that can be altered by the user later. + Compute reconstruction geometry and mode-dependent default parameters. + + For CPM, missing probe and spectral parameters are initialized from the current reconstruction geometry. + For FPM, the sample-plane sampling and pupil geometry are derived from the microscope magnification and numerical + aperture. + + The object-array size is determined from the range of reconstruction + positions with additional space for the probe. + + Notes: + This method may update both the ``Reconstruction`` instance and its + associated ``ExperimentalData`` object. + + If ``No`` has not been defined, a temporary value is assigned first so + that pixel positions can be evaluated before the final object size is + determined. """ if self.data.operationMode == "CPM": @@ -204,8 +441,26 @@ def computeParameters(self): self.No = np.max([self.Np, range_pixels]) def make_alignment_plot(self, saveit=False): - import time + """ + Create diagnostic plots for position alignment and autofocus history. + + The main plot compares the initial reconstruction positions with the + current corrected positions. Additional plots are included when autofocus + or total-variation history is available. + Args: + saveit (bool, optional): + If True, save the diagnostic plots to + ``plots/alignment.html``. Defaults to False. + + Returns: + bokeh.layouts.LayoutDOM: + Bokeh layout containing the available diagnostic plots. + + Notes: + Position coordinates are derived from ``positions`` and ``positions0``. + These quantities are expressed in reconstruction pixels. + """ t0 = time.time() p_new = self.positions.T p_old = self.positions0.T @@ -230,11 +485,11 @@ def make_alignment_plot(self, saveit=False): title=f'alignment (updated {time.strftime("%Y%h%d, %H:%M:%S")})', ) p.match_aspect = True - square = p.square( + p.square( p_old[0], p_old[1], fill_color="yellow", size=5, legend_label="original" ) # add a circle renderer for the new points - circle = p.circle( + p.circle( p_new[0], p_new[1], fill_color="red", size=5, legend_label="new" ) @@ -301,19 +556,26 @@ def make_alignment_plot(self, saveit=False): def initializeSettings(self): """ - Initialize the attributes that have to do with a reconstruction - or experimentalData fields which will become "reconstruction" + Initialize the default reconstruction model and initialization settings. - This method just sets the settings. It sets the what kind of initial guess should be used for initialObject - and initialProbe but it does not compute them yet. That will be done by calling initializeObjectProbe() + This method sets the number of wavelength, object, probe, and slice modes, + initializes purity-related state ``purityProbe`` and ``purityObject``, stores the initial reconstruction + positions, and selects the default object and probe initialization methods. - :return: + The object and probe arrays are not created by this method. They are + initialized later by ``initializeObjectProbe()``. + + Notes: + The default reconstruction model uses one wavelength, one object mode, + one probe mode, and one object slice. + + CPM initializes the object with ``"ones"`` and the probe with + ``"circ"``, while FPM uses ``"upsampled"`` for the object and + ``"circ"`` for the probe. """ - # create a 6D object where which allows to have: - # 1. polychromatic = nlambda - # 2. mixed state object - nosm - # 3. mixed state probe - npsm - # 4. multislice object (thick) - nslice + # Configure the reconstruction model dimensions. + # These support multiple wavelengths, mixed object/probe states, + # and multislice reconstruction. self.nlambda = 1 self.nosm = 1 self.npsm = 1 @@ -337,14 +599,35 @@ def initializeSettings(self): self.initialObject = "ones" def prepare_probe(self, i): - """Replace probe with the i-th TSVD estimate. + """ + Replace the current probe with a selected TSVD probe estimate. + + This method is intended for OPRP implementations and must be overridden + by a reconstruction class that provides the corresponding probe estimates. - This function is used in OPRP + Args: + i (int): + Index of the TSVD probe estimate to use. + + Raises: + NotImplementedError: + Always raised by the base ``Reconstruction`` implementation. """ raise NotImplementedError() def initializeObjectProbe(self, force=True): + """ + Initialize the object and probe used for reconstruction. + Initial object and probe estimates are generated using + ``initializeObject()`` and ``initializeProbe()``, then copied to + ``self.object`` and ``self.probe`` as the mutable reconstruction state. + + Args: + force (bool, optional): + Forwarded to the object and probe initialization methods. + Defaults to True. + """ # initialize object and probe self.initializeObject(force=force) self.initializeProbe(force=force) @@ -354,6 +637,32 @@ def initializeObjectProbe(self, force=True): self.probe = self.initialGuessProbe.copy() def initializeObject(self, type_of_init=None, force=True): + """ + Initialize the object estimate used for reconstruction. + + The object shape is determined from the configured wavelength, object-mode, + slice, and spatial dimensions. The initial object is either generated using + the selected initialization method or loaded from a previous reconstruction. + + Args: + type_of_init (str, optional): + Object initialization method. If provided, this overrides + ``self.initialObject``. If None, the currently configured + initialization method is used. + + force (bool, optional): + Whether to force object initialization. Defaults to True. + The current implementation does not support ``False``. + + Raises: + NotImplementedError: + If ``force`` is False. + + Notes: + The initialized object has shape + ``(nlambda, nosm, 1, nslice, No, No)`` and is stored as + ``complex64`` when generated by ``initialProbeOrObject()``. + """ if not force: raise NotImplementedError() if type_of_init is not None: @@ -368,7 +677,7 @@ def initializeObject(self, type_of_init=None, force=True): self.No, ) if self.initialObject == 'recon': - # Load the object from an existing reconstruction + # Load the object from an existing reconstruction. Confusing filename, but it contains both object and probe. self.initialGuessObject = self.loadResults(self.initialProbe_filename, datatype='object') else: self.initialGuessObject = initialProbeOrObject(self.shape_O, self.initialObject, self, self.logger).astype(np.complex64) @@ -378,13 +687,45 @@ def initializeObject(self, type_of_init=None, force=True): @staticmethod def loadResults(fileName, datatype='probe'): ''' - Loads data from a ptylab reconstruction file. + Load an object or probe from a saved PtyLab reconstruction. + + Args: + fileName (str or Path): + Path to the reconstruction HDF5 file. + + datatype (str, optional): + Name of the dataset to load, typically ``"probe"`` or + ``"object"``. Defaults to ``"probe"``. + + Returns: + np.ndarray: + Copy of the requested reconstruction dataset. ''' with h5py.File(fileName) as archive: data = np.copy(np.array(archive[datatype])) return data def initializeProbe(self, force=False): + """ + Initialize the probe estimate used for reconstruction. + + The probe shape is determined from the configured wavelength, probe-mode, + slice, and spatial dimensions. The initial probe is either generated using + the selected initialization method or loaded from a previous reconstruction. + + Args: + force (bool, optional): + Whether to reset the existing initial probe before generating a new + estimate. Defaults to False. + + Notes: + The initialized probe has shape + ``(nlambda, 1, npsm, nslice, Np, Np)`` and is stored as + ``complex64`` when generated by ``initialProbeOrObject()``. + + If ``entrancePupilDiameter`` is not available, it is set to one third + of the probe field of view before initialization. + """ if self.data.entrancePupilDiameter is None: # if it is not set, set it to something reasonable self.logger.warning( @@ -414,23 +755,31 @@ def initializeProbe(self, force=False): # initialize momentum, called in specific engines with momentum accelaration def initializeObjectMomentum(self): + """Initialize the object momentum buffer with zeros.""" self.objectMomentum = np.zeros_like(self.initialGuessObject) def initializeProbeMomentum(self): + """Initialize the probe momentum buffer with zeros.""" self.probeMomentum = np.zeros_like(self.initialGuessProbe) def load_object(self, filename): """ - Load the object from a previous reconstruction + Load an object from a previous reconstruction. + + The saved object is truncated to the dimensions required by the current + reconstruction and assigned to ``self.object``. - Parameters - ---------- - filename: .hdf5 file - Filenamne of the reconstruction whose object should be loaded. + Args: + filename (str or Path): + Path to a PtyLab reconstruction HDF5 file containing an + ``"object"`` dataset. - Returns - ------- + Raises: + RuntimeError: + If the loaded object cannot be matched to ``self.shape_O``. + Notes: + ``shape_O`` must already be defined before calling this method. """ with h5py.File(filename, "r") as archive: obj = np.array(archive["object"]) @@ -446,20 +795,36 @@ def load_object(self, filename): self.object = obj else: raise RuntimeError( - f'Shape of saved probe cannot be extended to shape of required probe. File: {archive["object"].shape}. Need: {self.shape_O}' + f'Shape of saved object cannot be extended to shape of required object. File: {archive["object"].shape}. Need: {self.shape_O}' ) def load_probe(self, filename, expand_npsm=False, center_phase=False): """ - Load the probe from a previous reconstruction. + Load a probe from a previous reconstruction. + + The saved probe is center-cropped to the current probe size and truncated + to the wavelength, probe-mode, and slice dimensions required by the + current reconstruction. + + Args: + filename (str or Path): + Path to a PtyLab reconstruction HDF5 file containing a + ``"probe"`` dataset. + + expand_npsm (bool, optional): + Reserved for probe-mode expansion. This argument is currently not + used by the implementation. Defaults to False. + + center_phase (bool, optional): + If True, center the probe propagation angle after loading. + Defaults to False. - Parameters - ---------- - filename: .hdf5 file - The filename of the reconstruction whose probe should be loaded. + Raises: + RuntimeError: + If the loaded probe cannot be matched to ``self.shape_P``. - Returns - ------- + Notes: + ``shape_P`` must already be defined before calling this method. """ with h5py.File(filename, "r") as archive: @@ -484,7 +849,12 @@ def load_probe(self, filename, expand_npsm=False, center_phase=False): self._center_probe_angle() def _center_probe_angle(self): - """ Center the angle of propagation for the probe. """ + """ + Remove the global propagation-angle offset from the probe. + + The offset is estimated from the first probe mode and corrected by + applying a compensating phase factor. + """ from skimage.registration import phase_cross_correlation from scipy.ndimage import fourier_shift p0 = np.squeeze(self.probe)[0] @@ -493,7 +863,25 @@ def _center_probe_angle(self): self.probe *= phexp def load(self, filename): - """Load the results given by saveResults.""" + """ + Load a previously saved reconstruction state. + + This method restores the reconstructed object and probe together with + selected reconstruction metadata from a PtyLab HDF5 result file. + + Args: + filename (str or Path): + Path to a reconstruction HDF5 file produced by + ``saveResults(type="all")``. + + Notes: + The current implementation expects only CPM-style result fields and does + not restore all saved reconstruction state, such as + ``encoder_corrected``. + + Unlike ``load_object()`` and ``load_probe()``, this method does not + adapt or validate the loaded object and probe shapes. + """ with h5py.File(filename, "r") as archive: self.probe = np.array(archive["probe"]) @@ -509,18 +897,30 @@ def load(self, filename): def saveResults(self, fileName="recent", type="all", squeeze=False): """ - Save reconstruction results. + Save reconstruction results to an HDF5 file. + + Args: + fileName (str or Path, optional): + Output filename. Defaults to ``"recent"``. + type (str, optional): + Type of reconstruction data to save. Supported values are + ``"all"``, ``"object"``, ``"probe"``, and ``"probe_stack"``. + Defaults to ``"all"``. - Parameters - ---------- - fileName - type - squeeze + squeeze (bool, optional): + If True, remove singleton dimensions when saving only the object + or probe. This option does not affect ``type="all"``. + Defaults to False. + Raises: + NotImplementedError: + If an unsupported save type is requested. - Returns - ------- + Notes: + The datasets saved by ``type="all"`` depend on the operation mode. + CPM and FPM reconstruction files currently contain different sets of + reconstruction metadata. """ @@ -690,17 +1090,15 @@ def Yo(self): # scan positions in pixel @property def positions(self): - """estimated positions in pixel numbers(real space for CPM, Fourier space for FPM) - note: Positions are given in row-column order and refer to the - pixel in the upper left corner of the respective data matrix; - -1st example: suppose the 2nd row of positions0 is [3, 4] and the - operation mode is 'CPM'. That implies that the second intensity - in the spectrogram updates an object patch that has - its left uppper corner pixel at the pixel coordinates [3, 4] - -2nd example: suppose the 2nd row of positions0 is [3, 4] and the - operation mode is 'FPM'. That implies that the second intensity - in the spectrogram is updates a patch which has pixel coordinates - [3,4] in the high-resolution Fourier transform + """ + Reconstruction positions in pixel coordinates. + + Positions are stored in row-column order and refer to the upper-left + corner of the reconstructed patch associated with each measurement. + + For CPM, the positions identify patches in the real-space object array. + For FPM, they identify patches in the high-resolution Fourier-space + object representation. """ if self.data.operationMode == "FPM": conv = -(1 / self.wavelength) * self.dxo * self.Np @@ -728,13 +1126,13 @@ def positions(self): # system property list @property def NAd(self): - """Detection NA""" + """Effective detection numerical aperture.""" NAd = self.Ld / (2 * self.zo) return NAd @property def DoF(self): - """expected Depth of field""" + """Estimated depth of field in meters.""" DoF = self.wavelength / self.NAd**2 # self.Dof2 = 5.2 *self.dxp**2 /self.wavelength return DoF @@ -750,6 +1148,17 @@ def _move_data_to_gpu(self): transfer_fields_to_gpu(self, self.possible_GPU_fields, self.logger) def describe_reconstruction(self): + """ + Print a summary of the reconstruction parameters and derived quantities. + The summary includes experimental-data dimensions and sampling, reconstruction-grid parameters, propagation geometry, and derived + quantities such as detector numerical aperture and depth of field. + + The summary is also written to the reconstruction logger. + + Returns: + str: + Formatted reconstruction summary. + """ minmax_tv = '' try: minmax_tv = f'(min: {self.params.TV_autofocus_min_z*1e3}, max: {self.params.TV_autofocus_max_z*1e3}.)' @@ -757,8 +1166,8 @@ def describe_reconstruction(self): pass info = f""" Experimental data: - - Number of ptychograms: {self.data.ptychogram.shape} - - Number of pixels ptychogram: {self.data.Nd} + - Ptychogram shape: {self.data.ptychogram.shape} + - Ptychogram size[px]: {self.data.Nd} - Ptychogram size: {self.data.Ld*1e3} mm - Pixel pitch: {self.data.dxd*1e6} um - Scan size: {1e3*(self.data.encoder.max(axis=0) - self.data.encoder.min(axis=0))} mm @@ -773,7 +1182,7 @@ def describe_reconstruction(self): Derived parameters: - NA detector: {self.NAd} - - DOF: {self.DoF*1e6} um + - Depth of field: {self.DoF*1e6} um """ self.logger.info(info) @@ -781,26 +1190,53 @@ def describe_reconstruction(self): @property def quadraticPhase(self): - """These functions are cached internally in Python and therefore no longer required.""" + """Deprecated property; quadratic phase is no longer cached.""" raise NotImplementedError("Quadratic phase is no longer cached. ") @property def transferFunction(self): - raise NotImplementedError("Quad phase is not longer cached") + """Deprecated property; transfer function is no longer cached.""" + raise NotImplementedError("Transfer function is not longer cached") @property def Q1(self): + """Deprecated property; Q1 is no longer available.""" raise NotImplementedError("Q1 is no longer available") @property def Q2(self): + """Deprecated property; Q2 is no longer available.""" raise NotImplementedError("Q2 is no longer available") def TV_autofocus(self, params: Params, loop): - """Perform an autofocusing step based on optimizing the total variation. - - If not required, returns none. Otherwise, returns the value of the TV at the current z0.""" + """ + Perform one autofocus update by optimizing a propagated-field metric. + + The selected object or probe field is propagated over a range of axial + positions around the current propagation distance. A focus metric is + evaluated at each plane and used to compute a momentum-based update of + ``zo``. + + Args: + params (Params): + Reconstruction parameters controlling autofocus range, metric, + update frequency, momentum, and axial bounds. + + loop (int or None): + Current reconstruction iteration. If provided, autofocus is only + run according to ``TV_autofocus_run_every``. + + Returns: + tuple: + Normalized metric at the current plane, selected propagated fields, + and autofocus score information. Returns ``(None, None, None)`` + when no autofocus update is required. + + Raises: + NotImplementedError: + If used with FPM or with an unsupported autofocus target. + """ start_time = time.time() if self.data.operationMode == "FPM": @@ -883,10 +1319,14 @@ def TV_autofocus(self, params: Params, loop): return merit[nplanes//2] / asNumpyArray(abs(self.object[..., sy, sx]).mean()), np.hstack(OEs), (scores, self.zo) def reset_TV_autofocus(self): - """Reset the settings of TV autofocus. Can be useful to reset the memory effect if the steps are getting really large.""" + """ + Reset the autofocus momentum. + + This clears the accumulated ``zMomentum`` used by TV autofocus and can be useful when the autofocus updates become excessively large. + """ self.zMomentum = 0 @property def TV(self): - """Return the TV of the object""" + """Total-variation metric of the current reconstructed object.""" return TV(self.object, 1e-2) diff --git a/docs/api/reconstruction.md b/docs/api/reconstruction.md index 99017cd..2913eb8 100644 --- a/docs/api/reconstruction.md +++ b/docs/api/reconstruction.md @@ -1,5 +1,20 @@ # Reconstruction -::: PtyLab.Reconstruction +::: PtyLab.Reconstruction.Reconstruction.Reconstruction options: - show_submodules: true + show_root_heading: true + show_root_full_path: false + members: + - copyAttributesFromExperiment + - computeParameters + - initializeObjectProbe + - initializeObject + - initializeProbe + - load_object + - load_probe + - saveResults + - describe_reconstruction + - make_alignment_plot + - reset_positioncorrection + - TV_autofocus + - reset_TV_autofocus \ No newline at end of file From 2a55dcd2b6d7761a39478c74c5b69a0fc7af52c4 Mon Sep 17 00:00:00 2001 From: Chang Liu Date: Thu, 24 Sep 2026 12:21:30 +0200 Subject: [PATCH 04/14] Minor change: move back FPMcalibration in API. Not sorted yet. I am not faimliar with FPM. --- docs/api/reconstruction.md | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/docs/api/reconstruction.md b/docs/api/reconstruction.md index 2913eb8..5208159 100644 --- a/docs/api/reconstruction.md +++ b/docs/api/reconstruction.md @@ -17,4 +17,11 @@ - make_alignment_plot - reset_positioncorrection - TV_autofocus - - reset_TV_autofocus \ No newline at end of file + - reset_TV_autofocus + +## FPM Illumination Calibration + +::: PtyLab.Reconstruction.CalibrationFPM.IlluminationCalibration + options: + show_root_heading: true + show_root_full_path: false \ No newline at end of file From 09291a4915ca1653ecaacb4aa72d251e51c3c82d Mon Sep 17 00:00:00 2001 From: Chang Liu Date: Fri, 25 Sep 2026 16:56:55 +0200 Subject: [PATCH 05/14] Minor change: more comprehensive param description in ExperimentalData.py --- PtyLab/ExperimentalData/ExperimentalData.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/PtyLab/ExperimentalData/ExperimentalData.py b/PtyLab/ExperimentalData/ExperimentalData.py index ee356a0..38c10b3 100644 --- a/PtyLab/ExperimentalData/ExperimentalData.py +++ b/PtyLab/ExperimentalData/ExperimentalData.py @@ -77,19 +77,19 @@ class ExperimentalData: Reference illumination or probe image. Optional for CPM datasets, not used in FPM. zled (float): - LED-to-sample distance in meters. Available for FPM datasets. + LED-to-sample distance in meters. Available for FPM datasets, and used to determine the illumination angle and corresponding Fourier-space shift for each measurement. magnification (float): - Microscope magnification. Available for FPM datasets. + 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 intensity of each measurement frame. + 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. maxProbePower (float): - Maximum integrated-amplitude scale derived from the ptychogram. + Probe-amplitude scale derived from the brightest diffraction frame, defined as the square root of its integrated intensity. Used to rescale the initial probe when probe-power correction is enabled. Raises: ValueError: From e66eee4355e2fbbfa4bac2736169032764b45963 Mon Sep 17 00:00:00 2001 From: Chang Liu Date: Tue, 29 Sep 2026 16:37:59 +0200 Subject: [PATCH 06/14] API documentation: BaseEngine.py - finished. Fixed some small - implement dtype matching used by convert2single() - return linear and area overlap metrics from getOverlap() - use normalized probe-amplitude autocorrelation for area-overlap estimation --- PtyLab/Engines/BaseEngine.py | 1020 +++++++++++++++++++++++--- docs/api/engines.md | 102 ++- tests/Engines/test_convert2single.py | 70 ++ 3 files changed, 1088 insertions(+), 104 deletions(-) create mode 100644 tests/Engines/test_convert2single.py diff --git a/PtyLab/Engines/BaseEngine.py b/PtyLab/Engines/BaseEngine.py index 3d55f83..b84b2a2 100644 --- a/PtyLab/Engines/BaseEngine.py +++ b/PtyLab/Engines/BaseEngine.py @@ -35,16 +35,34 @@ def smooth_amplitude( field: np.ndarray, width: float, aleph: float, amplitude_only: bool = True ): """ - Smooth the amplitude of a field. Optional phase can be smoothed as well. - Parameters - ---------- - field - width - aleph - amplitude_only + Smooth a complex field using a Gaussian filter. - Returns - ------- + By default, the Gaussian smoothing is applied only to the field + amplitude while preserving the original phase. If ``amplitude_only`` + is False, the complex field itself is smoothed. + + Args: + field (np.ndarray): + Complex-valued field to smooth. NumPy and CuPy arrays are + supported. + + width (float): + Width of the Gaussian smoothing filter applied along the two + spatial dimensions. + + aleph (float): + Weight of the smoothed field in the returned result. A value of + zero leaves the field unchanged, while a value of one returns + the fully smoothed field. + + amplitude_only (bool, optional): + If True, smooth only the amplitude and preserve the original + phase. If False, smooth the full complex field. + Defaults to True. + + Returns: + np.ndarray: + Smoothed field with the same shape as the input. """ xp = getArrayModule(field) @@ -68,10 +86,42 @@ def smooth_amplitude( class BaseEngine(object): """ - Common properties that are common for all reconstruction Engines are defined here. + Base class providing shared functionality for PtyLab reconstruction engines. + + ``BaseEngine`` coordinates operations that are common to different + reconstruction algorithms, including wave propagation, intensity + projection, error evaluation, reconstruction constraints, position + correction, monitoring, and CPU/GPU data transfer. + + The engine operates on shared ``Reconstruction``, ``ExperimentalData``, + ``Params``, and ``Monitor`` objects. These objects are referenced directly + rather than copied, so updates performed by the engine are reflected in + the associated reconstruction state. + + Specific reconstruction engines inherit from this class and implement + their algorithm-specific reconstruction and update steps. ``BaseEngine`` + is therefore generally not intended to be instantiated directly. - Unless you are testing the code, there's hardly any need to create this object. For your own implementation, - inherit from this object + Args: + reconstruction (Reconstruction): + Mutable reconstruction state containing the current object, probe, + geometry, and reconstruction results. + + experimentalData (ExperimentalData): + Experimental diffraction data and acquisition geometry. + + params (Params): + Reconstruction parameters controlling propagation, constraints, + correction methods, GPU usage, and other algorithm settings. + + monitor (Monitor): + Monitor used to visualize and report reconstruction progress. + + Attributes: + + betaObject (float): + Object update weight used by shared object-update routines. + Initialized to ``0.25``. """ @@ -82,21 +132,28 @@ def __init__( params: Params, monitor: Monitor, ): - # These statements don't copy any data, they just keep a reference to the object + # Keep references to the shared PtyLab components; no data are copied. self.betaObject = 0.25 self.reconstruction: Reconstruction = reconstruction self.experimentalData = experimentalData self.params = params self.monitor = monitor - self.monitor.reconstruction = reconstruction + self.monitor.reconstruction = reconstruction # share reconstruction state with monitor # datalogger self.logger = logging.getLogger("BaseEngine") def _prepareReconstruction(self): """ - Initialize everything that depends on user changeable attributes. - :return: + Prepare the reconstruction state before iterative updates. + + This method validates user-configurable settings and initializes shared + engine state, including FFT conventions, probe constraints, error arrays, + monitoring regions, position-correction parameters, and CPU/GPU data + placement. + + GPU setup is performed last because the preceding initialization steps + operate on host-side reconstruction data. """ # check miscellaneous quantities specific for certain Engines self._checkMISC() @@ -114,10 +171,13 @@ def _prepareReconstruction(self): def _setCPSC(self): """ - set constrained-pixel-sum constraint: - -save measured diffraction patterns into ptychograpmDownsampled - -pad the probe (useful when having a pre-calibrated probe) - -update the coordinates + Configure the constrained-pixel-sum (CPSC) model. + + CPSC reconstructs the diffraction field on a finer computational detector + grid while constraining the summed intensity within each group of subpixels + to match the corresponding measured detector pixel. The finer grid preserves + the detector field of view and therefore the real-space sampling, while + increasing the reconstruction field of view. """ # save the measured ptychogram into ptychograpmDownsampled @@ -170,12 +230,30 @@ def _setCPSC(self): self.logger.info("CPSCswitch is on, coordinates(dxd,dxp,dxo) have been updated") def update_data(self, experimentalData, reconstruction=None): - """Update the experimentalData if necessary""" + """ + Update the data objects referenced by the engine. + + Args: + experimentalData (ExperimentalData): + Experimental dataset to use for subsequent reconstruction steps. + + reconstruction (Reconstruction, optional): + Replacement reconstruction state. If None, the current + reconstruction object is retained. + """ self.experimentalData = experimentalData if reconstruction is not None: self.reconstruction = reconstruction def _initializePCParameters(self): + """ + Initialize parameters and state for pcPIE correction. + + When position correction is enabled, this method initializes the + feedback and momentum factors, the per-position correction vectors, + the candidate pixel shifts used for local correlation searches, and + the iteration threshold for starting position updates. + """ if self.params.positionCorrectionSwitch: # additional pcPIE parameters as they appear in Matlab self.daleth = 0.5 # feedback @@ -198,10 +276,12 @@ def _initializePCParameters(self): def _initializeErrors(self): """ - initialize all kinds of errors: - detectorError is a matrix calculated at each iteration (numFrames,Nd,Nd); - errorAtPos sums over detectorError at each iteration, (numFrames,1); - reconstruction.error sums over errorAtPos, one number at each iteration; + Initialize reconstruction error storage. + + Depending on ``saveMemory``, detector-plane errors are either stored for + every scan position or reduced immediately to per-position error values. + The method also initializes the per-position error array and the global + reconstruction error history. """ # initialize detector error matrices if self.params.saveMemory: @@ -225,6 +305,30 @@ def _initializeErrors(self): self.reconstruction.error = [] def _initialProbePowerCorrection(self): + r""" + Scale the initial probe to the measured diffraction power. + + When `probePowerCorrectionSwitch` is enabled, rescale the complex probe + $P$ using the measured amplitude scale $P_{\max}$: + + $$ + P_{\mathrm{new}} = \frac{P}{\sqrt{\sum |P|^2}} P_{\max} + $$ + + The sum in the normalization includes every element of the probe array. + The amplitude scale is derived from the brightest measured diffraction + pattern: + + $$ + P_{\max} = \sqrt{\max_j \sum_{x,y} I_j(x,y)} + $$ + + Here, $I_j(x,y)$ is the measured intensity at detector pixel $(x,y)$ + in frame $j$. The correction preserves the probe shape and phase while + placing its total power on the scale of the experimental data, reducing + large amplitude corrections at the beginning of reconstruction. + + """ if self.params.probePowerCorrectionSwitch: self.reconstruction.probe = ( self.reconstruction.probe @@ -235,6 +339,27 @@ def _initialProbePowerCorrection(self): ) def _probeWindow(self): + r""" + Create the spatial window used by probe-boundary constraints. + + For `absorbingProbeBoundary`, a smooth super-Gaussian window is defined as + + $$ + W(x,y) = \exp\left[-\left(\frac{x^2+y^2}{2\sigma^2}\right)^{10}\right] + $$ + + where + + $$ + \sigma = \frac{3}{4}\frac{N_p dx_p}{2.355} + $$ + + For `probeBoundary`, a circular support window is generated from the + entrance pupil diameter. + + The generated window is stored in `probeWindow` and applied later by + `applyConstraints()`. + """ # absorbing probe boundary: filter probe with super-gaussian window function if not self.params.saveMemory or self.params.absorbingProbeBoundary: self.probeWindow = np.exp( @@ -267,7 +392,20 @@ def _probeWindow(self): def _setObjectProbeROI(self, update=False): """ - Set object/probe ROI for monitoring + Set the object and probe regions of interest used for monitoring. + + The object ROI is derived from the scan-position extent and probe size, + scaled by ``monitor.objectZoom``. The probe ROI is centered on the probe + grid and derived from the entrance pupil diameter and + ``monitor.probeZoom``. + + If the corresponding zoom value is ``"full"`` or ``None``, the complete + object or probe is displayed. + + Args: + update (bool, optional): + If True, recompute existing ROIs. Otherwise, ROIs are only created + when they are not already defined. Defaults to False. """ if not hasattr(self.monitor, "objectROI") or update: if self.monitor.objectZoom == "full" or self.monitor.objectZoom is None: @@ -331,6 +469,14 @@ def _setObjectProbeROI(self, update=False): ] def _showInitialGuesses(self): + """ + Display the initial object and probe estimates in the reconstruction monitor. + + The object and probe are cropped to the monitoring regions defined by + ``objectROI`` and ``probeROI`` before being passed to the monitor together + with the current reconstruction error, propagation distance, mode purities, + and scan positions. + """ self.monitor.initializeMonitors() objectEstimate = np.squeeze( self.reconstruction.object[ @@ -357,7 +503,17 @@ def _showInitialGuesses(self): def _checkMISC(self): """ - checks miscellaneous quantities specific certain Engines + Initialize auxiliary reconstruction state and validate special settings. + + This method prepares additional variables required by selected intensity + constraints or background reconstruction, checks incompatible parameter + combinations, and initializes the constrained-pixel-sum configuration + when enabled. + + Raises: + ValueError: + If incompatible reconstruction options are enabled or required + parameters for a selected constraint are missing. """ if self.params.backgroundModeSwitch: self.reconstruction.background = 1e-1 * np.ones( @@ -404,7 +560,15 @@ def _checkMISC(self): def _checkFFT(self): """ - shift arrays to accelerate fft + Synchronize detector-domain arrays with the selected FFT convention. + + When ``fftshiftSwitch`` is enabled, detector-side quantities are shifted + to the FFT-native ordering using ``ifftshift``. When the switch is + disabled after a previous shift, the arrays are restored using + ``fftshift``. + + ``fftshiftFlag`` tracks the current data ordering to avoid applying the + shift repeatedly. """ if self.params.fftshiftSwitch: if self.params.fftshiftFlag == 0: @@ -418,7 +582,7 @@ def _checkFFT(self): self.experimentalData.ptychogramDownsampled = np.fft.ifftshift( self.experimentalData.ptychogramDownsampled, axes=(-1, -2) ) - if hasattr(self.experimentalData, "w"): + if hasattr(self.experimentalData, "W"): if self.experimentalData.W is not None: self.experimentalData.W = np.fft.ifftshift( self.experimentalData.W, axes=(-1, -2) @@ -452,12 +616,20 @@ def _checkFFT(self): self.experimentalData.emptyBeam = np.fft.fftshift( self.experimentalData.emptyBeam, axes=(-1, -2) ) + if hasattr(self.experimentalData, "PSD"): + if self.experimentalData.PSD is not None: + self.experimentalData.PSD = np.fft.fftshift( + self.experimentalData.PSD, axes=(-1, -2) + ) self.params.fftshiftFlag = 0 def _move_data_to_gpu(self): """ - Move the data to the GPU, called when the gpuSwitch is on. - :return: + Move reconstruction data required by the engine to the GPU. + + Reconstruction and experimental-data fields are transferred by their + respective container classes. Engine-specific fields, including the probe + window and additional aPIE data when required, are transferred separately. """ self.reconstruction._move_data_to_gpu() @@ -542,8 +714,11 @@ def _move_data_to_gpu(self): def _move_data_to_cpu(self): """ - Move the data to the CPU, called when the gpuSwitch is off. - :return: + Move reconstruction data required by the engine to the CPU. + + Reconstruction and experimental-data fields are transferred by their + respective container classes, together with engine-specific fields such + as the probe window. """ # reconstruction parameters @@ -578,7 +753,7 @@ def _move_data_to_cpu(self): # if self.params.aPIEflag: # self.theta = self.theta.get() # - fields_to_transfer = ["theta", "probeWindow"] + # fields_to_transfer = ["theta", "probeWindow"] # self.probeWindow = self.probeWindow.get() # non-reconstruction parameters @@ -611,6 +786,21 @@ def _move_data_to_cpu(self): # self.reconstruction.reference = self.reconstruction.reference.get() def _checkGPU(self): + """ + Synchronize reconstruction data with the selected computation device. + + If GPU execution is enabled, required reconstruction, experimental-data, + and engine fields are transferred to the GPU. If GPU execution is + disabled, the corresponding data are transferred back to the CPU. + + ``gpuFlag`` tracks the current device state, while repeated transfers + ensure that fields created after a previous device switch are also + synchronized. + + Raises: + ImportError: + If GPU execution is requested but CuPy is not available. + """ if not hasattr(self.params, "gpuFlag"): self.params.gpuFlag = 0 @@ -635,6 +825,24 @@ def _checkGPU(self): self.params.gpuFlag = 0 def setPositionOrder(self): + """ + Set the order in which scan positions are processed. + + The ordering is controlled by ``params.positionOrder``: + + - ``"sequential"`` processes frames in their original order. + - ``"random"`` uses the original order during the first two iterations + and randomly shuffles the positions afterwards. + - ``"NA"`` sorts positions by their distance from the scan center, + processing central positions first. This ordering is intended for FPM, + where central illumination angles correspond to bright-field data. + + The resulting frame indices are stored in ``positionIndices``. + + Raises: + ValueError: + If ``positionOrder`` is not one of the supported options. + """ if self.params.positionOrder == "sequential": self.positionIndices = np.arange(self.experimentalData.numFrames) @@ -665,14 +873,34 @@ def setPositionOrder(self): raise ValueError("position order not properly set") def changeExperimentalData(self, experimentalData: ExperimentalData): + """ + Replace the experimental-data object referenced by the engine. + + Args: + experimentalData (ExperimentalData): + Experimental dataset to use in subsequent reconstruction steps. + Raises: + TypeError: + If ``experimentalData`` is not an ``ExperimentalData`` instance. + """ if experimentalData is not None: if not isinstance(experimentalData, ExperimentalData): raise TypeError("Experimental data should be of class ExperimentalData") self.experimentalData = experimentalData def changeOptimizable(self, optimizable: Reconstruction): + """ + Replace the reconstruction object referenced by the engine. + Args: + optimizable (Reconstruction): + Reconstruction state to use in subsequent engine operations. + + Raises: + TypeError: + If ``optimizable`` is not a ``Reconstruction`` instance. + """ if optimizable is not None: if not isinstance(optimizable, Reconstruction): raise TypeError( @@ -682,8 +910,12 @@ def changeOptimizable(self, optimizable: Reconstruction): def convert2single(self): """ - Convert the datasets to single precision. Matches: convert2single.m - :return: + Configure single-precision data types for reconstruction arrays. + + This method sets the target complex and real data types to + ``numpy.complex64`` and ``numpy.float32`` and delegates the actual + conversion to dtype-matching helpers. + """ self.dtype_complex = np.complex64 self.dtype_real = np.float32 @@ -691,15 +923,67 @@ def convert2single(self): self._match_dtypes_real() def _match_dtypes_complex(self): - raise NotImplementedError() + """ + Convert complex-valued arrays to the configured complex dtype. + + Complex arrays stored by the engine, reconstruction, and experimental + data containers are converted to ``self.dtype_complex``. Non-array + attributes and non-complex arrays are left unchanged. + """ + array_types = (np.ndarray,) + if cp is not None: + array_types += (cp.ndarray,) + + for container in (self, self.reconstruction, self.experimentalData): + for name, value in vars(container).items(): + if not isinstance(value, array_types): + continue + + if value.dtype.kind == "c": + setattr( + container, + name, + value.astype(self.dtype_complex, copy=False), + ) def _match_dtypes_real(self): - raise NotImplementedError() + """ + Convert floating-point arrays to the configured real dtype. + + Floating-point arrays stored by the engine, reconstruction, and + experimental data containers are converted to ``self.dtype_real``. + Integer, boolean, complex, and non-array attributes are left unchanged. + """ + array_types = (np.ndarray,) + if cp is not None: + array_types += (cp.ndarray,) + + for container in (self, self.reconstruction, self.experimentalData): + for name, value in vars(container).items(): + if not isinstance(value, array_types): + continue + + if value.dtype.kind == "f": + setattr( + container, + name, + value.astype(self.dtype_real, copy=False), + ) + def object2detector(self, esw=None): """ - Implements object2detector.m. Modifies esw in-place - :return: + Propagate the exit surface wave from the object plane to the detector plane. + + If ``esw`` is not provided, ``reconstruction.esw`` is used. The propagation + is performed by the operator selected through the reconstruction + parameters, and the propagated detector-plane field is stored in + ``reconstruction.ESW``. + + Args: + esw (ndarray, optional): + Object-plane exit surface wave. If None, use + ``reconstruction.esw``. """ if esw is None: # todo: check this, it seems weird to store it in self.esw @@ -710,10 +994,17 @@ def object2detector(self, esw=None): def detector2object(self, ESW=None): """ - Propagate the ESW to the object plane (in-place). + Propagate the detector-plane field back to the object plane. - Matches: detector2object.m - :return: + If ``ESW`` is not provided, ``reconstruction.ESW`` is used. The + back-propagated exit surface wave is stored in ``reconstruction.esw``, + together with the corresponding update field in + ``reconstruction.eswUpdate``. + + Args: + ESW (ndarray, optional): + Detector-plane wavefield. If None, use + ``reconstruction.ESW``. """ if ESW is None: ESW = self.reconstruction.ESW @@ -742,9 +1033,43 @@ def ifft2s(self): ) def getBeamWidth(self): - """ - Calculate probe beam width (Full width half maximum) - :return: + r""" + Estimate the probe beam width from the second moment of its intensity. + + The probe intensity is summed over the non-spatial dimensions and + normalized as + + $$ + \tilde{P}(x,y) = \frac{P(x,y)}{\sum_{x,y} P(x,y)} + $$ + + The intensity-weighted centroid and variance are then calculated, for + example along $x$ as + + $$ + \langle x \rangle = \sum_{x,y} x\tilde{P}(x,y) + $$ + + $$ + \sigma_x^2 = \sum_{x,y}(x-\langle x\rangle)^2\tilde{P}(x,y) + $$ + + and converted to a Gaussian-equivalent full width at half maximum: + + $$ + \mathrm{FWHM}_x = 2\sqrt{2\ln 2}\sigma_x + $$ + + The same calculation is applied along $y$. + + Returns: + tuple: + ``(beamWidthY, beamWidthX)`` in meters. + + Notes: + For non-Gaussian probe + profiles, the returned values should be interpreted as second-moment + Gaussian-equivalent beam widths rather than direct half-maximum widths. """ xp = getArrayModule(self.reconstruction.probe) P = xp.sum( @@ -768,8 +1093,64 @@ def getBeamWidth(self): return self.reconstruction.beamWidthY, self.reconstruction.beamWidthX def getOverlap(self, ind1, ind2): - """ - Calculate linear and area overlap between two scan positions indexed ind1 and ind2 + r""" + Estimate the probe overlap between two scan positions. + + The physical displacement between the two positions is calculated from + their pixel-coordinate difference and the probe-plane pixel size: + + $$ + s_x = |x_2-x_1|dx_p,\qquad s_y = |y_2-y_1|dx_p + $$ + + The linear overlap is estimated from the radial scan displacement and + the smaller of the reconstructed probe widths: + + $$ + O_{\mathrm{linear}} = \max\left(1-\frac{\sqrt{s_x^2+s_y^2}}{\min(w_x,w_y)},0\right) + $$ + + where $w_x$ and $w_y$ are the Gaussian-equivalent probe widths returned + by `getBeamWidth()`. + + The area overlap is calculated from the normalized autocorrelation of the + probe amplitude. Let + + $$ + P(x,y) = |\mathrm{probe}(x,y)| + $$ + + and + + $$ + Q(f_x,f_y) = \mathcal{F}\{P(x,y)\}. + $$ + + Using the Fourier correlation theorem, the normalized area overlap is + + $$ + O_{\mathrm{area}} = \frac{1}{N_\lambda}\sum_\lambda\frac{\left|\sum_{f_x,f_y}|Q_\lambda(f_x,f_y)|^2\exp[-i2\pi(f_xs_x+f_ys_y)]\right|}{\sum_{f_x,f_y}|Q_\lambda(f_x,f_y)|^2}. + $$ + + This definition is invariant to the absolute centering of the probe and + yields an overlap of one for zero displacement. + + The current implementation uses the first probe mode and the last slice + when evaluating the area overlap. + + Args: + ind1 (int): + Index of the first scan position. + ind2 (int): + Index of the second scan position. + + Returns: + tuple: + ``(linearOverlap, areaOverlap)``. + + Notes: + The calculated values are also stored in + ``reconstruction.linearOverlap`` and ``reconstruction.areaOverlap``. """ sy = ( abs( @@ -808,18 +1189,56 @@ def getOverlap(self, ind1, ind2): self.reconstruction.areaOverlap = np.mean( abs( np.sum( - Q**2 * np.exp(-1.0j * 2 * np.pi * (Fx * sx + Fy * sy)), + abs(Q)**2 * np.exp(-1.0j * 2 * np.pi * (Fx * sx + Fy * sy)), axis=(-1, -2), ) ) / np.sum(abs(Q) ** 2, axis=(-1, -2)), axis=0, ) + return ( + self.reconstruction.linearOverlap, + self.reconstruction.areaOverlap, + ) def getErrorMetrics(self): - """ - matches getErrorMetrics.m - :return: + r""" + Compute the normalized reconstruction error for the current iteration. + + For each scan position $j$, the detector-domain error is summed over all + detector pixels: + + $$ + e_j = \sum_{x,y} E_j(x,y) + $$ + + where $E_j(x,y)$ is the absolute difference between measured and estimated + detector intensities. If `FourierMaskSwitch` is enabled, the detector error + is weighted by the Fourier mask $W$: + + $$ + e_j = \sum_{x,y} E_j(x,y)W(x,y) + $$ + + The error at each scan position is normalized by the measured diffraction + energy: + + $$ + \tilde{e}_j = \frac{e_j}{E_j^{\mathrm{meas}} + 10^{-20}} + $$ + + The total error for the current iteration is then + + $$ + e_{\mathrm{iter}} = \sum_j \tilde{e}_j + $$ + + and is appended to `reconstruction.error`. + + Notes: + If `saveMemory` is enabled, the per-position errors are accumulated + during the reconstruction loop instead of storing the full detector + error array. """ if not self.params.saveMemory: # Calculate mean error for all positions (make separate function for all of that) @@ -841,10 +1260,44 @@ def getErrorMetrics(self): self.reconstruction.error = np.append(self.reconstruction.error, eAverage) def getRMSD(self, positionIndex): - """ - Root mean square deviation between ptychogram and intensity estimate - :param positionIndex: - :return: + r""" + Compute the detector-domain intensity error for one scan position. + + The pixel-wise detector error is defined as + + $$ + E(x,y) = \left|I_{\mathrm{measured}}(x,y) - I_{\mathrm{estimated}}(x,y)\right|. + $$ + + If `saveMemory` is disabled, the full detector-error map is stored in + `reconstruction.detectorError`. + + If `saveMemory` is enabled, the detector error is reduced immediately to + a per-position scalar: + + $$ + e_j = \sum_{x,y} E_j(x,y). + $$ + + When `FourierMaskSwitch` is enabled, the masked error is used instead: + + $$ + e_j = \sum_{x,y} E_j(x,y)W(x,y). + $$ + + Args: + positionIndex (int): + Index of the current scan position. + + Raises: + NotImplementedError: + If `saveMemory`, `FourierMaskSwitch`, and `CPSCswitch` are all + enabled simultaneously. + + Notes: + Despite its name, this method does not currently compute a root mean + square deviation. It computes an absolute detector-intensity + difference. """ # find out wether or not to use the GPU xp = getArrayModule(self.reconstruction.Iestimated) @@ -868,8 +1321,48 @@ def getRMSD(self, positionIndex): self.reconstruction.detectorError[positionIndex] = self.currentDetectorError def intensityProjection(self, positionIndex): - """Compute the projected intensity. - Barebones, need to implement other methods + r""" + Apply the detector-plane intensity constraint for one scan position. + + The current exit surface wave is first propagated to the detector plane. + The estimated intensity is calculated from the propagated field as + + $$ + I_{\mathrm{estimated}}(x,y)=\sum_m |\Psi_m(x,y)|^2. + $$ + + For the standard intensity constraint, the detector-plane field is scaled + by + + $$ + f(x,y)=\sqrt{\frac{I_{\mathrm{measured}}(x,y)}{I_{\mathrm{estimated}}(x,y)+\epsilon}}, + $$ + + so that + + $$ + \Psi_{\mathrm{updated}}(x,y)=\Psi(x,y)f(x,y). + $$ + + Alternative projection rules are selected through + ``params.intensityConstraint``. The current implementation supports + ``"standard"``, ``"fluctuation"``, ``"exponential"``, + ``"poisson"``, and ``"interferometric"``. + + Depending on the active parameters, the method may additionally apply + constrained-pixel-sum decompression, adaptive denoising, Fourier masking, + background estimation, or an interferometric reference update. + + The constrained detector-plane field is finally propagated back to the + object plane. + + Args: + positionIndex (int): + Index of the current diffraction frame. + + Raises: + ValueError: + If ``intensityConstraint`` is not a supported value. """ # figure out whether or not to use the GPU xp = getArrayModule(self.reconstruction.esw) @@ -946,7 +1439,7 @@ def intensityProjection(self, positionIndex): ) frac = W * frac + (1 - W) - elif self.params.intensityConstraint == "poission": + elif self.params.intensityConstraint == "poisson": frac = self.reconstruction.Imeasured / ( self.reconstruction.Iestimated + gimmel ) @@ -1012,10 +1505,47 @@ def intensityProjection(self, positionIndex): self.detector2object() def decompressionProjection(self, positionIndex): - """ - calculate the upsampled Imeasured from downsampled Imeasured that is actually measured. - :param positionIndex: index for scan positions - :return: + r""" + Construct a high-resolution measured intensity for the + constrained-pixel-sum projection. + + The current high-resolution estimated intensity is divided into blocks + of size $s \times s$, where $s$ is `CPSCupsamplingFactor`. + + For each detector block, the predicted low-resolution intensity is + + $$ + \hat{I}^{\mathrm{LR}}_{mn} = \sum_{i,j \in \mathrm{block}_{mn}} I_{\mathrm{estimated}}(i,j) + $$ + + A block-wise correction factor is calculated from the measured + low-resolution diffraction intensity: + + $$ + S_{mn} = \frac{I^{\mathrm{LR}}_{\mathrm{measured},mn}}{\hat{I}^{\mathrm{LR}}_{mn} + \epsilon} + $$ + + The factor is expanded over the corresponding high-resolution block + and used to construct an effective high-resolution measured intensity: + + $$ + I^{\mathrm{HR}}_{\mathrm{measured}}(i,j) = I_{\mathrm{estimated}}(i,j) S_{mn} + $$ + + This preserves the current estimate of the sub-pixel intensity + distribution while enforcing that the summed intensity in each + high-resolution block matches the experimentally measured + low-resolution pixel intensity. + + Args: + positionIndex (int): + Index of the current diffraction frame. + + Notes: + The generated high-resolution intensity is stored in + `reconstruction.Imeasured`. When `FourierMaskSwitch` is enabled, + the CPSC correction is restricted by `experimentalData.W` after + the first five reconstruction iterations. """ # overwrite the measured intensity (just to have same dimensions as Iestimated) xp = getArrayModule(self.reconstruction.Iestimated) @@ -1044,9 +1574,23 @@ def decompressionProjection(self, positionIndex): def showReconstruction(self, loop): """ - Show the reconstruction process. - :param loop: the iteration number - :return: + Update reconstruction monitoring and optional iteration output. + + The object, probe, reconstruction error, propagation distance, mode + purities, scan positions, and beam width are sent to the active monitor + at intervals defined by ``monitor.figureUpdateFrequency``. + + For Fourier ptychography, the object is transformed to real space before + visualization. When the monitor verbosity is set to ``"high"``, the + current measured and estimated diffraction intensities, reconstruction + error, and estimated scan overlap are also displayed. + + If object dumping is enabled, the current reconstructed object is written + to disk for each iteration. + + Args: + loop (int): + Current reconstruction iteration. """ if np.mod(loop, self.monitor.figureUpdateFrequency) == 0: if self.experimentalData.operationMode == "FPM": @@ -1235,17 +1779,46 @@ def showReconstruction(self, loop): file_path = os.path.join(folder_path, filename) with h5py.File(file_path, "w") as hdf: - obj = self.reconstruction.object.get() + obj = asNumpyArray(self.reconstruction.object) hdf.create_dataset("Object", data=obj) def positionCorrection(self, objectPatch, positionIndex, sy, sx): - """ - Modified from pcPIE. Position correction is done by using positionCorrection and positionCorrectionUpdate - :param objectPatch: - :param positionIndex: - :param sy: - :param sx: - :return: + r""" + Estimate the position correction for one scan position. + + The current object patch is compared with the corresponding region of the + reconstructed object using cross-correlation. For small search radii, + shifted object patches are evaluated directly. For larger search radii, + the cross-correlation is calculated in the Fourier domain. + + A correlation-weighted displacement estimate is obtained from the tested + shifts. Conceptually, + + $$ + g_x = \beta \sum_k \frac{C_k-\bar{C}}{\|O_{\mathrm{patch}}\|_2^2}\Delta x_k + $$ + + and similarly for $g_y$, where $C_k$ is the correlation obtained for the + candidate shift $(\Delta y_k,\Delta x_k)$. + + The estimated correction is scaled by the position-correction feedback + factor and accumulated into the position search direction used by + `positionCorrectionUpdate()`. + + Args: + objectPatch (ndarray): + Current reconstructed object patch for the scan position. + positionIndex (int): + Index of the current scan position. + sy (slice): + Row slice locating the current object patch in the full object. + sx (slice): + Column slice locating the current object patch in the full object. + + Returns: + ndarray: + Estimated position correction ``[delta_y, delta_x]`` in pixels. + Returns zeros before position correction becomes active. """ xp = getArrayModule(objectPatch) @@ -1324,8 +1897,39 @@ def positionCorrection(self, objectPatch, positionIndex, sy, sx): return np.zeros(2) def position_update_to_change_in_z(self, loop): - """ - Update the z based on the position updates. + r""" + Map the global scaling of corrected scan positions to an update of the + propagation distance. + + The corrected and original encoder positions are centered and their + relative spatial scale is estimated as + + $$ + s = \frac{\sigma_{\mathrm{corrected}}}{\sigma_{\mathrm{original}}} + $$ + + A corresponding target propagation distance is estimated as + + $$ + 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, + 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. + + Args: + loop (int): + Current reconstruction iteration. + + 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 @@ -1398,6 +2002,42 @@ def position_update_to_change_in_z(self, loop): ) def positionCorrectionUpdate(self): + r""" + Apply the accumulated position corrections to the scan coordinates. + + Position corrections estimated by `positionCorrection()` are stored in + `D` in pixel units and are applied after the position-correction warm-up. + + For conventional ptychography, the corrected encoder coordinates are + updated as + + $$ + \mathbf{r}_j^{\mathrm{new}} = \mathbf{r}_j^{\mathrm{old}} - \alpha D_j d_{xo} + $$ + + where $\alpha$ is the adaptive step size and $d_{xo}$ is the object-plane + pixel size. The corrected scan grid is then recentered to the mean of the + original encoder positions. + + For Fourier ptychography, the corrected Fourier-space positions are + converted back to illumination coordinates using the inverse FPM + illumination geometry: + + $$ + \mathbf{r} = \operatorname{sign}(c)\frac{\mathbf{k}z_{\mathrm{LED}}}{\sqrt{c^2-k_x^2-k_y^2}} + $$ + + with + + $$ + c = -\frac{N_p d_{xo}}{\lambda} + $$ + + Notes: + The method updates `reconstruction.encoder_corrected` in place and is + called when position correction is enabled. + """ + # fit the scaling out, to put in the z if len(self.reconstruction.error) > self.startAtIteration: self.logger.info("Updating positions") @@ -1439,9 +2079,26 @@ def positionCorrectionUpdate(self): def applyConstraints(self, loop): """ - Apply constraints. - :param loop: loop number - :return: + Apply enabled reconstruction constraints after an iteration. + + This method acts as the central dispatcher for object, probe, position, + and autofocus constraints selected through `Params`. Depending on the + active switches, it may apply regularization, probe normalization, + modal orthogonalization, probe-boundary constraints, smoothing, + amplitude constraints, spectral coupling, position correction, or + autofocus updates. + + Constraints are applied sequentially in the order defined by this method, + so enabling multiple constraints may cause later operations to act on the + result of earlier ones. + + Args: + loop (int): + Current reconstruction iteration. + + Raises: + NotImplementedError: + If PSD estimation is enabled. """ # dirks additions, untested if self.params.l2reg: @@ -1612,9 +2269,47 @@ def applyConstraints(self, loop): # self.reconstruction.probe_storage.tsvd() def orthogonalization(self): - """ - Perform orthogonalization - :return: + r""" + Orthogonalize mixed-state probe or object modes. + + For multiple probe modes, the reconstructed modes are transformed to an + orthogonal modal basis independently for each wavelength and slice. + Conceptually, the transformed modes are + + $$ + P'_k(x,y) = \sum_j U_{kj}P_j(x,y) + $$ + + where $U$ is the modal transformation returned by `orthogonalizeModes()`. + + The normalized modal eigenvalues are used to calculate the probe purity: + + $$ + \mu = \sqrt{\sum_k \tilde{\lambda}_k^2} + $$ + + where + + $$ + \tilde{\lambda}_k = \frac{\lambda_k}{\sum_j \lambda_j} + $$ + + are the normalized modal eigenvalues. A purity of one + corresponds to a single dominant mode, while lower values indicate a + stronger mixed-state contribution. + + If momentum acceleration is enabled, the same modal transformation is + applied to the corresponding momentum and buffer arrays so that all + reconstruction state variables remain expressed in the same modal basis. + + If only multiple object modes are present, the analogous procedure is + applied to the object modes and `purityObject` is updated. + + Notes: + Probe modes are orthogonalized independently for each wavelength and + slice. The resulting probe purity is stored in + `reconstruction.purityProbe` and appended to + `reconstruction.purityProbeHist`. """ xp = getArrayModule(self.reconstruction.probe) if self.reconstruction.npsm > 1: @@ -1722,9 +2417,42 @@ def orthogonalization(self): pass def comStabilization(self): - """ - Perform center of mass stabilization (center the probe) - :return: + r""" + Stabilize the probe center of mass to suppress translational drift. + + Ptychographic reconstruction contains a translational ambiguity that can + allow the reconstructed probe and object to drift within their numerical + arrays without strongly affecting the data consistency. This method + estimates the probe center relative to the center of the reconstruction + window and shifts the reconstruction when the displacement exceeds + approximately one pixel. + + The current probe center is estimated as + + $$ + x_c = \frac{\sum_{x,y}X_p(x,y)A(x,y)}{d_{xp}\sum_{x,y}A(x,y)} + $$ + + and + + $$ + y_c = \frac{\sum_{x,y}Y_p(x,y)A(x,y)}{d_{xp}\sum_{x,y}A(x,y)} + $$ + + where $A(x,y)$ is the probe amplitude and $d_{xp}$ is the probe-plane + pixel size. The resulting coordinates are rounded to integer pixel + shifts. + + If the probe center is displaced by more than approximately one pixel, + both the probe and object are shifted by `(-yc, -xc)` to recenter the + reconstruction while preserving their relative spatial registration. + + When momentum acceleration is enabled, the corresponding probe and object + momentum and buffer arrays are shifted by the same amount. + + Notes: + For multislice reconstruction, the last probe slice is used to + estimate the center. """ self.logger.info("Doing probe com stabilization") xp = getArrayModule(self.reconstruction.probe) @@ -1732,7 +2460,7 @@ def comStabilization(self): P2 = xp.sum( abs(self.reconstruction.probe[:, :, :, -1, ...]) ** 2, axis=(0, 1, 2) ) - P2 = abs(self.reconstruction.probe[0, 0, 0, -1]) + P2 = abs(self.reconstruction.probe[0, 0, 0, -1]) ** 2 demon = xp.sum(P2) * self.reconstruction.dxp xc = int( xp.around(xp.sum(xp.array(self.reconstruction.Xp, xp.float32) * P2) / demon) @@ -1776,17 +2504,57 @@ def comStabilization(self): ) def modulusEnforcedProbe(self): - # propagate probe to detector + r""" + Constrain the reconstructed probe using a measured empty-beam intensity. + + The current probe is propagated to the detector plane, where its + estimated intensity is compared with `experimentalData.emptyBeam`. + + The detector-plane probe is scaled by + + $$ + f(x,y) = \sqrt{\frac{I_{\mathrm{empty}}(x,y)}{I_{\mathrm{probe}}(x,y)+\epsilon}} + $$ + + and updated as + + $$ + \Psi_{\mathrm{updated}}(x,y) = \Psi(x,y)f(x,y) + $$ + + so that the measured detector-plane amplitude is enforced while the + current phase estimate is retained. + + The constrained field is then propagated back to the probe plane and + stored in `reconstruction.probe`. + + If `FourierMaskSwitch` is enabled, the modulus constraint is applied only + inside the active Fourier mask, while the detector-plane field outside + the mask is left unchanged. + + Notes: + This constraint requires `experimentalData.emptyBeam`, containing a + measured detector intensity of the illumination without the sample. + """ xp = getArrayModule(self.reconstruction.esw) self.reconstruction.esw = self.reconstruction.probe self.object2detector() if self.params.FourierMaskSwitch: - self.reconstruction.ESW = self.reconstruction.ESW * xp.sqrt( - self.experimentalData.emptyBeam / 1e-10 - + xp.sum(xp.abs(self.reconstruction.ESW) ** 2, axis=(0, 1, 2, 3)) - ) * self.experimentalData.W + self.reconstruction.ESW * ( - 1 - self.experimentalData.W + self.reconstruction.ESW = ( + self.reconstruction.ESW + * xp.sqrt( + self.experimentalData.emptyBeam + / ( + 1e-10 + + xp.sum( + xp.abs(self.reconstruction.ESW) ** 2, + axis=(0, 1, 2, 3), + ) + ) + ) + * self.experimentalData.W + + self.reconstruction.ESW * (1 - self.experimentalData.W) ) else: self.reconstruction.ESW = self.reconstruction.ESW * np.sqrt( @@ -1803,11 +2571,34 @@ def modulusEnforcedProbe(self): self.reconstruction.probe = self.reconstruction.esw def adaptiveDenoising(self): - """ - Use the difference of mean intensities between the low-resolution - object estimate and the low-resolution raw data to estimate the - noise floor to be clipped. - :return: + r""" + Apply an adaptive amplitude-domain noise-floor correction. + + The measured and estimated detector intensities are converted to + amplitudes. + + A global noise level is estimated from their mean amplitude difference: + + $$ + n = \left|\left\langle A_{\mathrm{meas}} - A_{\mathrm{est}} \right\rangle\right| + $$ + + The estimated noise floor is subtracted from the measured amplitude and + negative values are clipped to zero: + + $$ + A'_{\mathrm{meas}} = \max(A_{\mathrm{meas}} - n, 0) + $$ + + The corrected measured intensity is then reconstructed as + + $$ + I'_{\mathrm{meas}} = \left(A'_{\mathrm{meas}}\right)^2 + $$ + + Notes: + This method modifies `reconstruction.Imeasured` in place and is + applied before the detector-plane intensity constraint. """ # figure out wether or not to use the GPU xp = getArrayModule(self.reconstruction.esw) @@ -1830,12 +2621,37 @@ def z_update(self, stepsize=0.01, roi_bounds=[0.3, 0.7], d=10): self.reconstruction.TV_autofocus() def objectPatchUpdate_TV(self, objectPatch: np.ndarray, DELTA: np.ndarray): - """ - Update the object patch with a TV regularization. + r""" + Update an object patch with an additional total-variation regularization + term. - :param objectPatch: - :param DELTA: - :return: + The data-driven object update is weighted by the conjugated probe: + + $$ + W_P = \frac{P^*}{\max\left(\sum |P|^2\right)} + $$ + + and the object patch is updated as + + $$ + O_{\mathrm{new}} = O + \beta_O \sum W_P\Delta\Psi + \lambda\beta_O G_{\mathrm{TV}}(O) + $$ + + where $\Delta\Psi$ is the exit-wave correction, $\beta_O$ is the object + update step size, $\lambda$ is `objectTVregStepSize`, and + $G_{\mathrm{TV}}(O)$ is the TV update returned by `grad_TV()`. + + Args: + objectPatch (ndarray): + Current object patch. + DELTA (ndarray): + Exit-wave correction used for the object update. + + Returns: + ndarray: + Updated object patch including the TV regularization term. + Notes: + This function is only called by engines supporting TV-regularized object update. """ xp = getArrayModule(objectPatch) diff --git a/docs/api/engines.md b/docs/api/engines.md index 8383fbc..33e3082 100644 --- a/docs/api/engines.md +++ b/docs/api/engines.md @@ -1,5 +1,103 @@ # Engines -::: PtyLab.Engines +::: PtyLab.Engines.BaseEngine.BaseEngine options: - show_submodules: true + show_root_heading: true + show_root_full_path: false + inherited_members: false + members: + - _initialProbePowerCorrection + - _probeWindow + - update_data + - setPositionOrder + - convert2single + - object2detector + - detector2object + - getBeamWidth + - getOverlap + - getErrorMetrics + - getRMSD + - intensityProjection + - decompressionProjection + - showReconstruction + - positionCorrection + - position_update_to_change_in_z + - positionCorrectionUpdate + - applyConstraints + - orthogonalization + - comStabilization + - modulusEnforcedProbe + - adaptiveDenoising + - objectPatchUpdate_TV + +::: PtyLab.Engines.aPIE.aPIE + options: + show_root_heading: true + show_root_full_path: false + inherited_members: false + +::: PtyLab.Engines.e3PIE.e3PIE + options: + show_root_heading: true + show_root_full_path: false + inherited_members: false + +::: PtyLab.Engines.ePIE_TV.ePIE_TV + options: + show_root_heading: true + show_root_full_path: false + inherited_members: false + +::: PtyLab.Engines.ePIE.ePIE + options: + show_root_heading: true + show_root_full_path: false + inherited_members: false + +::: PtyLab.Engines.mPIE_tv.mPIE_tv + options: + show_root_heading: true + show_root_full_path: false + inherited_members: false + +::: PtyLab.Engines.mPIE.mPIE + options: + show_root_heading: true + show_root_full_path: false + inherited_members: false + +::: PtyLab.Engines.mqNewton.mqNewton + options: + show_root_heading: true + show_root_full_path: false + inherited_members: false + +::: PtyLab.Engines.multiPIE.multiPIE + options: + show_root_heading: true + show_root_full_path: false + inherited_members: false + +::: PtyLab.Engines.OPR.OPR + options: + show_root_heading: true + show_root_full_path: false + inherited_members: false + +::: PtyLab.Engines.pcPIE.pcPIE + options: + show_root_heading: true + show_root_full_path: false + inherited_members: false + +::: PtyLab.Engines.qNewton.qNewton + options: + show_root_heading: true + show_root_full_path: false + inherited_members: false + +::: PtyLab.Engines.zPIE.zPIE + options: + show_root_heading: true + show_root_full_path: false + inherited_members: false \ No newline at end of file diff --git a/tests/Engines/test_convert2single.py b/tests/Engines/test_convert2single.py new file mode 100644 index 0000000..fa3bbd7 --- /dev/null +++ b/tests/Engines/test_convert2single.py @@ -0,0 +1,70 @@ +import numpy as np +import pytest +from numpy.testing import assert_allclose, assert_array_equal + +import PtyLab +from PtyLab import Engines + + +@pytest.fixture +def engine(generate_simu_hdf5): + """Initialize a fresh engine with double-precision arrays for each test.""" + experimental_data, reconstruction, _, _, engine = PtyLab.easyInitialize( + "example:simulation_cpm", + engine=Engines.ePIE, + operationMode="CPM", + dummyMonitor=True, + ) + reconstruction.probe = reconstruction.probe.astype(np.complex128) + reconstruction.object = reconstruction.object.astype(np.complex128) + experimental_data.ptychogram = experimental_data.ptychogram.astype(np.float64) + engine.test_complex_array = np.full((4, 4), 1.25 + 2.5j, dtype=np.complex128) + engine.test_real_array = np.full((4, 4), 1.25, dtype=np.float64) + return engine + + +@pytest.mark.parametrize( + "container_name,attribute,expected_dtype", + [ + ("reconstruction", "probe", np.complex64), + ("reconstruction", "object", np.complex64), + ("experimentalData", "ptychogram", np.float32), + (None, "test_complex_array", np.complex64), + (None, "test_real_array", np.float32), + ], +) +def test_convert2single_converts_arrays(engine, container_name, attribute, expected_dtype): + container = engine if container_name is None else getattr(engine, container_name) + original = getattr(container, attribute).copy() + + engine.convert2single() + + converted = getattr(container, attribute) + assert converted.dtype == expected_dtype + assert converted.shape == original.shape + assert_allclose(converted, original, rtol=1e-6, atol=1e-7) + + +def test_convert2single_preserves_other_attributes(engine): + integers = np.arange(10, dtype=np.int64) + booleans = np.array([True, False], dtype=np.bool_) + engine.test_integer_array = integers.copy() + engine.test_boolean_array = booleans.copy() + engine.test_string = "PtyLab" + engine.test_boolean = True + + engine.convert2single() + + assert engine.test_integer_array.dtype == np.int64 + assert_array_equal(engine.test_integer_array, integers) + assert engine.test_boolean_array.dtype == np.bool_ + assert_array_equal(engine.test_boolean_array, booleans) + assert engine.test_string == "PtyLab" + assert engine.test_boolean is True + + +def test_convert2single_sets_target_dtypes(engine): + engine.convert2single() + + assert engine.dtype_complex == np.complex64 + assert engine.dtype_real == np.float32 From eea4f048de8396623a943ef7a071457fc553efb2 Mon Sep 17 00:00:00 2001 From: Chang Liu Date: Tue, 29 Sep 2026 16:51:56 +0200 Subject: [PATCH 07/14] Minor: add 'footnote' extension for reference render in mkdocs. --- mkdocs.yml | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/mkdocs.yml b/mkdocs.yml index 4a9deef..fa8a5ee 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -92,7 +92,9 @@ markdown_extensions: - tables - toc: permalink: true - + - jetblack_markdown.latex2mathml + - footnotes + plugins: - search - mkdocstrings: From baf91135b3b456579b5ccc38b8cd0c6401a0f93f Mon Sep 17 00:00:00 2001 From: Chang Liu Date: Tue, 29 Sep 2026 17:38:56 +0200 Subject: [PATCH 08/14] API documentation: ePIE.py - finished. --- PtyLab/Engines/BaseEngine.py | 5 + PtyLab/Engines/ePIE.py | 217 +++++++++++++++++++++++++++++++---- docs/api/engines.md | 24 ++-- 3 files changed, 214 insertions(+), 32 deletions(-) diff --git a/PtyLab/Engines/BaseEngine.py b/PtyLab/Engines/BaseEngine.py index b84b2a2..8fc012e 100644 --- a/PtyLab/Engines/BaseEngine.py +++ b/PtyLab/Engines/BaseEngine.py @@ -102,6 +102,11 @@ class BaseEngine(object): their algorithm-specific reconstruction and update steps. ``BaseEngine`` is therefore generally not intended to be instantiated directly. + Many methods in `BaseEngine` implement functionality controlled by user-facing + `Params` switches. The `Params` API documentation describes when and why these + options are used, while the corresponding `BaseEngine` methods document their + implementation details. + Args: reconstruction (Reconstruction): Mutable reconstruction state containing the current object, probe, diff --git a/PtyLab/Engines/ePIE.py b/PtyLab/Engines/ePIE.py index 922ca00..50e947e 100644 --- a/PtyLab/Engines/ePIE.py +++ b/PtyLab/Engines/ePIE.py @@ -17,6 +17,55 @@ class ePIE(BaseEngine): + r""" + Extended Ptychographical Iterative Engine (ePIE). + + ePIE jointly reconstructs the complex object and illumination probe by + iterating over overlapping scan positions.[^maiden2009] For each position $j$, the + exit surface wave is formed as + + $$ + \Psi_j = O_j P + $$ + + where $O_j$ is the object patch illuminated by the probe $P$. + + After propagation to the detector plane and application of the measured + intensity constraint, the corrected exit wave $\Psi'_j$ is propagated + back to the object plane. The resulting exit-wave difference is + + $$ + \Delta\Psi_j = \Psi'_j - \Psi_j + $$ + + In the classical single-mode ePIE formulation, the object and probe are + updated according to + + $$ + O'_j = O_j + \beta_O \frac{P^*}{\max |P|^2}\Delta\Psi_j + $$ + + and + + $$ + P' = P + \beta_P \frac{O_j^*}{\max |O_j|^2}\Delta\Psi_j + $$ + + where $\beta_O$ and $\beta_P$ control the object and probe update step + sizes. + + The PtyLab implementation generalizes these updates to its multidimensional + reconstruction representation by summing the corresponding contributions + over the relevant wavelength, mode, and slice dimensions. + + The default ePIE settings are `betaObject = 0.25`, + `betaProbe = 0.25`, and `numIterations = 50`. + + [^maiden2009]: A. M. Maiden and J. M. Rodenburg, + "An improved ptychographical phase retrieval algorithm for diffractive + imaging," Ultramicroscopy 109, 1256-1262 (2009). + https://doi.org/10.1016/j.ultramic.2009.05.012 + """ def __init__( self, reconstruction: Reconstruction, @@ -24,8 +73,21 @@ def __init__( params: Params, monitor: Monitor, ): - # This contains reconstruction parameters that are specific to the reconstruction - # but not necessarily to ePIE reconstruction + """ + Initialize the ePIE reconstruction engine. + + Args: + reconstruction (Reconstruction): + Reconstruction state containing the current object, probe, and + geometry. + experimentalData (ExperimentalData): + Experimental diffraction data and acquisition parameters. + params (Params): + Shared reconstruction parameters and constraint settings. + monitor (Monitor): + Monitor used for reconstruction visualization and progress + reporting. + """ super().__init__(reconstruction, experimentalData, params, monitor) self.logger = logging.getLogger("ePIE") self.logger.info("Sucesfully created ePIE ePIE_engine") @@ -34,27 +96,86 @@ def __init__( def initializeReconstructionParams(self): """ - Set parameters that are specific to the ePIE settings. - :return: + Initialize ePIE-specific reconstruction parameters. + + Defaults: + betaObject (float): + Object update step size. Default is 0.25. + betaProbe (float): + Probe update step size. Default is 0.25. + numIterations (int): + Number of reconstruction iterations. Default is 50. """ self.betaProbe = 0.25 self.betaObject = 0.25 self.numIterations = 50 def reconstruct(self, experimentalData: ExperimentalData = None): - """Run the reconstruction to completion. + """ + Run the ePIE reconstruction to completion. - Use :meth:`reconstruct_stepwise` instead if you want to interleave your - own work between scan positions. + This method consumes the generator returned by `reconstruct_stepwise()` + until all reconstruction iterations and scan positions have been + processed. + + Use `reconstruct_stepwise()` when custom operations need to be inserted + between individual scan-position updates. + + Args: + experimentalData (ExperimentalData, optional): + Experimental dataset to use for the reconstruction. If provided, + it replaces the currently attached experimental data. """ for _ in self.reconstruct_stepwise(experimentalData): pass def reconstruct_stepwise(self, experimentalData: ExperimentalData = None): - """Generator variant of :meth:`reconstruct`. + r""" + Run the ePIE reconstruction one scan-position update at a time. + + For each reconstruction iteration, the scan positions are visited in the + order selected by `params.positionOrder`. At each position $j$, the + corresponding object patch is extracted and combined with the current + probe to form the exit surface wave: + + $$ + \Psi_j = O_j P + $$ + + The exit wave is propagated to the detector plane, constrained by the + measured diffraction intensity through `intensityProjection()`, and + propagated back to obtain an updated exit wave $\Psi'_j$. + + The exit-wave correction is - Yields ``(iteration, positionLoop)`` after every scan position. Nothing - happens until the generator is consumed. + $$ + \Delta\Psi_j = \Psi'_j - \Psi_j + $$ + + and is subsequently used by `objectPatchUpdate()` and `probeUpdate()` to + update the object and probe. + + If `params.OPRP` is enabled, position-dependent probe estimates are + retrieved from `reconstruction.probe_storage` before each scan-position + update and stored again after the probe update. Without OPRP, the same + probe estimate is shared and updated sequentially across all scan + positions. + + After all scan positions in an iteration have been processed, + `getErrorMetrics()` evaluates the reconstruction error and + `applyConstraints()` applies the enabled reconstruction constraints. + + The method yields after every scan-position update, allowing custom code + to be interleaved with the reconstruction. + + Args: + experimentalData (ExperimentalData, optional): + Experimental dataset to use for the reconstruction. If provided, + it replaces the currently attached experimental data. + + Yields: + tuple: + ``(iteration, positionLoop)`` after each scan-position update. """ if experimentalData is not None: self.reconstruction.data = experimentalData @@ -126,11 +247,39 @@ def reconstruct_stepwise(self, experimentalData: ExperimentalData = None): self.params.gpuFlag = 0 def objectPatchUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray): - """ - Todo add docstring - :param objectPatch: - :param DELTA: - :return: + r""" + Update the object patch using the ePIE object-update rule. + + For the classical single-mode case, the probe weighting is + + $$ + W_P(x,y) = \frac{P^*(x,y)}{\max_{x,y}|P(x,y)|^2} + $$ + + and the object patch is updated according to + + $$ + O'_j = O_j + \beta_O W_P\Delta\Psi_j + $$ + + where $O_j$ is the current object patch, $\Delta\Psi_j$ is the + exit-wave correction obtained from the detector-plane intensity + constraint, and $\beta_O$ is `betaObject`. + + In the multidimensional PtyLab representation, contributions from the + relevant wavelength, probe-mode, and slice dimensions are summed before + updating the object patch. + + Args: + objectPatch (ndarray): + Current object patch at the active scan position. + DELTA (ndarray): + Exit-wave correction + `reconstruction.eswUpdate - reconstruction.esw`. + + Returns: + ndarray: + Updated object patch. """ # find out which array module to use, numpy or cupy (or other...) xp = getArrayModule(objectPatch) @@ -143,11 +292,39 @@ def objectPatchUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray): ) def probeUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray): - """ - Todo add docstring - :param objectPatch: - :param DELTA: - :return: + r""" + Update the probe using the ePIE probe-update rule. + + For the classical single-mode case, the object weighting is + + $$ + W_O(x,y) = \frac{O_j^*(x,y)}{\max_{x,y}|O_j(x,y)|^2} + $$ + + and the probe is updated according to + + $$ + P' = P + \beta_P W_O\Delta\Psi_j + $$ + + where $O_j$ is the current object patch, $\Delta\Psi_j$ is the + exit-wave correction obtained from the detector-plane intensity + constraint, and $\beta_P$ is `betaProbe`. + + In the multidimensional PtyLab representation, the implemented update + sums the corresponding correction over axes `(0, 1, 3)` while preserving + the probe-mode dimension. + + Args: + objectPatch (ndarray): + Current object patch at the active scan position. + DELTA (ndarray): + Exit-wave correction + `reconstruction.eswUpdate - reconstruction.esw`. + + Returns: + ndarray: + Updated probe. """ # find out which array module to use, numpy or cupy (or other...) xp = getArrayModule(objectPatch) diff --git a/docs/api/engines.md b/docs/api/engines.md index 33e3082..e71fa7c 100644 --- a/docs/api/engines.md +++ b/docs/api/engines.md @@ -30,61 +30,61 @@ - adaptiveDenoising - objectPatchUpdate_TV -::: PtyLab.Engines.aPIE.aPIE +::: PtyLab.Engines.ePIE.ePIE options: show_root_heading: true show_root_full_path: false inherited_members: false -::: PtyLab.Engines.e3PIE.e3PIE +::: PtyLab.Engines.ePIE_TV.ePIE_TV options: show_root_heading: true show_root_full_path: false inherited_members: false -::: PtyLab.Engines.ePIE_TV.ePIE_TV +::: PtyLab.Engines.mPIE.mPIE options: show_root_heading: true show_root_full_path: false inherited_members: false -::: PtyLab.Engines.ePIE.ePIE +::: PtyLab.Engines.mPIE_tv.mPIE_tv options: show_root_heading: true show_root_full_path: false inherited_members: false -::: PtyLab.Engines.mPIE_tv.mPIE_tv +::: PtyLab.Engines.pcPIE.pcPIE options: show_root_heading: true show_root_full_path: false inherited_members: false -::: PtyLab.Engines.mPIE.mPIE +::: PtyLab.Engines.zPIE.zPIE options: show_root_heading: true show_root_full_path: false inherited_members: false -::: PtyLab.Engines.mqNewton.mqNewton +::: PtyLab.Engines.aPIE.aPIE options: show_root_heading: true show_root_full_path: false inherited_members: false -::: PtyLab.Engines.multiPIE.multiPIE +::: PtyLab.Engines.e3PIE.e3PIE options: show_root_heading: true show_root_full_path: false inherited_members: false -::: PtyLab.Engines.OPR.OPR +::: PtyLab.Engines.multiPIE.multiPIE options: show_root_heading: true show_root_full_path: false inherited_members: false -::: PtyLab.Engines.pcPIE.pcPIE +::: PtyLab.Engines.OPR.OPR options: show_root_heading: true show_root_full_path: false @@ -95,8 +95,8 @@ show_root_heading: true show_root_full_path: false inherited_members: false - -::: PtyLab.Engines.zPIE.zPIE + +::: PtyLab.Engines.mqNewton.mqNewton options: show_root_heading: true show_root_full_path: false From cde8067938026e49e2a88d106825148b5b7d35e8 Mon Sep 17 00:00:00 2001 From: Chang Liu Date: Wed, 30 Sep 2026 14:33:46 +0200 Subject: [PATCH 09/14] API documentation: mPIE.py --- PtyLab/Engines/mPIE.py | 123 ++++++++++++++++++++++++++++++++++++++--- docs/api/engines.md | 12 ---- 2 files changed, 114 insertions(+), 21 deletions(-) diff --git a/PtyLab/Engines/mPIE.py b/PtyLab/Engines/mPIE.py index 1238d90..4704a89 100644 --- a/PtyLab/Engines/mPIE.py +++ b/PtyLab/Engines/mPIE.py @@ -26,6 +26,77 @@ class mPIE(BaseEngine): + r""" + Momentum-accelerated ptychographic iterative engine (mPIE). + + mPIE extends the standard ePIE reconstruction by combining + regularized PIE (rPIE) object/probe updates with momentum + acceleration.[^maiden2017] + + As in ePIE, the exit-wave correction at scan position $j$ is + + $$ + \Delta\Psi_j = \Psi'_j - \Psi_j + $$ + + For conventional ptychography, the object update is regularized as + + $$ + O'_j = O_j + \beta_O \frac{P^*}{\alpha_O P_{\max} + (1-\alpha_O)|P|^2}\Delta\Psi_j + $$ + + where $P_{\max}=\max |P|^2$, $\beta_O$ controls the object update + step size, and $\alpha_O$ controls the spatial regularization. + + The probe is updated analogously: + + $$ + P' = P + \beta_P \frac{O_j^*}{\alpha_P O_{\max} + (1-\alpha_P)|O_j|^2}\Delta\Psi_j + $$ + + where $O_{\max}=\max |O_j|^2$, and $\alpha_P$ and $\beta_P$ control + probe regularization and update strength. + + Compared with ePIE, these denominators retain stronger updates in + moderately illuminated regions while suppressing unstable updates where + the corresponding probe or object intensity is small. + + Momentum acceleration is periodically applied to both object and probe. + In the PtyLab implementation, the momentum state is updated as + + $$ + M^{(n)} = G^{(n)} + \eta M^{(n-1)} + $$ + + followed by + + $$ + X^{(n+1)} = X^{(n)} - \gamma M^{(n)} + $$ + + where $X$ denotes the object or probe, $\eta$ is `frictionM`, and + $\gamma$ is `feedbackM`. The current implementation applies these + momentum updates stochastically during the scan-position loop. + + The default mPIE parameters are `betaObject = 0.25`, + `betaProbe = 0.25`, `alphaObject = 0.1`, `alphaProbe = 0.1`, + `feedbackM = 0.3`, and `frictionM = 0.7`. + + [^maiden2017]: A. M. Maiden, D. Johnson, and P. Li, + "Further improvements to the ptychographical iterative engine," + Optica 4, 736-745 (2017). + https://doi.org/10.1364/OPTICA.4.000736 + + Attributes: + keepPatches (bool): + If enabled, store the reconstructed object patch associated with each + scan position for debugging or detailed analysis. This option can + require substantial additional memory. + See Also: + `ePIE` + Baseline ePIE reconstruction without rPIE regularization or + momentum acceleration. + """ def __init__( self, reconstruction: Reconstruction, @@ -33,11 +104,32 @@ def __init__( params: Params, monitor: Monitor, ): - # This contains reconstruction parameters that are specific to the reconstruction - # but not necessarily to ePIE reconstruction + """ + Initialize the mPIE reconstruction engine. + + Shared reconstruction state is initialized through `BaseEngine`, followed + by the mPIE-specific reconstruction parameters and momentum buffers. + + Momentum acceleration is enabled through + `params.momentumAcceleration`, allowing shared BaseEngine operations such + as modal orthogonalization to keep the corresponding momentum and buffer + arrays consistent with the reconstructed object and probe. + + Args: + reconstruction (Reconstruction): + Reconstruction state containing the current object, probe, and + geometry. + experimentalData (ExperimentalData): + Experimental diffraction data and acquisition parameters. + params (Params): + Shared reconstruction parameters and constraint settings. + monitor (Monitor): + Monitor used for reconstruction visualization and progress + reporting. + """ super().__init__(reconstruction, experimentalData, params, monitor) self.logger = logging.getLogger("mPIE") - self.logger.info("Sucesfully created mPIE mPIE_engine") + self.logger.info("Successfully created mPIE engine") self.logger.info("Wavelength attribute: %s", self.reconstruction.wavelength) # initialize mPIE Params self.initializeReconstructionParams() @@ -46,10 +138,12 @@ def __init__( @property def keepPatches(self): - """Wether or not to keep track of the individual object update patches. + """ + Whether to store the reconstructed object patch for every scan position. - This strongly increases the amount of memory required, only use when absolutely required. + Enable with `engine.keepPatches = True` and disable with `engine.keepPatches = False`. + This option is intended for debugging or detailed analysis and may require a large amount of additional memory. """ return hasattr(self, "patches") @@ -72,8 +166,20 @@ def keepPatches(self, keep_them): def initializeReconstructionParams(self): """ - Set parameters that are specific to the mPIE settings. - :return: + Initialize mPIE-specific reconstruction parameters and momentum state. + + The default mPIE parameters are: + + - `betaObject = 0.25`: object update step size. + - `betaProbe = 0.25`: probe update step size. + - `alphaObject = 0.1`: object-update regularization parameter. + - `alphaProbe = 0.1`: probe-update regularization parameter. + - `feedbackM = 0.3`: momentum feedback strength. + - `frictionM = 0.7`: momentum memory coefficient. + - `numIterations = 50`: number of reconstruction iterations. + + Object and probe momentum arrays are initialized together with corresponding + buffers that store the reconstruction state used by the momentum updates. """ # self.eswUpdate = self.reconstruction.esw.copy() self.betaProbe = 0.25 @@ -144,12 +250,11 @@ def reconstruct(self, experimentalData=None, reconstruction=None, vis_after_each else: object_patch = self.objectPatchUpdate(objectPatch, DELTA) + self.reconstruction.object[..., sy, sx] = object_patch if self.keepPatches: self.patches[positionIndex, ..., sy, sx] = asNumpyArray( abs(object_patch) ** 2 ) - else: - self.reconstruction.object[..., sy, sx] = object_patch # probe update weight = 1 diff --git a/docs/api/engines.md b/docs/api/engines.md index e71fa7c..659958a 100644 --- a/docs/api/engines.md +++ b/docs/api/engines.md @@ -36,24 +36,12 @@ show_root_full_path: false inherited_members: false -::: PtyLab.Engines.ePIE_TV.ePIE_TV - options: - show_root_heading: true - show_root_full_path: false - inherited_members: false - ::: PtyLab.Engines.mPIE.mPIE options: show_root_heading: true show_root_full_path: false inherited_members: false -::: PtyLab.Engines.mPIE_tv.mPIE_tv - options: - show_root_heading: true - show_root_full_path: false - inherited_members: false - ::: PtyLab.Engines.pcPIE.pcPIE options: show_root_heading: true From b0251df5ca6312517e6860d9800702f05aa55c5b Mon Sep 17 00:00:00 2001 From: Chang Liu Date: Wed, 30 Sep 2026 16:00:22 +0200 Subject: [PATCH 10/14] API documetation - ePIE and mPIE finished. --- PtyLab/Engines/BaseEngine.py | 42 ++++--- PtyLab/Engines/ePIE.py | 19 ++- PtyLab/Engines/mPIE.py | 220 ++++++++++++++++++++++++++++++++--- mkdocs.yml | 6 + pyproject.toml | 6 + scripts/docs_notebooks.py | 27 +++++ 6 files changed, 282 insertions(+), 38 deletions(-) create mode 100644 scripts/docs_notebooks.py diff --git a/PtyLab/Engines/BaseEngine.py b/PtyLab/Engines/BaseEngine.py index 3e1525e..60335a6 100644 --- a/PtyLab/Engines/BaseEngine.py +++ b/PtyLab/Engines/BaseEngine.py @@ -2627,39 +2627,51 @@ def z_update(self, stepsize=0.01, roi_bounds=[0.3, 0.7], d=10): def objectPatchUpdate_TV(self, objectPatch: np.ndarray, DELTA: np.ndarray): r""" - Update an object patch with an additional total-variation regularization - term. + Apply total-variation regularization to the engine-specific object update. - The data-driven object update is weighted by the conjugated probe: + The data-driven object update is first calculated using the current + engine's `objectPatchUpdate()` implementation: $$ - W_P = \frac{P^*}{\max\left(\sum |P|^2\right)} + O_{\mathrm{data}} = U_{\mathrm{engine}}(O_j,\Delta\Psi_j) $$ - and the object patch is updated as + where $U_{\mathrm{engine}}$ denotes the object-update rule implemented by + the active reconstruction engine. For example, `ePIE` uses the standard + ePIE update, while `mPIE` uses its regularized mPIE/rPIE update. + + A total-variation regularization term is then added: $$ - O_{\mathrm{new}} = O + \beta_O \sum W_P\Delta\Psi + \lambda\beta_O G_{\mathrm{TV}}(O) + O'_j = O_{\mathrm{data}} + \lambda\beta_O G_{\mathrm{TV}}(O_j) $$ - where $\Delta\Psi$ is the exit-wave correction, $\beta_O$ is the object - update step size, $\lambda$ is `objectTVregStepSize`, and - $G_{\mathrm{TV}}(O)$ is the TV update returned by `grad_TV()`. + where $G_{\mathrm{TV}}(O_j)$ is calculated by `grad_TV()` using + `epsilon=1e-2`, $\beta_O$ is the engine's object update step size, and + $\lambda$ is controlled by `params.objectTVregStepSize`. + + This design keeps the TV regularization independent of the underlying + reconstruction engine, allowing different engines to retain their native + object-update rules while sharing the same TV regularizer. + + The activation and application frequency of this update are controlled by + `params.objectTVregSwitch` and `params.objectTVfreq` in the reconstruction + loop. Args: objectPatch (ndarray): - Current object patch. + Current object patch at the active scan position. DELTA (ndarray): - Exit-wave correction used for the object update. + Exit-wave correction, typically + `reconstruction.eswUpdate - reconstruction.esw`. Returns: ndarray: - Updated object patch including the TV regularization term. - Notes: - This function is only called by engines supporting TV-regularized object update. + Engine-specific object update with the additional TV + regularization term. """ - xp = getArrayModule(objectPatch) + #xp = getArrayModule(objectPatch) #frac = self.reconstruction.probe.conj() / xp.max( # xp.sum(xp.abs(self.reconstruction.probe) ** 2, axis=(0, 1, 2, 3)) #) diff --git a/PtyLab/Engines/ePIE.py b/PtyLab/Engines/ePIE.py index 3dd22ee..2a04794 100644 --- a/PtyLab/Engines/ePIE.py +++ b/PtyLab/Engines/ePIE.py @@ -152,14 +152,21 @@ def reconstruct_stepwise(self, experimentalData: ExperimentalData = None): \Delta\Psi_j = \Psi'_j - \Psi_j $$ - and is subsequently used by `objectPatchUpdate()` and `probeUpdate()` to - update the object and probe. + By default, the object is updated using the standard ePIE rule implemented + by `objectPatchUpdate()`. + + If `params.objectTVregSwitch` is enabled, the TV-regularized update + `objectPatchUpdate_TV()` is used every `params.objectTVfreq` iterations. + The standard ePIE object update is retained and an additional TV + regularization term is added with strength controlled by + `params.objectTVregStepSize`. + + The probe is updated using `probeUpdate()` after each object update. If `params.OPRP` is enabled, position-dependent probe estimates are retrieved from `reconstruction.probe_storage` before each scan-position - update and stored again after the probe update. Without OPRP, the same - probe estimate is shared and updated sequentially across all scan - positions. + update and stored again after the probe update. Without OPRP, a shared + probe estimate is updated sequentially across all scan positions. After all scan positions in an iteration have been processed, `getErrorMetrics()` evaluates the reconstruction error and @@ -175,7 +182,7 @@ def reconstruct_stepwise(self, experimentalData: ExperimentalData = None): Yields: tuple: - ``(iteration, positionLoop)`` after each scan-position update. + `(iteration, positionLoop)` after each scan-position update. """ if experimentalData is not None: self.reconstruction.data = experimentalData diff --git a/PtyLab/Engines/mPIE.py b/PtyLab/Engines/mPIE.py index 4704a89..ac5215f 100644 --- a/PtyLab/Engines/mPIE.py +++ b/PtyLab/Engines/mPIE.py @@ -200,8 +200,70 @@ def initializeReconstructionParams(self): self.reconstruction.probeWindow = np.abs(self.reconstruction.probe) def reconstruct(self, experimentalData=None, reconstruction=None, vis_after_each_iteration=None): - """Reconstruct object. If experimentalData is given, it replaces the current data. Idem for reconstruction.""" + r""" + Run the mPIE reconstruction to completion. + The reconstruction follows the standard ptychographic position loop. + At each scan position $j$, the current object patch and probe form the + exit surface wave + + $$ + \Psi_j = O_j P + $$ + + which is propagated to the detector plane and constrained by the measured + diffraction intensity through `intensityProjection()`. The corresponding + exit-wave correction is + + $$ + \Delta\Psi_j = \Psi'_j - \Psi_j + $$ + + The object and probe are then updated using the mPIE/rPIE update rules + implemented by `objectPatchUpdate()` and `probeUpdate()`. + + If `params.objectTVregSwitch` is enabled, the TV-regularized object update + `objectPatchUpdate_TV()` is used every `params.objectTVfreq` iterations. + The native mPIE object update is retained and an additional TV + regularization term is added with strength controlled by + `params.objectTVregStepSize`. + + If `params.weigh_probe_updates_by_intensity` is enabled, the probe update + is scaled by the relative intensity of the current diffraction frame. + + If `params.positionCorrectionSwitch` is enabled, scan-position correction + is applied after the object and probe updates. + + Momentum acceleration is applied stochastically during the scan-position + loop. In the current implementation, each position update has + approximately a 5% probability of triggering `objectMomentumUpdate()` and + `probeMomentumUpdate()`. + + If `keepPatches` is enabled, the reconstructed patch associated with each + scan position is additionally stored in `self.patches` for diagnostics or + further analysis. + + After all scan positions in an iteration have been processed, + `getErrorMetrics()` evaluates the reconstruction error, + `applyConstraints()` applies the enabled reconstruction constraints, and + `showReconstruction()` updates the reconstruction monitor. + + Args: + experimentalData (ExperimentalData, optional): + Experimental dataset to use for the reconstruction. If provided, + it replaces the currently attached experimental data. + reconstruction (Reconstruction, optional): + Reconstruction state to optimize. If provided, it replaces the + currently attached reconstruction object. + vis_after_each_iteration (callable, optional): + Callback executed after each reconstruction iteration as + `vis_after_each_iteration(loop, reconstruction)`. + + Notes: + The object and probe momentum buffers are reset after + `_prepareReconstruction()` so that they remain synchronized with any + initialization changes applied before reconstruction starts. + """ self.changeExperimentalData(experimentalData) self.changeOptimizable(reconstruction) @@ -299,9 +361,38 @@ def reconstruct(self, experimentalData=None, reconstruction=None, vis_after_each # todo clearMemory implementation def objectMomentumUpdate(self): - """ - momentum update object, save updated objectMomentum and objectBuffer. - :return: + r""" + Apply the mPIE momentum update to the reconstructed object. + + The change in the object since the previous momentum update is estimated + from the stored object buffer: + + $$ + G_O^{(n)} = O_{\mathrm{buf}}^{(n)} - O^{(n)} + $$ + + The object momentum is updated according to + + $$ + M_O^{(n)} = G_O^{(n)} + \eta M_O^{(n-1)} + $$ + + where $\eta$ is `frictionM`. + + The accumulated momentum is then fed back into the object estimate: + + $$ + O^{(n+1)} = O^{(n)} - \gamma M_O^{(n)} + $$ + + where $\gamma$ is `feedbackM`. + + After the momentum correction, `objectBuffer` is updated with the current + object estimate for the next momentum step. + + Notes: + This update is triggered stochastically from `reconstruct()` rather + than after every scan-position update. """ gradient = self.reconstruction.objectBuffer - self.reconstruction.object self.reconstruction.objectMomentum = ( @@ -314,9 +405,12 @@ def objectMomentumUpdate(self): self.reconstruction.objectBuffer = self.reconstruction.object.copy() def probeMomentumUpdate(self): - """ - momentum update probe, save updated probeMomentum and probeBuffer. - :return: + r""" + Apply the mPIE momentum update to the reconstructed probe. Similar to `objectMomentumUpdate()`. + + See Also: + `objectMomentumUpdate` + Equivalent momentum update applied to the reconstructed object. """ gradient = self.reconstruction.probeBuffer - self.reconstruction.probe self.reconstruction.probeMomentum = ( @@ -329,11 +423,56 @@ def probeMomentumUpdate(self): self.reconstruction.probeBuffer = self.reconstruction.probe.copy() def objectPatchUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray): - """ - Todo add docstring - :param objectPatch: - :param DELTA: - :return: + r""" + Update the object patch using the regularized mPIE/rPIE object-update rule. + + The probe intensity is first evaluated and its maximum value is used as a + global normalization scale: + + $$ + P_{\max} = \max_{x,y}\sum |P(x,y)|^2 + $$ + + For conventional ptychography, the probe weighting is + + $$ + W_P =\frac{P^*}{\alpha_O P_{\max} + (1-\alpha_O)|P|^2} + $$ + + and the object patch is updated according to + + $$ + O'_j =O_j + \beta_O \sum W_P\Delta\Psi_j + $$ + + where $\Delta\Psi_j$ is the exit-wave correction, $\beta_O$ is + `betaObject`, and $\alpha_O$ is `alphaObject`. + + The parameter `alphaObject` controls the balance between global + normalization by the maximum probe intensity and local normalization by + the spatially varying probe intensity. + + For Fourier ptychography (`operationMode == "FPM"`), an additional + probe-amplitude weighting is applied: + + $$ + W_P^{\mathrm{FPM}} =\frac{|P|}{P_{\max}}\frac{P^*}{\alpha_O P_{\max} + (1-\alpha_O)|P|^2} + $$ + + In the multidimensional PtyLab representation, the object correction is + summed over the probe-mode axis before being added to the current object + patch. + + Args: + objectPatch (ndarray): + Current object patch at the active scan position. + DELTA (ndarray): + Exit-wave correction + `reconstruction.eswUpdate - reconstruction.esw`. + + Returns: + ndarray: + Updated object patch. """ # find out which array module to use, numpy or cupy (or other...) xp = getArrayModule(objectPatch) @@ -356,11 +495,58 @@ def objectPatchUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray): ) def probeUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray, weight: float): - """ - Todo add docstring - :param objectPatch: - :param DELTA: - :return: + r""" + Update the probe using the regularized mPIE/rPIE probe-update rule. + + The object intensity is first evaluated and its maximum value is used as + a global normalization scale: + + $$ + O_{\max} = \max_{x,y}\sum |O_j(x,y)|^2 + $$ + + The object weighting is then calculated as + + $$ + W_O = \frac{O_j^*}{\alpha_P O_{\max} + (1-\alpha_P)|O_j|^2} + $$ + + and the probe is updated according to + + $$ + P' = P + w\beta_P\sum W_O\Delta\Psi_j + $$ + + where $\Delta\Psi_j$ is the exit-wave correction, $\beta_P$ is + `betaProbe`, $\alpha_P$ is `alphaProbe`, and $w$ is the optional + intensity-dependent update weight. + + The parameter `alphaProbe` controls the balance between global + normalization by the maximum object intensity and local normalization by + the spatially varying object intensity. + + By default, $w=1$. If `params.weigh_probe_updates_by_intensity` is + enabled in `reconstruct()`, $w$ is set to the relative intensity of the + current diffraction frame. + + In the current multidimensional PtyLab representation, the probe + correction is summed over axis `1` before being added to the current + probe estimate. + + Args: + objectPatch (ndarray): + Current object patch at the active scan position. + DELTA (ndarray): + Exit-wave correction + `reconstruction.eswUpdate - reconstruction.esw`. + weight (float): + Multiplicative weight applied to the probe update. Typically `1`, + or the relative intensity of the current diffraction frame when + intensity-weighted probe updates are enabled. + + Returns: + ndarray: + Updated probe. """ # find out which array module to use, numpy or cupy (or other...) xp = getArrayModule(objectPatch) diff --git a/mkdocs.yml b/mkdocs.yml index fa8a5ee..9970650 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -5,6 +5,12 @@ repo_url: https://github.com/PtyLab/PtyLab.py repo_name: PtyLab/PtyLab.py dev_addr: '127.0.0.1:8000' +hooks: + - scripts/docs_notebooks.py + +watch: + - jupyter_tutorials + theme: name: material logo: assets/ptylab_logo.svg diff --git a/pyproject.toml b/pyproject.toml index 627d764..e727830 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,6 +34,7 @@ dev = [ gpu = ["cupy-cuda12x[ctk]"] tensorflow = ["tensorflow~=2.14"] docs = [ + "jetblack-markdown", "mkdocs>=1.6.1", "mkdocs-material>=9.7.6", "mkdocs-jupyter>=0.24", @@ -53,3 +54,8 @@ build-backend = "uv_build" [tool.uv.build-backend] module-name = "PtyLab" module-root = "" + +[tool.uv] +required-environments = [ + "sys_platform == 'win32' and platform_machine == 'AMD64'", +] diff --git a/scripts/docs_notebooks.py b/scripts/docs_notebooks.py new file mode 100644 index 0000000..6a120c1 --- /dev/null +++ b/scripts/docs_notebooks.py @@ -0,0 +1,27 @@ +"""Resolve notebook symlink placeholders produced by Git on Windows.""" + +from pathlib import Path + +from mkdocs.exceptions import PluginError +from mkdocs.plugins import event_priority + + +@event_priority(-100) +def on_files(files, config): + # Run after mkdocs-jupyter has created its NotebookFile instances, retaining + # their documentation URLs while pointing conversion at the real notebooks. + notebook_root = (Path(config["docs_dir"]).parent / "jupyter_tutorials").resolve() + for file in files: + if not file.src_uri.endswith(".ipynb") or not file.abs_src_path: + continue + source = Path(file.abs_src_path) + if source.is_symlink() or source.stat().st_size > 1024: + continue + reference = source.read_text(encoding="utf-8").strip() + if not reference.startswith("../../jupyter_tutorials/"): + continue + target = (source.parent / reference).resolve() + if not target.is_relative_to(notebook_root) or not target.is_file(): + raise PluginError(f"Notebook link {source} has an invalid target: {reference}") + file.abs_src_path = str(target) + return files From a990ec647896f1bd1431c688948525a6531339f3 Mon Sep 17 00:00:00 2001 From: Chang Liu Date: Wed, 30 Sep 2026 17:19:27 +0200 Subject: [PATCH 11/14] Remove pcPIE.py but keep the same interface, functions already included in mPIE --- PtyLab/Engines/__init__.py | 14 +- PtyLab/Engines/mPIE.py | 392 ++++++++++++++++++++-- tests/Engines/test_remove_pcPIE_engine.py | 43 +++ 3 files changed, 418 insertions(+), 31 deletions(-) create mode 100644 tests/Engines/test_remove_pcPIE_engine.py diff --git a/PtyLab/Engines/__init__.py b/PtyLab/Engines/__init__.py index 2f87ce9..17ed5a0 100644 --- a/PtyLab/Engines/__init__.py +++ b/PtyLab/Engines/__init__.py @@ -14,16 +14,24 @@ "mqNewton", "multiPIE", "OPR", - "pcPIE", "qNewton", "zPIE", ] from .e3PIE import e3PIE from .ePIE import ePIE -from .mPIE import mPIE +from .mPIE import mPIE, pcPIE from .mqNewton import mqNewton from .multiPIE import multiPIE from .OPR import OPR -from .pcPIE import pcPIE from .qNewton import qNewton from .zPIE import zPIE + + +import warnings + +warnings.warn( + "`pcPIE` is deprecated. Use `mPIE` with " + "`params.positionCorrectionSwitch = True` instead.", + DeprecationWarning, + stacklevel=2, +) \ No newline at end of file diff --git a/PtyLab/Engines/mPIE.py b/PtyLab/Engines/mPIE.py index 1238d90..a2a93c0 100644 --- a/PtyLab/Engines/mPIE.py +++ b/PtyLab/Engines/mPIE.py @@ -26,6 +26,77 @@ class mPIE(BaseEngine): + r""" + Momentum-accelerated ptychographic iterative engine (mPIE). + + mPIE extends the standard ePIE reconstruction by combining + regularized PIE (rPIE) object/probe updates with momentum + acceleration.[^maiden2017] + + As in ePIE, the exit-wave correction at scan position $j$ is + + $$ + \Delta\Psi_j = \Psi'_j - \Psi_j + $$ + + For conventional ptychography, the object update is regularized as + + $$ + O'_j = O_j + \beta_O \frac{P^*}{\alpha_O P_{\max} + (1-\alpha_O)|P|^2}\Delta\Psi_j + $$ + + where $P_{\max}=\max |P|^2$, $\beta_O$ controls the object update + step size, and $\alpha_O$ controls the spatial regularization. + + The probe is updated analogously: + + $$ + P' = P + \beta_P \frac{O_j^*}{\alpha_P O_{\max} + (1-\alpha_P)|O_j|^2}\Delta\Psi_j + $$ + + where $O_{\max}=\max |O_j|^2$, and $\alpha_P$ and $\beta_P$ control + probe regularization and update strength. + + Compared with ePIE, these denominators retain stronger updates in + moderately illuminated regions while suppressing unstable updates where + the corresponding probe or object intensity is small. + + Momentum acceleration is periodically applied to both object and probe. + In the PtyLab implementation, the momentum state is updated as + + $$ + M^{(n)} = G^{(n)} + \eta M^{(n-1)} + $$ + + followed by + + $$ + X^{(n+1)} = X^{(n)} - \gamma M^{(n)} + $$ + + where $X$ denotes the object or probe, $\eta$ is `frictionM`, and + $\gamma$ is `feedbackM`. The current implementation applies these + momentum updates stochastically during the scan-position loop. + + The default mPIE parameters are `betaObject = 0.25`, + `betaProbe = 0.25`, `alphaObject = 0.1`, `alphaProbe = 0.1`, + `feedbackM = 0.3`, and `frictionM = 0.7`. + + [^maiden2017]: A. M. Maiden, D. Johnson, and P. Li, + "Further improvements to the ptychographical iterative engine," + Optica 4, 736-745 (2017). + https://doi.org/10.1364/OPTICA.4.000736 + + Attributes: + keepPatches (bool): + If enabled, store the reconstructed object patch associated with each + scan position for debugging or detailed analysis. This option can + require substantial additional memory. + See Also: + `ePIE` + Baseline ePIE reconstruction without rPIE regularization or + momentum acceleration. + """ def __init__( self, reconstruction: Reconstruction, @@ -33,11 +104,32 @@ def __init__( params: Params, monitor: Monitor, ): - # This contains reconstruction parameters that are specific to the reconstruction - # but not necessarily to ePIE reconstruction + """ + Initialize the mPIE reconstruction engine. + + Shared reconstruction state is initialized through `BaseEngine`, followed + by the mPIE-specific reconstruction parameters and momentum buffers. + + Momentum acceleration is enabled through + `params.momentumAcceleration`, allowing shared BaseEngine operations such + as modal orthogonalization to keep the corresponding momentum and buffer + arrays consistent with the reconstructed object and probe. + + Args: + reconstruction (Reconstruction): + Reconstruction state containing the current object, probe, and + geometry. + experimentalData (ExperimentalData): + Experimental diffraction data and acquisition parameters. + params (Params): + Shared reconstruction parameters and constraint settings. + monitor (Monitor): + Monitor used for reconstruction visualization and progress + reporting. + """ super().__init__(reconstruction, experimentalData, params, monitor) self.logger = logging.getLogger("mPIE") - self.logger.info("Sucesfully created mPIE mPIE_engine") + self.logger.info("Successfully created mPIE engine") self.logger.info("Wavelength attribute: %s", self.reconstruction.wavelength) # initialize mPIE Params self.initializeReconstructionParams() @@ -46,10 +138,12 @@ def __init__( @property def keepPatches(self): - """Wether or not to keep track of the individual object update patches. + """ + Whether to store the reconstructed object patch for every scan position. - This strongly increases the amount of memory required, only use when absolutely required. + Enable with `engine.keepPatches = True` and disable with `engine.keepPatches = False`. + This option is intended for debugging or detailed analysis and may require a large amount of additional memory. """ return hasattr(self, "patches") @@ -72,8 +166,20 @@ def keepPatches(self, keep_them): def initializeReconstructionParams(self): """ - Set parameters that are specific to the mPIE settings. - :return: + Initialize mPIE-specific reconstruction parameters and momentum state. + + The default mPIE parameters are: + + - `betaObject = 0.25`: object update step size. + - `betaProbe = 0.25`: probe update step size. + - `alphaObject = 0.1`: object-update regularization parameter. + - `alphaProbe = 0.1`: probe-update regularization parameter. + - `feedbackM = 0.3`: momentum feedback strength. + - `frictionM = 0.7`: momentum memory coefficient. + - `numIterations = 50`: number of reconstruction iterations. + + Object and probe momentum arrays are initialized together with corresponding + buffers that store the reconstruction state used by the momentum updates. """ # self.eswUpdate = self.reconstruction.esw.copy() self.betaProbe = 0.25 @@ -94,8 +200,70 @@ def initializeReconstructionParams(self): self.reconstruction.probeWindow = np.abs(self.reconstruction.probe) def reconstruct(self, experimentalData=None, reconstruction=None, vis_after_each_iteration=None): - """Reconstruct object. If experimentalData is given, it replaces the current data. Idem for reconstruction.""" - + r""" + Run the mPIE reconstruction to completion. + + The reconstruction follows the standard ptychographic position loop. + At each scan position $j$, the current object patch and probe form the + exit surface wave + + $$ + \Psi_j = O_j P + $$ + + which is propagated to the detector plane and constrained by the measured + diffraction intensity through `intensityProjection()`. The corresponding + exit-wave correction is + + $$ + \Delta\Psi_j = \Psi'_j - \Psi_j + $$ + + The object and probe are then updated using the mPIE/rPIE update rules + implemented by `objectPatchUpdate()` and `probeUpdate()`. + + If `params.objectTVregSwitch` is enabled, the TV-regularized object update + `objectPatchUpdate_TV()` is used every `params.objectTVfreq` iterations. + The native mPIE object update is retained and an additional TV + regularization term is added with strength controlled by + `params.objectTVregStepSize`. + + If `params.weigh_probe_updates_by_intensity` is enabled, the probe update + is scaled by the relative intensity of the current diffraction frame. + + If `params.positionCorrectionSwitch` is enabled, scan-position correction + is applied after the object and probe updates. + + Momentum acceleration is applied stochastically during the scan-position + loop. In the current implementation, each position update has + approximately a 5% probability of triggering `objectMomentumUpdate()` and + `probeMomentumUpdate()`. + + If `keepPatches` is enabled, the reconstructed patch associated with each + scan position is additionally stored in `self.patches` for diagnostics or + further analysis. + + After all scan positions in an iteration have been processed, + `getErrorMetrics()` evaluates the reconstruction error, + `applyConstraints()` applies the enabled reconstruction constraints, and + `showReconstruction()` updates the reconstruction monitor. + + Args: + experimentalData (ExperimentalData, optional): + Experimental dataset to use for the reconstruction. If provided, + it replaces the currently attached experimental data. + reconstruction (Reconstruction, optional): + Reconstruction state to optimize. If provided, it replaces the + currently attached reconstruction object. + vis_after_each_iteration (callable, optional): + Callback executed after each reconstruction iteration as + `vis_after_each_iteration(loop, reconstruction)`. + + Notes: + The object and probe momentum buffers are reset after + `_prepareReconstruction()` so that they remain synchronized with any + initialization changes applied before reconstruction starts. + """ self.changeExperimentalData(experimentalData) self.changeOptimizable(reconstruction) @@ -144,12 +312,11 @@ def reconstruct(self, experimentalData=None, reconstruction=None, vis_after_each else: object_patch = self.objectPatchUpdate(objectPatch, DELTA) + self.reconstruction.object[..., sy, sx] = object_patch if self.keepPatches: self.patches[positionIndex, ..., sy, sx] = asNumpyArray( abs(object_patch) ** 2 ) - else: - self.reconstruction.object[..., sy, sx] = object_patch # probe update weight = 1 @@ -161,7 +328,7 @@ def reconstruct(self, experimentalData=None, reconstruction=None, vis_after_each # self.reconstruction.push_probe_update(self.reconstruction.probe, positionIndex, self.experimentalData.ptychogram.shape[0]) if self.params.positionCorrectionSwitch: - shifter = self.positionCorrection( + self.positionCorrection( objectPatch, positionIndex, sy, sx ) # self.pbar_pos.write(f'Corr: {shifter[0]*1e6:.2f} um x {shifter[1]*1e6:.2f} um') @@ -194,9 +361,38 @@ def reconstruct(self, experimentalData=None, reconstruction=None, vis_after_each # todo clearMemory implementation def objectMomentumUpdate(self): - """ - momentum update object, save updated objectMomentum and objectBuffer. - :return: + r""" + Apply the mPIE momentum update to the reconstructed object. + + The change in the object since the previous momentum update is estimated + from the stored object buffer: + + $$ + G_O^{(n)} = O_{\mathrm{buf}}^{(n)} - O^{(n)} + $$ + + The object momentum is updated according to + + $$ + M_O^{(n)} = G_O^{(n)} + \eta M_O^{(n-1)} + $$ + + where $\eta$ is `frictionM`. + + The accumulated momentum is then fed back into the object estimate: + + $$ + O^{(n+1)} = O^{(n)} - \gamma M_O^{(n)} + $$ + + where $\gamma$ is `feedbackM`. + + After the momentum correction, `objectBuffer` is updated with the current + object estimate for the next momentum step. + + Notes: + This update is triggered stochastically from `reconstruct()` rather + than after every scan-position update. """ gradient = self.reconstruction.objectBuffer - self.reconstruction.object self.reconstruction.objectMomentum = ( @@ -209,9 +405,12 @@ def objectMomentumUpdate(self): self.reconstruction.objectBuffer = self.reconstruction.object.copy() def probeMomentumUpdate(self): - """ - momentum update probe, save updated probeMomentum and probeBuffer. - :return: + r""" + Apply the mPIE momentum update to the reconstructed probe. Similar to `objectMomentumUpdate()`. + + See Also: + `objectMomentumUpdate` + Equivalent momentum update applied to the reconstructed object. """ gradient = self.reconstruction.probeBuffer - self.reconstruction.probe self.reconstruction.probeMomentum = ( @@ -224,11 +423,56 @@ def probeMomentumUpdate(self): self.reconstruction.probeBuffer = self.reconstruction.probe.copy() def objectPatchUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray): - """ - Todo add docstring - :param objectPatch: - :param DELTA: - :return: + r""" + Update the object patch using the regularized mPIE/rPIE object-update rule. + + The probe intensity is first evaluated and its maximum value is used as a + global normalization scale: + + $$ + P_{\max} = \max_{x,y}\sum |P(x,y)|^2 + $$ + + For conventional ptychography, the probe weighting is + + $$ + W_P =\frac{P^*}{\alpha_O P_{\max} + (1-\alpha_O)|P|^2} + $$ + + and the object patch is updated according to + + $$ + O'_j =O_j + \beta_O \sum W_P\Delta\Psi_j + $$ + + where $\Delta\Psi_j$ is the exit-wave correction, $\beta_O$ is + `betaObject`, and $\alpha_O$ is `alphaObject`. + + The parameter `alphaObject` controls the balance between global + normalization by the maximum probe intensity and local normalization by + the spatially varying probe intensity. + + For Fourier ptychography (`operationMode == "FPM"`), an additional + probe-amplitude weighting is applied: + + $$ + W_P^{\mathrm{FPM}} =\frac{|P|}{P_{\max}}\frac{P^*}{\alpha_O P_{\max} + (1-\alpha_O)|P|^2} + $$ + + In the multidimensional PtyLab representation, the object correction is + summed over the probe-mode axis before being added to the current object + patch. + + Args: + objectPatch (ndarray): + Current object patch at the active scan position. + DELTA (ndarray): + Exit-wave correction + `reconstruction.eswUpdate - reconstruction.esw`. + + Returns: + ndarray: + Updated object patch. """ # find out which array module to use, numpy or cupy (or other...) xp = getArrayModule(objectPatch) @@ -251,11 +495,58 @@ def objectPatchUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray): ) def probeUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray, weight: float): - """ - Todo add docstring - :param objectPatch: - :param DELTA: - :return: + r""" + Update the probe using the regularized mPIE/rPIE probe-update rule. + + The object intensity is first evaluated and its maximum value is used as + a global normalization scale: + + $$ + O_{\max} = \max_{x,y}\sum |O_j(x,y)|^2 + $$ + + The object weighting is then calculated as + + $$ + W_O = \frac{O_j^*}{\alpha_P O_{\max} + (1-\alpha_P)|O_j|^2} + $$ + + and the probe is updated according to + + $$ + P' = P + w\beta_P\sum W_O\Delta\Psi_j + $$ + + where $\Delta\Psi_j$ is the exit-wave correction, $\beta_P$ is + `betaProbe`, $\alpha_P$ is `alphaProbe`, and $w$ is the optional + intensity-dependent update weight. + + The parameter `alphaProbe` controls the balance between global + normalization by the maximum object intensity and local normalization by + the spatially varying object intensity. + + By default, $w=1$. If `params.weigh_probe_updates_by_intensity` is + enabled in `reconstruct()`, $w$ is set to the relative intensity of the + current diffraction frame. + + In the current multidimensional PtyLab representation, the probe + correction is summed over axis `1` before being added to the current + probe estimate. + + Args: + objectPatch (ndarray): + Current object patch at the active scan position. + DELTA (ndarray): + Exit-wave correction + `reconstruction.eswUpdate - reconstruction.esw`. + weight (float): + Multiplicative weight applied to the probe update. Typically `1`, + or the relative intensity of the current diffraction frame when + intensity-weighted probe updates are enabled. + + Returns: + ndarray: + Updated probe. """ # find out which array module to use, numpy or cupy (or other...) xp = getArrayModule(objectPatch) @@ -268,3 +559,48 @@ def probeUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray, weight: float) frac * DELTA, axis=1, keepdims=True ) return r + + +class pcPIE(mPIE): + """ + Backward-compatible wrapper for :class:`mPIE`. + + Position correction is now provided by `mPIE` through + `params.positionCorrectionSwitch`. This class is retained only for + compatibility with existing code using `Engines.pcPIE`. + """ + + def __init__( + self, + reconstruction: Reconstruction, + experimentalData: ExperimentalData, + params: Params, + monitor: Monitor, + ): + super().__init__( + reconstruction, + experimentalData, + params, + monitor, + ) + + self.name = "pcPIE" + self.logger = logging.getLogger("pcPIE") + + @property + def betaM(self): + """Deprecated alias for `feedbackM`.""" + return self.feedbackM + + @betaM.setter + def betaM(self, value): + self.feedbackM = value + + @property + def stepM(self): + """Deprecated alias for `frictionM`.""" + return self.frictionM + + @stepM.setter + def stepM(self, value): + self.frictionM = value \ No newline at end of file diff --git a/tests/Engines/test_remove_pcPIE_engine.py b/tests/Engines/test_remove_pcPIE_engine.py new file mode 100644 index 0000000..c60e5a8 --- /dev/null +++ b/tests/Engines/test_remove_pcPIE_engine.py @@ -0,0 +1,43 @@ +import pytest + +from PtyLab import Engines +from PtyLab.Engines.mPIE import mPIE, pcPIE + + +@pytest.fixture +def engine(): + """Create a wrapper instance without initializing reconstruction data.""" + return object.__new__(pcPIE) + + +def test_pcpie_is_exported(): + """Engines.pcPIE should remain available for backward compatibility.""" + assert hasattr(Engines, "pcPIE") + assert Engines.pcPIE is pcPIE + + +def test_pcpie_inherits_mpie(): + """pcPIE should reuse the mPIE implementation.""" + assert issubclass(pcPIE, mPIE) + + +@pytest.mark.parametrize( + "legacy_name,current_name,initial_value,updated_value", + [ + ("betaM", "feedbackM", 0.3, 0.5), + ("stepM", "frictionM", 0.7, 0.8), + ], + ids=["betaM-feedbackM", "stepM-frictionM"], +) +def test_legacy_momentum_parameter_aliases( + engine, legacy_name, current_name, initial_value, updated_value +): + """Legacy pcPIE momentum names should map to the mPIE parameters.""" + setattr(engine, current_name, initial_value) + + # Old pcPIE names should read from the new mPIE attributes + assert getattr(engine, legacy_name) == initial_value + + # Setting the old names should update the new attributes + setattr(engine, legacy_name, updated_value) + assert getattr(engine, current_name) == updated_value From 81d1127092c63351fe0c5cd7c9f7a950371795c2 Mon Sep 17 00:00:00 2001 From: Chang Liu Date: Wed, 30 Sep 2026 17:22:24 +0200 Subject: [PATCH 12/14] delete pcPIE.py --- PtyLab/Engines/pcPIE.py | 202 ---------------------------------------- 1 file changed, 202 deletions(-) delete mode 100644 PtyLab/Engines/pcPIE.py diff --git a/PtyLab/Engines/pcPIE.py b/PtyLab/Engines/pcPIE.py deleted file mode 100644 index 0802a0e..0000000 --- a/PtyLab/Engines/pcPIE.py +++ /dev/null @@ -1,202 +0,0 @@ -import numpy as np -from matplotlib import pyplot as plt - -from PtyLab.utils import gpuUtils - -try: - import cupy as cp -except ImportError: - # print("Cupy not available, will not be able to run GPU based computation") - # Still define the name, we'll take care of it later but in this way it's still possible - # to see that gPIE exists for example. - cp = None - -import logging -import sys - -import tqdm - -from PtyLab.Engines.BaseEngine import BaseEngine -from PtyLab.ExperimentalData.ExperimentalData import ExperimentalData -from PtyLab.Monitor.Monitor import Monitor -from PtyLab.Params.Params import Params - -# PtyLab imports -from PtyLab.Reconstruction.Reconstruction import Reconstruction -from PtyLab.utils.gpuUtils import asNumpyArray, getArrayModule -from PtyLab.utils.utils import fft2c, ifft2c - - -class pcPIE(BaseEngine): - def __init__( - self, - reconstruction: Reconstruction, - experimentalData: ExperimentalData, - params: Params, - monitor: Monitor, - ): - # This contains reconstruction parameters that are specific to the reconstruction - # but not necessarily to ePIE reconstruction - super().__init__(reconstruction, experimentalData, params, monitor) - self.logger = logging.getLogger("pcPIE") - self.logger.info("Successfully created pcPIE pcPIE_engine") - self.logger.info("Wavelength attribute: %s", self.reconstruction.wavelength) - # initialize pcPIE Params - self.initializeReconstructionParams() - # initialize momentum - self.reconstruction.initializeObjectMomentum() - self.reconstruction.initializeProbeMomentum() - # set object and probe buffers - self.reconstruction.objectBuffer = self.reconstruction.object.copy() - self.reconstruction.probeBuffer = self.reconstruction.probe.copy() - - self.params.momentumAcceleration = True - - def initializeReconstructionParams(self): - """ - Set parameters that are specific to the pcPIE settings. - :return: - """ - # these are same as mPIE - # self.eswUpdate = self.reconstruction.esw.copy() - self.betaProbe = 0.25 - self.betaObject = 0.25 - self.alphaProbe = 0.1 # probe regularization - self.alphaObject = 0.1 # object regularization - self.betaM = 0.3 # feedback - self.stepM = 0.7 # friction - # self.probeWindow = np.abs(self.reconstruction.probe) - self.numIterations = 50 - - def reconstruct(self): - self._prepareReconstruction() - - # actual reconstruction ePIE_engine - - self.pbar = tqdm.trange( - self.numIterations, desc="pcPIE", file=sys.stdout, leave=True - ) # in order to change description to the tqdm progress bar - for loop in self.pbar: - # set position order - self.setPositionOrder() - - for positionLoop, positionIndex in enumerate(self.positionIndices): - # get object patch - row, col = self.reconstruction.positions[positionIndex] - sy = slice(row, row + self.reconstruction.Np) - sx = slice(col, col + self.reconstruction.Np) - # note that object patch has size of probe array - objectPatch = self.reconstruction.object[..., sy, sx].copy() - - # make exit surface wave - self.reconstruction.esw = objectPatch * self.reconstruction.probe - - # propagate to camera, intensityProjection, propagate back to object - self.intensityProjection(positionIndex) - - # difference term - DELTA = self.reconstruction.eswUpdate - self.reconstruction.esw - - # object update - self.reconstruction.object[..., sy, sx] = self.objectPatchUpdate( - objectPatch, DELTA - ) - - # probe update - self.reconstruction.probe = self.probeUpdate(objectPatch, DELTA) - if self.params.positionCorrectionSwitch: - self.positionCorrection(objectPatch, positionIndex, sy, sx) - - # momentum updates - if np.random.rand(1) > 0.95: - self.objectMomentumUpdate() - self.probeMomentumUpdate() - - # get error metric - self.getErrorMetrics() - - # apply Constraints - self.applyConstraints(loop) - - # show reconstruction - self.showReconstruction(loop) - - # todo clearMemory implementation - - if self.params.gpuFlag: - self.logger.info("switch to cpu") - self._move_data_to_cpu() - self.params.gpuFlag = 0 - - def objectMomentumUpdate(self): - """ - momentum update object, save updated objectMomentum and objectBuffer. - :return: - """ - gradient = self.reconstruction.objectBuffer - self.reconstruction.object - self.reconstruction.objectMomentum = ( - gradient + self.stepM * self.reconstruction.objectMomentum - ) - self.reconstruction.object = ( - self.reconstruction.object - self.betaM * self.reconstruction.objectMomentum - ) - self.reconstruction.objectBuffer = self.reconstruction.object.copy() - - def probeMomentumUpdate(self): - """ - momentum update probe, save updated probeMomentum and probeBuffer. - :return: - """ - gradient = self.reconstruction.probeBuffer - self.reconstruction.probe - self.reconstruction.probeMomentum = ( - gradient + self.stepM * self.reconstruction.probeMomentum - ) - self.reconstruction.probe = ( - self.reconstruction.probe - self.betaM * self.reconstruction.probeMomentum - ) - self.reconstruction.probeBuffer = self.reconstruction.probe.copy() - - def objectPatchUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray): - """ - Todo add docstring - :param objectPatch: - :param DELTA: - :return: - """ - # find out which array module to use, numpy or cupy (or other...) - xp = getArrayModule(objectPatch) - absP2 = xp.abs(self.reconstruction.probe) ** 2 - Pmax = xp.max(xp.sum(absP2, axis=(0, 1, 2, 3)), axis=(-1, -2)) - if self.experimentalData.operationMode == "FPM": - frac = ( - abs(self.reconstruction.probe) - / Pmax - * self.reconstruction.probe.conj() - / (self.alphaObject * Pmax + (1 - self.alphaObject) * absP2) - ) - else: - frac = self.reconstruction.probe.conj() / ( - self.alphaObject * Pmax + (1 - self.alphaObject) * absP2 - ) - return objectPatch + self.betaObject * xp.sum( - frac * DELTA, axis=2, keepdims=True - ) - - def probeUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray): - """ - Todo add docstring - :param objectPatch: - :param DELTA: - :return: - """ - # find out which array module to use, numpy or cupy (or other...) - xp = getArrayModule(objectPatch) - absO2 = xp.abs(objectPatch) ** 2 - Omax = xp.max(xp.sum(absO2, axis=(0, 1, 2, 3)), axis=(-1, -2)) - frac = objectPatch.conj() / ( - self.alphaProbe * Omax + (1 - self.alphaProbe) * absO2 - ) - r = self.reconstruction.probe + self.betaProbe * xp.sum( - frac * DELTA, axis=1, keepdims=True - ) - return r From c327b547a6413ada5ada97630b87a7ae26e712f4 Mon Sep 17 00:00:00 2001 From: Chang Liu Date: Wed, 30 Sep 2026 17:42:44 +0200 Subject: [PATCH 13/14] modify GPU requirement --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index e727830..c8e3cb4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,7 +31,7 @@ dev = [ "pytest~=9.0", "imageio~=2.0", ] -gpu = ["cupy-cuda12x[ctk]"] +gpu = ["cupy-cuda12x[ctk]; sys_platform != 'darwin'"] tensorflow = ["tensorflow~=2.14"] docs = [ "jetblack-markdown", From 58abc7e0a3b83b82e7ca3bbd70789c1b6a5a5186 Mon Sep 17 00:00:00 2001 From: Chang Liu Date: Wed, 30 Sep 2026 17:44:38 +0200 Subject: [PATCH 14/14] modified tool.uv requirement --- pyproject.toml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index c8e3cb4..b1a5cc3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -57,5 +57,7 @@ module-root = "" [tool.uv] required-environments = [ + "sys_platform == 'linux' and platform_machine == 'x86_64'", "sys_platform == 'win32' and platform_machine == 'AMD64'", + "sys_platform == 'darwin' and platform_machine == 'arm64'", ]