From 7eafb2b210389f2ca55e1898f57291b008a85f32 Mon Sep 17 00:00:00 2001 From: Chang Liu Date: Wed, 30 Sep 2026 17:54:52 +0200 Subject: [PATCH 01/11] API Documentation - zPIE- started --- PtyLab/Engines/zPIE.py | 71 ++++++++++++++++++++++++++++++++++++++++++ docs/api/engines.md | 6 ---- 2 files changed, 71 insertions(+), 6 deletions(-) diff --git a/PtyLab/Engines/zPIE.py b/PtyLab/Engines/zPIE.py index 1d03c685..3dff0f7d 100644 --- a/PtyLab/Engines/zPIE.py +++ b/PtyLab/Engines/zPIE.py @@ -25,6 +25,77 @@ class zPIE(BaseEngine): + r""" + Ptychographic iterative engine with axial-distance refinement. + + zPIE extends the standard PIE reconstruction by jointly refining the + object/probe estimate and the axial propagation distance $z$.[^loetgering2020] + + During reconstruction, a set of candidate axial offsets is generated + around the current propagation distance: + + $$ + \Delta z_k \in [-d\,\mathrm{DoF},\, d\,\mathrm{DoF}] + $$ + + For each candidate offset, either the reconstructed object or probe is + propagated to the corresponding defocus plane using angular-spectrum + propagation. The quantity used for axial optimization is selected through + `focusObject`. + + A total-variation-based focus metric is evaluated from the propagated + complex field: + + $$ + M_k = \sum \sqrt{|\nabla_x U_k|^2 + |\nabla_y U_k|^2 + \epsilon} + $$ + + where $U_k$ is the propagated field at axial offset $\Delta z_k$ and + $\epsilon$ is a small numerical stabilization term. + + The axial feedback is calculated from the merit values as + + $$ + f_z = \frac{\sum_k \Delta z_k M_k}{\sum_k M_k} + $$ + + and accumulated using a momentum-like update: + + $$ + m_z^{(n)} = \gamma m_z^{(n-1)} + \eta f_z + $$ + + followed by + + $$ + z^{(n+1)} = z^{(n)} + m_z^{(n)} + $$ + + where $\gamma$ is `zPIEfriction` and $\eta$ is + `zPIEgradientStepSize`. + + The updated axial distance is stored in `reconstruction.zo`. For + propagation schemes whose sampling depends explicitly on distance, the + corresponding propagation coordinates are updated after each axial + correction. + + The object and probe are subsequently updated using the standard PIE + object and probe update rules. + + Parameters specific to zPIE include: + + - `DoF`: depth-of-field scale used to define the axial search range. + - `zPIEgradientStepSize`: strength of the axial feedback update. + - `zPIEfriction`: momentum retention factor for axial refinement. + - `focusObject`: if `True`, evaluate the focus metric on the object; + otherwise evaluate it on the probe. + - `zMomentun`: accumulated axial momentum. + + [^loetgering2020]: L. Loetgering, M. Du, K. S. E. Eikema, and S. Witte, + "zPIE: an autofocusing algorithm for ptychography," + Optics Letters 45, 2030-2033 (2020). + https://doi.org/10.1364/OL.389492 + """ def __init__( self, reconstruction: Reconstruction, diff --git a/docs/api/engines.md b/docs/api/engines.md index 659958a8..59e5825c 100644 --- a/docs/api/engines.md +++ b/docs/api/engines.md @@ -42,12 +42,6 @@ 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.zPIE.zPIE options: show_root_heading: true From 3dd53118f6ec2a68b21a7514f42eaf49f59ef5f1 Mon Sep 17 00:00:00 2001 From: Chang Liu Date: Fri, 2 Oct 2026 22:38:20 +0200 Subject: [PATCH 02/11] Implemented purityPIE for purity-based axial-distance calibration. Issues: external initialization --- PtyLab/Engines/BaseEngine.py | 4 +- PtyLab/Engines/__init__.py | 2 + PtyLab/Engines/purityPIE.py | 917 ++++++++++++++++++ PtyLab/Params/Params.py | 6 + example_scripts/exampleReconstructionCPM.py | 2 +- .../example_purity_Zscan_real_Data.py | 302 ++++++ 6 files changed, 1230 insertions(+), 3 deletions(-) create mode 100644 PtyLab/Engines/purityPIE.py create mode 100644 example_scripts/example_purity_Zscan_real_Data.py diff --git a/PtyLab/Engines/BaseEngine.py b/PtyLab/Engines/BaseEngine.py index 60335a69..20be5a7f 100644 --- a/PtyLab/Engines/BaseEngine.py +++ b/PtyLab/Engines/BaseEngine.py @@ -1577,7 +1577,7 @@ def decompressionProjection(self, positionIndex): axis=-2, ) - def showReconstruction(self, loop): + def showReconstruction(self, loop, force=False): """ Update reconstruction monitoring and optional iteration output. @@ -1597,7 +1597,7 @@ def showReconstruction(self, loop): loop (int): Current reconstruction iteration. """ - if np.mod(loop, self.monitor.figureUpdateFrequency) == 0: + if force or np.mod(loop, self.monitor.figureUpdateFrequency) == 0: if self.experimentalData.operationMode == "FPM": object_estimate = np.squeeze( asNumpyArray( diff --git a/PtyLab/Engines/__init__.py b/PtyLab/Engines/__init__.py index 17ed5a03..d5e3bef2 100644 --- a/PtyLab/Engines/__init__.py +++ b/PtyLab/Engines/__init__.py @@ -16,6 +16,7 @@ "OPR", "qNewton", "zPIE", + "purityPIE", ] from .e3PIE import e3PIE from .ePIE import ePIE @@ -25,6 +26,7 @@ from .OPR import OPR from .qNewton import qNewton from .zPIE import zPIE +from .purityPIE import purityPIE import warnings diff --git a/PtyLab/Engines/purityPIE.py b/PtyLab/Engines/purityPIE.py new file mode 100644 index 00000000..4824d61a --- /dev/null +++ b/PtyLab/Engines/purityPIE.py @@ -0,0 +1,917 @@ +import numpy as np +from matplotlib import pyplot as plt + +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 purityPIE(BaseEngine): + ''' + purityPIE extends mPIE-style reconstruction with purity-based axial calibration. + ''' + def __init__( + self, + reconstruction: Reconstruction, + experimentalData: ExperimentalData, + params: Params, + monitor: Monitor, + ): + """ + 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("purityPIE") + self.logger.info("Successfully created purityPIE engine") + self.logger.info("Wavelength attribute: %s", self.reconstruction.wavelength) + + self.initializeReconstructionParams() + self.params.momentumAcceleration = True + self.name = "purityPIE" + + @property + def keepPatches(self): + """ + Whether to store the reconstructed object patch for every scan position. + + 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") + + @keepPatches.setter + def keepPatches(self, keep_them): + + if keep_them: + self.logger.info("Keeping patches!") + self.patches = np.zeros( + ( + self.experimentalData.ptychogram.shape[0], + *self.reconstruction.shape_O, + ), + np.complex64, + ) + else: + self.logger.info("Not keeping patches") + if hasattr(self, "patches"): + del self.patches + + def initializeReconstructionParams(self): + """ + 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 + self.betaObject = 0.25 + self.alphaProbe = 0.1 # probe regularization + self.alphaObject = 0.1 # object regularization + self.feedbackM = 0.3 # feedback + self.frictionM = 0.7 # friction + self.numIterations = 50 + + # 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.reconstruction.probeWindow = np.abs(self.reconstruction.probe) + + def reconstruct( + self, + experimentalData=None, + reconstruction=None, + ): + """ + Run purity-based axial calibration. + + Each candidate propagation distance is reconstructed from the same + initial state using `numIterations` mPIE-style iterations. The distance + yielding the highest probe purity is selected, and the corresponding + reconstruction state is retained. + + If `params.purityZScanPlot` is enabled, the resulting purity-versus-z + curve is displayed after calibration. + + Args: + experimentalData (ExperimentalData, optional): + Experimental dataset to use. If provided, it replaces the + currently attached dataset. + reconstruction (Reconstruction, optional): + Reconstruction state to calibrate. If provided, it replaces the + currently attached reconstruction. + + Returns: + tuple: + `(best_z, best_purity)`, where `best_z` is the selected axial + propagation distance and `best_purity` is the corresponding + probe purity. + """ + + self.changeExperimentalData(experimentalData) + self.changeOptimizable(reconstruction) + + self._prepareReconstruction() + + # Reset momentum buffers after preparation + self.reconstruction.objectBuffer = self.reconstruction.object.copy() + self.reconstruction.probeBuffer = self.reconstruction.probe.copy() + + # Purity-based axial calibration + best_z, best_purity = self.purityZScan() + + # Optional scan visualization + if self.params.purityZScanPlot: + self.plotPurityZScan() + + # Return reconstruction data to CPU + if self.params.gpuFlag: + self.logger.info("switch to cpu") + self._move_data_to_cpu() + self.params.gpuFlag = 0 + + return best_z, best_purity + + def purityZScan(self): + """ + Scan a fixed axial range and select the distance that maximizes + reconstructed probe purity. + + Every candidate z is evaluated from the same initial reconstruction + state using `numIterations` mPIE iterations. + """ + + if self.reconstruction.npsm < 2: + raise ValueError( + "Purity-based z scanning requires at least two probe modes " + "(reconstruction.npsm >= 2)." + ) + + # Initial z guess defines the center of the scan + z0 = self.reconstruction.zo + self.purityZInitialGuess = z0 + + z_values = np.linspace( + z0 - self.params.purityZScanRange, + z0 + self.params.purityZScanRange, + self.params.purityZScanPoints, + ) + + # Save the common starting state + object0 = self.reconstruction.object.copy() + probe0 = self.reconstruction.probe.copy() + + objectMomentum0 = self.reconstruction.objectMomentum.copy() + probeMomentum0 = self.reconstruction.probeMomentum.copy() + + objectBuffer0 = self.reconstruction.objectBuffer.copy() + probeBuffer0 = self.reconstruction.probeBuffer.copy() + + purity_values = [] + + best_purity = -np.inf + best_z = z0 + best_state = None + + # Do not allow orthogonalization during the internal mPIE iterations. + # Purity is evaluated once after each candidate reconstruction. + #orthogonalization_switch = self.params.orthogonalizationSwitch + #self.params.orthogonalizationSwitch = False + + tqdm.tqdm.write("") + tqdm.tqdm.write("Starting purity-based z scan") + tqdm.tqdm.write( + f"Initial z = {z0 * 1e3:.6f} mm | " + f"range = ±{self.params.purityZScanRange * 1e6:.1f} µm | " + f"points = {len(z_values)} | " + f"iterations/z = {self.numIterations}" + ) + tqdm.tqdm.write("") + + #try: + for z_index, z in enumerate(z_values): + + # Restore exactly the same starting state for every candidate z + self.reconstruction.object = object0.copy() + self.reconstruction.probe = probe0.copy() + + self.reconstruction.objectMomentum = objectMomentum0.copy() + self.reconstruction.probeMomentum = probeMomentum0.copy() + + self.reconstruction.objectBuffer = objectBuffer0.copy() + self.reconstruction.probeBuffer = probeBuffer0.copy() + + # Set candidate propagation distance + self.reconstruction.zo = z + + dz_um = (z - z0) * 1e6 + + tqdm.tqdm.write( + f"[{z_index + 1}/{len(z_values)}] " + f"Reconstructing z = {z * 1e3:.6f} mm " + f"(Δz = {dz_um:+.2f} µm)" + ) + + # Run mPIE reconstruction at this z + for loop in range(self.numIterations): + self._run_single_iteration(loop) + #self.showReconstruction(loop) + + # Orthogonalize once and evaluate purity + self.orthogonalization() + + + purity = float( + np.asarray( + asNumpyArray(self.reconstruction.purityProbe) + ).squeeze() + ) + + purity_values.append(purity) + # Force monitor to display the final state and final purity + self.showReconstruction( + self.numIterations - 1, + force=True, + ) + tqdm.tqdm.write( + f" Probe purity = {purity:.6f}" + ) + + if purity > best_purity: + best_purity = purity + best_z = z + + best_state = { + "object": self.reconstruction.object.copy(), + "probe": self.reconstruction.probe.copy(), + "objectMomentum": self.reconstruction.objectMomentum.copy(), + "probeMomentum": self.reconstruction.probeMomentum.copy(), + "objectBuffer": self.reconstruction.objectBuffer.copy(), + "probeBuffer": self.reconstruction.probeBuffer.copy(), + } + + tqdm.tqdm.write( + f" NEW BEST: z = {best_z * 1e3:.6f} mm, " + f"purity = {best_purity:.6f}" + ) + + tqdm.tqdm.write("") + + #finally: + # self.params.orthogonalizationSwitch = orthogonalization_switch + + if best_state is None: + raise RuntimeError( + "Purity z scan did not produce a valid purity value." + ) + + # Restore best-z reconstruction state + self.reconstruction.object = best_state["object"] + self.reconstruction.probe = best_state["probe"] + + self.reconstruction.objectMomentum = best_state["objectMomentum"] + self.reconstruction.probeMomentum = best_state["probeMomentum"] + + self.reconstruction.objectBuffer = best_state["objectBuffer"] + self.reconstruction.probeBuffer = best_state["probeBuffer"] + + self.reconstruction.zo = best_z + + # Store scan results + self.purityZValues = z_values + self.purityZMetrics = np.asarray(purity_values) + self.purityBestZ = best_z + self.purityBestValue = best_purity + + # Final summary + tqdm.tqdm.write("=" * 60) + tqdm.tqdm.write("Purity-based z scan finished") + tqdm.tqdm.write( + f"Best z = {best_z * 1e3:.6f} mm" + ) + tqdm.tqdm.write( + f"Best Δz = {(best_z - z0) * 1e6:+.2f} µm" + ) + tqdm.tqdm.write( + f"Best purity = {best_purity:.6f}" + ) + tqdm.tqdm.write("=" * 60) + + return best_z, best_purity + + def purityZScan_X(self): + """ + Scan a fixed axial range and select the distance that maximizes + reconstructed probe purity. + + Every candidate z is evaluated from the same initial reconstruction + state using `numIterations` mPIE iterations. + """ + + if self.reconstruction.npsm < 2: + raise ValueError( + "Purity-based z scanning requires at least two probe modes " + "(reconstruction.npsm >= 2)." + ) + + # Initial z guess defines the center of the scan + z0 = self.reconstruction.zo + self.purityZInitialGuess = z0 + + z_values = np.linspace( + z0 - self.params.purityZScanRange, + z0 + self.params.purityZScanRange, + self.params.purityZScanPoints, + ) + + # Save the common starting state + object0 = self.reconstruction.object.copy() + probe0 = self.reconstruction.probe.copy() + + objectMomentum0 = self.reconstruction.objectMomentum.copy() + probeMomentum0 = self.reconstruction.probeMomentum.copy() + + objectBuffer0 = self.reconstruction.objectBuffer.copy() + probeBuffer0 = self.reconstruction.probeBuffer.copy() + + purity_values = [] + + best_purity = -np.inf + best_z = z0 + best_state = None + + # Do not allow orthogonalization during the internal mPIE iterations. + # Purity is evaluated once after each candidate reconstruction. + orthogonalization_switch = self.params.orthogonalizationSwitch + self.params.orthogonalizationSwitch = False + + tqdm.tqdm.write("") + tqdm.tqdm.write("Starting purity-based z scan") + tqdm.tqdm.write( + f"Initial z = {z0 * 1e3:.6f} mm | " + f"range = ±{self.params.purityZScanRange * 1e6:.1f} µm | " + f"points = {len(z_values)} | " + f"iterations/z = {self.numIterations}" + ) + tqdm.tqdm.write("") + + try: + for z_index, z in enumerate(z_values): + + self._resetCandidateState(z) + #np.random.seed(0) + + + dz_um = (z - z0) * 1e6 + + tqdm.tqdm.write( + f"[{z_index + 1}/{len(z_values)}] " + f"Reconstructing z = {z * 1e3:.6f} mm " + f"(Δz = {dz_um:+.2f} µm)" + ) + + # Run mPIE reconstruction at this z + for loop in range(self.numIterations): + self._run_single_iteration(loop) + self.showReconstruction(loop) + + # Orthogonalize once and evaluate purity + self.orthogonalization() + + + purity = float( + np.asarray( + asNumpyArray(self.reconstruction.purityProbe) + ).squeeze() + ) + self.showReconstruction(0) + + purity_values.append(purity) + + tqdm.tqdm.write( + f" Probe purity = {purity:.6f}" + ) + + if purity > best_purity: + best_purity = purity + best_z = z + + best_state = { + "object": self.reconstruction.object.copy(), + "probe": self.reconstruction.probe.copy(), + "objectMomentum": self.reconstruction.objectMomentum.copy(), + "probeMomentum": self.reconstruction.probeMomentum.copy(), + "objectBuffer": self.reconstruction.objectBuffer.copy(), + "probeBuffer": self.reconstruction.probeBuffer.copy(), + } + + tqdm.tqdm.write( + f" NEW BEST: z = {best_z * 1e3:.6f} mm, " + f"purity = {best_purity:.6f}" + ) + + tqdm.tqdm.write("") + + finally: + self.params.orthogonalizationSwitch = orthogonalization_switch + + if best_state is None: + raise RuntimeError( + "Purity z scan did not produce a valid purity value." + ) + + # Restore best-z reconstruction state + self.reconstruction.object = best_state["object"] + self.reconstruction.probe = best_state["probe"] + + self.reconstruction.objectMomentum = best_state["objectMomentum"] + self.reconstruction.probeMomentum = best_state["probeMomentum"] + + self.reconstruction.objectBuffer = best_state["objectBuffer"] + self.reconstruction.probeBuffer = best_state["probeBuffer"] + + self.reconstruction.zo = best_z + + # Store scan results + self.purityZValues = z_values + self.purityZMetrics = np.asarray(purity_values) + self.purityBestZ = best_z + self.purityBestValue = best_purity + + # Final summary + tqdm.tqdm.write("=" * 60) + tqdm.tqdm.write("Purity-based z scan finished") + tqdm.tqdm.write( + f"Best z = {best_z * 1e3:.6f} mm" + ) + tqdm.tqdm.write( + f"Best Δz = {(best_z - z0) * 1e6:+.2f} µm" + ) + tqdm.tqdm.write( + f"Best purity = {best_purity:.6f}" + ) + tqdm.tqdm.write("=" * 60) + + return best_z, best_purity + + def _resetCandidateState(self, z): + """ + Reset the reconstruction state for one candidate z. + + This implementation is intended for analytic initialization: + initialProbe = "circ" + initialObject = "ones" + + The object/probe are regenerated after setting the candidate z so that + their physical sampling is consistent with the candidate geometry. + Momentum and buffers are also reset. + """ + if self.reconstruction.initialProbe != "circ": + raise NotImplementedError( + "Candidate-z reinitialization is currently implemented " + "only for initialProbe='circ'." + ) + + if self.reconstruction.initialObject != "ones": + raise NotImplementedError( + "Candidate-z reinitialization is currently implemented " + "only for initialObject='ones'." + ) + + # Set candidate geometry first + self.reconstruction.zo = z + + # Rebuild object/probe on the candidate-z geometry + self.reconstruction.initializeObjectProbe() + + self._initialProbePowerCorrection() + + # Reset momentum + self.reconstruction.initializeObjectMomentum() + self.reconstruction.initializeProbeMomentum() + + # Reset buffers + self.reconstruction.objectBuffer = self.reconstruction.object.copy() + self.reconstruction.probeBuffer = self.reconstruction.probe.copy() + + self._prepareReconstruction() + + def plotPurityZScan(self, show=True): + """ + Plot probe purity as a function of axial distance. + + The axial coordinate is shown relative to the initial z guess used + for the purity scan. + + Args: + show (bool, optional): + Display the figure immediately. Default is `True`. + + Returns: + tuple: + Matplotlib `(figure, axis)` objects. + """ + + if not hasattr(self, "purityZValues"): + raise RuntimeError( + "No purity z-scan results are available. " + "Run purityZScan() first." + ) + + if not hasattr(self, "purityZInitialGuess"): + raise RuntimeError( + "Initial z guess is not available." + ) + + z_relative_um = ( + np.asarray(self.purityZValues) + - self.purityZInitialGuess + ) * 1e6 + + purity = np.asarray(self.purityZMetrics) + + best_index = np.argmax(purity) + + fig, ax = plt.subplots(figsize=(6, 4)) + + ax.plot( + z_relative_um, + purity, + marker="o", + label="Probe purity", + ) + + # Initial z guess + ax.axvline( + 0, + linestyle="--", + alpha=0.5, + label="Initial z guess", + ) + + # Best z + ax.plot( + z_relative_um[best_index], + purity[best_index], + marker="*", + markersize=12, + linestyle="None", + label=( + f"Best Δz = {z_relative_um[best_index]:.1f} µm" + ), + ) + + ax.set_xlabel("Axial offset from initial z (µm)") + ax.set_ylabel("Probe purity") + ax.set_title("Purity-based axial calibration") + + ax.grid(alpha=0.3) + ax.legend() + fig.tight_layout() + + if show: + plt.show() + + return fig, ax + + def _run_single_iteration(self, loop): + """ + Run one complete mPIE reconstruction iteration. + """ + + # set position order + self.setPositionOrder() + + self.pbar_pos = tqdm.tqdm( + self.positionIndices, + leave=False, + desc="ptychogram", + file=sys.stdout, + ) + + for positionLoop, positionIndex in enumerate(self.pbar_pos): + + row, col = self.reconstruction.positions[positionIndex] + sy = slice(row, row + self.reconstruction.Np) + sx = slice(col, col + self.reconstruction.Np) + + objectPatch = self.reconstruction.object[..., sy, sx].copy() + + # exit surface wave + self.reconstruction.esw = ( + objectPatch * self.reconstruction.probe + ) + + # intensity constraint + self.intensityProjection(positionIndex) + + DELTA = ( + self.reconstruction.eswUpdate + - self.reconstruction.esw + ) + + # object update + if ( + self.params.objectTVregSwitch + and loop % self.params.objectTVfreq == 0 + ): + object_patch = self.objectPatchUpdate_TV( + objectPatch, + DELTA, + ) + 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) + ) + + # probe update + weight = 1 + if self.params.weigh_probe_updates_by_intensity: + weight = self.experimentalData.relative_intensity( + positionIndex + ) + + self.reconstruction.probe = self.probeUpdate( + objectPatch, + DELTA, + weight, + ) + + # position correction + if self.params.positionCorrectionSwitch: + self.positionCorrection( + objectPatch, + positionIndex, + sy, + sx, + ) + + # momentum updates + if np.random.rand(1) > 0.95: + self.objectMomentumUpdate() + self.probeMomentumUpdate() + + self.getErrorMetrics() + self.applyConstraints(loop) + + def objectMomentumUpdate(self): + 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 = ( + gradient + self.frictionM * self.reconstruction.objectMomentum + ) + self.reconstruction.object = ( + self.reconstruction.object + - self.feedbackM * self.reconstruction.objectMomentum + ) + self.reconstruction.objectBuffer = self.reconstruction.object.copy() + + def probeMomentumUpdate(self): + 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 = ( + gradient + self.frictionM * self.reconstruction.probeMomentum + ) + self.reconstruction.probe = ( + self.reconstruction.probe + - self.feedbackM * self.reconstruction.probeMomentum + ) + self.reconstruction.probeBuffer = self.reconstruction.probe.copy() + + def objectPatchUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray): + 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) + 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, weight: float): + 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) + 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 + weight * self.betaProbe * xp.sum( + frac * DELTA, axis=1, keepdims=True + ) + return r diff --git a/PtyLab/Params/Params.py b/PtyLab/Params/Params.py index cc77de05..17d3bff8 100644 --- a/PtyLab/Params/Params.py +++ b/PtyLab/Params/Params.py @@ -153,6 +153,12 @@ def __init__(self): # map a change in positions to a change in z. Experimental, do not use self.map_position_to_z_change = False + # purity based Zscan + self.purityZScanSwitch = False + self.purityZScanRange = 100e-6 + self.purityZScanPoints = 11 + self.purityZScanPlot = True + self.OPRP = False # Default values of all OPR parameters diff --git a/example_scripts/exampleReconstructionCPM.py b/example_scripts/exampleReconstructionCPM.py index 280d9dc2..7b03e42c 100644 --- a/example_scripts/exampleReconstructionCPM.py +++ b/example_scripts/exampleReconstructionCPM.py @@ -148,7 +148,7 @@ # Params = Params() ## main parameters params.positionOrder = "random" # 'sequential' or 'random' -params.propagatorType = "ASP" # Fresnel'# 'Fresnel' #aunhofer' # Fraunhofer Fresnel ASP scaledASP polychromeASP scaledPolychromeASP +params.propagatorType = "Fraunhofer" # Fresnel'# 'Fresnel' #aunhofer' # Fraunhofer Fresnel ASP scaledASP polychromeASP scaledPolychromeASP params.positionCorrectionSwitch = False params.modulusEnforcedProbeSwitch = False diff --git a/example_scripts/example_purity_Zscan_real_Data.py b/example_scripts/example_purity_Zscan_real_Data.py new file mode 100644 index 00000000..8fc0d4f1 --- /dev/null +++ b/example_scripts/example_purity_Zscan_real_Data.py @@ -0,0 +1,302 @@ +import argparse +import logging + +import matplotlib + +matplotlib.use("qt5agg") + +import matplotlib.pyplot as plt +import numpy as np + +import PtyLab +from PtyLab import Engines +from PtyLab.io import getExampleDataFolder + + +logging.basicConfig(level=logging.INFO) + + +""" +Purity-based axial calibration example. + +Workflow +-------- +1. Load a conventional ptychography dataset. +2. Initialize a multimode reconstruction. +3. Run purityPIE over a fixed z range. +4. Select the z value that gives the highest reconstructed probe purity. +5. Optionally continue the reconstruction with mPIE from the best-z state. +""" + + +# ------------------------------------------------------------------------- +# Command-line arguments +# ------------------------------------------------------------------------- + +import logging +import os + +import matplotlib +matplotlib.use("QtAgg") +import matplotlib.pyplot as plt + +import PtyLab +from PtyLab import Engines + + +logging.basicConfig(level=logging.INFO) + +logging.basicConfig(level=logging.INFO) + + +# ------------------------------------------------------------------------- +# Data path +# ------------------------------------------------------------------------- + +PATH = r"\\dionysios.iap.uni-jena.de\faserlaser$\AG_Imaging/2023_XUV_Bio-imaging\XUV imaging\_RAWDATA" +subdir = "20240522/176" + +filePath = os.path.join( + PATH, + f"{subdir}/dp/dp_processed.hdf5", +) + +filePath_recon = os.path.join( + PATH, + f"{subdir}/recons/seed.hdf5", +) + + +# ------------------------------------------------------------------------- +# Experimental parameters +# ------------------------------------------------------------------------- + +initial_z = 32.0e-3 +wavelength = 13.5e-9 +bin_factor = 1 + +# ------------------------------------------------------------------------- +# Load data +# ------------------------------------------------------------------------- + +experimentalData, reconstruction, params, monitor, _ = ( + PtyLab.easyInitialize( + filePath, + operationMode="CPM", + ) +) + +experimentalData.setOrientation(4) + +print("Loaded file:", experimentalData.filename) +print("Ptychogram shape:", experimentalData.ptychogram.shape) +print("Initial z:", reconstruction.zo) +print("Wavelength:", reconstruction.wavelength) + +# ------------------------------------------------------------------------- +# Optional initial z offset for testing +# ------------------------------------------------------------------------- + +# Uncomment this to deliberately start from an incorrect distance. +# +# experimentalData.zo += 100e-6 +# reconstruction.zo += 100e-6 + + + +experimentalData.zo = initial_z +reconstruction.zo = initial_z + +print(f"Initial z guess: {reconstruction.zo * 1e3:.6f} mm") + + +# ------------------------------------------------------------------------- +# Reconstruction dimensions +# ------------------------------------------------------------------------- + +# Purity-based calibration requires multiple probe modes. +reconstruction.npsm = 4 + +reconstruction.nosm = 1 +reconstruction.nlambda = 1 +reconstruction.nslice = 1 + + +# ------------------------------------------------------------------------- +# Object and probe initialization +# ------------------------------------------------------------------------- + +reconstruction.initialProbe = "recon" +reconstruction.initialProbe_filename = filePath_recon +reconstruction.initialObject = "ones" + +reconstruction.initializeObjectProbe() + + +# Optional quadratic phase initialization, if required for your dataset. +# +# reconstruction.probe *= np.exp( +# 1.0j +# * 2 +# * np.pi +# / (reconstruction.wavelength * reconstruction.zo * 2) +# * (reconstruction.Xp**2 + reconstruction.Yp**2) +# / 2 +# ) + + +reconstruction.describe_reconstruction() + + +# ------------------------------------------------------------------------- +# Monitor settings +# ------------------------------------------------------------------------- + +monitor.figureUpdateFrequency = 5 +monitor.objectPlot = "complex" +monitor.verboseLevel = "low" + +monitor.objectZoom = None +monitor.probeZoom = None + + +# ------------------------------------------------------------------------- +# General reconstruction parameters +# ------------------------------------------------------------------------- + +params.positionOrder = "random" + +# Start with the propagator used for your actual dataset. +params.propagatorType = "Fraunhofer" + +params.gpuSwitch = True +params.fftshiftSwitch = False + +params.intensityConstraint = "standard" + +params.positionCorrectionSwitch = False + +params.modulusEnforcedProbeSwitch = False +params.probePowerCorrectionSwitch = True + +params.probeSmoothenessSwitch = True +params.probeSmoothnessAleph = 1e-2 +params.probeSmoothenessWidth = 10 + +params.comStabilizationSwitch = 10 + +# purityPIE performs an explicit orthogonalization after each candidate-z +# reconstruction. Automatic orthogonalization during the internal +# reconstruction is temporarily disabled by purityZScan(). +params.orthogonalizationSwitch = True +params.orthogonalizationFrequency = 10 + +params.absorbingProbeBoundary = False +params.objectContrastSwitch = False +params.absObjectSwitch = False +params.backgroundModeSwitch = False + +params.couplingSwitch = True +params.couplingAleph = 1 + +params.TV_autofocus = False +params.l2reg = False + + +# ------------------------------------------------------------------------- +# Purity-based z calibration parameters +# ------------------------------------------------------------------------- + +params.purityZScanSwitch = True + +# Half-range around the current reconstruction.zo. +# +# Example: +# reconstruction.zo = 50 mm +# purityZScanRange = 100 µm +# +# scan range: +# 49.9 mm ... 50.1 mm +params.purityZScanRange = 2e-3 + +# Number of candidate z values. +params.purityZScanPoints = 10 + +# Plot purity versus z after the scan. +params.purityZScanPlot = True + + +monitor.describe_parameters(params) + + +# ------------------------------------------------------------------------- +# Run purity-based axial calibration +# ------------------------------------------------------------------------- + +if params.purityZScanSwitch: + + purity_engine = Engines.purityPIE( + reconstruction, + experimentalData, + params, + monitor, + ) + + # Number of mPIE iterations used to evaluate EACH candidate z. + purity_engine.numIterations = 5 + + best_z, best_purity = purity_engine.reconstruct() + + print() + print("Purity-based axial calibration finished") + print("---------------------------------------") + print(f"Initial z : {purity_engine.purityZInitialGuess * 1e3:.6f} mm") + print(f"Best z : {best_z * 1e3:.6f} mm") + print( + f"Delta z : " + f"{(best_z - purity_engine.purityZInitialGuess) * 1e6:.3f} um" + ) + print(f"Best purity: {best_purity:.6f}") + + print() + print("z scan:") + for z, purity in zip( + purity_engine.purityZValues, + purity_engine.purityZMetrics, + ): + print( + f"z = {z * 1e3:.6f} mm, " + f"delta z = {(z - purity_engine.purityZInitialGuess) * 1e6:+.2f} um, " + f"purity = {purity:.6f}" + ) + +# ------------------------------------------------------------------------- +# Optional: continue reconstruction using the best-z state +# ------------------------------------------------------------------------- + +continue_with_mPIE = False + +if continue_with_mPIE: + + mPIE_engine = Engines.mPIE( + reconstruction, + experimentalData, + params, + monitor, + ) + + mPIE_engine.numIterations = 50 + + mPIE_engine.reconstruct() + + +# ------------------------------------------------------------------------- +# Save / display +# ------------------------------------------------------------------------- + +# reconstruction.saveResults( +# f"{getExampleDataFolder()}/purity_calibrated_recon.hdf5", +# squeeze=False, +# ) + +plt.show(block=True) \ No newline at end of file From b99f4e874cca0c4c74a11af54fda4851036551cf Mon Sep 17 00:00:00 2001 From: Chang Liu Date: Sat, 3 Oct 2026 00:31:41 +0200 Subject: [PATCH 03/11] Change the backend of purityPIE to directly call mPIE, this make sure exact same initialization of each z. --- PtyLab/Engines/purityPIE.py | 705 ++++++++++---------------------- tests/Engines/test_purityPIE.py | 174 ++++++++ 2 files changed, 396 insertions(+), 483 deletions(-) create mode 100644 tests/Engines/test_purityPIE.py diff --git a/PtyLab/Engines/purityPIE.py b/PtyLab/Engines/purityPIE.py index 4824d61a..a637139a 100644 --- a/PtyLab/Engines/purityPIE.py +++ b/PtyLab/Engines/purityPIE.py @@ -24,8 +24,12 @@ from PtyLab.utils.gpuUtils import asNumpyArray, getArrayModule from PtyLab.utils.utils import fft2c, ifft2c +from copy import deepcopy -class purityPIE(BaseEngine): +from PtyLab.Monitor.Monitor import DummyMonitor + +from PtyLab.Engines.mPIE import mPIE +class purityPIE(mPIE): ''' purityPIE extends mPIE-style reconstruction with purity-based axial calibration. ''' @@ -36,101 +40,22 @@ def __init__( params: Params, monitor: Monitor, ): - """ - 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. + super().__init__( + reconstruction, + experimentalData, + params, + monitor, + ) - 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("purityPIE") self.logger.info("Successfully created purityPIE engine") - self.logger.info("Wavelength attribute: %s", self.reconstruction.wavelength) + self.logger.info( + "Wavelength attribute: %s", + self.reconstruction.wavelength, + ) - self.initializeReconstructionParams() - self.params.momentumAcceleration = True self.name = "purityPIE" - @property - def keepPatches(self): - """ - Whether to store the reconstructed object patch for every scan position. - - 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") - - @keepPatches.setter - def keepPatches(self, keep_them): - - if keep_them: - self.logger.info("Keeping patches!") - self.patches = np.zeros( - ( - self.experimentalData.ptychogram.shape[0], - *self.reconstruction.shape_O, - ), - np.complex64, - ) - else: - self.logger.info("Not keeping patches") - if hasattr(self, "patches"): - del self.patches - - def initializeReconstructionParams(self): - """ - 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 - self.betaObject = 0.25 - self.alphaProbe = 0.1 # probe regularization - self.alphaObject = 0.1 # object regularization - self.feedbackM = 0.3 # feedback - self.frictionM = 0.7 # friction - self.numIterations = 50 - - # 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.reconstruction.probeWindow = np.abs(self.reconstruction.probe) - def reconstruct( self, experimentalData=None, @@ -186,13 +111,63 @@ def reconstruct( return best_z, best_purity + def _createCandidateReconstruction(self, z): + """ + Create an independent reconstruction state for one candidate z. + + The candidate inherits the user's reconstruction settings but is built + from a fresh geometry at the requested propagation distance. + """ + + # Independent copies so one candidate cannot modify another + candidate_data = deepcopy(self.experimentalData) + candidate_params = deepcopy(self.params) + + # Set candidate geometry before creating Reconstruction + candidate_data.zo = z + + candidate_reconstruction = Reconstruction( + candidate_data, + candidate_params, + ) + + # Copy user-defined reconstruction settings + candidate_reconstruction.npsm = self.reconstruction.npsm + candidate_reconstruction.nosm = self.reconstruction.nosm + candidate_reconstruction.nlambda = self.reconstruction.nlambda + candidate_reconstruction.nslice = self.reconstruction.nslice + + #candidate_reconstruction.No = self.reconstruction.No + + candidate_reconstruction.initialProbe = self.reconstruction.initialProbe + candidate_reconstruction.initialObject = self.reconstruction.initialObject + + if hasattr(self.reconstruction, "initialProbe_filename"): + candidate_reconstruction.initialProbe_filename = ( + self.reconstruction.initialProbe_filename + ) + + if hasattr(self.reconstruction, "initialObject_filename"): + candidate_reconstruction.initialObject_filename = ( + self.reconstruction.initialObject_filename + ) + + # Fresh object/probe for this candidate geometry + candidate_reconstruction.initializeObjectProbe() + + return ( + candidate_data, + candidate_reconstruction, + candidate_params, + ) + def purityZScan(self): """ Scan a fixed axial range and select the distance that maximizes reconstructed probe purity. - Every candidate z is evaluated from the same initial reconstruction - state using `numIterations` mPIE iterations. + Each candidate z is evaluated using an independent reconstruction + state and a standard mPIE reconstruction. """ if self.reconstruction.npsm < 2: @@ -201,7 +176,10 @@ def purityZScan(self): "(reconstruction.npsm >= 2)." ) - # Initial z guess defines the center of the scan + # ------------------------------------------------------------------ + # Scan grid + # ------------------------------------------------------------------ + z0 = self.reconstruction.zo self.purityZInitialGuess = z0 @@ -211,26 +189,16 @@ def purityZScan(self): self.params.purityZScanPoints, ) - # Save the common starting state - object0 = self.reconstruction.object.copy() - probe0 = self.reconstruction.probe.copy() - - objectMomentum0 = self.reconstruction.objectMomentum.copy() - probeMomentum0 = self.reconstruction.probeMomentum.copy() - - objectBuffer0 = self.reconstruction.objectBuffer.copy() - probeBuffer0 = self.reconstruction.probeBuffer.copy() - purity_values = [] best_purity = -np.inf best_z = z0 - best_state = None - # Do not allow orthogonalization during the internal mPIE iterations. - # Purity is evaluated once after each candidate reconstruction. - #orthogonalization_switch = self.params.orthogonalizationSwitch - #self.params.orthogonalizationSwitch = False + best_reconstruction = None + + # ------------------------------------------------------------------ + # Scan information + # ------------------------------------------------------------------ tqdm.tqdm.write("") tqdm.tqdm.write("Starting purity-based z scan") @@ -242,21 +210,11 @@ def purityZScan(self): ) tqdm.tqdm.write("") - #try: - for z_index, z in enumerate(z_values): - - # Restore exactly the same starting state for every candidate z - self.reconstruction.object = object0.copy() - self.reconstruction.probe = probe0.copy() - - self.reconstruction.objectMomentum = objectMomentum0.copy() - self.reconstruction.probeMomentum = probeMomentum0.copy() + # ------------------------------------------------------------------ + # Candidate-z loop + # ------------------------------------------------------------------ - self.reconstruction.objectBuffer = objectBuffer0.copy() - self.reconstruction.probeBuffer = probeBuffer0.copy() - - # Set candidate propagation distance - self.reconstruction.zo = z + for z_index, z in enumerate(z_values): dz_um = (z - z0) * 1e6 @@ -266,43 +224,99 @@ def purityZScan(self): f"(Δz = {dz_um:+.2f} µm)" ) - # Run mPIE reconstruction at this z - for loop in range(self.numIterations): - self._run_single_iteration(loop) - #self.showReconstruction(loop) + # -------------------------------------------------------------- + # Create a completely independent reconstruction for this z + # -------------------------------------------------------------- + #np.random.seed(0) + ( + candidate_data, + candidate_reconstruction, + candidate_params, + ) = self._createCandidateReconstruction(z) + + # Candidate reconstructions do not need their own GUI monitor + candidate_monitor = DummyMonitor() + + # -------------------------------------------------------------- + # Run standard mPIE + # -------------------------------------------------------------- + + candidate_engine = mPIE( + candidate_reconstruction, + candidate_data, + candidate_params, + candidate_monitor, + ) + + # Match the purityPIE reconstruction settings + candidate_engine.numIterations = self.numIterations + + candidate_engine.betaProbe = self.betaProbe + candidate_engine.betaObject = self.betaObject + + candidate_engine.alphaProbe = self.alphaProbe + candidate_engine.alphaObject = self.alphaObject - # Orthogonalize once and evaluate purity - self.orthogonalization() - + candidate_engine.feedbackM = self.feedbackM + candidate_engine.frictionM = self.frictionM + + candidate_engine.reconstruct() + + # -------------------------------------------------------------- + # Final modal decomposition and purity evaluation + # -------------------------------------------------------------- + + candidate_engine.orthogonalization() purity = float( np.asarray( - asNumpyArray(self.reconstruction.purityProbe) + asNumpyArray(candidate_reconstruction.purityProbe) ).squeeze() ) purity_values.append(purity) - # Force monitor to display the final state and final purity + + tqdm.tqdm.write( + f" Probe purity = {purity:.6f}" + ) + + # -------------------------------------------------------------- + # Update main monitor with the completed candidate state + # -------------------------------------------------------------- + + self.reconstruction.zo = z + self.reconstruction.object = ( + asNumpyArray(candidate_reconstruction.object).copy() + ) + self.reconstruction.probe = ( + asNumpyArray(candidate_reconstruction.probe).copy() + ) + + self.reconstruction.purityProbe = purity + + if hasattr(candidate_reconstruction, "error"): + self.reconstruction.error = np.asarray( + asNumpyArray(candidate_reconstruction.error) + ).copy() + + # Display only completed candidate reconstructions self.showReconstruction( self.numIterations - 1, force=True, ) - tqdm.tqdm.write( - f" Probe purity = {purity:.6f}" - ) + + # -------------------------------------------------------------- + # Keep best candidate + # -------------------------------------------------------------- if purity > best_purity: + best_purity = purity best_z = z - best_state = { - "object": self.reconstruction.object.copy(), - "probe": self.reconstruction.probe.copy(), - "objectMomentum": self.reconstruction.objectMomentum.copy(), - "probeMomentum": self.reconstruction.probeMomentum.copy(), - "objectBuffer": self.reconstruction.objectBuffer.copy(), - "probeBuffer": self.reconstruction.probeBuffer.copy(), - } + # Candidate objects are independent, so keeping this reference + # preserves the complete best reconstruction state. + best_reconstruction = candidate_reconstruction tqdm.tqdm.write( f" NEW BEST: z = {best_z * 1e3:.6f} mm, " @@ -311,188 +325,94 @@ def purityZScan(self): tqdm.tqdm.write("") - #finally: - # self.params.orthogonalizationSwitch = orthogonalization_switch + # ------------------------------------------------------------------ + # Safety check + # ------------------------------------------------------------------ - if best_state is None: + if best_reconstruction is None: raise RuntimeError( "Purity z scan did not produce a valid purity value." ) - # Restore best-z reconstruction state - self.reconstruction.object = best_state["object"] - self.reconstruction.probe = best_state["probe"] - - self.reconstruction.objectMomentum = best_state["objectMomentum"] - self.reconstruction.probeMomentum = best_state["probeMomentum"] - - self.reconstruction.objectBuffer = best_state["objectBuffer"] - self.reconstruction.probeBuffer = best_state["probeBuffer"] + # ------------------------------------------------------------------ + # Restore the best candidate into the original Reconstruction object + # + # Keep the original object identity because external scripts may still + # hold a reference to self.reconstruction. + # ------------------------------------------------------------------ self.reconstruction.zo = best_z - # Store scan results - self.purityZValues = z_values - self.purityZMetrics = np.asarray(purity_values) - self.purityBestZ = best_z - self.purityBestValue = best_purity - - # Final summary - tqdm.tqdm.write("=" * 60) - tqdm.tqdm.write("Purity-based z scan finished") - tqdm.tqdm.write( - f"Best z = {best_z * 1e3:.6f} mm" - ) - tqdm.tqdm.write( - f"Best Δz = {(best_z - z0) * 1e6:+.2f} µm" - ) - tqdm.tqdm.write( - f"Best purity = {best_purity:.6f}" - ) - tqdm.tqdm.write("=" * 60) - - return best_z, best_purity - - def purityZScan_X(self): - """ - Scan a fixed axial range and select the distance that maximizes - reconstructed probe purity. - - Every candidate z is evaluated from the same initial reconstruction - state using `numIterations` mPIE iterations. - """ - - if self.reconstruction.npsm < 2: - raise ValueError( - "Purity-based z scanning requires at least two probe modes " - "(reconstruction.npsm >= 2)." - ) - - # Initial z guess defines the center of the scan - z0 = self.reconstruction.zo - self.purityZInitialGuess = z0 - - z_values = np.linspace( - z0 - self.params.purityZScanRange, - z0 + self.params.purityZScanRange, - self.params.purityZScanPoints, + self.reconstruction.object = ( + asNumpyArray(best_reconstruction.object).copy() ) - # Save the common starting state - object0 = self.reconstruction.object.copy() - probe0 = self.reconstruction.probe.copy() - - objectMomentum0 = self.reconstruction.objectMomentum.copy() - probeMomentum0 = self.reconstruction.probeMomentum.copy() - - objectBuffer0 = self.reconstruction.objectBuffer.copy() - probeBuffer0 = self.reconstruction.probeBuffer.copy() - - purity_values = [] - - best_purity = -np.inf - best_z = z0 - best_state = None - - # Do not allow orthogonalization during the internal mPIE iterations. - # Purity is evaluated once after each candidate reconstruction. - orthogonalization_switch = self.params.orthogonalizationSwitch - self.params.orthogonalizationSwitch = False - - tqdm.tqdm.write("") - tqdm.tqdm.write("Starting purity-based z scan") - tqdm.tqdm.write( - f"Initial z = {z0 * 1e3:.6f} mm | " - f"range = ±{self.params.purityZScanRange * 1e6:.1f} µm | " - f"points = {len(z_values)} | " - f"iterations/z = {self.numIterations}" + self.reconstruction.probe = ( + asNumpyArray(best_reconstruction.probe).copy() ) - tqdm.tqdm.write("") - - try: - for z_index, z in enumerate(z_values): - - self._resetCandidateState(z) - #np.random.seed(0) - - - dz_um = (z - z0) * 1e6 - - tqdm.tqdm.write( - f"[{z_index + 1}/{len(z_values)}] " - f"Reconstructing z = {z * 1e3:.6f} mm " - f"(Δz = {dz_um:+.2f} µm)" - ) - - # Run mPIE reconstruction at this z - for loop in range(self.numIterations): - self._run_single_iteration(loop) - self.showReconstruction(loop) - - # Orthogonalize once and evaluate purity - self.orthogonalization() - - - purity = float( - np.asarray( - asNumpyArray(self.reconstruction.purityProbe) - ).squeeze() - ) - self.showReconstruction(0) - - purity_values.append(purity) - tqdm.tqdm.write( - f" Probe purity = {purity:.6f}" - ) - - if purity > best_purity: - best_purity = purity - best_z = z - - best_state = { - "object": self.reconstruction.object.copy(), - "probe": self.reconstruction.probe.copy(), - "objectMomentum": self.reconstruction.objectMomentum.copy(), - "probeMomentum": self.reconstruction.probeMomentum.copy(), - "objectBuffer": self.reconstruction.objectBuffer.copy(), - "probeBuffer": self.reconstruction.probeBuffer.copy(), - } + self.reconstruction.purityProbe = best_purity - tqdm.tqdm.write( - f" NEW BEST: z = {best_z * 1e3:.6f} mm, " - f"purity = {best_purity:.6f}" - ) - - tqdm.tqdm.write("") - - finally: - self.params.orthogonalizationSwitch = orthogonalization_switch + if hasattr(best_reconstruction, "encoder_corrected"): + self.reconstruction.encoder_corrected = ( + asNumpyArray( + best_reconstruction.encoder_corrected + ).copy() + ) - if best_state is None: - raise RuntimeError( - "Purity z scan did not produce a valid purity value." + if hasattr(best_reconstruction, "objectMomentum"): + self.reconstruction.objectMomentum = ( + asNumpyArray( + best_reconstruction.objectMomentum + ).copy() ) - # Restore best-z reconstruction state - self.reconstruction.object = best_state["object"] - self.reconstruction.probe = best_state["probe"] + if hasattr(best_reconstruction, "probeMomentum"): + self.reconstruction.probeMomentum = ( + asNumpyArray( + best_reconstruction.probeMomentum + ).copy() + ) - self.reconstruction.objectMomentum = best_state["objectMomentum"] - self.reconstruction.probeMomentum = best_state["probeMomentum"] + if hasattr(best_reconstruction, "objectBuffer"): + self.reconstruction.objectBuffer = ( + asNumpyArray( + best_reconstruction.objectBuffer + ).copy() + ) - self.reconstruction.objectBuffer = best_state["objectBuffer"] - self.reconstruction.probeBuffer = best_state["probeBuffer"] + if hasattr(best_reconstruction, "probeBuffer"): + self.reconstruction.probeBuffer = ( + asNumpyArray( + best_reconstruction.probeBuffer + ).copy() + ) - self.reconstruction.zo = best_z + if hasattr(best_reconstruction, "error"): + self.reconstruction.error = np.asarray( + asNumpyArray(best_reconstruction.error) + ).copy() + # ------------------------------------------------------------------ # Store scan results - self.purityZValues = z_values + # ------------------------------------------------------------------ + + self.purityZValues = np.asarray(z_values) self.purityZMetrics = np.asarray(purity_values) + self.purityBestZ = best_z self.purityBestValue = best_purity - # Final summary + # Final monitor update with the selected best candidate + self.showReconstruction( + self.numIterations - 1, + force=True, + ) + + # ------------------------------------------------------------------ + # Summary + # ------------------------------------------------------------------ + tqdm.tqdm.write("=" * 60) tqdm.tqdm.write("Purity-based z scan finished") tqdm.tqdm.write( @@ -508,7 +428,7 @@ def purityZScan_X(self): return best_z, best_purity - def _resetCandidateState(self, z): + """ Reset the reconstruction state for one candidate z. @@ -628,7 +548,7 @@ def plotPurityZScan(self, show=True): return fig, ax - def _run_single_iteration(self, loop): + """ Run one complete mPIE reconstruction iteration. """ @@ -716,51 +636,6 @@ def _run_single_iteration(self, loop): self.getErrorMetrics() self.applyConstraints(loop) - def objectMomentumUpdate(self): - 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 = ( - gradient + self.frictionM * self.reconstruction.objectMomentum - ) - self.reconstruction.object = ( - self.reconstruction.object - - self.feedbackM * self.reconstruction.objectMomentum - ) - self.reconstruction.objectBuffer = self.reconstruction.object.copy() - - def probeMomentumUpdate(self): r""" Apply the mPIE momentum update to the reconstructed probe. Similar to `objectMomentumUpdate()`. @@ -778,140 +653,4 @@ def probeMomentumUpdate(self): ) self.reconstruction.probeBuffer = self.reconstruction.probe.copy() - def objectPatchUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray): - 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) - 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, weight: float): - 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) - 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 + weight * self.betaProbe * xp.sum( - frac * DELTA, axis=1, keepdims=True - ) - return r + \ No newline at end of file diff --git a/tests/Engines/test_purityPIE.py b/tests/Engines/test_purityPIE.py new file mode 100644 index 00000000..d0e20b8e --- /dev/null +++ b/tests/Engines/test_purityPIE.py @@ -0,0 +1,174 @@ +from copy import deepcopy + +import numpy as np +import pytest +from numpy.testing import assert_allclose + +import PtyLab +from PtyLab import Engines +from PtyLab.Reconstruction.Reconstruction import Reconstruction + + +Z0 = 50e-3 +SCAN_RANGE = 2e-3 +SCAN_POINTS = 3 +N_ITERATIONS = 5 +RNG_SEED = 0 + + +def configure_reconstruction(reconstruction): + reconstruction.npsm = 4 + reconstruction.nosm = 1 + reconstruction.nlambda = 1 + reconstruction.nslice = 1 + + reconstruction.initialProbe = "circ" + reconstruction.initialObject = "ones" + + reconstruction.initializeObjectProbe() + + +def configure_params(params): + params.positionOrder = "random" + params.propagatorType = "Fraunhofer" + + params.gpuSwitch = False + params.fftshiftSwitch = False + + params.intensityConstraint = "standard" + + params.positionCorrectionSwitch = False + params.modulusEnforcedProbeSwitch = False + params.probePowerCorrectionSwitch = True + + params.probeSmoothenessSwitch = True + params.probeSmoothnessAleph = 1e-2 + params.probeSmoothenessWidth = 10 + + params.comStabilizationSwitch = 10 + + # Keep normal mPIE behavior + params.orthogonalizationSwitch = True + params.orthogonalizationFrequency = 10 + + params.absorbingProbeBoundary = False + + params.objectContrastSwitch = False + params.absObjectSwitch = False + params.backgroundModeSwitch = False + + params.couplingSwitch = True + params.couplingAleph = 1 + + params.TV_autofocus = False + params.l2reg = False + + +def create_fresh_reference( + source_data, + source_reconstruction, + source_params, + z, +): + """ + Independently create one fresh reconstruction at candidate z. + This deliberately does NOT use purityPIE's helper. + """ + + candidate_data = deepcopy(source_data) + candidate_params = deepcopy(source_params) + + candidate_data.zo = z + + candidate_reconstruction = Reconstruction( + candidate_data, + candidate_params, + ) + + candidate_reconstruction.npsm = source_reconstruction.npsm + candidate_reconstruction.nosm = source_reconstruction.nosm + candidate_reconstruction.nlambda = source_reconstruction.nlambda + candidate_reconstruction.nslice = source_reconstruction.nslice + + candidate_reconstruction.initialProbe = ( + source_reconstruction.initialProbe + ) + candidate_reconstruction.initialObject = ( + source_reconstruction.initialObject + ) + + if hasattr(source_reconstruction, "initialProbe_filename"): + candidate_reconstruction.initialProbe_filename = ( + source_reconstruction.initialProbe_filename + ) + + if hasattr(source_reconstruction, "initialObject_filename"): + candidate_reconstruction.initialObject_filename = ( + source_reconstruction.initialObject_filename + ) + + candidate_reconstruction.initializeObjectProbe() + + return ( + candidate_data, + candidate_reconstruction, + candidate_params, + ) + + +@pytest.fixture +def scan_setup(generate_simu_hdf5): + """Create a CPU reconstruction and restore global RNG state after the test.""" + random_state = np.random.get_state() + try: + np.random.seed(RNG_SEED) + data, reconstruction, params, monitor, _ = PtyLab.easyInitialize( + generate_simu_hdf5, + operationMode="CPM", + dummyMonitor=True, + ) + reconstruction.zo = Z0 + configure_reconstruction(reconstruction) + configure_params(params) + params.purityZScanRange = SCAN_RANGE + params.purityZScanPoints = SCAN_POINTS + params.purityZScanPlot = False + yield data, reconstruction, params, monitor + finally: + np.random.set_state(random_state) + + +def test_fresh_curve_matches_new_purity_scan(scan_setup): + """The scan must match independent mPIE runs and select their best distance.""" + data, reconstruction, params, monitor = scan_setup + z_values = np.linspace(Z0 - SCAN_RANGE, Z0 + SCAN_RANGE, SCAN_POINTS) + + # Seed once per complete scan; candidates consume the RNG sequence naturally. + np.random.seed(RNG_SEED) + fresh_purities = [] + for z in z_values: + fresh_data, fresh_reconstruction, fresh_params = create_fresh_reference( + data, reconstruction, params, z + ) + fresh_engine = Engines.mPIE( + fresh_reconstruction, fresh_data, fresh_params, monitor + ) + fresh_engine.numIterations = N_ITERATIONS + fresh_engine.reconstruct() + fresh_engine.orthogonalization() + fresh_purities.append(float(np.asarray(fresh_reconstruction.purityProbe).squeeze())) + + fresh_purities = np.asarray(fresh_purities) + + # Use exactly the same random sequence for the purityPIE scan. + np.random.seed(RNG_SEED) + purity_engine = Engines.purityPIE(reconstruction, data, params, monitor) + purity_engine.numIterations = N_ITERATIONS + best_z, best_purity = purity_engine.reconstruct() + + assert_allclose(purity_engine.purityZValues, z_values) + assert_allclose( + purity_engine.purityZMetrics, fresh_purities, rtol=1e-6, atol=1e-7 + ) + assert best_z == pytest.approx(z_values[np.argmax(fresh_purities)]) + assert best_purity == pytest.approx(np.max(fresh_purities), rel=1e-6, abs=1e-7) From c8ad57318a3d81e475d9a536940ea01d00954237 Mon Sep 17 00:00:00 2001 From: Chang Liu Date: Sat, 3 Oct 2026 12:04:47 +0200 Subject: [PATCH 04/11] clarified and sort the example script --- ...ta.py => example_purity_Zscan_yourData.py} | 118 ++++-------------- 1 file changed, 24 insertions(+), 94 deletions(-) rename example_scripts/{example_purity_Zscan_real_Data.py => example_purity_Zscan_yourData.py} (67%) diff --git a/example_scripts/example_purity_Zscan_real_Data.py b/example_scripts/example_purity_Zscan_yourData.py similarity index 67% rename from example_scripts/example_purity_Zscan_real_Data.py rename to example_scripts/example_purity_Zscan_yourData.py index 8fc0d4f1..bbbb0027 100644 --- a/example_scripts/example_purity_Zscan_real_Data.py +++ b/example_scripts/example_purity_Zscan_yourData.py @@ -1,80 +1,52 @@ -import argparse import logging +from tkinter import Tk, filedialog import matplotlib - -matplotlib.use("qt5agg") +matplotlib.use("QtAgg") import matplotlib.pyplot as plt import numpy as np import PtyLab from PtyLab import Engines -from PtyLab.io import getExampleDataFolder logging.basicConfig(level=logging.INFO) -""" -Purity-based axial calibration example. - -Workflow --------- -1. Load a conventional ptychography dataset. -2. Initialize a multimode reconstruction. -3. Run purityPIE over a fixed z range. -4. Select the z value that gives the highest reconstructed probe purity. -5. Optionally continue the reconstruction with mPIE from the best-z state. -""" +def select_hdf5_file(title): + root = Tk() + root.withdraw() + root.attributes("-topmost", True) + file_path = filedialog.askopenfilename( + title=title, + filetypes=[ + ("HDF5 files", "*.hdf5"), + ("All files", "*.*"), + ], + ) -# ------------------------------------------------------------------------- -# Command-line arguments -# ------------------------------------------------------------------------- - -import logging -import os - -import matplotlib -matplotlib.use("QtAgg") -import matplotlib.pyplot as plt - -import PtyLab -from PtyLab import Engines - + root.destroy() -logging.basicConfig(level=logging.INFO) + if not file_path: + raise RuntimeError(f"No file selected: {title}") -logging.basicConfig(level=logging.INFO) + return file_path # ------------------------------------------------------------------------- -# Data path +# Select input files # ------------------------------------------------------------------------- -PATH = r"\\dionysios.iap.uni-jena.de\faserlaser$\AG_Imaging/2023_XUV_Bio-imaging\XUV imaging\_RAWDATA" -subdir = "20240522/176" - -filePath = os.path.join( - PATH, - f"{subdir}/dp/dp_processed.hdf5", +filePath = select_hdf5_file( + "Select ptychography dataset" ) -filePath_recon = os.path.join( - PATH, - f"{subdir}/recons/seed.hdf5", +filePath_recon = select_hdf5_file( + "Select seed reconstruction" ) - -# ------------------------------------------------------------------------- -# Experimental parameters -# ------------------------------------------------------------------------- - -initial_z = 32.0e-3 -wavelength = 13.5e-9 -bin_factor = 1 - # ------------------------------------------------------------------------- # Load data # ------------------------------------------------------------------------- @@ -86,26 +58,15 @@ ) ) -experimentalData.setOrientation(4) +#experimentalData.setOrientation(0) print("Loaded file:", experimentalData.filename) print("Ptychogram shape:", experimentalData.ptychogram.shape) print("Initial z:", reconstruction.zo) print("Wavelength:", reconstruction.wavelength) -# ------------------------------------------------------------------------- -# Optional initial z offset for testing -# ------------------------------------------------------------------------- - -# Uncomment this to deliberately start from an incorrect distance. -# -# experimentalData.zo += 100e-6 -# reconstruction.zo += 100e-6 - - -experimentalData.zo = initial_z -reconstruction.zo = initial_z +initial_z = reconstruction.zo print(f"Initial z guess: {reconstruction.zo * 1e3:.6f} mm") @@ -132,19 +93,6 @@ reconstruction.initializeObjectProbe() - -# Optional quadratic phase initialization, if required for your dataset. -# -# reconstruction.probe *= np.exp( -# 1.0j -# * 2 -# * np.pi -# / (reconstruction.wavelength * reconstruction.zo * 2) -# * (reconstruction.Xp**2 + reconstruction.Yp**2) -# / 2 -# ) - - reconstruction.describe_reconstruction() @@ -170,38 +118,20 @@ params.propagatorType = "Fraunhofer" params.gpuSwitch = True -params.fftshiftSwitch = False params.intensityConstraint = "standard" -params.positionCorrectionSwitch = False - -params.modulusEnforcedProbeSwitch = False params.probePowerCorrectionSwitch = True params.probeSmoothenessSwitch = True params.probeSmoothnessAleph = 1e-2 params.probeSmoothenessWidth = 10 - params.comStabilizationSwitch = 10 - -# purityPIE performs an explicit orthogonalization after each candidate-z -# reconstruction. Automatic orthogonalization during the internal -# reconstruction is temporarily disabled by purityZScan(). params.orthogonalizationSwitch = True params.orthogonalizationFrequency = 10 - -params.absorbingProbeBoundary = False -params.objectContrastSwitch = False -params.absObjectSwitch = False -params.backgroundModeSwitch = False - params.couplingSwitch = True params.couplingAleph = 1 -params.TV_autofocus = False -params.l2reg = False - # ------------------------------------------------------------------------- # Purity-based z calibration parameters From 264177913f5dd5e8e3899c51a583128fc9c631d0 Mon Sep 17 00:00:00 2001 From: Chang Liu Date: Sat, 3 Oct 2026 20:55:11 +0200 Subject: [PATCH 05/11] Add adaptive search function for purity-based axial calibratoin --- PtyLab/Engines/purityPIE.py | 717 +++++++++++++++++- PtyLab/Params/Params.py | 12 +- .../example_purity_Zscan_yourData.py | 152 +++- tests/Engines/test_purity_adaptive_scan.py | 417 ++++++++++ 4 files changed, 1274 insertions(+), 24 deletions(-) create mode 100644 tests/Engines/test_purity_adaptive_scan.py diff --git a/PtyLab/Engines/purityPIE.py b/PtyLab/Engines/purityPIE.py index a637139a..ffecdc0f 100644 --- a/PtyLab/Engines/purityPIE.py +++ b/PtyLab/Engines/purityPIE.py @@ -95,9 +95,12 @@ def reconstruct( # Reset momentum buffers after preparation self.reconstruction.objectBuffer = self.reconstruction.object.copy() self.reconstruction.probeBuffer = self.reconstruction.probe.copy() - + # Purity-based axial calibration - best_z, best_purity = self.purityZScan() + if self.params.purityZAdaptive: + best_z, best_purity = self.adaptivePurityZScan() + else: + best_z, best_purity = self.purityZScan() # Optional scan visualization if self.params.purityZScanPlot: @@ -161,7 +164,7 @@ def _createCandidateReconstruction(self, z): candidate_params, ) - def purityZScan(self): + """ Scan a fixed axial range and select the distance that maximizes reconstructed probe purity. @@ -470,6 +473,712 @@ def purityZScan(self): self._prepareReconstruction() + def _evaluatePurityAtZ(self, z): + """ + Evaluate probe purity for one candidate propagation distance. + + A fresh reconstruction is created at the requested z value and + reconstructed using the standard mPIE engine. Probe modes are + orthogonalized after reconstruction and the resulting probe purity + is returned together with the candidate reconstruction state. + + Parameters + ---------- + z : float + Candidate propagation distance in meters. + + Returns + ------- + purity : float + Reconstructed probe purity at the candidate distance. + + candidate_reconstruction : Reconstruction + Reconstruction state obtained at the candidate distance. + """ + + # ------------------------------------------------------------------ + # Create independent candidate reconstruction + # ------------------------------------------------------------------ + + ( + candidate_data, + candidate_reconstruction, + candidate_params, + ) = self._createCandidateReconstruction(z) + + candidate_monitor = DummyMonitor() + + # ------------------------------------------------------------------ + # Run standard mPIE + # ------------------------------------------------------------------ + + candidate_engine = mPIE( + candidate_reconstruction, + candidate_data, + candidate_params, + candidate_monitor, + ) + + # Match the purityPIE engine settings + candidate_engine.numIterations = self.numIterations + + candidate_engine.betaProbe = self.betaProbe + candidate_engine.betaObject = self.betaObject + + candidate_engine.alphaProbe = self.alphaProbe + candidate_engine.alphaObject = self.alphaObject + + candidate_engine.feedbackM = self.feedbackM + candidate_engine.frictionM = self.frictionM + + candidate_engine.reconstruct() + + # ------------------------------------------------------------------ + # Final modal decomposition and purity evaluation + # ------------------------------------------------------------------ + + candidate_engine.orthogonalization() + + purity = float( + np.asarray( + asNumpyArray( + candidate_reconstruction.purityProbe + ) + ).squeeze() + ) + + return purity, candidate_reconstruction + + def purityZScan(self): + """ + Scan a fixed axial range and select the distance that maximizes + reconstructed probe purity. + + Each candidate z is evaluated using an independent reconstruction + state and a standard mPIE reconstruction. + """ + + if self.reconstruction.npsm < 2: + raise ValueError( + "Purity-based z scanning requires at least two probe modes " + "(reconstruction.npsm >= 2)." + ) + + # ------------------------------------------------------------------ + # Scan grid + # ------------------------------------------------------------------ + + z0 = self.reconstruction.zo + self.purityZInitialGuess = z0 + + z_values = np.linspace( + z0 - self.params.purityZScanRange, + z0 + self.params.purityZScanRange, + self.params.purityZScanPoints, + ) + + purity_values = [] + + best_purity = -np.inf + best_z = z0 + best_state = None + + # ------------------------------------------------------------------ + # Scan information + # ------------------------------------------------------------------ + + tqdm.tqdm.write("") + tqdm.tqdm.write("Starting purity-based z scan") + tqdm.tqdm.write( + f"Initial z = {z0 * 1e3:.6f} mm | " + f"range = ±{self.params.purityZScanRange * 1e6:.1f} µm | " + f"points = {len(z_values)} | " + f"iterations/z = {self.numIterations}" + ) + tqdm.tqdm.write("") + + # ------------------------------------------------------------------ + # Candidate-z loop + # ------------------------------------------------------------------ + + for z_index, z in enumerate(z_values): + + dz_um = (z - z0) * 1e6 + + tqdm.tqdm.write( + f"[{z_index + 1}/{len(z_values)}] " + f"Reconstructing z = {z * 1e3:.6f} mm " + f"(Δz = {dz_um:+.2f} µm)" + ) + + # -------------------------------------------------------------- + # Evaluate candidate + # -------------------------------------------------------------- + + purity, candidate_reconstruction = ( + self._evaluatePurityAtZ(z) + ) + + purity_values.append(purity) + + tqdm.tqdm.write( + f" Probe purity = {purity:.6f}" + ) + + # -------------------------------------------------------------- + # Update main monitor with completed candidate + # -------------------------------------------------------------- + + self.reconstruction.zo = z + + self.reconstruction.object = ( + asNumpyArray( + candidate_reconstruction.object + ).copy() + ) + + self.reconstruction.probe = ( + asNumpyArray( + candidate_reconstruction.probe + ).copy() + ) + + self.reconstruction.purityProbe = purity + + if hasattr(candidate_reconstruction, "error"): + self.reconstruction.error = np.asarray( + asNumpyArray( + candidate_reconstruction.error + ) + ).copy() + + self.showReconstruction( + self.numIterations - 1, + force=True, + ) + + # -------------------------------------------------------------- + # Keep only the state required from the best candidate + # -------------------------------------------------------------- + + if purity > best_purity: + + best_purity = purity + best_z = z + + best_state = { + "object": asNumpyArray( + candidate_reconstruction.object + ).copy(), + + "probe": asNumpyArray( + candidate_reconstruction.probe + ).copy(), + + "purityProbe": purity, + } + + if hasattr( + candidate_reconstruction, + "encoder_corrected", + ): + best_state["encoder_corrected"] = ( + asNumpyArray( + candidate_reconstruction.encoder_corrected + ).copy() + ) + + if hasattr( + candidate_reconstruction, + "objectMomentum", + ): + best_state["objectMomentum"] = ( + asNumpyArray( + candidate_reconstruction.objectMomentum + ).copy() + ) + + if hasattr( + candidate_reconstruction, + "probeMomentum", + ): + best_state["probeMomentum"] = ( + asNumpyArray( + candidate_reconstruction.probeMomentum + ).copy() + ) + + if hasattr( + candidate_reconstruction, + "objectBuffer", + ): + best_state["objectBuffer"] = ( + asNumpyArray( + candidate_reconstruction.objectBuffer + ).copy() + ) + + if hasattr( + candidate_reconstruction, + "probeBuffer", + ): + best_state["probeBuffer"] = ( + asNumpyArray( + candidate_reconstruction.probeBuffer + ).copy() + ) + + if hasattr( + candidate_reconstruction, + "error", + ): + best_state["error"] = np.asarray( + asNumpyArray( + candidate_reconstruction.error + ) + ).copy() + + tqdm.tqdm.write( + f" NEW BEST: z = {best_z * 1e3:.6f} mm, " + f"purity = {best_purity:.6f}" + ) + + tqdm.tqdm.write("") + + # ------------------------------------------------------------------ + # Safety check + # ------------------------------------------------------------------ + + if best_state is None: + raise RuntimeError( + "Purity z scan did not produce a valid purity value." + ) + + # ------------------------------------------------------------------ + # Restore best candidate state + # ------------------------------------------------------------------ + + self.reconstruction.zo = best_z + + self.reconstruction.object = ( + best_state["object"] + ) + + self.reconstruction.probe = ( + best_state["probe"] + ) + + self.reconstruction.purityProbe = ( + best_state["purityProbe"] + ) + + if "encoder_corrected" in best_state: + self.reconstruction.encoder_corrected = ( + best_state["encoder_corrected"] + ) + + if "objectMomentum" in best_state: + self.reconstruction.objectMomentum = ( + best_state["objectMomentum"] + ) + + if "probeMomentum" in best_state: + self.reconstruction.probeMomentum = ( + best_state["probeMomentum"] + ) + + if "objectBuffer" in best_state: + self.reconstruction.objectBuffer = ( + best_state["objectBuffer"] + ) + + if "probeBuffer" in best_state: + self.reconstruction.probeBuffer = ( + best_state["probeBuffer"] + ) + + if "error" in best_state: + self.reconstruction.error = ( + best_state["error"] + ) + + # ------------------------------------------------------------------ + # Store scan results + # ------------------------------------------------------------------ + + self.purityZValues = np.asarray( + z_values + ) + + self.purityZMetrics = np.asarray( + purity_values + ) + + self.purityBestZ = best_z + self.purityBestValue = best_purity + + # Final monitor update + self.showReconstruction( + self.numIterations - 1, + force=True, + ) + + # ------------------------------------------------------------------ + # Summary + # ------------------------------------------------------------------ + + tqdm.tqdm.write("=" * 60) + tqdm.tqdm.write("Purity-based z scan finished") + + tqdm.tqdm.write( + f"Best z = {best_z * 1e3:.6f} mm" + ) + + tqdm.tqdm.write( + f"Best Δz = " + f"{(best_z - z0) * 1e6:+.2f} µm" + ) + + tqdm.tqdm.write( + f"Best purity = {best_purity:.6f}" + ) + + tqdm.tqdm.write("=" * 60) + + return best_z, best_purity + + def adaptivePurityZScan(self): + """ + Adaptively search for the propagation distance that maximizes + reconstructed probe purity. + + The search first determines the uphill direction, then expands the + axial step while purity keeps increasing. Once the maximum is bracketed, + the step size is reduced and the search is refined around the current + best position. + """ + + if self.reconstruction.npsm < 2: + raise ValueError( + "Purity-based z scanning requires at least two probe modes " + "(reconstruction.npsm >= 2)." + ) + + z0 = self.reconstruction.zo + self.purityZInitialGuess = z0 + + step = self.params.purityZInitialStep + growth = self.params.purityZStepGrowth + shrink = self.params.purityZStepShrink + min_step = self.params.purityZMinStep + purity_tol = self.params.purityZPurityTolerance + max_evaluations = self.params.purityZMaxEvaluations + + evaluated = {} + reconstruction_states = {} + + best_z = None + best_purity = -np.inf + best_state = None + + z_history = [] + purity_history = [] + + def evaluate(z): + nonlocal best_z + nonlocal best_purity + nonlocal best_state + + # Reuse already evaluated points + key = float(z) + + if key in evaluated: + return evaluated[key] + + purity, candidate_reconstruction = ( + self._evaluatePurityAtZ(z) + ) + + evaluated[key] = purity + + z_history.append(z) + purity_history.append(purity) + + tqdm.tqdm.write( + f"z = {z * 1e3:.6f} mm | " + f"purity = {purity:.6f}" + ) + + # Store best candidate state + if purity > best_purity: + + best_purity = purity + best_z = z + + best_state = { + "object": asNumpyArray( + candidate_reconstruction.object + ).copy(), + + "probe": asNumpyArray( + candidate_reconstruction.probe + ).copy(), + + "purityProbe": purity, + } + + if hasattr( + candidate_reconstruction, + "encoder_corrected", + ): + best_state["encoder_corrected"] = ( + asNumpyArray( + candidate_reconstruction.encoder_corrected + ).copy() + ) + + if hasattr( + candidate_reconstruction, + "objectMomentum", + ): + best_state["objectMomentum"] = ( + asNumpyArray( + candidate_reconstruction.objectMomentum + ).copy() + ) + + if hasattr( + candidate_reconstruction, + "probeMomentum", + ): + best_state["probeMomentum"] = ( + asNumpyArray( + candidate_reconstruction.probeMomentum + ).copy() + ) + + if hasattr( + candidate_reconstruction, + "objectBuffer", + ): + best_state["objectBuffer"] = ( + asNumpyArray( + candidate_reconstruction.objectBuffer + ).copy() + ) + + if hasattr( + candidate_reconstruction, + "probeBuffer", + ): + best_state["probeBuffer"] = ( + asNumpyArray( + candidate_reconstruction.probeBuffer + ).copy() + ) + + if hasattr( + candidate_reconstruction, + "error", + ): + best_state["error"] = np.asarray( + asNumpyArray( + candidate_reconstruction.error + ) + ).copy() + + tqdm.tqdm.write( + f"NEW BEST: z = {best_z * 1e3:.6f} mm | " + f"purity = {best_purity:.6f}" + ) + + return purity + + tqdm.tqdm.write("") + tqdm.tqdm.write("Starting adaptive purity-based z scan") + tqdm.tqdm.write( + f"Initial z = {z0 * 1e3:.6f} mm | " + f"initial step = {step * 1e6:.1f} µm" + ) + tqdm.tqdm.write("") + + # -------------------------------------------------------------- + # 1. Initial three-point evaluation + # -------------------------------------------------------------- + + p0 = evaluate(z0) + + z_left = z0 - step + z_right = z0 + step + + p_left = evaluate(z_left) + p_right = evaluate(z_right) + + evaluations = len(evaluated) + + # -------------------------------------------------------------- + # 2. Determine initial uphill direction + # -------------------------------------------------------------- + + if p_left > p0 and p_left >= p_right: + direction = -1 + + elif p_right > p0 and p_right > p_left: + direction = +1 + + else: + direction = 0 + + # -------------------------------------------------------------- + # 3. Expansion phase + # -------------------------------------------------------------- + + if direction != 0: + + current_z = z0 + direction * step + current_purity = evaluate(current_z) + + while evaluations < max_evaluations: + + step *= growth + + next_z = current_z + direction * step + next_purity = evaluate(next_z) + + evaluations = len(evaluated) + + improvement = next_purity - current_purity + + if improvement > purity_tol: + + current_z = next_z + current_purity = next_purity + continue + + # Purity stopped increasing: + # maximum should now be bracketed. + break + + # -------------------------------------------------------------- + # 4. Refinement phase + # -------------------------------------------------------------- + + step *= shrink + + while ( + step >= min_step + and len(evaluated) < max_evaluations + ): + + left_z = best_z - step + right_z = best_z + step + + left_purity = evaluate(left_z) + + if len(evaluated) >= max_evaluations: + break + + right_purity = evaluate(right_z) + + # If neither side significantly improves the best result, + # reduce the search step. + if ( + left_purity <= best_purity + purity_tol + and right_purity <= best_purity + purity_tol + ): + step *= shrink + continue + + # A better point was found. Keep the current step and + # continue refining around the new best position. + + # -------------------------------------------------------------- + # Safety check + # -------------------------------------------------------------- + + if best_state is None: + raise RuntimeError( + "Adaptive purity z scan did not produce a valid result." + ) + + # -------------------------------------------------------------- + # Restore best candidate + # -------------------------------------------------------------- + + self.reconstruction.zo = best_z + self.reconstruction.object = best_state["object"] + self.reconstruction.probe = best_state["probe"] + self.reconstruction.purityProbe = best_state["purityProbe"] + + if "encoder_corrected" in best_state: + self.reconstruction.encoder_corrected = ( + best_state["encoder_corrected"] + ) + + if "objectMomentum" in best_state: + self.reconstruction.objectMomentum = ( + best_state["objectMomentum"] + ) + + if "probeMomentum" in best_state: + self.reconstruction.probeMomentum = ( + best_state["probeMomentum"] + ) + + if "objectBuffer" in best_state: + self.reconstruction.objectBuffer = ( + best_state["objectBuffer"] + ) + + if "probeBuffer" in best_state: + self.reconstruction.probeBuffer = ( + best_state["probeBuffer"] + ) + + if "error" in best_state: + self.reconstruction.error = ( + best_state["error"] + ) + + # -------------------------------------------------------------- + # Store results + # -------------------------------------------------------------- + + order = np.argsort(z_history) + + self.purityZValues = np.asarray(z_history)[order] + self.purityZMetrics = np.asarray(purity_history)[order] + + self.purityBestZ = best_z + self.purityBestValue = best_purity + + self.showReconstruction( + self.numIterations - 1, + force=True, + ) + + # -------------------------------------------------------------- + # Summary + # -------------------------------------------------------------- + + tqdm.tqdm.write("=" * 60) + tqdm.tqdm.write("Adaptive purity-based z scan finished") + tqdm.tqdm.write( + f"Best z = {best_z * 1e3:.6f} mm" + ) + tqdm.tqdm.write( + f"Best Δz = {(best_z - z0) * 1e6:+.2f} µm" + ) + tqdm.tqdm.write( + f"Best purity = {best_purity:.6f}" + ) + tqdm.tqdm.write( + f"Evaluations = {len(evaluated)}" + ) + tqdm.tqdm.write("=" * 60) + + return best_z, best_purity + def plotPurityZScan(self, show=True): """ Plot probe purity as a function of axial distance. @@ -653,4 +1362,4 @@ def plotPurityZScan(self, show=True): ) self.reconstruction.probeBuffer = self.reconstruction.probe.copy() - \ No newline at end of file + \ No newline at end of file diff --git a/PtyLab/Params/Params.py b/PtyLab/Params/Params.py index 17d3bff8..bcd8ff30 100644 --- a/PtyLab/Params/Params.py +++ b/PtyLab/Params/Params.py @@ -154,11 +154,21 @@ def __init__(self): self.map_position_to_z_change = False # purity based Zscan - self.purityZScanSwitch = False self.purityZScanRange = 100e-6 self.purityZScanPoints = 11 + self.purityZScanPlot = True + self.purityZAdaptive = False + + self.purityZInitialStep = 200e-6 + self.purityZStepGrowth = 1.5 + self.purityZStepShrink = 0.5 + + self.purityZMinStep = 20e-6 + self.purityZPurityTolerance = 1e-4 + self.purityZMaxEvaluations = 20 + self.OPRP = False # Default values of all OPR parameters diff --git a/example_scripts/example_purity_Zscan_yourData.py b/example_scripts/example_purity_Zscan_yourData.py index bbbb0027..0c3a5d5c 100644 --- a/example_scripts/example_purity_Zscan_yourData.py +++ b/example_scripts/example_purity_Zscan_yourData.py @@ -14,6 +14,10 @@ logging.basicConfig(level=logging.INFO) +# ============================================================================= +# File selection +# ============================================================================= + def select_hdf5_file(title): root = Tk() root.withdraw() @@ -29,24 +33,80 @@ def select_hdf5_file(title): root.destroy() - if not file_path: - raise RuntimeError(f"No file selected: {title}") - return file_path -# ------------------------------------------------------------------------- -# Select input files -# ------------------------------------------------------------------------- +print() +print("=" * 70) +print("Purity-based axial calibration") +print("=" * 70) + +print() +print( + "Step 1: Select the ptychography dataset to calibrate.\n" + "This should be the HDF5 file containing the measured diffraction data." +) filePath = select_hdf5_file( "Select ptychography dataset" ) -filePath_recon = select_hdf5_file( - "Select seed reconstruction" +if not filePath: + raise RuntimeError( + "No ptychography dataset was selected." + ) + + +print() +print("Selected dataset:") +print(filePath) + + +print() +print( + "Step 2: Choose the initial probe.\n" + "If you already have a previous reconstruction, it can be used as a " + "seed probe.\n" + "Otherwise, purityPIE can start from a circular probe estimate." ) +use_seed = input( + "Use an existing reconstruction as seed probe? [y/N]: " +).strip().lower() in ("y", "yes") + + +filePath_recon = None + +if use_seed: + + print() + print( + "Select the HDF5 reconstruction file containing the seed probe." + ) + + filePath_recon = select_hdf5_file( + "Select seed reconstruction" + ) + + if not filePath_recon: + raise RuntimeError( + "Seed probe was requested, but no reconstruction file " + "was selected." + ) + + print() + print("Selected seed reconstruction:") + print(filePath_recon) + +else: + + print() + print( + "No seed reconstruction selected.\n" + "The reconstruction will start from a circular probe " + "and a uniform object." + ) + # ------------------------------------------------------------------------- # Load data # ------------------------------------------------------------------------- @@ -66,9 +126,16 @@ def select_hdf5_file(title): print("Wavelength:", reconstruction.wavelength) -initial_z = reconstruction.zo +# Use the z value stored in the dataset by default. +initial_z = reconstruction.zo + +# Uncomment to set the initial z guess manually. +# initial_z = 32.5e-3 -print(f"Initial z guess: {reconstruction.zo * 1e3:.6f} mm") +experimentalData.zo = initial_z +reconstruction.zo = initial_z + +print(f"Initial z guess: {initial_z * 1e3:.6f} mm") # ------------------------------------------------------------------------- @@ -83,17 +150,55 @@ def select_hdf5_file(title): reconstruction.nslice = 1 -# ------------------------------------------------------------------------- +# ============================================================================= # Object and probe initialization -# ------------------------------------------------------------------------- +# ============================================================================= -reconstruction.initialProbe = "recon" -reconstruction.initialProbe_filename = filePath_recon reconstruction.initialObject = "ones" -reconstruction.initializeObjectProbe() +if filePath_recon is not None: + + reconstruction.initialProbe = "recon" + reconstruction.initialProbe_filename = ( + filePath_recon + ) + + print() + print( + "Initial probe: loaded from seed reconstruction" + ) + +else: + + reconstruction.initialProbe = "circ" + + print() + print( + "Initial probe: circular estimate" + ) + + # Circular initialization requires an entrance pupil diameter. + pupil_diameter = getattr( + experimentalData, + "entrancePupilDiameter", + None, + ) + + if pupil_diameter is None: + raise RuntimeError( + "Circular probe initialization requires " + "`experimentalData.entrancePupilDiameter`.\n" + "Define the probe / entrance pupil diameter before calling " + "`initializeObjectProbe()`, or provide a seed reconstruction." + ) + + print( + f"Entrance pupil diameter: " + f"{pupil_diameter * 1e6:.2f} um" + ) -reconstruction.describe_reconstruction() + +reconstruction.initializeObjectProbe() # ------------------------------------------------------------------------- @@ -138,7 +243,6 @@ def select_hdf5_file(title): # ------------------------------------------------------------------------- params.purityZScanSwitch = True - # Half-range around the current reconstruction.zo. # # Example: @@ -152,10 +256,18 @@ def select_hdf5_file(title): # Number of candidate z values. params.purityZScanPoints = 10 -# Plot purity versus z after the scan. -params.purityZScanPlot = True +params.purityZAdaptive = True + +params.purityZInitialStep = 200e-6 +params.purityZStepGrowth = 1.5 +params.purityZStepShrink = 0.5 +params.purityZMinStep = 20e-6 +params.purityZPurityTolerance = 1e-4 +params.purityZMaxEvaluations = 20 +# Plot purity versus z after the scan. +params.purityZScanPlot = True monitor.describe_parameters(params) @@ -200,6 +312,8 @@ def select_hdf5_file(title): f"purity = {purity:.6f}" ) + + # ------------------------------------------------------------------------- # Optional: continue reconstruction using the best-z state # ------------------------------------------------------------------------- diff --git a/tests/Engines/test_purity_adaptive_scan.py b/tests/Engines/test_purity_adaptive_scan.py new file mode 100644 index 00000000..05ca34d0 --- /dev/null +++ b/tests/Engines/test_purity_adaptive_scan.py @@ -0,0 +1,417 @@ +import numpy as np +import pytest + +def adaptive_search( + objective, + z0, + initial_step, + growth=1.5, + shrink=0.5, + min_step=20e-6, + tolerance=1e-8, + max_evaluations=30, +): + """ + Standalone version of the adaptive z-search logic. + + Parameters + ---------- + objective : callable + Function returning the quantity to maximize. + z0 : float + Initial z position. + initial_step : float + Initial axial step. + growth : float + Step expansion factor. + shrink : float + Step refinement factor. + min_step : float + Minimum allowed step. + tolerance : float + Minimum significant improvement. + max_evaluations : int + Maximum number of objective evaluations. + + Returns + ------- + best_z : float + Best position found. + best_value : float + Best objective value found. + z_values : ndarray + Evaluated z positions. + values : ndarray + Corresponding objective values. + """ + + step = initial_step + + evaluated = {} + + best_z = None + best_value = -np.inf + + z_history = [] + value_history = [] + + def evaluate(z): + nonlocal best_z, best_value + + key = float(z) + + if key in evaluated: + return evaluated[key] + + value = float(objective(z)) + + evaluated[key] = value + z_history.append(z) + value_history.append(value) + + if value > best_value: + best_value = value + best_z = z + + return value + + # -------------------------------------------------------------- + # Initial three-point sampling + # -------------------------------------------------------------- + + p0 = evaluate(z0) + + z_left = z0 - step + z_right = z0 + step + + p_left = evaluate(z_left) + p_right = evaluate(z_right) + + # -------------------------------------------------------------- + # Determine uphill direction + # -------------------------------------------------------------- + + if p_left > p0 and p_left >= p_right: + direction = -1 + + elif p_right > p0 and p_right > p_left: + direction = +1 + + else: + direction = 0 + + # -------------------------------------------------------------- + # Expansion phase + # -------------------------------------------------------------- + + if direction != 0: + + current_z = z0 + direction * step + current_value = evaluate(current_z) + + while len(evaluated) < max_evaluations: + + step *= growth + + next_z = current_z + direction * step + next_value = evaluate(next_z) + + improvement = next_value - current_value + + if improvement > tolerance: + current_z = next_z + current_value = next_value + continue + + break + + # -------------------------------------------------------------- + # Refinement phase + # -------------------------------------------------------------- + + step *= shrink + + while ( + step >= min_step + and len(evaluated) < max_evaluations + ): + + left_z = best_z - step + right_z = best_z + step + + left_value = evaluate(left_z) + + if len(evaluated) >= max_evaluations: + break + + right_value = evaluate(right_z) + + if ( + left_value <= best_value + tolerance + and right_value <= best_value + tolerance + ): + step *= shrink + + order = np.argsort(z_history) + + return ( + best_z, + best_value, + np.asarray(z_history)[order], + np.asarray(value_history)[order], + ) + +def test_adaptive_search_finds_parabolic_peak(): + + true_z = 50e-3 + + def purity_function(z): + """ + Synthetic purity curve with a maximum at 50 mm. + """ + return 1.0 - ((z - true_z) / 1e-3) ** 2 + + best_z, best_purity, z_values, purity_values = adaptive_search( + objective=purity_function, + z0=47e-3, + initial_step=200e-6, + growth=1.5, + shrink=0.5, + min_step=10e-6, + tolerance=1e-12, + max_evaluations=30, + ) + + print() + print("=" * 70) + print("Adaptive search test") + print("=" * 70) + + for z, purity in zip(z_values, purity_values): + print( + f"z = {z * 1e3:8.4f} mm | " + f"purity = {purity:.8f}" + ) + + print("-" * 70) + + print( + f"true z = {true_z * 1e3:.6f} mm" + ) + + print( + f"best z = {best_z * 1e3:.6f} mm" + ) + + print( + f"error = " + f"{abs(best_z - true_z) * 1e6:.3f} um" + ) + + print( + f"best purity = {best_purity:.8f}" + ) + + print( + f"evaluations = {len(z_values)}" + ) + + print("=" * 70) + + assert abs(best_z - true_z) <= 20e-6 + + assert np.isclose( + best_purity, + 1.0, + atol=1e-3, + ) + + assert len(z_values) <= 30 + +def test_adaptive_search_with_noise(): + + true_z = 50e-3 + + rng = np.random.default_rng(0) + + def purity_function(z): + """ + Synthetic purity curve with small reconstruction-like noise. + """ + clean = 1.0 - ((z - true_z) / 1e-3) ** 2 + + noise = rng.normal( + loc=0.0, + scale=0.002, + ) + + return clean + noise + + best_z, best_purity, z_values, purity_values = adaptive_search( + objective=purity_function, + z0=47e-3, + initial_step=200e-6, + growth=1.5, + shrink=0.5, + min_step=20e-6, + + # Larger than the typical point-to-point numerical noise + tolerance=5e-4, + + max_evaluations=30, + ) + + error_um = abs(best_z - true_z) * 1e6 + + print() + print("=" * 70) + print("Adaptive search with noise") + print("=" * 70) + + for z, purity in zip(z_values, purity_values): + print( + f"z = {z * 1e3:8.4f} mm | " + f"purity = {purity:.8f}" + ) + + print("-" * 70) + + print( + f"true z = {true_z * 1e3:.6f} mm" + ) + + print( + f"best z = {best_z * 1e3:.6f} mm" + ) + + print( + f"error = {error_um:.3f} um" + ) + + print( + f"best purity = {best_purity:.8f}" + ) + + print( + f"evaluations = {len(z_values)}" + ) + + print("=" * 70) + + # The noisy search does not need to hit the exact analytical maximum, + # but it should remain close to it. + assert error_um <= 100 + + assert len(z_values) <= 30 + +@pytest.mark.parametrize("seed", range(20)) +def test_adaptive_search_with_noise_multiple_seeds(seed): + + true_z = 50e-3 + + rng = np.random.default_rng(seed) + + def purity_function(z): + clean = 1.0 - ((z - true_z) / 1e-3) ** 2 + + noise = rng.normal( + loc=0.0, + scale=0.002, + ) + + return clean + noise + + best_z, best_purity, z_values, purity_values = adaptive_search( + objective=purity_function, + z0=47e-3, + initial_step=200e-6, + growth=1.5, + shrink=0.5, + min_step=20e-6, + tolerance=5e-4, + max_evaluations=30, + ) + + error_um = abs(best_z - true_z) * 1e6 + + assert error_um <= 100 + + assert len(z_values) <= 30 + + +def test_adaptive_search_avoids_local_bump(): + + true_z = 50e-3 + + def purity_function(z): + """ + Synthetic purity curve with: + - global maximum near 50 mm + - smaller local bump near 48.5 mm + """ + + # Main global peak + main_peak = 1.0 - ((z - true_z) / 1e-3) ** 2 + + # Local bump centered at 48.5 mm + local_bump = 0.35 * np.exp( + -((z - 48.5e-3) / 0.20e-3) ** 2 + ) + + return main_peak + local_bump + + best_z, best_purity, z_values, purity_values = adaptive_search( + objective=purity_function, + z0=47e-3, + initial_step=200e-6, + growth=1.5, + shrink=0.5, + min_step=20e-6, + tolerance=5e-4, + max_evaluations=30, + ) + + error_um = abs(best_z - true_z) * 1e6 + + print() + print("=" * 70) + print("Adaptive search with local bump") + print("=" * 70) + + for z, purity in zip(z_values, purity_values): + print( + f"z = {z * 1e3:8.4f} mm | " + f"purity = {purity:.8f}" + ) + + print("-" * 70) + + print( + f"true z = {true_z * 1e3:.6f} mm" + ) + + print( + f"best z = {best_z * 1e3:.6f} mm" + ) + + print( + f"error = {error_um:.3f} um" + ) + + print( + f"best purity = {best_purity:.8f}" + ) + + print( + f"evaluations = {len(z_values)}" + ) + + print("=" * 70) + + assert error_um <= 100 + assert len(z_values) <= 30 + +if __name__ == "__main__": + test_adaptive_search_finds_parabolic_peak() + test_adaptive_search_with_noise() + test_adaptive_search_avoids_local_bump() \ No newline at end of file From d7c3b7caea4677bdc04786921450c893bca8d81d Mon Sep 17 00:00:00 2001 From: Chang Liu Date: Sat, 3 Oct 2026 21:03:17 +0200 Subject: [PATCH 06/11] minor change: remove switch --- .../example_purity_Zscan_yourData.py | 62 +++++++++---------- 1 file changed, 30 insertions(+), 32 deletions(-) diff --git a/example_scripts/example_purity_Zscan_yourData.py b/example_scripts/example_purity_Zscan_yourData.py index 0c3a5d5c..c57b7173 100644 --- a/example_scripts/example_purity_Zscan_yourData.py +++ b/example_scripts/example_purity_Zscan_yourData.py @@ -241,8 +241,6 @@ def select_hdf5_file(title): # ------------------------------------------------------------------------- # Purity-based z calibration parameters # ------------------------------------------------------------------------- - -params.purityZScanSwitch = True # Half-range around the current reconstruction.zo. # # Example: @@ -275,42 +273,42 @@ def select_hdf5_file(title): # Run purity-based axial calibration # ------------------------------------------------------------------------- -if params.purityZScanSwitch: - purity_engine = Engines.purityPIE( - reconstruction, - experimentalData, - params, - monitor, - ) - # Number of mPIE iterations used to evaluate EACH candidate z. - purity_engine.numIterations = 5 +purity_engine = Engines.purityPIE( + reconstruction, + experimentalData, + params, + monitor, +) + +# Number of mPIE iterations used to evaluate EACH candidate z. +purity_engine.numIterations = 5 - best_z, best_purity = purity_engine.reconstruct() +best_z, best_purity = purity_engine.reconstruct() - print() - print("Purity-based axial calibration finished") - print("---------------------------------------") - print(f"Initial z : {purity_engine.purityZInitialGuess * 1e3:.6f} mm") - print(f"Best z : {best_z * 1e3:.6f} mm") +print() +print("Purity-based axial calibration finished") +print("---------------------------------------") +print(f"Initial z : {purity_engine.purityZInitialGuess * 1e3:.6f} mm") +print(f"Best z : {best_z * 1e3:.6f} mm") +print( + f"Delta z : " + f"{(best_z - purity_engine.purityZInitialGuess) * 1e6:.3f} um" +) +print(f"Best purity: {best_purity:.6f}") + +print() +print("z scan:") +for z, purity in zip( + purity_engine.purityZValues, + purity_engine.purityZMetrics, +): print( - f"Delta z : " - f"{(best_z - purity_engine.purityZInitialGuess) * 1e6:.3f} um" + f"z = {z * 1e3:.6f} mm, " + f"delta z = {(z - purity_engine.purityZInitialGuess) * 1e6:+.2f} um, " + f"purity = {purity:.6f}" ) - print(f"Best purity: {best_purity:.6f}") - - print() - print("z scan:") - for z, purity in zip( - purity_engine.purityZValues, - purity_engine.purityZMetrics, - ): - print( - f"z = {z * 1e3:.6f} mm, " - f"delta z = {(z - purity_engine.purityZInitialGuess) * 1e6:+.2f} um, " - f"purity = {purity:.6f}" - ) From ef2193301a12ad55a071198e602be5591e0b84c0 Mon Sep 17 00:00:00 2001 From: Chang Liu Date: Mon, 5 Oct 2026 11:09:26 +0200 Subject: [PATCH 07/11] API doc for purityPIE and zPIE engines --- PtyLab/Engines/ePIE.py | 15 - PtyLab/Engines/mPIE.py | 25 +- PtyLab/Engines/purityPIE.py | 763 ++++++++++++++---------------------- PtyLab/Engines/zPIE.py | 184 ++++++++- docs/api/engines.md | 6 + 5 files changed, 481 insertions(+), 512 deletions(-) diff --git a/PtyLab/Engines/ePIE.py b/PtyLab/Engines/ePIE.py index 2a047940..c88c495c 100644 --- a/PtyLab/Engines/ePIE.py +++ b/PtyLab/Engines/ePIE.py @@ -73,21 +73,6 @@ def __init__( params: Params, monitor: Monitor, ): - """ - 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") diff --git a/PtyLab/Engines/mPIE.py b/PtyLab/Engines/mPIE.py index b677a321..583693ae 100644 --- a/PtyLab/Engines/mPIE.py +++ b/PtyLab/Engines/mPIE.py @@ -105,29 +105,6 @@ def __init__( params: Params, monitor: Monitor, ): - """ - 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("Successfully created mPIE engine") @@ -564,7 +541,7 @@ def probeUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray, weight: float) class pcPIE(mPIE): """ - Backward-compatible wrapper for :class:`mPIE`. + Backward-compatible wrapper for `mPIE`. Position correction is now provided by `mPIE` through `params.positionCorrectionSwitch`. This class is retained only for diff --git a/PtyLab/Engines/purityPIE.py b/PtyLab/Engines/purityPIE.py index ffecdc0f..deeb7d3b 100644 --- a/PtyLab/Engines/purityPIE.py +++ b/PtyLab/Engines/purityPIE.py @@ -14,7 +14,7 @@ import tqdm -from PtyLab.Engines.BaseEngine import BaseEngine +#from PtyLab.Engines.BaseEngine import BaseEngine from PtyLab.ExperimentalData.ExperimentalData import ExperimentalData from PtyLab.Monitor.Monitor import Monitor from PtyLab.Params.Params import Params @@ -30,9 +30,89 @@ from PtyLab.Engines.mPIE import mPIE class purityPIE(mPIE): - ''' - purityPIE extends mPIE-style reconstruction with purity-based axial calibration. - ''' + r""" + Purity-based axial self-calibration for ptychographic reconstruction. + + `purityPIE` extends `mPIE` by estimating the sample-to-detector + axial distance from the reconstructed probe purity.[^liu2025] + + In mixed-state ptychography, the measured diffraction intensity at + scan position $j$ is modeled as an incoherent sum over multiple probe + modes: + + $$ + I_j(\mathbf{q}) = \sum_m \left| \mathcal{F}\left[P_m(\mathbf{r}) O(\mathbf{r}-\mathbf{r}_j)\right] \right|^2 + $$ + + where $P_m$ denotes probe mode $m$, $O$ is the object transmission + function, and $\mathcal{F}$ denotes propagation to the detector plane. + + The mutual intensity of the reconstructed probe is + + $$ + J(\mathbf{r}_1,\mathbf{r}_2) = \sum_m P_m^*(\mathbf{r}_1)P_m(\mathbf{r}_2) + $$ + + and can be decomposed into orthonormal coherent modes as + + $$ + J(\mathbf{r}_1,\mathbf{r}_2) = \sum_m \lambda_m \phi_m^*(\mathbf{r}_1)\phi_m(\mathbf{r}_2) + $$ + + where $\lambda_m$ represents the intensity contribution, or relative + power content, of the corresponding orthonormal mode. + + The probe purity is defined as + + $$ + \mu = \frac{\sqrt{\sum_m \lambda_m^2}}{\sum_m \lambda_m} + $$ + + with $\mu \in [0,1]$. + + If the assumed axial distance is incorrect, a mismatch is introduced + in the diffraction integral connecting real and reciprocal space. + This produces stretching deformation and resolution degradation in + the reconstructed object and probe, which appears as a reduction in + reconstructed purity. + + The calibrated axial distance is obtained by maximizing the + reconstructed probe purity: + + $$ + z_{\mathrm{opt}} = \operatorname*{arg\,max}_z \mu(z) + $$ + + For each candidate axial distance, `purityPIE` creates an independent + reconstruction initialized with the corresponding geometry, performs + an `mPIE` reconstruction, orthogonalizes the reconstructed probe + modes, and evaluates the resulting purity. The candidate with the + highest purity is retained as the calibrated reconstruction. + + Two search strategies are available: + + - A fixed-grid scan evaluates uniformly spaced axial positions defined + by `params.purityZScanRange` and `params.purityZScanPoints`. + + - An adaptive scan expands the axial step while purity improves and + subsequently refines the search around the best candidate with + progressively smaller steps. It is enabled with + `params.purityZAdaptive = True`. + + Purity-based axial calibration requires at least two reconstructed + probe modes, i.e. `reconstruction.npsm >= 2`. + + [^liu2025]: C. Liu, W. Eschen, D. S. Penagos Molina, + L. Licht, M. Abdelaal, and J. Rothhardt, + "Purity-based self-calibration in ptychography," + Opt. Lett. 50, 1581-1584 (2025). + https://doi.org/10.1364/OL.543865 + + See Also: + `mPIE` + Momentum-accelerated reconstruction engine used for each + candidate axial distance. + """ def __init__( self, reconstruction: Reconstruction, @@ -62,29 +142,41 @@ def reconstruct( reconstruction=None, ): """ - Run purity-based axial calibration. + Run purity-based axial self-calibration. - Each candidate propagation distance is reconstructed from the same - initial state using `numIterations` mPIE-style iterations. The distance - yielding the highest probe purity is selected, and the corresponding - reconstruction state is retained. + The reconstruction first prepares the current object, probe, geometry, + constraints, and optional GPU state using the standard `mPIE` setup. + + The axial calibration is then performed using one of two search strategies: + + - If `params.purityZAdaptive = False`, `purityZScan()` evaluates a fixed + grid of candidate axial distances. + + - If `params.purityZAdaptive = True`, `adaptivePurityZScan()` searches the + axial distance using adaptive step-size expansion and refinement. + + Each candidate distance is evaluated using an independent `mPIE` + reconstruction. The candidate yielding the highest reconstructed probe + purity is selected, and its reconstruction state is restored to the + original `Reconstruction` object. If `params.purityZScanPlot` is enabled, the resulting purity-versus-z curve is displayed after calibration. Args: experimentalData (ExperimentalData, optional): - Experimental dataset to use. If provided, it replaces the - currently attached dataset. + Experimental dataset to use for the calibration. If provided, it + replaces the dataset currently attached to the engine. + reconstruction (Reconstruction, optional): Reconstruction state to calibrate. If provided, it replaces the - currently attached reconstruction. + reconstruction currently attached to the engine. Returns: tuple: - `(best_z, best_purity)`, where `best_z` is the selected axial - propagation distance and `best_purity` is the corresponding - probe purity. + `(best_z, best_purity)`, where `best_z` is the calibrated axial + propagation distance in meters and `best_purity` is the + corresponding reconstructed probe purity. """ self.changeExperimentalData(experimentalData) @@ -116,10 +208,33 @@ def reconstruct( def _createCandidateReconstruction(self, z): """ - Create an independent reconstruction state for one candidate z. + Create an independent reconstruction for one candidate axial distance. + + A deep copy of the experimental data and reconstruction parameters is + created so that evaluating one candidate distance cannot modify the state + of another candidate. + + The candidate axial distance is assigned before constructing the new + `Reconstruction` object. This ensures that all geometry-dependent + quantities, including sampling and scan positions, are initialized + consistently for the requested distance. + + User-defined reconstruction settings such as the number of probe modes, + object modes, wavelengths, slices, and object/probe initialization method + are copied from the parent reconstruction. + + The object and probe are then initialized independently for the candidate + geometry. + + Args: + z (float): + Candidate sample-to-detector propagation distance in meters. - The candidate inherits the user's reconstruction settings but is built - from a fresh geometry at the requested propagation distance. + Returns: + tuple: + `(candidate_data, candidate_reconstruction, candidate_params)`, + containing the independent experimental data, reconstruction state, + and parameter set for the candidate distance. """ # Independent copies so one candidate cannot modify another @@ -164,336 +279,30 @@ def _createCandidateReconstruction(self, z): candidate_params, ) - - """ - Scan a fixed axial range and select the distance that maximizes - reconstructed probe purity. - - Each candidate z is evaluated using an independent reconstruction - state and a standard mPIE reconstruction. - """ - - if self.reconstruction.npsm < 2: - raise ValueError( - "Purity-based z scanning requires at least two probe modes " - "(reconstruction.npsm >= 2)." - ) - - # ------------------------------------------------------------------ - # Scan grid - # ------------------------------------------------------------------ - - z0 = self.reconstruction.zo - self.purityZInitialGuess = z0 - - z_values = np.linspace( - z0 - self.params.purityZScanRange, - z0 + self.params.purityZScanRange, - self.params.purityZScanPoints, - ) - - purity_values = [] - - best_purity = -np.inf - best_z = z0 - - best_reconstruction = None - - # ------------------------------------------------------------------ - # Scan information - # ------------------------------------------------------------------ - - tqdm.tqdm.write("") - tqdm.tqdm.write("Starting purity-based z scan") - tqdm.tqdm.write( - f"Initial z = {z0 * 1e3:.6f} mm | " - f"range = ±{self.params.purityZScanRange * 1e6:.1f} µm | " - f"points = {len(z_values)} | " - f"iterations/z = {self.numIterations}" - ) - tqdm.tqdm.write("") - - # ------------------------------------------------------------------ - # Candidate-z loop - # ------------------------------------------------------------------ - - for z_index, z in enumerate(z_values): - - dz_um = (z - z0) * 1e6 - - tqdm.tqdm.write( - f"[{z_index + 1}/{len(z_values)}] " - f"Reconstructing z = {z * 1e3:.6f} mm " - f"(Δz = {dz_um:+.2f} µm)" - ) - - # -------------------------------------------------------------- - # Create a completely independent reconstruction for this z - # -------------------------------------------------------------- - #np.random.seed(0) - ( - candidate_data, - candidate_reconstruction, - candidate_params, - ) = self._createCandidateReconstruction(z) - - # Candidate reconstructions do not need their own GUI monitor - candidate_monitor = DummyMonitor() - - # -------------------------------------------------------------- - # Run standard mPIE - # -------------------------------------------------------------- - - candidate_engine = mPIE( - candidate_reconstruction, - candidate_data, - candidate_params, - candidate_monitor, - ) - - # Match the purityPIE reconstruction settings - candidate_engine.numIterations = self.numIterations - - candidate_engine.betaProbe = self.betaProbe - candidate_engine.betaObject = self.betaObject - - candidate_engine.alphaProbe = self.alphaProbe - candidate_engine.alphaObject = self.alphaObject - - candidate_engine.feedbackM = self.feedbackM - candidate_engine.frictionM = self.frictionM - - candidate_engine.reconstruct() - - # -------------------------------------------------------------- - # Final modal decomposition and purity evaluation - # -------------------------------------------------------------- - - candidate_engine.orthogonalization() - - purity = float( - np.asarray( - asNumpyArray(candidate_reconstruction.purityProbe) - ).squeeze() - ) - - purity_values.append(purity) - - tqdm.tqdm.write( - f" Probe purity = {purity:.6f}" - ) - - # -------------------------------------------------------------- - # Update main monitor with the completed candidate state - # -------------------------------------------------------------- - - self.reconstruction.zo = z - self.reconstruction.object = ( - asNumpyArray(candidate_reconstruction.object).copy() - ) - self.reconstruction.probe = ( - asNumpyArray(candidate_reconstruction.probe).copy() - ) - - self.reconstruction.purityProbe = purity - - if hasattr(candidate_reconstruction, "error"): - self.reconstruction.error = np.asarray( - asNumpyArray(candidate_reconstruction.error) - ).copy() - - # Display only completed candidate reconstructions - self.showReconstruction( - self.numIterations - 1, - force=True, - ) - - # -------------------------------------------------------------- - # Keep best candidate - # -------------------------------------------------------------- - - if purity > best_purity: - - best_purity = purity - best_z = z - - # Candidate objects are independent, so keeping this reference - # preserves the complete best reconstruction state. - best_reconstruction = candidate_reconstruction - - tqdm.tqdm.write( - f" NEW BEST: z = {best_z * 1e3:.6f} mm, " - f"purity = {best_purity:.6f}" - ) - - tqdm.tqdm.write("") - - # ------------------------------------------------------------------ - # Safety check - # ------------------------------------------------------------------ - - if best_reconstruction is None: - raise RuntimeError( - "Purity z scan did not produce a valid purity value." - ) - - # ------------------------------------------------------------------ - # Restore the best candidate into the original Reconstruction object - # - # Keep the original object identity because external scripts may still - # hold a reference to self.reconstruction. - # ------------------------------------------------------------------ - - self.reconstruction.zo = best_z - - self.reconstruction.object = ( - asNumpyArray(best_reconstruction.object).copy() - ) - - self.reconstruction.probe = ( - asNumpyArray(best_reconstruction.probe).copy() - ) - - self.reconstruction.purityProbe = best_purity - - if hasattr(best_reconstruction, "encoder_corrected"): - self.reconstruction.encoder_corrected = ( - asNumpyArray( - best_reconstruction.encoder_corrected - ).copy() - ) - - if hasattr(best_reconstruction, "objectMomentum"): - self.reconstruction.objectMomentum = ( - asNumpyArray( - best_reconstruction.objectMomentum - ).copy() - ) - - if hasattr(best_reconstruction, "probeMomentum"): - self.reconstruction.probeMomentum = ( - asNumpyArray( - best_reconstruction.probeMomentum - ).copy() - ) - - if hasattr(best_reconstruction, "objectBuffer"): - self.reconstruction.objectBuffer = ( - asNumpyArray( - best_reconstruction.objectBuffer - ).copy() - ) - - if hasattr(best_reconstruction, "probeBuffer"): - self.reconstruction.probeBuffer = ( - asNumpyArray( - best_reconstruction.probeBuffer - ).copy() - ) - - if hasattr(best_reconstruction, "error"): - self.reconstruction.error = np.asarray( - asNumpyArray(best_reconstruction.error) - ).copy() - - # ------------------------------------------------------------------ - # Store scan results - # ------------------------------------------------------------------ - - self.purityZValues = np.asarray(z_values) - self.purityZMetrics = np.asarray(purity_values) - - self.purityBestZ = best_z - self.purityBestValue = best_purity - - # Final monitor update with the selected best candidate - self.showReconstruction( - self.numIterations - 1, - force=True, - ) - - # ------------------------------------------------------------------ - # Summary - # ------------------------------------------------------------------ - - tqdm.tqdm.write("=" * 60) - tqdm.tqdm.write("Purity-based z scan finished") - tqdm.tqdm.write( - f"Best z = {best_z * 1e3:.6f} mm" - ) - tqdm.tqdm.write( - f"Best Δz = {(best_z - z0) * 1e6:+.2f} µm" - ) - tqdm.tqdm.write( - f"Best purity = {best_purity:.6f}" - ) - tqdm.tqdm.write("=" * 60) - - return best_z, best_purity - - - """ - Reset the reconstruction state for one candidate z. - - This implementation is intended for analytic initialization: - initialProbe = "circ" - initialObject = "ones" - - The object/probe are regenerated after setting the candidate z so that - their physical sampling is consistent with the candidate geometry. - Momentum and buffers are also reset. - """ - if self.reconstruction.initialProbe != "circ": - raise NotImplementedError( - "Candidate-z reinitialization is currently implemented " - "only for initialProbe='circ'." - ) - - if self.reconstruction.initialObject != "ones": - raise NotImplementedError( - "Candidate-z reinitialization is currently implemented " - "only for initialObject='ones'." - ) - - # Set candidate geometry first - self.reconstruction.zo = z - - # Rebuild object/probe on the candidate-z geometry - self.reconstruction.initializeObjectProbe() - - self._initialProbePowerCorrection() - - # Reset momentum - self.reconstruction.initializeObjectMomentum() - self.reconstruction.initializeProbeMomentum() - - # Reset buffers - self.reconstruction.objectBuffer = self.reconstruction.object.copy() - self.reconstruction.probeBuffer = self.reconstruction.probe.copy() - - self._prepareReconstruction() - def _evaluatePurityAtZ(self, z): """ - Evaluate probe purity for one candidate propagation distance. + Evaluate reconstructed probe purity at one candidate axial distance. - A fresh reconstruction is created at the requested z value and - reconstructed using the standard mPIE engine. Probe modes are - orthogonalized after reconstruction and the resulting probe purity - is returned together with the candidate reconstruction state. + An independent reconstruction is first created for the requested distance + using `_createCandidateReconstruction()`. - Parameters - ---------- - z : float - Candidate propagation distance in meters. + A standard `mPIE` reconstruction is then performed using the same iteration + count and update parameters as the parent `purityPIE` engine. After the + reconstruction, the probe modes are orthogonalized and the resulting probe + purity is extracted. - Returns - ------- - purity : float - Reconstructed probe purity at the candidate distance. + This helper contains the common candidate-evaluation procedure shared by + both the fixed-grid and adaptive axial scans. - candidate_reconstruction : Reconstruction - Reconstruction state obtained at the candidate distance. + Args: + z (float): + Candidate sample-to-detector propagation distance in meters. + + Returns: + tuple: + `(purity, candidate_reconstruction)`, where `purity` is the + reconstructed probe purity and `candidate_reconstruction` is the + completed reconstruction state for the candidate distance. """ # ------------------------------------------------------------------ @@ -550,12 +359,58 @@ def _evaluatePurityAtZ(self, z): return purity, candidate_reconstruction def purityZScan(self): - """ - Scan a fixed axial range and select the distance that maximizes - reconstructed probe purity. + r""" + Perform a fixed-grid purity-based axial calibration. + + The scan is centered on the current propagation distance + `reconstruction.zo`. Candidate axial positions are generated uniformly + within the interval + + $$ + z \in [z_0-\Delta z,\; z_0+\Delta z] + $$ + + where $z_0$ is the current axial-distance estimate and + `params.purityZScanRange` defines the half-width $\Delta z$ of the scan. + + The number of evaluated candidate positions is controlled by + `params.purityZScanPoints`. + + Each candidate distance is evaluated independently using + `_evaluatePurityAtZ()`, which creates a fresh reconstruction geometry, + performs an `mPIE` reconstruction, orthogonalizes the reconstructed probe + modes, and evaluates the resulting probe purity. + + The candidate with the highest purity is selected as the calibrated axial + distance. Its object, probe, momentum state, reconstruction buffers, and + error history are restored to the original `Reconstruction` object. + + The scan results are stored in: + + - `purityZValues`: evaluated axial distances in meters. + - `purityZMetrics`: reconstructed probe purity at each distance. + - `purityBestZ`: axial distance corresponding to the highest purity. + - `purityBestValue`: maximum reconstructed probe purity. + - `purityZInitialGuess`: axial distance used as the center of the scan. + + Returns: + tuple: + `(best_z, best_purity)`, where `best_z` is the calibrated axial + propagation distance in meters and `best_purity` is the maximum + reconstructed probe purity. + + Raises: + ValueError: + If fewer than two probe modes are reconstructed + (`reconstruction.npsm < 2`). + + RuntimeError: + If no valid candidate reconstruction produces a purity value. + + See Also: + `adaptivePurityZScan` + Adaptive alternative using variable axial step sizes. - Each candidate z is evaluated using an independent reconstruction - state and a standard mPIE reconstruction. """ if self.reconstruction.npsm < 2: @@ -848,16 +703,105 @@ def purityZScan(self): return best_z, best_purity def adaptivePurityZScan(self): - """ - Adaptively search for the propagation distance that maximizes - reconstructed probe purity. + r""" + Perform adaptive purity-based axial calibration. - The search first determines the uphill direction, then expands the - axial step while purity keeps increasing. Once the maximum is bracketed, - the step size is reduced and the search is refined around the current - best position. - """ + The adaptive scan searches for the axial distance that maximizes the + reconstructed probe purity while reducing the number of candidate + reconstructions compared with a dense fixed-grid scan. + + The search starts from the current propagation distance + `reconstruction.zo` and evaluates three initial positions: + + $$ + z_0-\Delta z,\quad z_0,\quad z_0+\Delta z + $$ + + where the initial step size $\Delta z$ is defined by + `params.purityZInitialStep`. + + If the purity increases toward one side, the search continues in that + direction and progressively expands the axial step: + + $$ + \Delta z_{k+1} = g\,\Delta z_k + $$ + + where the expansion factor $g$ is set by `params.purityZStepGrowth`. + Expansion continues while the increase in purity between consecutive + candidate positions exceeds `params.purityZPurityTolerance`. + + Once the purity stops increasing, the search enters a refinement phase. + The axial step is progressively reduced according to + + $$ + \Delta z_{k+1} = s\,\Delta z_k + $$ + + where the reduction factor $s$ is set by `params.purityZStepShrink`. + Candidate positions on both sides of the current best distance are then + evaluated until the requested axial resolution is reached. + + The adaptive search is controlled by the following parameters: + + - `params.purityZInitialStep`: + Initial axial step size used for the three-point search around the + starting distance. + + - `params.purityZStepGrowth`: + Multiplicative factor used to increase the step size while moving + toward higher purity. + + - `params.purityZStepShrink`: + Multiplicative factor used to reduce the step size during refinement. + + - `params.purityZMinStep`: + Minimum axial step size. Refinement stops once the current step becomes + smaller than this value. + + - `params.purityZPurityTolerance`: + Minimum purity improvement considered significant when determining + whether the search should continue in the current direction. + + - `params.purityZMaxEvaluations`: + Maximum number of candidate axial distances that may be reconstructed + during one adaptive scan. + + Candidate distances that have already been evaluated are cached and reused, + avoiding repeated `mPIE` reconstructions at the same axial position. + + Each new candidate is evaluated independently using + `_evaluatePurityAtZ()`. The candidate with the highest reconstructed probe + purity is retained as the calibrated axial distance, and its reconstruction + state is restored to the original `Reconstruction` object. + + The adaptive scan results are stored in: + + - `purityZValues`: evaluated axial distances in meters, sorted by z. + - `purityZMetrics`: reconstructed probe purity corresponding to each + evaluated distance. + - `purityBestZ`: axial distance corresponding to the highest purity. + - `purityBestValue`: maximum reconstructed probe purity. + - `purityZInitialGuess`: axial distance used as the starting point. + + Returns: + tuple: + `(best_z, best_purity)`, where `best_z` is the calibrated axial + propagation distance in meters and `best_purity` is the maximum + reconstructed probe purity. + + Raises: + ValueError: + If fewer than two probe modes are reconstructed + (`reconstruction.npsm < 2`). + RuntimeError: + If no valid candidate reconstruction produces a purity value. + + See Also: + `purityZScan` + Fixed-grid purity-based axial calibration. + """ if self.reconstruction.npsm < 2: raise ValueError( "Purity-based z scanning requires at least two probe modes " @@ -1255,111 +1199,4 @@ def plotPurityZScan(self, show=True): if show: plt.show() - return fig, ax - - - """ - Run one complete mPIE reconstruction iteration. - """ - - # set position order - self.setPositionOrder() - - self.pbar_pos = tqdm.tqdm( - self.positionIndices, - leave=False, - desc="ptychogram", - file=sys.stdout, - ) - - for positionLoop, positionIndex in enumerate(self.pbar_pos): - - row, col = self.reconstruction.positions[positionIndex] - sy = slice(row, row + self.reconstruction.Np) - sx = slice(col, col + self.reconstruction.Np) - - objectPatch = self.reconstruction.object[..., sy, sx].copy() - - # exit surface wave - self.reconstruction.esw = ( - objectPatch * self.reconstruction.probe - ) - - # intensity constraint - self.intensityProjection(positionIndex) - - DELTA = ( - self.reconstruction.eswUpdate - - self.reconstruction.esw - ) - - # object update - if ( - self.params.objectTVregSwitch - and loop % self.params.objectTVfreq == 0 - ): - object_patch = self.objectPatchUpdate_TV( - objectPatch, - DELTA, - ) - 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) - ) - - # probe update - weight = 1 - if self.params.weigh_probe_updates_by_intensity: - weight = self.experimentalData.relative_intensity( - positionIndex - ) - - self.reconstruction.probe = self.probeUpdate( - objectPatch, - DELTA, - weight, - ) - - # position correction - if self.params.positionCorrectionSwitch: - self.positionCorrection( - objectPatch, - positionIndex, - sy, - sx, - ) - - # momentum updates - if np.random.rand(1) > 0.95: - self.objectMomentumUpdate() - self.probeMomentumUpdate() - - self.getErrorMetrics() - self.applyConstraints(loop) - - 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 = ( - gradient + self.frictionM * self.reconstruction.probeMomentum - ) - self.reconstruction.probe = ( - self.reconstruction.probe - - self.feedbackM * self.reconstruction.probeMomentum - ) - self.reconstruction.probeBuffer = self.reconstruction.probe.copy() - - \ No newline at end of file + return fig, ax \ No newline at end of file diff --git a/PtyLab/Engines/zPIE.py b/PtyLab/Engines/zPIE.py index 3dff0f7d..7a7beebf 100644 --- a/PtyLab/Engines/zPIE.py +++ b/PtyLab/Engines/zPIE.py @@ -114,8 +114,38 @@ def __init__( def initializeReconstructionParams(self): """ - Set parameters that are specific to the ePIE settings. - :return: + Initialize zPIE-specific reconstruction parameters. + + The default parameters are: + + - `betaProbe = 0.25`: + Probe update step size used during the PIE reconstruction. + + - `betaObject = 0.25`: + Object update step size used during the PIE reconstruction. + + - `numIterations = 50`: + Number of reconstruction iterations. + + - `DoF = reconstruction.DoF`: + Depth-of-field scale used to define the axial search interval. + + - `zPIEgradientStepSize = 100`: + Scaling factor applied to the axial feedback term when updating the + sample-to-detector distance. + + - `zPIEfriction = 0.7`: + Momentum retention factor used for axial-distance refinement. + + - `focusObject = True`: + If `True`, evaluate the axial focus metric on the reconstructed object. + If `False`, evaluate it on the reconstructed probe. + + - `zMomentun = 0`: + Initial accumulated axial momentum used for iterative z refinement. + + These values can be modified after creating the zPIE engine and before + calling `reconstruct()`. """ self.betaProbe = 0.25 self.betaObject = 0.25 @@ -127,6 +157,48 @@ def initializeReconstructionParams(self): self.zMomentun = 0 def show_defocus(self, viewer=None, scanrange_times_dof=1000, N_points=10): + r""" + Visualize a defocus stack of the reconstructed object. + + A set of axial offsets is generated symmetrically around the current + reconstruction plane: + + $$ + \Delta z_k \in [-s\,\mathrm{DoF},\; s\,\mathrm{DoF}] + $$ + + where `s` is given by `scanrange_times_dof` and `DoF` is taken from + `reconstruction.DoF`. + + For each axial offset, the current reconstructed object is propagated + using angular-spectrum propagation (`aspw`). The propagated intensity + is calculated as the squared magnitude of the complex field and collected + into a defocus stack. + + The resulting stack is displayed in a napari viewer. If no viewer is + provided, a new napari viewer is created automatically. + + Args: + viewer (napari.Viewer, optional): + Existing napari viewer to which the defocus stack is added. + If `None`, a new viewer is created. + + scanrange_times_dof (float, optional): + Half-width of the axial scan expressed in units of the + reconstruction depth of field. Default is `1000`. + + N_points (int, optional): + Number of equally spaced axial positions in the defocus stack. + Default is `10`. + + Raises: + ImportError: + If `viewer` is `None` and napari is not installed. + + Notes: + This function is intended as a diagnostic visualization tool and does + not modify the reconstructed object or the current axial distance. + """ z = np.linspace(-1, 1, N_points) * scanrange_times_dof * self.reconstruction.DoF from PtyLab.Operators.Operators import aspw @@ -161,6 +233,92 @@ def show_defocus(self, viewer=None, scanrange_times_dof=1000, N_points=10): viewer.add_image(defocii) def reconstruct(self, experimentalData=None, reconstruction=None): + r""" + Run the zPIE reconstruction with iterative axial-distance refinement. + + The reconstruction jointly updates the object, probe, and + sample-to-detector propagation distance. Before the iterative loop starts, + the experimental data and reconstruction state are updated if new objects + are supplied, followed by the standard reconstruction preparation. + + During each zPIE iteration, a set of candidate axial offsets is generated + around the current propagation distance: + + $$ + \Delta z_k \in [-10\,\mathrm{DoF},\;10\,\mathrm{DoF}] + $$ + + using 11 equally spaced candidate positions. + + Depending on `focusObject`, either the reconstructed object or probe is + propagated to each candidate defocus plane using angular-spectrum + propagation. A total-variation-based merit value is calculated for each + propagated field. + + The axial feedback is then obtained from the weighted candidate offsets: + + $$ + f_z = \frac{\sum_k \Delta z_k M_k}{\sum_k M_k} + $$ + + where $M_k$ is the focus merit evaluated at candidate offset + $\Delta z_k$. + + The feedback is accumulated through the axial momentum update + + $$ + m_z^{(n)} = \gamma m_z^{(n-1)} + \eta f_z + $$ + + where $\gamma$ is `zPIEfriction` and $\eta$ is + `zPIEgradientStepSize`. The propagation distance is then updated as + + $$ + z^{(n+1)} = z^{(n)} + m_z^{(n)} + $$ + + After the axial update, a standard PIE reconstruction loop is performed + over all scan positions. The exit surface wave is calculated from the + current object patch and probe, followed by the intensity constraint and + object/probe updates. + + The current axial distance is stored during reconstruction in + `reconstruction.zHistory`. The total-variation merit values and candidate + axial offsets are also stored in `reconstruction.merit` and + `reconstruction.dz` for diagnostic use. + + For propagation schemes whose object-plane sampling depends explicitly on + the propagation distance, the sampling `reconstruction.dxp` is updated + after each axial correction. + + At the end of each iteration, reconstruction error metrics and enabled + constraints are evaluated, followed by an update of the reconstruction + monitor. + + Args: + experimentalData (ExperimentalData, optional): + Experimental dataset to use. If provided, it replaces the dataset + currently attached to the engine. + + reconstruction (Reconstruction, optional): + Reconstruction state to optimize. If provided, it replaces the + reconstruction currently attached to the engine. + + Returns: + None + + See Also: + `show_defocus` + Visualize the reconstructed object over a range of axial defocus + positions. + + `objectPatchUpdate` + Update the reconstructed object patch during the PIE loop. + + `probeUpdate` + Update the reconstructed probe during the PIE loop. + """ + self.changeExperimentalData(experimentalData) self.changeOptimizable(reconstruction) self._prepareReconstruction() @@ -383,10 +541,13 @@ def reconstruct(self, experimentalData=None, reconstruction=None): def objectPatchUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray): """ - Todo add docstring - :param objectPatch: - :param DELTA: - :return: + Update the object patch using the standard ePIE object-update rule. + + This implementation is identical to `ePIE.objectPatchUpdate()`. + + See Also: + `ePIE.objectPatchUpdate` + Standard ePIE object update. """ # find out which array module to use, numpy or cupy (or other...) xp = getArrayModule(objectPatch) @@ -400,10 +561,13 @@ 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: + Update the probe using the standard ePIE probe-update rule. + + This implementation is identical to `ePIE.probeUpdate()`. + + See Also: + `ePIE.probeUpdate` + Standard ePIE probe update. """ # 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 bc68dbe4..541dbafe 100644 --- a/docs/api/engines.md +++ b/docs/api/engines.md @@ -48,6 +48,12 @@ show_root_full_path: false inherited_members: false +::: PtyLab.Engines.purityPIE.purityPIE + options: + show_root_heading: true + show_root_full_path: false + inherited_members: false + ::: PtyLab.Engines.zPIE.zPIE options: show_root_heading: true From fb2b1ff6d496fdfa73ce7f17da23a15fbcf851a1 Mon Sep 17 00:00:00 2001 From: Chang Liu Date: Mon, 5 Oct 2026 11:13:25 +0200 Subject: [PATCH 08/11] minor change --- PtyLab/Engines/zPIE.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/PtyLab/Engines/zPIE.py b/PtyLab/Engines/zPIE.py index 7a7beebf..7861cb10 100644 --- a/PtyLab/Engines/zPIE.py +++ b/PtyLab/Engines/zPIE.py @@ -304,9 +304,6 @@ def reconstruct(self, experimentalData=None, reconstruction=None): Reconstruction state to optimize. If provided, it replaces the reconstruction currently attached to the engine. - Returns: - None - See Also: `show_defocus` Visualize the reconstructed object over a range of axial defocus From b26292b5d3ff9f185a421828021996513ee66509 Mon Sep 17 00:00:00 2001 From: Chang Liu Date: Tue, 6 Oct 2026 11:44:28 +0200 Subject: [PATCH 09/11] minor change: delete old code. --- PtyLab/Engines/purityPIE.py | 412 ------------------------------------ 1 file changed, 412 deletions(-) diff --git a/PtyLab/Engines/purityPIE.py b/PtyLab/Engines/purityPIE.py index ffecdc0f..ec6a24b7 100644 --- a/PtyLab/Engines/purityPIE.py +++ b/PtyLab/Engines/purityPIE.py @@ -164,315 +164,6 @@ def _createCandidateReconstruction(self, z): candidate_params, ) - - """ - Scan a fixed axial range and select the distance that maximizes - reconstructed probe purity. - - Each candidate z is evaluated using an independent reconstruction - state and a standard mPIE reconstruction. - """ - - if self.reconstruction.npsm < 2: - raise ValueError( - "Purity-based z scanning requires at least two probe modes " - "(reconstruction.npsm >= 2)." - ) - - # ------------------------------------------------------------------ - # Scan grid - # ------------------------------------------------------------------ - - z0 = self.reconstruction.zo - self.purityZInitialGuess = z0 - - z_values = np.linspace( - z0 - self.params.purityZScanRange, - z0 + self.params.purityZScanRange, - self.params.purityZScanPoints, - ) - - purity_values = [] - - best_purity = -np.inf - best_z = z0 - - best_reconstruction = None - - # ------------------------------------------------------------------ - # Scan information - # ------------------------------------------------------------------ - - tqdm.tqdm.write("") - tqdm.tqdm.write("Starting purity-based z scan") - tqdm.tqdm.write( - f"Initial z = {z0 * 1e3:.6f} mm | " - f"range = ±{self.params.purityZScanRange * 1e6:.1f} µm | " - f"points = {len(z_values)} | " - f"iterations/z = {self.numIterations}" - ) - tqdm.tqdm.write("") - - # ------------------------------------------------------------------ - # Candidate-z loop - # ------------------------------------------------------------------ - - for z_index, z in enumerate(z_values): - - dz_um = (z - z0) * 1e6 - - tqdm.tqdm.write( - f"[{z_index + 1}/{len(z_values)}] " - f"Reconstructing z = {z * 1e3:.6f} mm " - f"(Δz = {dz_um:+.2f} µm)" - ) - - # -------------------------------------------------------------- - # Create a completely independent reconstruction for this z - # -------------------------------------------------------------- - #np.random.seed(0) - ( - candidate_data, - candidate_reconstruction, - candidate_params, - ) = self._createCandidateReconstruction(z) - - # Candidate reconstructions do not need their own GUI monitor - candidate_monitor = DummyMonitor() - - # -------------------------------------------------------------- - # Run standard mPIE - # -------------------------------------------------------------- - - candidate_engine = mPIE( - candidate_reconstruction, - candidate_data, - candidate_params, - candidate_monitor, - ) - - # Match the purityPIE reconstruction settings - candidate_engine.numIterations = self.numIterations - - candidate_engine.betaProbe = self.betaProbe - candidate_engine.betaObject = self.betaObject - - candidate_engine.alphaProbe = self.alphaProbe - candidate_engine.alphaObject = self.alphaObject - - candidate_engine.feedbackM = self.feedbackM - candidate_engine.frictionM = self.frictionM - - candidate_engine.reconstruct() - - # -------------------------------------------------------------- - # Final modal decomposition and purity evaluation - # -------------------------------------------------------------- - - candidate_engine.orthogonalization() - - purity = float( - np.asarray( - asNumpyArray(candidate_reconstruction.purityProbe) - ).squeeze() - ) - - purity_values.append(purity) - - tqdm.tqdm.write( - f" Probe purity = {purity:.6f}" - ) - - # -------------------------------------------------------------- - # Update main monitor with the completed candidate state - # -------------------------------------------------------------- - - self.reconstruction.zo = z - self.reconstruction.object = ( - asNumpyArray(candidate_reconstruction.object).copy() - ) - self.reconstruction.probe = ( - asNumpyArray(candidate_reconstruction.probe).copy() - ) - - self.reconstruction.purityProbe = purity - - if hasattr(candidate_reconstruction, "error"): - self.reconstruction.error = np.asarray( - asNumpyArray(candidate_reconstruction.error) - ).copy() - - # Display only completed candidate reconstructions - self.showReconstruction( - self.numIterations - 1, - force=True, - ) - - # -------------------------------------------------------------- - # Keep best candidate - # -------------------------------------------------------------- - - if purity > best_purity: - - best_purity = purity - best_z = z - - # Candidate objects are independent, so keeping this reference - # preserves the complete best reconstruction state. - best_reconstruction = candidate_reconstruction - - tqdm.tqdm.write( - f" NEW BEST: z = {best_z * 1e3:.6f} mm, " - f"purity = {best_purity:.6f}" - ) - - tqdm.tqdm.write("") - - # ------------------------------------------------------------------ - # Safety check - # ------------------------------------------------------------------ - - if best_reconstruction is None: - raise RuntimeError( - "Purity z scan did not produce a valid purity value." - ) - - # ------------------------------------------------------------------ - # Restore the best candidate into the original Reconstruction object - # - # Keep the original object identity because external scripts may still - # hold a reference to self.reconstruction. - # ------------------------------------------------------------------ - - self.reconstruction.zo = best_z - - self.reconstruction.object = ( - asNumpyArray(best_reconstruction.object).copy() - ) - - self.reconstruction.probe = ( - asNumpyArray(best_reconstruction.probe).copy() - ) - - self.reconstruction.purityProbe = best_purity - - if hasattr(best_reconstruction, "encoder_corrected"): - self.reconstruction.encoder_corrected = ( - asNumpyArray( - best_reconstruction.encoder_corrected - ).copy() - ) - - if hasattr(best_reconstruction, "objectMomentum"): - self.reconstruction.objectMomentum = ( - asNumpyArray( - best_reconstruction.objectMomentum - ).copy() - ) - - if hasattr(best_reconstruction, "probeMomentum"): - self.reconstruction.probeMomentum = ( - asNumpyArray( - best_reconstruction.probeMomentum - ).copy() - ) - - if hasattr(best_reconstruction, "objectBuffer"): - self.reconstruction.objectBuffer = ( - asNumpyArray( - best_reconstruction.objectBuffer - ).copy() - ) - - if hasattr(best_reconstruction, "probeBuffer"): - self.reconstruction.probeBuffer = ( - asNumpyArray( - best_reconstruction.probeBuffer - ).copy() - ) - - if hasattr(best_reconstruction, "error"): - self.reconstruction.error = np.asarray( - asNumpyArray(best_reconstruction.error) - ).copy() - - # ------------------------------------------------------------------ - # Store scan results - # ------------------------------------------------------------------ - - self.purityZValues = np.asarray(z_values) - self.purityZMetrics = np.asarray(purity_values) - - self.purityBestZ = best_z - self.purityBestValue = best_purity - - # Final monitor update with the selected best candidate - self.showReconstruction( - self.numIterations - 1, - force=True, - ) - - # ------------------------------------------------------------------ - # Summary - # ------------------------------------------------------------------ - - tqdm.tqdm.write("=" * 60) - tqdm.tqdm.write("Purity-based z scan finished") - tqdm.tqdm.write( - f"Best z = {best_z * 1e3:.6f} mm" - ) - tqdm.tqdm.write( - f"Best Δz = {(best_z - z0) * 1e6:+.2f} µm" - ) - tqdm.tqdm.write( - f"Best purity = {best_purity:.6f}" - ) - tqdm.tqdm.write("=" * 60) - - return best_z, best_purity - - - """ - Reset the reconstruction state for one candidate z. - - This implementation is intended for analytic initialization: - initialProbe = "circ" - initialObject = "ones" - - The object/probe are regenerated after setting the candidate z so that - their physical sampling is consistent with the candidate geometry. - Momentum and buffers are also reset. - """ - if self.reconstruction.initialProbe != "circ": - raise NotImplementedError( - "Candidate-z reinitialization is currently implemented " - "only for initialProbe='circ'." - ) - - if self.reconstruction.initialObject != "ones": - raise NotImplementedError( - "Candidate-z reinitialization is currently implemented " - "only for initialObject='ones'." - ) - - # Set candidate geometry first - self.reconstruction.zo = z - - # Rebuild object/probe on the candidate-z geometry - self.reconstruction.initializeObjectProbe() - - self._initialProbePowerCorrection() - - # Reset momentum - self.reconstruction.initializeObjectMomentum() - self.reconstruction.initializeProbeMomentum() - - # Reset buffers - self.reconstruction.objectBuffer = self.reconstruction.object.copy() - self.reconstruction.probeBuffer = self.reconstruction.probe.copy() - - self._prepareReconstruction() - def _evaluatePurityAtZ(self, z): """ Evaluate probe purity for one candidate propagation distance. @@ -1258,108 +949,5 @@ def plotPurityZScan(self, show=True): return fig, ax - """ - Run one complete mPIE reconstruction iteration. - """ - - # set position order - self.setPositionOrder() - - self.pbar_pos = tqdm.tqdm( - self.positionIndices, - leave=False, - desc="ptychogram", - file=sys.stdout, - ) - - for positionLoop, positionIndex in enumerate(self.pbar_pos): - - row, col = self.reconstruction.positions[positionIndex] - sy = slice(row, row + self.reconstruction.Np) - sx = slice(col, col + self.reconstruction.Np) - - objectPatch = self.reconstruction.object[..., sy, sx].copy() - - # exit surface wave - self.reconstruction.esw = ( - objectPatch * self.reconstruction.probe - ) - - # intensity constraint - self.intensityProjection(positionIndex) - - DELTA = ( - self.reconstruction.eswUpdate - - self.reconstruction.esw - ) - - # object update - if ( - self.params.objectTVregSwitch - and loop % self.params.objectTVfreq == 0 - ): - object_patch = self.objectPatchUpdate_TV( - objectPatch, - DELTA, - ) - 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) - ) - - # probe update - weight = 1 - if self.params.weigh_probe_updates_by_intensity: - weight = self.experimentalData.relative_intensity( - positionIndex - ) - - self.reconstruction.probe = self.probeUpdate( - objectPatch, - DELTA, - weight, - ) - - # position correction - if self.params.positionCorrectionSwitch: - self.positionCorrection( - objectPatch, - positionIndex, - sy, - sx, - ) - - # momentum updates - if np.random.rand(1) > 0.95: - self.objectMomentumUpdate() - self.probeMomentumUpdate() - - self.getErrorMetrics() - self.applyConstraints(loop) - - 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 = ( - gradient + self.frictionM * self.reconstruction.probeMomentum - ) - self.reconstruction.probe = ( - self.reconstruction.probe - - self.feedbackM * self.reconstruction.probeMomentum - ) - self.reconstruction.probeBuffer = self.reconstruction.probe.copy() \ No newline at end of file From 8379a6ebbb69f0a612c882143a914e3a773c3456 Mon Sep 17 00:00:00 2001 From: Chang Liu Date: Wed, 7 Oct 2026 11:44:23 +0200 Subject: [PATCH 10/11] Extend purityPIE with wavelength scans, fitted optima, and error comparison - add fixed and adaptive wavelength calibration - add local quadratic fitting of purity scan maxima - store and plot reconstruction error alongside probe purity - unify scan plotting for z and wavelength calibration - update the example script for selectable z or wavelength calibration --- PtyLab/Engines/purityPIE.py | 2061 +++++++++++++++-- PtyLab/Params/Params.py | 24 +- .../example_purity_Zscan_yourData.py | 343 ++- 3 files changed, 2202 insertions(+), 226 deletions(-) diff --git a/PtyLab/Engines/purityPIE.py b/PtyLab/Engines/purityPIE.py index ec6a24b7..1ce7d51a 100644 --- a/PtyLab/Engines/purityPIE.py +++ b/PtyLab/Engines/purityPIE.py @@ -93,69 +93,156 @@ def reconstruct( self._prepareReconstruction() # Reset momentum buffers after preparation - self.reconstruction.objectBuffer = self.reconstruction.object.copy() - self.reconstruction.probeBuffer = self.reconstruction.probe.copy() - - # Purity-based axial calibration - if self.params.purityZAdaptive: - best_z, best_purity = self.adaptivePurityZScan() + self.reconstruction.objectBuffer = ( + self.reconstruction.object.copy() + ) + + self.reconstruction.probeBuffer = ( + self.reconstruction.probe.copy() + ) + + target = self.params.purityCalibrationTarget.lower() + + # ------------------------------------------------------------------ + # Axial-distance calibration + # ------------------------------------------------------------------ + + if target == "z": + + if self.params.purityAdaptive: + best_value, best_purity = ( + self.adaptivePurityZScan() + ) + else: + best_value, best_purity = ( + self.purityZScan() + ) + + # ------------------------------------------------------------------ + # Wavelength calibration + # ------------------------------------------------------------------ + + elif target == "wavelength": + + if self.params.purityAdaptive: + best_value, best_purity = ( + self.adaptivePurityWavelengthScan() + ) + else: + best_value, best_purity = ( + self.purityWavelengthScan() + ) + + # ------------------------------------------------------------------ + # Unsupported target + # ------------------------------------------------------------------ + else: - best_z, best_purity = self.purityZScan() - # Optional scan visualization - if self.params.purityZScanPlot: - self.plotPurityZScan() + raise ValueError( + "Unsupported purity calibration target: " + f"{self.params.purityCalibrationTarget!r}. " + "Supported targets are 'z' and 'wavelength'." + ) + + if self.params.purityScanPlot: + self.plotPurityScan() + # ------------------------------------------------------------------ # Return reconstruction data to CPU + # ------------------------------------------------------------------ + if self.params.gpuFlag: self.logger.info("switch to cpu") self._move_data_to_cpu() self.params.gpuFlag = 0 - return best_z, best_purity + return best_value, best_purity - def _createCandidateReconstruction(self, z): + def _createCandidateReconstruction( + self, + target, + value, + ): """ - Create an independent reconstruction state for one candidate z. - - The candidate inherits the user's reconstruction settings but is built - from a fresh geometry at the requested propagation distance. + Create an independent reconstruction for one candidate calibration value. """ - # Independent copies so one candidate cannot modify another - candidate_data = deepcopy(self.experimentalData) - candidate_params = deepcopy(self.params) + candidate_data = deepcopy( + self.experimentalData + ) + + candidate_params = deepcopy( + self.params + ) + + # -------------------------------------------------------------- + # Set candidate parameter before creating Reconstruction + # -------------------------------------------------------------- - # Set candidate geometry before creating Reconstruction - candidate_data.zo = z + if target == "z": + candidate_data.zo = value + + elif target == "wavelength": + candidate_data.wavelength = value + + if self.reconstruction.nlambda == 1: + candidate_data.spectralDensity = np.atleast_1d(value) + else: + raise ValueError( + f"Unsupported purity calibration target: {target!r}" + ) + # Fresh geometry for this candidate candidate_reconstruction = Reconstruction( candidate_data, candidate_params, ) - # Copy user-defined reconstruction settings - candidate_reconstruction.npsm = self.reconstruction.npsm - candidate_reconstruction.nosm = self.reconstruction.nosm - candidate_reconstruction.nlambda = self.reconstruction.nlambda - candidate_reconstruction.nslice = self.reconstruction.nslice + # -------------------------------------------------------------- + # Copy reconstruction settings + # -------------------------------------------------------------- + + candidate_reconstruction.npsm = ( + self.reconstruction.npsm + ) + + candidate_reconstruction.nosm = ( + self.reconstruction.nosm + ) + + candidate_reconstruction.nlambda = ( + self.reconstruction.nlambda + ) + + candidate_reconstruction.nslice = ( + self.reconstruction.nslice + ) - #candidate_reconstruction.No = self.reconstruction.No + candidate_reconstruction.initialProbe = ( + self.reconstruction.initialProbe + ) - candidate_reconstruction.initialProbe = self.reconstruction.initialProbe - candidate_reconstruction.initialObject = self.reconstruction.initialObject + candidate_reconstruction.initialObject = ( + self.reconstruction.initialObject + ) - if hasattr(self.reconstruction, "initialProbe_filename"): + if hasattr( + self.reconstruction, + "initialProbe_filename", + ): candidate_reconstruction.initialProbe_filename = ( self.reconstruction.initialProbe_filename ) - if hasattr(self.reconstruction, "initialObject_filename"): + if hasattr( + self.reconstruction, + "initialObject_filename", + ): candidate_reconstruction.initialObject_filename = ( self.reconstruction.initialObject_filename ) - # Fresh object/probe for this candidate geometry candidate_reconstruction.initializeObjectProbe() return ( @@ -165,44 +252,40 @@ def _createCandidateReconstruction(self, z): ) def _evaluatePurityAtZ(self, z): - """ - Evaluate probe purity for one candidate propagation distance. - - A fresh reconstruction is created at the requested z value and - reconstructed using the standard mPIE engine. Probe modes are - orthogonalized after reconstruction and the resulting probe purity - is returned together with the candidate reconstruction state. - - Parameters - ---------- - z : float - Candidate propagation distance in meters. + return self._evaluatePurityAtParameter( + "z", + z, + ) - Returns - ------- - purity : float - Reconstructed probe purity at the candidate distance. + def _evaluatePurityAtWavelength( + self, + wavelength, + ): + return self._evaluatePurityAtParameter( + "wavelength", + wavelength, + ) - candidate_reconstruction : Reconstruction - Reconstruction state obtained at the candidate distance. + def _evaluatePurityAtParameter( + self, + target, + value, + ): + """ + Evaluate reconstructed probe purity for one candidate calibration value. """ - - # ------------------------------------------------------------------ - # Create independent candidate reconstruction - # ------------------------------------------------------------------ ( candidate_data, candidate_reconstruction, candidate_params, - ) = self._createCandidateReconstruction(z) + ) = self._createCandidateReconstruction( + target, + value, + ) candidate_monitor = DummyMonitor() - # ------------------------------------------------------------------ - # Run standard mPIE - # ------------------------------------------------------------------ - candidate_engine = mPIE( candidate_reconstruction, candidate_data, @@ -210,23 +293,35 @@ def _evaluatePurityAtZ(self, z): candidate_monitor, ) - # Match the purityPIE engine settings - candidate_engine.numIterations = self.numIterations + candidate_engine.numIterations = ( + self.numIterations + ) - candidate_engine.betaProbe = self.betaProbe - candidate_engine.betaObject = self.betaObject + candidate_engine.betaProbe = ( + self.betaProbe + ) - candidate_engine.alphaProbe = self.alphaProbe - candidate_engine.alphaObject = self.alphaObject + candidate_engine.betaObject = ( + self.betaObject + ) - candidate_engine.feedbackM = self.feedbackM - candidate_engine.frictionM = self.frictionM + candidate_engine.alphaProbe = ( + self.alphaProbe + ) - candidate_engine.reconstruct() + candidate_engine.alphaObject = ( + self.alphaObject + ) - # ------------------------------------------------------------------ - # Final modal decomposition and purity evaluation - # ------------------------------------------------------------------ + candidate_engine.feedbackM = ( + self.feedbackM + ) + + candidate_engine.frictionM = ( + self.frictionM + ) + + candidate_engine.reconstruct() candidate_engine.orthogonalization() @@ -238,7 +333,10 @@ def _evaluatePurityAtZ(self, z): ).squeeze() ) - return purity, candidate_reconstruction + return ( + purity, + candidate_reconstruction, + ) def purityZScan(self): """ @@ -269,6 +367,7 @@ def purityZScan(self): ) purity_values = [] + error_values = [] best_purity = -np.inf best_z = z0 @@ -310,10 +409,20 @@ def purityZScan(self): self._evaluatePurityAtZ(z) ) + candidate_error = float( + np.asarray( + asNumpyArray( + candidate_reconstruction.error + ) + ).reshape(-1)[-1] + ) + purity_values.append(purity) + error_values.append(candidate_error) tqdm.tqdm.write( - f" Probe purity = {purity:.6f}" + f" Probe purity = {purity:.6f} | " + f"reconstruction error = {candidate_error:.6e}" ) # -------------------------------------------------------------- @@ -505,6 +614,10 @@ def purityZScan(self): purity_values ) + self.purityZErrors = np.asarray( + error_values + ) + self.purityBestZ = best_z self.purityBestValue = best_purity @@ -563,10 +676,9 @@ def adaptivePurityZScan(self): shrink = self.params.purityZStepShrink min_step = self.params.purityZMinStep purity_tol = self.params.purityZPurityTolerance - max_evaluations = self.params.purityZMaxEvaluations + max_evaluations = self.params.purityMaxEvaluations evaluated = {} - reconstruction_states = {} best_z = None best_purity = -np.inf @@ -574,6 +686,7 @@ def adaptivePurityZScan(self): z_history = [] purity_history = [] + error_history = [] def evaluate(z): nonlocal best_z @@ -589,15 +702,23 @@ def evaluate(z): purity, candidate_reconstruction = ( self._evaluatePurityAtZ(z) ) - + candidate_error = float( + np.asarray( + asNumpyArray( + candidate_reconstruction.error + ) + ).reshape(-1)[-1] + ) + evaluated[key] = purity z_history.append(z) purity_history.append(purity) - + error_history.append(candidate_error) tqdm.tqdm.write( f"z = {z * 1e3:.6f} mm | " - f"purity = {purity:.6f}" + f"purity = {purity:.6f} | " + f"error = {candidate_error:.6e}" ) # Store best candidate state @@ -839,6 +960,7 @@ def evaluate(z): self.purityZValues = np.asarray(z_history)[order] self.purityZMetrics = np.asarray(purity_history)[order] + self.purityZErrors = np.asarray(error_history)[order] self.purityBestZ = best_z self.purityBestValue = best_purity @@ -869,85 +991,1746 @@ def evaluate(z): tqdm.tqdm.write("=" * 60) return best_z, best_purity + + def purityWavelengthScan(self): - def plotPurityZScan(self, show=True): - """ - Plot probe purity as a function of axial distance. - - The axial coordinate is shown relative to the initial z guess used - for the purity scan. - Args: - show (bool, optional): - Display the figure immediately. Default is `True`. + """ + Scan a fixed wavelength range and select the wavelength that maximizes + reconstructed probe purity. - Returns: - tuple: - Matplotlib `(figure, axis)` objects. + Each candidate wavelength is evaluated using an independent reconstruction + state and a standard mPIE reconstruction. """ - if not hasattr(self, "purityZValues"): - raise RuntimeError( - "No purity z-scan results are available. " - "Run purityZScan() first." + if self.reconstruction.npsm < 2: + raise ValueError( + "Purity-based wavelength scanning requires at least two probe modes " + "(reconstruction.npsm >= 2)." ) - if not hasattr(self, "purityZInitialGuess"): - raise RuntimeError( - "Initial z guess is not available." - ) + wavelength0 = float( + self.reconstruction.wavelength + ) - z_relative_um = ( - np.asarray(self.purityZValues) - - self.purityZInitialGuess - ) * 1e6 + self.purityWavelengthInitialGuess = ( + wavelength0 + ) - purity = np.asarray(self.purityZMetrics) + wavelength_values = np.linspace( + wavelength0 + - self.params.purityWavelengthScanRange, + wavelength0 + + self.params.purityWavelengthScanRange, + self.params.purityWavelengthScanPoints, + ) - best_index = np.argmax(purity) + purity_values = [] + error_values = [] - fig, ax = plt.subplots(figsize=(6, 4)) + best_purity = -np.inf + best_wavelength = wavelength0 + best_state = None - ax.plot( - z_relative_um, - purity, - marker="o", - label="Probe purity", + tqdm.tqdm.write("") + tqdm.tqdm.write( + "Starting purity-based wavelength scan" ) - # Initial z guess - ax.axvline( - 0, - linestyle="--", - alpha=0.5, - label="Initial z guess", + tqdm.tqdm.write( + f"Initial wavelength = " + f"{wavelength0 * 1e9:.6f} nm | " + f"range = ±" + f"{self.params.purityWavelengthScanRange * 1e9:.4f} nm | " + f"points = {len(wavelength_values)} | " + f"iterations/wavelength = {self.numIterations}" ) - # Best z - ax.plot( - z_relative_um[best_index], - purity[best_index], - marker="*", - markersize=12, - linestyle="None", - label=( - f"Best Δz = {z_relative_um[best_index]:.1f} µm" - ), - ) + tqdm.tqdm.write("") - ax.set_xlabel("Axial offset from initial z (µm)") - ax.set_ylabel("Probe purity") - ax.set_title("Purity-based axial calibration") + for wavelength_index, wavelength in enumerate( + wavelength_values + ): - ax.grid(alpha=0.3) - ax.legend() - fig.tight_layout() + delta_wavelength_nm = ( + wavelength - wavelength0 + ) * 1e9 - if show: - plt.show() + tqdm.tqdm.write( + f"[{wavelength_index + 1}/" + f"{len(wavelength_values)}] " + f"Reconstructing wavelength = " + f"{wavelength * 1e9:.6f} nm " + f"(Δλ = {delta_wavelength_nm:+.4f} nm)" + ) + + ( + purity, + candidate_reconstruction, + ) = self._evaluatePurityAtWavelength( + wavelength + ) + + candidate_error = float( + np.asarray( + asNumpyArray( + candidate_reconstruction.error + ) + ).reshape(-1)[-1] + ) + + purity_values.append( + purity + ) + + error_values.append( + candidate_error + ) + + tqdm.tqdm.write( + f" Probe purity = {purity:.6f} | " + f"reconstruction error = {candidate_error:.6e}" + ) + + # Update monitor with current candidate + self.reconstruction.wavelength = ( + wavelength + ) + + self.reconstruction.dxp = ( + candidate_reconstruction.dxp + ) + + self.reconstruction.object = ( + asNumpyArray( + candidate_reconstruction.object + ).copy() + ) + + self.reconstruction.probe = ( + asNumpyArray( + candidate_reconstruction.probe + ).copy() + ) + + self.reconstruction.purityProbe = ( + purity + ) + + if hasattr( + candidate_reconstruction, + "error", + ): + self.reconstruction.error = ( + np.asarray( + asNumpyArray( + candidate_reconstruction.error + ) + ).copy() + ) + + self.showReconstruction( + self.numIterations - 1, + force=True, + ) + + if purity > best_purity: - return fig, ax - - + best_purity = purity + best_wavelength = wavelength + + best_state = { + "object": asNumpyArray( + candidate_reconstruction.object + ).copy(), + + "probe": asNumpyArray( + candidate_reconstruction.probe + ).copy(), + + "purityProbe": purity, + + "dxp": float( + candidate_reconstruction.dxp + ), + } + + if hasattr( + candidate_reconstruction, + "spectralDensity", + ): + best_state["spectralDensity"] = ( + np.asarray( + candidate_reconstruction.spectralDensity + ).copy() + ) + + if hasattr( + candidate_reconstruction, + "encoder_corrected", + ): + best_state["encoder_corrected"] = ( + asNumpyArray( + candidate_reconstruction.encoder_corrected + ).copy() + ) + + if hasattr( + candidate_reconstruction, + "objectMomentum", + ): + best_state["objectMomentum"] = ( + asNumpyArray( + candidate_reconstruction.objectMomentum + ).copy() + ) + + if hasattr( + candidate_reconstruction, + "probeMomentum", + ): + best_state["probeMomentum"] = ( + asNumpyArray( + candidate_reconstruction.probeMomentum + ).copy() + ) + + if hasattr( + candidate_reconstruction, + "objectBuffer", + ): + best_state["objectBuffer"] = ( + asNumpyArray( + candidate_reconstruction.objectBuffer + ).copy() + ) - \ No newline at end of file + if hasattr( + candidate_reconstruction, + "probeBuffer", + ): + best_state["probeBuffer"] = ( + asNumpyArray( + candidate_reconstruction.probeBuffer + ).copy() + ) + + if hasattr( + candidate_reconstruction, + "error", + ): + best_state["error"] = ( + np.asarray( + asNumpyArray( + candidate_reconstruction.error + ) + ).copy() + ) + + tqdm.tqdm.write( + f" NEW BEST: wavelength = " + f"{best_wavelength * 1e9:.6f} nm | " + f"purity = {best_purity:.6f}" + ) + + tqdm.tqdm.write("") + + if best_state is None: + raise RuntimeError( + "Purity wavelength scan did not produce " + "a valid purity value." + ) + + # Restore best candidate + self.reconstruction.wavelength = ( + best_wavelength + ) + + self.reconstruction.dxp = ( + best_state["dxp"] + ) + + if "spectralDensity" in best_state: + self.reconstruction.spectralDensity = ( + best_state["spectralDensity"] + ) + + self.reconstruction.object = ( + best_state["object"] + ) + + self.reconstruction.probe = ( + best_state["probe"] + ) + + self.reconstruction.purityProbe = ( + best_state["purityProbe"] + ) + + if "encoder_corrected" in best_state: + self.reconstruction.encoder_corrected = ( + best_state["encoder_corrected"] + ) + + if "objectMomentum" in best_state: + self.reconstruction.objectMomentum = ( + best_state["objectMomentum"] + ) + + if "probeMomentum" in best_state: + self.reconstruction.probeMomentum = ( + best_state["probeMomentum"] + ) + + if "objectBuffer" in best_state: + self.reconstruction.objectBuffer = ( + best_state["objectBuffer"] + ) + + if "probeBuffer" in best_state: + self.reconstruction.probeBuffer = ( + best_state["probeBuffer"] + ) + + if "error" in best_state: + self.reconstruction.error = ( + best_state["error"] + ) + + # Store scan results + self.purityWavelengthValues = ( + np.asarray( + wavelength_values + ) + ) + + self.purityWavelengthMetrics = ( + np.asarray( + purity_values + ) + ) + + self.purityWavelengthErrors = ( + np.asarray( + error_values + ) + ) + + self.purityBestWavelength = ( + best_wavelength + ) + + self.purityBestWavelengthValue = ( + best_purity + ) + + self.showReconstruction( + self.numIterations - 1, + force=True, + ) + + tqdm.tqdm.write("=" * 60) + tqdm.tqdm.write( + "Purity-based wavelength scan finished" + ) + + tqdm.tqdm.write( + f"Best wavelength = " + f"{best_wavelength * 1e9:.6f} nm" + ) + + tqdm.tqdm.write( + f"Best Δλ = " + f"{(best_wavelength - wavelength0) * 1e9:+.4f} nm" + ) + + tqdm.tqdm.write( + f"Best purity = " + f"{best_purity:.6f}" + ) + + tqdm.tqdm.write("=" * 60) + + return ( + best_wavelength, + best_purity, + ) + + def adaptivePurityWavelengthScan(self): + + """ + Adaptively search for the wavelength that maximizes reconstructed + probe purity. + + The search first determines the uphill direction, then expands the + wavelength step while purity keeps increasing. Once the maximum is + bracketed, the step size is reduced and the search is refined around + the current best wavelength. + """ + + if self.reconstruction.npsm < 2: + raise ValueError( + "Purity-based wavelength scanning requires at least two probe modes " + "(reconstruction.npsm >= 2)." + ) + + wavelength0 = float( + self.reconstruction.wavelength + ) + + self.purityWavelengthInitialGuess = ( + wavelength0 + ) + + step = self.params.purityWavelengthInitialStep + growth = self.params.purityWavelengthStepGrowth + shrink = self.params.purityWavelengthStepShrink + min_step = self.params.purityWavelengthMinStep + purity_tol = self.params.purityWavelengthPurityTolerance + max_evaluations = self.params.purityMaxEvaluations + + evaluated = {} + + best_wavelength = None + best_purity = -np.inf + best_state = None + + wavelength_history = [] + purity_history = [] + error_history = [] + + def evaluate(wavelength): + + nonlocal best_wavelength + nonlocal best_purity + nonlocal best_state + + if wavelength <= 0: + return -np.inf + + key = float(wavelength) + + if key in evaluated: + return evaluated[key] + + ( + purity, + candidate_reconstruction, + ) = self._evaluatePurityAtWavelength( + wavelength + ) + + candidate_error = float( + np.asarray( + asNumpyArray( + candidate_reconstruction.error + ) + ).reshape(-1)[-1] + ) + + evaluated[key] = purity + + wavelength_history.append( + wavelength + ) + + purity_history.append( + purity + ) + + error_history.append( + candidate_error + ) + + tqdm.tqdm.write( + f"wavelength = {wavelength * 1e9:.6f} nm | " + f"purity = {purity:.6f} | " + f"error = {candidate_error:.6e}" + ) + + if purity > best_purity: + + best_purity = purity + best_wavelength = wavelength + + best_state = { + "object": asNumpyArray( + candidate_reconstruction.object + ).copy(), + + "probe": asNumpyArray( + candidate_reconstruction.probe + ).copy(), + + "purityProbe": purity, + + "dxp": float( + candidate_reconstruction.dxp + ), + } + + if hasattr( + candidate_reconstruction, + "spectralDensity", + ): + best_state["spectralDensity"] = ( + np.asarray( + candidate_reconstruction.spectralDensity + ).copy() + ) + + if hasattr( + candidate_reconstruction, + "encoder_corrected", + ): + best_state["encoder_corrected"] = ( + asNumpyArray( + candidate_reconstruction.encoder_corrected + ).copy() + ) + + if hasattr( + candidate_reconstruction, + "objectMomentum", + ): + best_state["objectMomentum"] = ( + asNumpyArray( + candidate_reconstruction.objectMomentum + ).copy() + ) + + if hasattr( + candidate_reconstruction, + "probeMomentum", + ): + best_state["probeMomentum"] = ( + asNumpyArray( + candidate_reconstruction.probeMomentum + ).copy() + ) + + if hasattr( + candidate_reconstruction, + "objectBuffer", + ): + best_state["objectBuffer"] = ( + asNumpyArray( + candidate_reconstruction.objectBuffer + ).copy() + ) + + if hasattr( + candidate_reconstruction, + "probeBuffer", + ): + best_state["probeBuffer"] = ( + asNumpyArray( + candidate_reconstruction.probeBuffer + ).copy() + ) + + if hasattr( + candidate_reconstruction, + "error", + ): + best_state["error"] = ( + np.asarray( + asNumpyArray( + candidate_reconstruction.error + ) + ).copy() + ) + + tqdm.tqdm.write( + f"NEW BEST: wavelength = " + f"{best_wavelength * 1e9:.6f} nm | " + f"purity = {best_purity:.6f}" + ) + + return purity + + tqdm.tqdm.write("") + tqdm.tqdm.write( + "Starting adaptive purity-based wavelength scan" + ) + + tqdm.tqdm.write( + f"Initial wavelength = " + f"{wavelength0 * 1e9:.6f} nm | " + f"initial step = {step * 1e9:.4f} nm" + ) + + tqdm.tqdm.write("") + + # ------------------------------------------------------------------ + # 1. Initial three-point evaluation + # ------------------------------------------------------------------ + + p0 = evaluate( + wavelength0 + ) + + wavelength_left = ( + wavelength0 - step + ) + + wavelength_right = ( + wavelength0 + step + ) + + p_left = evaluate( + wavelength_left + ) + + p_right = evaluate( + wavelength_right + ) + + # ------------------------------------------------------------------ + # 2. Determine uphill direction + # ------------------------------------------------------------------ + + if ( + p_left > p0 + and p_left >= p_right + ): + direction = -1 + + elif ( + p_right > p0 + and p_right > p_left + ): + direction = +1 + + else: + direction = 0 + + # ------------------------------------------------------------------ + # 3. Expansion phase + # ------------------------------------------------------------------ + + if direction != 0: + + current_wavelength = ( + wavelength0 + + direction * step + ) + + current_purity = evaluate( + current_wavelength + ) + + while ( + len(evaluated) + < max_evaluations + ): + + step *= growth + + next_wavelength = ( + current_wavelength + + direction * step + ) + + if next_wavelength <= 0: + break + + next_purity = evaluate( + next_wavelength + ) + + improvement = ( + next_purity + - current_purity + ) + + if improvement > purity_tol: + + current_wavelength = ( + next_wavelength + ) + + current_purity = ( + next_purity + ) + + continue + + break + + # ------------------------------------------------------------------ + # 4. Refinement phase + # ------------------------------------------------------------------ + + step *= shrink + + while ( + step >= min_step + and len(evaluated) + < max_evaluations + ): + + left_wavelength = ( + best_wavelength - step + ) + + right_wavelength = ( + best_wavelength + step + ) + + if left_wavelength > 0: + left_purity = evaluate( + left_wavelength + ) + else: + left_purity = -np.inf + + if ( + len(evaluated) + >= max_evaluations + ): + break + + right_purity = evaluate( + right_wavelength + ) + + if ( + left_purity + <= best_purity + purity_tol + and right_purity + <= best_purity + purity_tol + ): + step *= shrink + continue + + if best_state is None: + raise RuntimeError( + "Adaptive purity wavelength scan did not " + "produce a valid result." + ) + + # ------------------------------------------------------------------ + # Restore best candidate + # ------------------------------------------------------------------ + + self.reconstruction.wavelength = ( + best_wavelength + ) + + self.reconstruction.dxp = ( + best_state["dxp"] + ) + + if "spectralDensity" in best_state: + self.reconstruction.spectralDensity = ( + best_state["spectralDensity"] + ) + + self.reconstruction.object = ( + best_state["object"] + ) + + self.reconstruction.probe = ( + best_state["probe"] + ) + + self.reconstruction.purityProbe = ( + best_state["purityProbe"] + ) + + if "encoder_corrected" in best_state: + self.reconstruction.encoder_corrected = ( + best_state["encoder_corrected"] + ) + + if "objectMomentum" in best_state: + self.reconstruction.objectMomentum = ( + best_state["objectMomentum"] + ) + + if "probeMomentum" in best_state: + self.reconstruction.probeMomentum = ( + best_state["probeMomentum"] + ) + + if "objectBuffer" in best_state: + self.reconstruction.objectBuffer = ( + best_state["objectBuffer"] + ) + + if "probeBuffer" in best_state: + self.reconstruction.probeBuffer = ( + best_state["probeBuffer"] + ) + + if "error" in best_state: + self.reconstruction.error = ( + best_state["error"] + ) + + # ------------------------------------------------------------------ + # Store results + # ------------------------------------------------------------------ + + order = np.argsort( + wavelength_history + ) + + self.purityWavelengthValues = ( + np.asarray( + wavelength_history + )[order] + ) + + self.purityWavelengthMetrics = ( + np.asarray( + purity_history + )[order] + ) + + self.purityWavelengthErrors = ( + np.asarray( + error_history + )[order] + ) + + self.purityBestWavelength = ( + best_wavelength + ) + + self.purityBestWavelengthValue = ( + best_purity + ) + + self.showReconstruction( + self.numIterations - 1, + force=True, + ) + + tqdm.tqdm.write("=" * 60) + tqdm.tqdm.write( + "Adaptive purity-based wavelength scan finished" + ) + + tqdm.tqdm.write( + f"Best wavelength = " + f"{best_wavelength * 1e9:.6f} nm" + ) + + tqdm.tqdm.write( + f"Best Δλ = " + f"{(best_wavelength - wavelength0) * 1e9:+.4f} nm" + ) + + tqdm.tqdm.write( + f"Best purity = " + f"{best_purity:.6f}" + ) + + tqdm.tqdm.write( + f"Evaluations = " + f"{len(evaluated)}" + ) + + tqdm.tqdm.write("=" * 60) + + return ( + best_wavelength, + best_purity, + ) + + def plotPurityScan_X( + self, + show=True, + fit=True, + fit_points=5, +): + """ + Plot probe purity and reconstruction error for the selected + purity-calibration target. + + Optionally fit a quadratic curve locally around the measured purity + maximum and estimate the fitted calibration optimum. + + Args: + show (bool, optional): + Display the figure immediately. Default is `True`. + + fit (bool, optional): + If `True`, perform a local quadratic fit around the measured + purity maximum. Default is `True`. + + fit_points (int, optional): + Number of measured points used for the local quadratic fit. + Default is 5. + + Returns: + tuple: + Matplotlib `(figure, purity_axis, error_axis)` objects. + """ + + target = ( + self.params.purityCalibrationTarget.lower() + ) + + # ------------------------------------------------------------------ + # Select calibration results + # ------------------------------------------------------------------ + + if target == "z": + + if not hasattr( + self, + "purityZValues", + ): + raise RuntimeError( + "No purity z-scan results are available." + ) + + values = ( + np.asarray( + self.purityZValues + ) + * 1e3 + ) + + purity = np.asarray( + self.purityZMetrics + ) + + errors = ( + np.asarray( + self.purityZErrors + ) + if hasattr( + self, + "purityZErrors", + ) + else None + ) + + initial_value = ( + self.purityZInitialGuess + * 1e3 + ) + + xlabel = ( + "Sample-to-detector distance z (mm)" + ) + + title = ( + "Purity-based axial calibration" + ) + + value_symbol = "z" + value_unit = "mm" + value_format = ".3f" + + elif target == "wavelength": + + if not hasattr( + self, + "purityWavelengthValues", + ): + raise RuntimeError( + "No purity wavelength-scan results are available." + ) + + values = ( + np.asarray( + self.purityWavelengthValues + ) + * 1e9 + ) + + purity = np.asarray( + self.purityWavelengthMetrics + ) + + errors = ( + np.asarray( + self.purityWavelengthErrors + ) + if hasattr( + self, + "purityWavelengthErrors", + ) + else None + ) + + initial_value = ( + self.purityWavelengthInitialGuess + * 1e9 + ) + + xlabel = ( + "Wavelength λ (nm)" + ) + + title = ( + "Purity-based wavelength calibration" + ) + + value_symbol = "λ" + value_unit = "nm" + value_format = ".4f" + + else: + + raise ValueError( + "Unsupported purity calibration target: " + f"{self.params.purityCalibrationTarget!r}" + ) + + # ------------------------------------------------------------------ + # Best measured purity + # ------------------------------------------------------------------ + + best_purity_index = np.argmax( + purity + ) + + fig, ax_purity = plt.subplots( + figsize=(7, 4.5) + ) + + ax_purity.plot( + values, + purity, + marker="o", + label="Probe purity", + ) + + ax_purity.plot( + values[ + best_purity_index + ], + purity[ + best_purity_index + ], + marker="*", + markersize=12, + linestyle="None", + label=( + f"Measured maximum: " + f"{value_symbol} = " + f"{format(values[best_purity_index], value_format)} " + f"{value_unit}" + ), + ) + + # ------------------------------------------------------------------ + # Local quadratic fitting + # ------------------------------------------------------------------ + + if fit: + + if fit_points < 3: + raise ValueError( + "`fit_points` must be at least 3." + ) + + if fit_points > len(values): + fit_points = len(values) + + half_window = ( + fit_points // 2 + ) + + start = max( + 0, + best_purity_index + - half_window, + ) + + stop = min( + len(values), + start + fit_points, + ) + + start = max( + 0, + stop - fit_points, + ) + + fit_x = values[ + start:stop + ] + + fit_y = purity[ + start:stop + ] + + if len(fit_x) >= 3: + + coefficients = np.polyfit( + fit_x, + fit_y, + deg=2, + ) + + a, b, c = coefficients + + if a < 0: + + fitted_best_value = ( + -b / (2 * a) + ) + + fitted_best_purity = ( + np.polyval( + coefficients, + fitted_best_value, + ) + ) + + fit_curve_x = np.linspace( + fit_x.min(), + fit_x.max(), + 200, + ) + + fit_curve_y = np.polyval( + coefficients, + fit_curve_x, + ) + + ax_purity.plot( + fit_curve_x, + fit_curve_y, + linestyle="--", + label="Quadratic fit", + ) + + ax_purity.plot( + fitted_best_value, + fitted_best_purity, + marker="X", + markersize=9, + linestyle="None", + label=( + f"Fitted maximum: " + f"{value_symbol} = " + f"{format(fitted_best_value, value_format)} " + f"{value_unit}" + ), + ) + + if target == "z": + self.purityBestZFitted = ( + fitted_best_value + * 1e-3 + ) + + self.purityBestZFittedValue = ( + fitted_best_purity + ) + + elif target == "wavelength": + self.purityBestWavelengthFitted = ( + fitted_best_value + * 1e-9 + ) + + self.purityBestWavelengthFittedValue = ( + fitted_best_purity + ) + + else: + + self.logger.warning( + "Quadratic fit does not produce a local maximum " + "(fitted curvature is non-negative)." + ) + + ax_purity.set_xlabel( + xlabel + ) + + ax_purity.set_ylabel( + "Probe purity" + ) + + # ------------------------------------------------------------------ + # Reconstruction error + # ------------------------------------------------------------------ + + ax_error = ax_purity.twinx() + + if errors is not None: + + ax_error.plot( + values, + errors, + marker="s", + linestyle="--", + label="Reconstruction error", + ) + + best_error_index = np.argmin( + errors + ) + + ax_error.plot( + values[ + best_error_index + ], + errors[ + best_error_index + ], + marker="*", + markersize=12, + linestyle="None", + label=( + f"Minimum error: " + f"{value_symbol} = " + f"{format(values[best_error_index], value_format)} " + f"{value_unit}" + ), + ) + + ax_error.set_ylabel( + "Reconstruction error" + ) + + # ------------------------------------------------------------------ + # Initial guess + # ------------------------------------------------------------------ + + ax_purity.axvline( + initial_value, + linestyle="--", + alpha=0.5, + label=( + f"Initial {value_symbol}: " + f"{format(initial_value, value_format)} " + f"{value_unit}" + ), + ) + + # ------------------------------------------------------------------ + # Combined legend + # ------------------------------------------------------------------ + + purity_handles, purity_labels = ( + ax_purity.get_legend_handles_labels() + ) + + error_handles, error_labels = ( + ax_error.get_legend_handles_labels() + ) + + ax_purity.legend( + purity_handles + error_handles, + purity_labels + error_labels, + loc="best", + ) + + ax_purity.set_title( + title + ) + + ax_purity.grid( + alpha=0.3 + ) + + fig.tight_layout() + + if show: + plt.show() + + return ( + fig, + ax_purity, + ax_error, + ) + + + def plotPurityScan( + self, + show=True, + fit=True, + fit_range=None, + fit_fraction=0.25, + ): + """ + Plot probe purity and reconstruction error for the selected + purity-calibration target. + + A local quadratic fit can optionally be performed around the measured + purity maximum. The fitting window can either be specified explicitly + through `fit_range` or determined automatically as a fraction of the + scanned parameter range. + + Args: + show (bool, optional): + Display the figure immediately. Default is `True`. + + fit (bool, optional): + Perform a local quadratic fit around the measured purity maximum. + Default is `True`. + + fit_range (float, optional): + Half-width of the fitting window in the displayed unit + (mm for z and nm for wavelength). If `None`, the fitting window + is determined from `fit_fraction`. + + fit_fraction (float, optional): + Fraction of the total scanned parameter range used as the + fitting half-width when `fit_range` is `None`. + Default is `0.25`. + + Returns: + tuple: + Matplotlib `(figure, purity_axis, error_axis)` objects. + """ + + target = ( + self.params.purityCalibrationTarget.lower() + ) + + # ------------------------------------------------------------------ + # Select calibration results + # ------------------------------------------------------------------ + + if target == "z": + + if not hasattr( + self, + "purityZValues", + ): + raise RuntimeError( + "No purity z-scan results are available." + ) + + values = ( + np.asarray( + self.purityZValues + ) + * 1e3 + ) + + purity = np.asarray( + self.purityZMetrics + ) + + errors = ( + np.asarray( + self.purityZErrors + ) + if hasattr( + self, + "purityZErrors", + ) + else None + ) + + initial_value = ( + self.purityZInitialGuess + * 1e3 + ) + + xlabel = ( + "Sample-to-detector distance z (mm)" + ) + + title = ( + "Purity-based axial calibration" + ) + + value_symbol = "z" + value_unit = "mm" + value_format = ".3f" + + elif target == "wavelength": + + if not hasattr( + self, + "purityWavelengthValues", + ): + raise RuntimeError( + "No purity wavelength-scan results are available." + ) + + values = ( + np.asarray( + self.purityWavelengthValues + ) + * 1e9 + ) + + purity = np.asarray( + self.purityWavelengthMetrics + ) + + errors = ( + np.asarray( + self.purityWavelengthErrors + ) + if hasattr( + self, + "purityWavelengthErrors", + ) + else None + ) + + initial_value = ( + self.purityWavelengthInitialGuess + * 1e9 + ) + + xlabel = ( + "Wavelength λ (nm)" + ) + + title = ( + "Purity-based wavelength calibration" + ) + + value_symbol = "λ" + value_unit = "nm" + value_format = ".4f" + + else: + + raise ValueError( + "Unsupported purity calibration target: " + f"{self.params.purityCalibrationTarget!r}" + ) + + # ------------------------------------------------------------------ + # Measured maximum + # ------------------------------------------------------------------ + + best_purity_index = np.argmax( + purity + ) + + measured_best_value = ( + values[ + best_purity_index + ] + ) + + measured_best_purity = ( + purity[ + best_purity_index + ] + ) + + fig, ax_purity = plt.subplots( + figsize=(7, 4.5) + ) + + ax_purity.plot( + values, + purity, + marker="o", + label="Probe purity", + ) + + ax_purity.plot( + measured_best_value, + measured_best_purity, + marker="*", + markersize=12, + linestyle="None", + label=( + f"Measured maximum: " + f"{value_symbol} = " + f"{format(measured_best_value, value_format)} " + f"{value_unit}" + ), + ) + + # ------------------------------------------------------------------ + # Local quadratic fit + # ------------------------------------------------------------------ + + if fit: + + if fit_range is None: + + total_range = ( + values.max() + - values.min() + ) + + fit_half_range = ( + fit_fraction + * total_range + ) + + else: + + fit_half_range = ( + float(fit_range) + ) + + fit_mask = ( + np.abs( + values + - measured_best_value + ) + <= fit_half_range + ) + + fit_x = values[ + fit_mask + ] + + fit_y = purity[ + fit_mask + ] + + if len(fit_x) < 3: + + self.logger.warning( + "Quadratic fit skipped because fewer than " + "three scan points fall inside the fitting window." + ) + + else: + + coefficients = np.polyfit( + fit_x, + fit_y, + deg=2, + ) + + a, b, c = coefficients + + fitted_y = np.polyval( + coefficients, + fit_x, + ) + + residual_sum = np.sum( + ( + fit_y + - fitted_y + ) + ** 2 + ) + + total_sum = np.sum( + ( + fit_y + - np.mean(fit_y) + ) + ** 2 + ) + + if total_sum > 0: + + fit_r2 = ( + 1 + - residual_sum + / total_sum + ) + + else: + + fit_r2 = np.nan + + fitted_best_value = ( + -b + / (2 * a) + if a != 0 + else np.nan + ) + + valid_fit = ( + a < 0 + and np.isfinite( + fitted_best_value + ) + and fit_x.min() + <= fitted_best_value + <= fit_x.max() + ) + + if valid_fit: + + fitted_best_purity = ( + np.polyval( + coefficients, + fitted_best_value, + ) + ) + + fit_curve_x = np.linspace( + fit_x.min(), + fit_x.max(), + 300, + ) + + fit_curve_y = np.polyval( + coefficients, + fit_curve_x, + ) + + ax_purity.plot( + fit_curve_x, + fit_curve_y, + linestyle="--", + label=( + f"Quadratic fit " + f"($R^2$ = {fit_r2:.3f})" + ), + ) + + ax_purity.plot( + fitted_best_value, + fitted_best_purity, + marker="X", + markersize=9, + linestyle="None", + label=( + f"Fitted maximum: " + f"{value_symbol} = " + f"{format(fitted_best_value, value_format)} " + f"{value_unit}" + ), + ) + + if target == "z": + + self.purityBestZFitted = ( + fitted_best_value + * 1e-3 + ) + + self.purityBestZFittedValue = ( + fitted_best_purity + ) + + self.purityBestZFitR2 = ( + fit_r2 + ) + + elif target == "wavelength": + + self.purityBestWavelengthFitted = ( + fitted_best_value + * 1e-9 + ) + + self.purityBestWavelengthFittedValue = ( + fitted_best_purity + ) + + self.purityBestWavelengthFitR2 = ( + fit_r2 + ) + + else: + + self.logger.warning( + "Quadratic fit did not produce a valid local maximum. " + "The fitted curvature may be non-negative or the fitted " + "maximum may lie outside the fitting window." + ) + + ax_purity.set_xlabel( + xlabel + ) + + ax_purity.set_ylabel( + "Probe purity" + ) + + # ------------------------------------------------------------------ + # Reconstruction error + # ------------------------------------------------------------------ + + ax_error = ax_purity.twinx() + + if errors is not None: + + ax_error.plot( + values, + errors, + marker="s", + linestyle="--", + label="Reconstruction error", + ) + + best_error_index = np.argmin( + errors + ) + + ax_error.plot( + values[ + best_error_index + ], + errors[ + best_error_index + ], + marker="*", + markersize=12, + linestyle="None", + label=( + f"Minimum error: " + f"{value_symbol} = " + f"{format(values[best_error_index], value_format)} " + f"{value_unit}" + ), + ) + + ax_error.set_ylabel( + "Reconstruction error" + ) + + # ------------------------------------------------------------------ + # Initial guess + # ------------------------------------------------------------------ + + ax_purity.axvline( + initial_value, + linestyle="--", + alpha=0.5, + label=( + f"Initial {value_symbol}: " + f"{format(initial_value, value_format)} " + f"{value_unit}" + ), + ) + + # ------------------------------------------------------------------ + # Combined legend + # ------------------------------------------------------------------ + + purity_handles, purity_labels = ( + ax_purity.get_legend_handles_labels() + ) + + error_handles, error_labels = ( + ax_error.get_legend_handles_labels() + ) + + ax_purity.legend( + purity_handles + error_handles, + purity_labels + error_labels, + loc="best", + ) + + ax_purity.set_title( + title + ) + + ax_purity.grid( + alpha=0.3 + ) + + fig.tight_layout() + + if show: + plt.show() + + return ( + fig, + ax_purity, + ax_error, + ) \ No newline at end of file diff --git a/PtyLab/Params/Params.py b/PtyLab/Params/Params.py index bcd8ff30..dc764fa7 100644 --- a/PtyLab/Params/Params.py +++ b/PtyLab/Params/Params.py @@ -154,20 +154,34 @@ def __init__(self): self.map_position_to_z_change = False # purity based Zscan + self.purityCalibrationTarget = "z" # aternatively wavelength + + self.purityWavelengthScanRange = 1e-9 + self.purityWavelengthScanPoints = 11 + self.purityWavelengthScanPlot = True + self.purityZScanRange = 100e-6 self.purityZScanPoints = 11 - self.purityZScanPlot = True - self.purityZAdaptive = False + # Shared adaptive-scan settings + self.purityAdaptive = True + self.purityMaxEvaluations = 50 + # z-specific adaptive settings self.purityZInitialStep = 200e-6 self.purityZStepGrowth = 1.5 self.purityZStepShrink = 0.5 - self.purityZMinStep = 20e-6 - self.purityZPurityTolerance = 1e-4 - self.purityZMaxEvaluations = 20 + self.purityZPurityTolerance = 1e-5 + + # wavelength-specific adaptive settings + self.purityWavelengthInitialStep = 0.1e-9 + self.purityWavelengthStepGrowth = 1.5 + self.purityWavelengthStepShrink = 0.5 + self.purityWavelengthMinStep = 0.01e-9 + self.purityWavelengthPurityTolerance = 1e-5 + self.OPRP = False diff --git a/example_scripts/example_purity_Zscan_yourData.py b/example_scripts/example_purity_Zscan_yourData.py index c57b7173..63d7fc65 100644 --- a/example_scripts/example_purity_Zscan_yourData.py +++ b/example_scripts/example_purity_Zscan_yourData.py @@ -5,7 +5,6 @@ matplotlib.use("QtAgg") import matplotlib.pyplot as plt -import numpy as np import PtyLab from PtyLab import Engines @@ -14,6 +13,22 @@ logging.basicConfig(level=logging.INFO) +# ============================================================================= +# User settings +# ============================================================================= + +# Choose which experimental parameter to calibrate: +# +# "z" -> sample-to-detector distance +# "wavelength" -> illumination wavelength +CALIBRATION_TARGET = "z" # "z" or "wavelength" + +# Z / Wavelength calibration assumes that the wavelength / Z is known to a resonable degree of accuracy." +initial_z = 35e-3 # 32.5 mm + +initial_wavelength = 13.5e-9 # 13.5 nm + + # ============================================================================= # File selection # ============================================================================= @@ -38,7 +53,7 @@ def select_hdf5_file(title): print() print("=" * 70) -print("Purity-based axial calibration") +print("Purity-based parameter calibration") print("=" * 70) print() @@ -56,12 +71,15 @@ def select_hdf5_file(title): "No ptychography dataset was selected." ) - print() print("Selected dataset:") print(filePath) +# ============================================================================= +# Optional seed reconstruction +# ============================================================================= + print() print( "Step 2: Choose the initial probe.\n" @@ -107,9 +125,10 @@ def select_hdf5_file(title): "and a uniform object." ) -# ------------------------------------------------------------------------- + +# ============================================================================= # Load data -# ------------------------------------------------------------------------- +# ============================================================================= experimentalData, reconstruction, params, monitor, _ = ( PtyLab.easyInitialize( @@ -118,29 +137,130 @@ def select_hdf5_file(title): ) ) -#experimentalData.setOrientation(0) +# Change only if required by your dataset. +experimentalData.setOrientation(4) +print() print("Loaded file:", experimentalData.filename) print("Ptychogram shape:", experimentalData.ptychogram.shape) -print("Initial z:", reconstruction.zo) -print("Wavelength:", reconstruction.wavelength) +print(f"Dataset z: {reconstruction.zo * 1e3:.6f} mm") +print(f"Dataset wavelength: {reconstruction.wavelength * 1e9:.6f} nm") + + + +# ============================================================================= +# Purity-based calibration settings +# ============================================================================= + +params.purityCalibrationTarget = ( + CALIBRATION_TARGET +) + +# Shared settings +params.purityAdaptive = True +params.purityMaxEvaluations = 20 +params.purityScanPlot = True + +# ----------------------------------------------------------------------------- +# z calibration settings +# ----------------------------------------------------------------------------- + +# Used by the fixed-grid scan. +params.purityZScanRange = 2e-3 +params.purityZScanPoints = 10 + +# Used by the adaptive scan. +params.purityZInitialStep = 200e-6 +params.purityZStepGrowth = 1.5 +params.purityZStepShrink = 0.5 +params.purityZMinStep = 20e-6 +params.purityZPurityTolerance = 1e-4 + + +# ----------------------------------------------------------------------------- +# Wavelength calibration settings +# ----------------------------------------------------------------------------- +# Used by the fixed-grid scan. +params.purityWavelengthScanRange = 1.0e-9 +params.purityWavelengthScanPoints = 11 -# Use the z value stored in the dataset by default. -initial_z = reconstruction.zo +# Used by the adaptive scan. +params.purityWavelengthInitialStep = 0.1e-9 +params.purityWavelengthStepGrowth = 1.5 +params.purityWavelengthStepShrink = 0.5 +params.purityWavelengthMinStep = 0.01e-9 +params.purityWavelengthPurityTolerance = 1e-4 -# Uncomment to set the initial z guess manually. -# initial_z = 32.5e-3 +# ============================================================================= +# Initial calibration guess +# ============================================================================= -experimentalData.zo = initial_z -reconstruction.zo = initial_z +if CALIBRATION_TARGET == "z": -print(f"Initial z guess: {initial_z * 1e3:.6f} mm") + # Use the value stored in the dataset by default. + # initial_z = reconstruction.zo + # Uncomment to set the initial z guess manually. + # initial_z = 30e-3 -# ------------------------------------------------------------------------- + experimentalData.zo = initial_z + reconstruction.zo = initial_z + + print() + print( + f"Initial z guess: " + f"{reconstruction.zo * 1e3:.6f} mm" + ) + + +elif CALIBRATION_TARGET == "wavelength": + + reconstruction.zo = initial_z + + print( + f"Fixed z used for wavelength calibration: " + f"{reconstruction.zo * 1e3:.6f} mm" + ) + # Use the wavelength stored in the dataset by default. + # initial_wavelength = reconstruction.wavelength + + # Uncomment to set the initial wavelength manually. + # initial_wavelength = 13.0e-9 + + experimentalData.wavelength = initial_wavelength + reconstruction.wavelength = initial_wavelength + + # Keep the single-wavelength spectral-density entry consistent. + if reconstruction.nlambda == 1: + experimentalData.spectralDensity = [initial_wavelength] + reconstruction.spectralDensity = [initial_wavelength] + + # Update CPM sampling because wavelength has no property setter. + reconstruction.dxp = ( + reconstruction.wavelength + * reconstruction.zo + / reconstruction.Ld + ) + + print() + print( + f"Initial wavelength guess: " + f"{reconstruction.wavelength * 1e9:.6f} nm" + ) + + +else: + + raise ValueError( + "CALIBRATION_TARGET must be either " + "'z' or 'wavelength'." + ) + + +# ============================================================================= # Reconstruction dimensions -# ------------------------------------------------------------------------- +# ============================================================================= # Purity-based calibration requires multiple probe modes. reconstruction.npsm = 4 @@ -177,7 +297,6 @@ def select_hdf5_file(title): "Initial probe: circular estimate" ) - # Circular initialization requires an entrance pupil diameter. pupil_diameter = getattr( experimentalData, "entrancePupilDiameter", @@ -200,10 +319,12 @@ def select_hdf5_file(title): reconstruction.initializeObjectProbe() +reconstruction.describe_reconstruction() -# ------------------------------------------------------------------------- + +# ============================================================================= # Monitor settings -# ------------------------------------------------------------------------- +# ============================================================================= monitor.figureUpdateFrequency = 5 monitor.objectPlot = "complex" @@ -213,67 +334,48 @@ def select_hdf5_file(title): monitor.probeZoom = None -# ------------------------------------------------------------------------- +# ============================================================================= # General reconstruction parameters -# ------------------------------------------------------------------------- +# ============================================================================= params.positionOrder = "random" -# Start with the propagator used for your actual dataset. +# Use the propagator appropriate for your dataset. params.propagatorType = "Fraunhofer" params.gpuSwitch = True params.intensityConstraint = "standard" +params.positionCorrectionSwitch = False + +params.modulusEnforcedProbeSwitch = False params.probePowerCorrectionSwitch = True params.probeSmoothenessSwitch = True params.probeSmoothnessAleph = 1e-2 params.probeSmoothenessWidth = 10 + params.comStabilizationSwitch = 10 + params.orthogonalizationSwitch = True params.orthogonalizationFrequency = 10 -params.couplingSwitch = True -params.couplingAleph = 1 +params.absorbingProbeBoundary = False +params.objectContrastSwitch = False +params.absObjectSwitch = False +params.backgroundModeSwitch = False -# ------------------------------------------------------------------------- -# Purity-based z calibration parameters -# ------------------------------------------------------------------------- -# Half-range around the current reconstruction.zo. -# -# Example: -# reconstruction.zo = 50 mm -# purityZScanRange = 100 µm -# -# scan range: -# 49.9 mm ... 50.1 mm -params.purityZScanRange = 2e-3 - -# Number of candidate z values. -params.purityZScanPoints = 10 - -params.purityZAdaptive = True +params.couplingSwitch = True +params.couplingAleph = 1 -params.purityZInitialStep = 200e-6 -params.purityZStepGrowth = 1.5 -params.purityZStepShrink = 0.5 -params.purityZMinStep = 20e-6 -params.purityZPurityTolerance = 1e-4 -params.purityZMaxEvaluations = 20 - -# Plot purity versus z after the scan. -params.purityZScanPlot = True monitor.describe_parameters(params) -# ------------------------------------------------------------------------- -# Run purity-based axial calibration -# ------------------------------------------------------------------------- - - +# ============================================================================= +# Run purity-based calibration +# ============================================================================= purity_engine = Engines.purityPIE( reconstruction, @@ -282,39 +384,116 @@ def select_hdf5_file(title): monitor, ) -# Number of mPIE iterations used to evaluate EACH candidate z. +# Number of mPIE iterations used to evaluate EACH candidate. purity_engine.numIterations = 5 -best_z, best_purity = purity_engine.reconstruct() - -print() -print("Purity-based axial calibration finished") -print("---------------------------------------") -print(f"Initial z : {purity_engine.purityZInitialGuess * 1e3:.6f} mm") -print(f"Best z : {best_z * 1e3:.6f} mm") -print( - f"Delta z : " - f"{(best_z - purity_engine.purityZInitialGuess) * 1e6:.3f} um" +best_value, best_purity = ( + purity_engine.reconstruct() ) -print(f"Best purity: {best_purity:.6f}") + + +# ============================================================================= +# Print results +# ============================================================================= print() -print("z scan:") -for z, purity in zip( - purity_engine.purityZValues, - purity_engine.purityZMetrics, -): +print("=" * 70) +print("Purity-based calibration finished") +print("=" * 70) + + +if CALIBRATION_TARGET == "z": + print( - f"z = {z * 1e3:.6f} mm, " - f"delta z = {(z - purity_engine.purityZInitialGuess) * 1e6:+.2f} um, " - f"purity = {purity:.6f}" + f"Initial z : " + f"{purity_engine.purityZInitialGuess * 1e3:.6f} mm" ) + print( + f"Best z : " + f"{best_value * 1e3:.6f} mm" + ) + print( + f"Delta z : " + f"{( + best_value + - purity_engine.purityZInitialGuess + ) * 1e6:+.3f} um" + ) -# ------------------------------------------------------------------------- -# Optional: continue reconstruction using the best-z state -# ------------------------------------------------------------------------- + print( + f"Best purity: " + f"{best_purity:.6f}" + ) + + print() + print("z scan:") + + for z, purity in zip( + purity_engine.purityZValues, + purity_engine.purityZMetrics, + ): + + print( + f"z = {z * 1e3:.6f} mm | " + f"delta z = " + f"{( + z + - purity_engine.purityZInitialGuess + ) * 1e6:+.2f} um | " + f"purity = {purity:.6f}" + ) + + +elif CALIBRATION_TARGET == "wavelength": + + print( + f"Initial wavelength : " + f"{purity_engine.purityWavelengthInitialGuess * 1e9:.6f} nm" + ) + + print( + f"Best wavelength : " + f"{best_value * 1e9:.6f} nm" + ) + + print( + f"Delta wavelength : " + f"{( + best_value + - purity_engine.purityWavelengthInitialGuess + ) * 1e9:+.4f} nm" + ) + + print( + f"Best purity : " + f"{best_purity:.6f}" + ) + + print() + print("Wavelength scan:") + + for wavelength, purity in zip( + purity_engine.purityWavelengthValues, + purity_engine.purityWavelengthMetrics, + ): + + print( + f"wavelength = " + f"{wavelength * 1e9:.6f} nm | " + f"delta wavelength = " + f"{( + wavelength + - purity_engine.purityWavelengthInitialGuess + ) * 1e9:+.4f} nm | " + f"purity = {purity:.6f}" + ) + + +# ============================================================================= +# Optional: continue reconstruction using the calibrated state +# ============================================================================= continue_with_mPIE = False @@ -332,12 +511,12 @@ def select_hdf5_file(title): mPIE_engine.reconstruct() -# ------------------------------------------------------------------------- +# ============================================================================= # Save / display -# ------------------------------------------------------------------------- +# ============================================================================= # reconstruction.saveResults( -# f"{getExampleDataFolder()}/purity_calibrated_recon.hdf5", +# "purity_calibrated_recon.hdf5", # squeeze=False, # ) From 83347e02c545e9620009561258365a74f1ad6f74 Mon Sep 17 00:00:00 2001 From: Chang Liu Date: Wed, 7 Oct 2026 14:35:34 +0200 Subject: [PATCH 11/11] doc of purityPIE and aPIE --- PtyLab/Engines/aPIE.py | 88 ++++- PtyLab/Engines/purityPIE.py | 736 +++++++++--------------------------- 2 files changed, 271 insertions(+), 553 deletions(-) diff --git a/PtyLab/Engines/aPIE.py b/PtyLab/Engines/aPIE.py index 5c8fe6c5..40beb8a0 100644 --- a/PtyLab/Engines/aPIE.py +++ b/PtyLab/Engines/aPIE.py @@ -30,8 +30,46 @@ class aPIE(BaseEngine): - """ - aPIE: angle correction PIE: ePIE combined with Luus-Jaakola algorithm (the latter for angle correction) + momentum + r""" + Angle-corrected ptychographic iterative engine (aPIE). + + `aPIE` extends the standard PIE reconstruction by jointly refining the + object, probe, and detector tilt angle. The angular correction is based on + a Luus-Jaakola-type stochastic search combined with momentum + acceleration.[^debeurs2022] + + The detector tilt is described by the angle `reconstruction.theta`. For + each reconstruction iteration, two candidate angles are evaluated: the + current angle and a randomly perturbed angle within the current search + radius. The diffraction data are transformed into the corresponding + detector coordinates for each candidate, followed by a standard PIE + object/probe reconstruction. + + The candidate producing the lower reconstruction error is retained. The + angular update is then combined with a momentum term according to + + $$ + \theta_{\mathrm{mom}}^{(k+1)} = f\,\Delta\theta_k + \gamma\,\theta_{\mathrm{mom}}^{(k)} + $$ + + where `f = feedback`, `\gamma = aPIEfriction`, and `\Delta\theta_k` is the + accepted angular change. + + The angular search radius is gradually reduced from + `thetaSearchRadiusMax` to `thetaSearchRadiusMin` over the reconstruction + iterations, allowing coarse angular exploration at the beginning and + progressively finer correction toward convergence. + + The evolution of the estimated detector angle is stored in + `reconstruction.thetaHistory`. + + Angle correction requires an initial value of `reconstruction.theta`. + + [^debeurs2022]: A. de Beurs, L. Loetgering, M. Herczog, + M. Du, K. S. E. Eikema, and S. Witte, + "aPIE: an angle calibration algorithm for reflection ptychography," + Opt. Lett. 47, 1949-1952 (2022). + https://doi.org/10.1364/OL.453655 """ def __init__( @@ -41,8 +79,6 @@ def __init__( params: Params, monitor: Monitor, ): - # This contains reconstruction parameters that are specific to the reconstruction - # but not necessarily to aPIE reconstruction super().__init__(reconstruction, experimentalData, params, monitor) self.logger = logging.getLogger("aPIE") self.logger.info("Sucesfully created aPIE aPIE_engine") @@ -51,8 +87,20 @@ def __init__( def initializeReconstructionParams(self): """ - Set parameters that are specific to the ePIE settings. - :return: + Initialize parameters specific to the aPIE reconstruction. + + This method sets the object and probe update strengths, angular momentum + parameters, reconstruction length, and angular search-radius limits. It + also initializes the detector-angle momentum and history if they are not + already present. + + A copy of the original ptychogram is stored in + `self.ptychogramUntransformed` for repeated angle-dependent coordinate + transformations during reconstruction. + + Raises: + ValueError: + If `reconstruction.theta` has not been initialized. """ self.betaProbe = 0.25 self.betaObject = 0.25 @@ -74,6 +122,34 @@ def initializeReconstructionParams(self): raise ValueError("theta value is not given") def doReconstruction(self): + """ + Run the aPIE reconstruction with iterative detector-angle correction. + + The reconstruction starts from the current detector angle + `reconstruction.theta` and gradually reduces the angular search radius + from `thetaSearchRadiusMax` to `thetaSearchRadiusMin` over the requested + number of iterations. + + At each iteration, two candidate detector angles are evaluated: the + current angle and a randomly perturbed angle within the current search + radius, both including the accumulated angular momentum. + + For each candidate angle, the original ptychogram is transformed into the + corresponding detector coordinates. A standard PIE reconstruction is then + performed over all scan positions, including object and probe updates, + constraint application, and reconstruction-error evaluation. + + The candidate producing the lower reconstruction error is retained. The + accepted angular change is used to update `reconstruction.thetaMomentum`, + which contributes to the candidate angles in the following iteration. + + The accepted detector angle is appended to + `reconstruction.thetaHistory`, and the reconstruction monitor is updated + throughout the optimization. + + If GPU acceleration is enabled, the reconstruction data are returned to + CPU memory after completion. + """ self._prepareReconstruction() xp = getArrayModule(self.reconstruction.object) diff --git a/PtyLab/Engines/purityPIE.py b/PtyLab/Engines/purityPIE.py index aebdb3ed..839108af 100644 --- a/PtyLab/Engines/purityPIE.py +++ b/PtyLab/Engines/purityPIE.py @@ -31,14 +31,15 @@ from PtyLab.Engines.mPIE import mPIE class purityPIE(mPIE): r""" - Purity-based axial self-calibration for ptychographic reconstruction. + Purity-based axial-distance and wavelength self-calibration for + ptychographic reconstruction. - `purityPIE` extends `mPIE` by estimating the sample-to-detector - axial distance from the reconstructed probe purity.[^liu2025] + `purityPIE` extends `mPIE` by estimating either the sample-to-detector + axial distance or the illumination wavelength from the reconstructed + probe purity.[^liu2025] - In mixed-state ptychography, the measured diffraction intensity at - scan position $j$ is modeled as an incoherent sum over multiple probe - modes: + In mixed-state ptychography, the measured diffraction intensity at scan + position $j$ is modeled as an incoherent sum over multiple probe modes: $$ I_j(\mathbf{q}) = \sum_m \left| \mathcal{F}\left[P_m(\mathbf{r}) O(\mathbf{r}-\mathbf{r}_j)\right] \right|^2 @@ -59,48 +60,59 @@ class purityPIE(mPIE): J(\mathbf{r}_1,\mathbf{r}_2) = \sum_m \lambda_m \phi_m^*(\mathbf{r}_1)\phi_m(\mathbf{r}_2) $$ - where $\lambda_m$ represents the intensity contribution, or relative - power content, of the corresponding orthonormal mode. + where $\lambda_m$ represents the power contribution of the corresponding + orthonormal mode. The probe purity is defined as $$ \mu = \frac{\sqrt{\sum_m \lambda_m^2}}{\sum_m \lambda_m} $$ - with $\mu \in [0,1]$. - If the assumed axial distance is incorrect, a mismatch is introduced - in the diffraction integral connecting real and reciprocal space. - This produces stretching deformation and resolution degradation in - the reconstructed object and probe, which appears as a reduction in - reconstructed purity. + An incorrect axial distance or illumination wavelength introduces a + mismatch in the diffraction geometry connecting the object and detector + planes. This mismatch can produce geometric deformation and resolution + degradation in the reconstructed object and probe and is reflected by a + reduction in the reconstructed probe purity. - The calibrated axial distance is obtained by maximizing the - reconstructed probe purity: + The calibrated parameter is therefore obtained by maximizing the + reconstructed probe purity. - $$ - z_{\mathrm{opt}} = \operatorname*{arg\,max}_z \mu(z) - $$ + For each candidate value, `purityPIE` creates an independent + reconstruction with the corresponding geometry, performs an `mPIE` + reconstruction, orthogonalizes the reconstructed probe modes, and + evaluates the resulting probe purity. The candidate yielding the highest + purity is retained as the calibrated reconstruction. + + Two calibration targets are supported: - For each candidate axial distance, `purityPIE` creates an independent - reconstruction initialized with the corresponding geometry, performs - an `mPIE` reconstruction, orthogonalizes the reconstructed probe - modes, and evaluates the resulting purity. The candidate with the - highest purity is retained as the calibrated reconstruction. + - **Axial-distance calibration:** optimizes the sample-to-detector + distance `z`. - Two search strategies are available: + - **Wavelength calibration:** optimizes the illumination wavelength + `λ`. - - A fixed-grid scan evaluates uniformly spaced axial positions defined - by `params.purityZScanRange` and `params.purityZScanPoints`. + The calibration target is selected through + `params.purityCalibrationTarget`, which can be set to `"z"` or + `"wavelength"`. - - An adaptive scan expands the axial step while purity improves and - subsequently refines the search around the best candidate with - progressively smaller steps. It is enabled with - `params.purityZAdaptive = True`. + For either calibration target, two search strategies are available: - Purity-based axial calibration requires at least two reconstructed - probe modes, i.e. `reconstruction.npsm >= 2`. + - **Fixed-grid search:** evaluates uniformly spaced candidate values over + a predefined range. Axial-distance scans are controlled by + `params.purityZScanRange` and `params.purityZScanPoints`, while + wavelength scans are controlled by `params.purityWavelengthScanRange` + and `params.purityWavelengthScanPoints`. + + - **Adaptive search:** starts from the current parameter estimate, + determines the direction of increasing probe purity, expands the search + step while the purity continues to improve, and then refines the search + around the best candidate using progressively smaller steps. Adaptive + search is enabled with `params.purityAdaptive = True`. + + Purity-based calibration requires at least two reconstructed probe modes, + i.e. `reconstruction.npsm >= 2`. [^liu2025]: C. Liu, W. Eschen, D. S. Penagos Molina, L. Licht, M. Abdelaal, and J. Rothhardt, @@ -111,7 +123,7 @@ class purityPIE(mPIE): See Also: `mPIE` Momentum-accelerated reconstruction engine used for each - candidate axial distance. + candidate calibration value. """ def __init__( self, @@ -142,30 +154,25 @@ def reconstruct( reconstruction=None, ): """ - Run purity-based axial self-calibration. + Run purity-based self-calibration. The reconstruction first prepares the current object, probe, geometry, constraints, and optional GPU state using the standard `mPIE` setup. - The axial calibration is then performed using one of two search strategies: - - - If `params.purityZAdaptive = False`, `purityZScan()` evaluates a fixed - grid of candidate axial distances. + Each candidate parameter value is evaluated using an independent `mPIE` + reconstruction. The reconstructed probe modes are orthogonalized, the + corresponding probe purity is evaluated, and the candidate yielding the + highest purity is retained. - - If `params.purityZAdaptive = True`, `adaptivePurityZScan()` searches the - axial distance using adaptive step-size expansion and refinement. + After calibration, the reconstruction state corresponding to the optimal + candidate is restored to the original `Reconstruction` object. - Each candidate distance is evaluated using an independent `mPIE` - reconstruction. The candidate yielding the highest reconstructed probe - purity is selected, and its reconstruction state is restored to the - original `Reconstruction` object. - - If `params.purityZScanPlot` is enabled, the resulting purity-versus-z - curve is displayed after calibration. + If `params.purityScanPlot` is enabled, the purity and reconstruction error + as functions of the calibrated parameter are displayed after the scan. Args: experimentalData (ExperimentalData, optional): - Experimental dataset to use for the calibration. If provided, it + Experimental dataset to use for calibration. If provided, it replaces the dataset currently attached to the engine. reconstruction (Reconstruction, optional): @@ -174,9 +181,16 @@ def reconstruct( Returns: tuple: - `(best_z, best_purity)`, where `best_z` is the calibrated axial - propagation distance in meters and `best_purity` is the + `(best_value, best_purity)`, where `best_value` is the calibrated + axial distance in meters for `purityCalibrationTarget = "z"` or the + calibrated wavelength in meters for + `purityCalibrationTarget = "wavelength"`, and `best_purity` is the corresponding reconstructed probe purity. + + Raises: + ValueError: + If `params.purityCalibrationTarget` is neither `"z"` nor + `"wavelength"`. """ self.changeExperimentalData(experimentalData) @@ -257,7 +271,29 @@ def _createCandidateReconstruction( value, ): """ - Create an independent reconstruction for one candidate calibration value. + Create an independent reconstruction for one calibration candidate. + + The requested calibration parameter is applied to a deep copy of the + experimental data before a new `Reconstruction` object is created. + Relevant reconstruction settings and initial object/probe definitions are + then copied from the current reconstruction. + + Args: + target (str): + Calibration target. Supported values are `"z"` and `"wavelength"`. + + value (float): + Candidate axial distance or wavelength in meters. + + Returns: + tuple: + `(candidate_data, candidate_reconstruction, candidate_params)`, + containing the copied experimental data, newly initialized + reconstruction, and copied parameter set for the candidate. + + Raises: + ValueError: + If `target` is neither `"z"` nor `"wavelength"`. """ candidate_data = deepcopy( @@ -364,7 +400,25 @@ def _evaluatePurityAtParameter( value, ): """ - Evaluate reconstructed probe purity for one candidate calibration value. + Evaluate probe purity for one calibration candidate. + + An independent candidate reconstruction is created for the specified + axial distance or wavelength and reconstructed using `mPIE`. The resulting + probe modes are orthogonalized before the reconstructed probe purity is + evaluated. + + Args: + target (str): + Calibration target. Supported values are `"z"` and `"wavelength"`. + + value (float): + Candidate axial distance or wavelength in meters. + + Returns: + tuple: + `(purity, candidate_reconstruction)`, where `purity` is the + reconstructed probe purity and `candidate_reconstruction` contains + the reconstruction state obtained for the candidate. """ ( @@ -432,11 +486,39 @@ def _evaluatePurityAtParameter( def purityZScan(self): """ - Scan a fixed axial range and select the distance that maximizes - reconstructed probe purity. + Perform a fixed-grid purity-based axial calibration. + + The scan evaluates a uniformly spaced set of candidate sample-to-detector + distances centered on the current value `reconstruction.zo`. The scan range + and number of candidate positions are defined by + `params.purityZScanRange` and `params.purityZScanPoints`. + + Each candidate distance is evaluated using an independent `mPIE` + reconstruction. The reconstructed probe modes are orthogonalized, and the + candidate yielding the highest probe purity is retained as the calibrated + axial distance. + + The reconstruction error for each candidate is recorded together with the + purity metric. After the scan, the reconstruction state corresponding to + the highest-purity candidate is restored to the original `Reconstruction` + object. + + The evaluated distances, purity values, and reconstruction errors are stored + in `purityZValues`, `purityZMetrics`, and `purityZErrors`, respectively. - Each candidate z is evaluated using an independent reconstruction - state and a standard mPIE reconstruction. + Returns: + tuple: + `(best_z, best_purity)`, where `best_z` is the calibrated + sample-to-detector distance in meters and `best_purity` is the + corresponding reconstructed probe purity. + + Raises: + ValueError: + If fewer than two probe modes are reconstructed + (`reconstruction.npsm < 2`). + + RuntimeError: + If no valid candidate reconstruction produces a purity value. """ if self.reconstruction.npsm < 2: @@ -745,91 +827,50 @@ def purityZScan(self): def adaptivePurityZScan(self): r""" - Perform adaptive purity-based axial calibration. + Perform an adaptive purity-based axial calibration. - The adaptive scan searches for the axial distance that maximizes the - reconstructed probe purity while reducing the number of candidate - reconstructions compared with a dense fixed-grid scan. + The search starts from the current sample-to-detector distance + `reconstruction.zo` and evaluates the initial position together with two + neighboring candidates separated by `params.purityZInitialStep`. - The search starts from the current propagation distance - `reconstruction.zo` and evaluates three initial positions: - - $$ - z_0-\Delta z,\quad z_0,\quad z_0+\Delta z - $$ - - where the initial step size $\Delta z$ is defined by - `params.purityZInitialStep`. - - If the purity increases toward one side, the search continues in that - direction and progressively expands the axial step: + If the reconstructed probe purity increases toward either side, the search + continues in that direction while progressively increasing the step size + according to $$ \Delta z_{k+1} = g\,\Delta z_k $$ - where the expansion factor $g$ is set by `params.purityZStepGrowth`. - Expansion continues while the increase in purity between consecutive - candidate positions exceeds `params.purityZPurityTolerance`. + where `g = params.purityZStepGrowth`. - Once the purity stops increasing, the search enters a refinement phase. - The axial step is progressively reduced according to + Once the purity no longer improves significantly, the search switches to a + local refinement around the best candidate. The step size is progressively + reduced according to $$ \Delta z_{k+1} = s\,\Delta z_k $$ - where the reduction factor $s$ is set by `params.purityZStepShrink`. - Candidate positions on both sides of the current best distance are then - evaluated until the requested axial resolution is reached. - - The adaptive search is controlled by the following parameters: - - - `params.purityZInitialStep`: - Initial axial step size used for the three-point search around the - starting distance. - - - `params.purityZStepGrowth`: - Multiplicative factor used to increase the step size while moving - toward higher purity. - - - `params.purityZStepShrink`: - Multiplicative factor used to reduce the step size during refinement. + where `s = params.purityZStepShrink`. - - `params.purityZMinStep`: - Minimum axial step size. Refinement stops once the current step becomes - smaller than this value. + Refinement continues until the step size falls below + `params.purityZMinStep` or the maximum number of candidate evaluations + defined by `params.purityMaxEvaluations` is reached. Previously evaluated + axial positions are cached and reused. - - `params.purityZPurityTolerance`: - Minimum purity improvement considered significant when determining - whether the search should continue in the current direction. - - - `params.purityZMaxEvaluations`: - Maximum number of candidate axial distances that may be reconstructed - during one adaptive scan. - - Candidate distances that have already been evaluated are cached and reused, - avoiding repeated `mPIE` reconstructions at the same axial position. - - Each new candidate is evaluated independently using - `_evaluatePurityAtZ()`. The candidate with the highest reconstructed probe - purity is retained as the calibrated axial distance, and its reconstruction - state is restored to the original `Reconstruction` object. - - The adaptive scan results are stored in: + Each candidate distance is evaluated using an independent `mPIE` + reconstruction. The candidate yielding the highest reconstructed probe + purity is retained, and its reconstruction state is restored to the original + `Reconstruction` object. - - `purityZValues`: evaluated axial distances in meters, sorted by z. - - `purityZMetrics`: reconstructed probe purity corresponding to each - evaluated distance. - - `purityBestZ`: axial distance corresponding to the highest purity. - - `purityBestValue`: maximum reconstructed probe purity. - - `purityZInitialGuess`: axial distance used as the starting point. + The evaluated distances, purity values, and reconstruction errors are stored + in `purityZValues`, `purityZMetrics`, and `purityZErrors`, respectively. Returns: tuple: - `(best_z, best_purity)`, where `best_z` is the calibrated axial - propagation distance in meters and `best_purity` is the maximum - reconstructed probe purity. + `(best_z, best_purity)`, where `best_z` is the calibrated + sample-to-detector distance in meters and `best_purity` is the + corresponding reconstructed probe purity. Raises: ValueError: @@ -838,10 +879,6 @@ def adaptivePurityZScan(self): RuntimeError: If no valid candidate reconstruction produces a purity value. - - See Also: - `purityZScan` - Fixed-grid purity-based axial calibration. """ if self.reconstruction.npsm < 2: raise ValueError( @@ -1174,14 +1211,21 @@ def evaluate(z): return best_z, best_purity def purityWavelengthScan(self): - - """ - Scan a fixed wavelength range and select the wavelength that maximizes - reconstructed probe purity. + Perform a fixed-grid purity-based wavelength calibration. - Each candidate wavelength is evaluated using an independent reconstruction - state and a standard mPIE reconstruction. + Candidate wavelengths are evaluated over a uniformly spaced range defined + by `params.purityWavelengthScanRange` and + `params.purityWavelengthScanPoints`. The wavelength yielding the highest + reconstructed probe purity is retained. + + Returns: + tuple: + `(best_wavelength, best_purity)`. + + See Also: + `purityZScan` + Fixed-grid purity-based axial calibration. """ if self.reconstruction.npsm < 2: @@ -1538,15 +1582,20 @@ def purityWavelengthScan(self): ) def adaptivePurityWavelengthScan(self): - """ - Adaptively search for the wavelength that maximizes reconstructed - probe purity. + Perform an adaptive purity-based wavelength calibration. - The search first determines the uphill direction, then expands the - wavelength step while purity keeps increasing. Once the maximum is - bracketed, the step size is reduced and the search is refined around - the current best wavelength. + The search starts from the current wavelength, expands the step in the + direction of increasing probe purity, and then refines around the best + candidate using progressively smaller steps. + + Returns: + tuple: + `(best_wavelength, best_purity)`. + + See Also: + `adaptivePurityZScan` + adaptive stepsize purity-based axial distance calibration. """ if self.reconstruction.npsm < 2: @@ -2021,413 +2070,6 @@ def evaluate(wavelength): best_purity, ) - def plotPurityScan_X( - self, - show=True, - fit=True, - fit_points=5, -): - """ - Plot probe purity and reconstruction error for the selected - purity-calibration target. - - Optionally fit a quadratic curve locally around the measured purity - maximum and estimate the fitted calibration optimum. - - Args: - show (bool, optional): - Display the figure immediately. Default is `True`. - - fit (bool, optional): - If `True`, perform a local quadratic fit around the measured - purity maximum. Default is `True`. - - fit_points (int, optional): - Number of measured points used for the local quadratic fit. - Default is 5. - - Returns: - tuple: - Matplotlib `(figure, purity_axis, error_axis)` objects. - """ - - target = ( - self.params.purityCalibrationTarget.lower() - ) - - # ------------------------------------------------------------------ - # Select calibration results - # ------------------------------------------------------------------ - - if target == "z": - - if not hasattr( - self, - "purityZValues", - ): - raise RuntimeError( - "No purity z-scan results are available." - ) - - values = ( - np.asarray( - self.purityZValues - ) - * 1e3 - ) - - purity = np.asarray( - self.purityZMetrics - ) - - errors = ( - np.asarray( - self.purityZErrors - ) - if hasattr( - self, - "purityZErrors", - ) - else None - ) - - initial_value = ( - self.purityZInitialGuess - * 1e3 - ) - - xlabel = ( - "Sample-to-detector distance z (mm)" - ) - - title = ( - "Purity-based axial calibration" - ) - - value_symbol = "z" - value_unit = "mm" - value_format = ".3f" - - elif target == "wavelength": - - if not hasattr( - self, - "purityWavelengthValues", - ): - raise RuntimeError( - "No purity wavelength-scan results are available." - ) - - values = ( - np.asarray( - self.purityWavelengthValues - ) - * 1e9 - ) - - purity = np.asarray( - self.purityWavelengthMetrics - ) - - errors = ( - np.asarray( - self.purityWavelengthErrors - ) - if hasattr( - self, - "purityWavelengthErrors", - ) - else None - ) - - initial_value = ( - self.purityWavelengthInitialGuess - * 1e9 - ) - - xlabel = ( - "Wavelength λ (nm)" - ) - - title = ( - "Purity-based wavelength calibration" - ) - - value_symbol = "λ" - value_unit = "nm" - value_format = ".4f" - - else: - - raise ValueError( - "Unsupported purity calibration target: " - f"{self.params.purityCalibrationTarget!r}" - ) - - # ------------------------------------------------------------------ - # Best measured purity - # ------------------------------------------------------------------ - - best_purity_index = np.argmax( - purity - ) - - fig, ax_purity = plt.subplots( - figsize=(7, 4.5) - ) - - ax_purity.plot( - values, - purity, - marker="o", - label="Probe purity", - ) - - ax_purity.plot( - values[ - best_purity_index - ], - purity[ - best_purity_index - ], - marker="*", - markersize=12, - linestyle="None", - label=( - f"Measured maximum: " - f"{value_symbol} = " - f"{format(values[best_purity_index], value_format)} " - f"{value_unit}" - ), - ) - - # ------------------------------------------------------------------ - # Local quadratic fitting - # ------------------------------------------------------------------ - - if fit: - - if fit_points < 3: - raise ValueError( - "`fit_points` must be at least 3." - ) - - if fit_points > len(values): - fit_points = len(values) - - half_window = ( - fit_points // 2 - ) - - start = max( - 0, - best_purity_index - - half_window, - ) - - stop = min( - len(values), - start + fit_points, - ) - - start = max( - 0, - stop - fit_points, - ) - - fit_x = values[ - start:stop - ] - - fit_y = purity[ - start:stop - ] - - if len(fit_x) >= 3: - - coefficients = np.polyfit( - fit_x, - fit_y, - deg=2, - ) - - a, b, c = coefficients - - if a < 0: - - fitted_best_value = ( - -b / (2 * a) - ) - - fitted_best_purity = ( - np.polyval( - coefficients, - fitted_best_value, - ) - ) - - fit_curve_x = np.linspace( - fit_x.min(), - fit_x.max(), - 200, - ) - - fit_curve_y = np.polyval( - coefficients, - fit_curve_x, - ) - - ax_purity.plot( - fit_curve_x, - fit_curve_y, - linestyle="--", - label="Quadratic fit", - ) - - ax_purity.plot( - fitted_best_value, - fitted_best_purity, - marker="X", - markersize=9, - linestyle="None", - label=( - f"Fitted maximum: " - f"{value_symbol} = " - f"{format(fitted_best_value, value_format)} " - f"{value_unit}" - ), - ) - - if target == "z": - self.purityBestZFitted = ( - fitted_best_value - * 1e-3 - ) - - self.purityBestZFittedValue = ( - fitted_best_purity - ) - - elif target == "wavelength": - self.purityBestWavelengthFitted = ( - fitted_best_value - * 1e-9 - ) - - self.purityBestWavelengthFittedValue = ( - fitted_best_purity - ) - - else: - - self.logger.warning( - "Quadratic fit does not produce a local maximum " - "(fitted curvature is non-negative)." - ) - - ax_purity.set_xlabel( - xlabel - ) - - ax_purity.set_ylabel( - "Probe purity" - ) - - # ------------------------------------------------------------------ - # Reconstruction error - # ------------------------------------------------------------------ - - ax_error = ax_purity.twinx() - - if errors is not None: - - ax_error.plot( - values, - errors, - marker="s", - linestyle="--", - label="Reconstruction error", - ) - - best_error_index = np.argmin( - errors - ) - - ax_error.plot( - values[ - best_error_index - ], - errors[ - best_error_index - ], - marker="*", - markersize=12, - linestyle="None", - label=( - f"Minimum error: " - f"{value_symbol} = " - f"{format(values[best_error_index], value_format)} " - f"{value_unit}" - ), - ) - - ax_error.set_ylabel( - "Reconstruction error" - ) - - # ------------------------------------------------------------------ - # Initial guess - # ------------------------------------------------------------------ - - ax_purity.axvline( - initial_value, - linestyle="--", - alpha=0.5, - label=( - f"Initial {value_symbol}: " - f"{format(initial_value, value_format)} " - f"{value_unit}" - ), - ) - - # ------------------------------------------------------------------ - # Combined legend - # ------------------------------------------------------------------ - - purity_handles, purity_labels = ( - ax_purity.get_legend_handles_labels() - ) - - error_handles, error_labels = ( - ax_error.get_legend_handles_labels() - ) - - ax_purity.legend( - purity_handles + error_handles, - purity_labels + error_labels, - loc="best", - ) - - ax_purity.set_title( - title - ) - - ax_purity.grid( - alpha=0.3 - ) - - fig.tight_layout() - - if show: - plt.show() - - return ( - fig, - ax_purity, - ax_error, - ) - - def plotPurityScan( self, show=True,