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 c4c6c0b2..17f66597 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,3 +26,14 @@ from .OPR import OPR from .qNewton import qNewton from .zPIE import zPIE +from .purityPIE import purityPIE + + +import warnings + +warnings.warn( + "`pcPIE` is deprecated. Use `mPIE` with " + "`params.positionCorrectionSwitch = True` instead.", + DeprecationWarning, + stacklevel=2, +) 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/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 new file mode 100644 index 00000000..839108af --- /dev/null +++ b/PtyLab/Engines/purityPIE.py @@ -0,0 +1,2559 @@ +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 + +from copy import deepcopy + +from PtyLab.Monitor.Monitor import DummyMonitor + +from PtyLab.Engines.mPIE import mPIE +class purityPIE(mPIE): + r""" + Purity-based axial-distance and wavelength self-calibration for + ptychographic reconstruction. + + `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: + + $$ + 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 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]$. + + 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 parameter is therefore obtained by maximizing the + reconstructed probe purity. + + 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: + + - **Axial-distance calibration:** optimizes the sample-to-detector + distance `z`. + + - **Wavelength calibration:** optimizes the illumination wavelength + `λ`. + + The calibration target is selected through + `params.purityCalibrationTarget`, which can be set to `"z"` or + `"wavelength"`. + + For either calibration target, two search strategies are available: + + - **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, + "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 calibration value. + """ + def __init__( + self, + reconstruction: Reconstruction, + experimentalData: ExperimentalData, + params: Params, + monitor: Monitor, + ): + 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.name = "purityPIE" + + def reconstruct( + self, + experimentalData=None, + reconstruction=None, + ): + """ + Run purity-based self-calibration. + + The reconstruction first prepares the current object, probe, geometry, + constraints, and optional GPU state using the standard `mPIE` setup. + + 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. + + After calibration, the reconstruction state corresponding to the optimal + candidate is restored to the original `Reconstruction` object. + + 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 calibration. If provided, it + replaces the dataset currently attached to the engine. + + reconstruction (Reconstruction, optional): + Reconstruction state to calibrate. If provided, it replaces the + reconstruction currently attached to the engine. + + Returns: + tuple: + `(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) + self.changeOptimizable(reconstruction) + + self._prepareReconstruction() + + # Reset momentum buffers after preparation + 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: + + 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_value, best_purity + + def _createCandidateReconstruction( + self, + target, + 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( + self.experimentalData + ) + + candidate_params = deepcopy( + self.params + ) + + # -------------------------------------------------------------- + # Set candidate parameter before creating Reconstruction + # -------------------------------------------------------------- + + 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 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.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 + ) + + candidate_reconstruction.initializeObjectProbe() + + return ( + candidate_data, + candidate_reconstruction, + candidate_params, + ) + + def _evaluatePurityAtZ(self, z): + return self._evaluatePurityAtParameter( + "z", + z, + ) + + def _evaluatePurityAtWavelength( + self, + wavelength, + ): + return self._evaluatePurityAtParameter( + "wavelength", + wavelength, + ) + + def _evaluatePurityAtParameter( + self, + target, + 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. + """ + + ( + candidate_data, + candidate_reconstruction, + candidate_params, + ) = self._createCandidateReconstruction( + target, + value, + ) + + candidate_monitor = DummyMonitor() + + candidate_engine = mPIE( + candidate_reconstruction, + candidate_data, + candidate_params, + candidate_monitor, + ) + + 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() + + candidate_engine.orthogonalization() + + purity = float( + np.asarray( + asNumpyArray( + candidate_reconstruction.purityProbe + ) + ).squeeze() + ) + + return ( + purity, + candidate_reconstruction, + ) + + def purityZScan(self): + """ + 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. + + 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: + 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 = [] + error_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) + ) + + 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 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.purityZErrors = np.asarray( + error_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): + r""" + Perform an adaptive purity-based axial calibration. + + 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`. + + 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 `g = params.purityZStepGrowth`. + + 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 `s = params.purityZStepShrink`. + + 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. + + 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. + + 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 + 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: + 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.purityMaxEvaluations + + evaluated = {} + + best_z = None + best_purity = -np.inf + best_state = None + + z_history = [] + purity_history = [] + error_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) + ) + 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"error = {candidate_error:.6e}" + ) + + # 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.purityZErrors = np.asarray(error_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 purityWavelengthScan(self): + """ + Perform a fixed-grid purity-based wavelength calibration. + + 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: + raise ValueError( + "Purity-based wavelength scanning requires at least two probe modes " + "(reconstruction.npsm >= 2)." + ) + + wavelength0 = float( + self.reconstruction.wavelength + ) + + self.purityWavelengthInitialGuess = ( + wavelength0 + ) + + wavelength_values = np.linspace( + wavelength0 + - self.params.purityWavelengthScanRange, + wavelength0 + + self.params.purityWavelengthScanRange, + self.params.purityWavelengthScanPoints, + ) + + purity_values = [] + error_values = [] + + best_purity = -np.inf + best_wavelength = wavelength0 + best_state = None + + tqdm.tqdm.write("") + tqdm.tqdm.write( + "Starting purity-based wavelength scan" + ) + + 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}" + ) + + tqdm.tqdm.write("") + + for wavelength_index, wavelength in enumerate( + wavelength_values + ): + + delta_wavelength_nm = ( + wavelength - wavelength0 + ) * 1e9 + + 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: + + 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}" + ) + + 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): + """ + Perform an adaptive purity-based wavelength calibration. + + 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: + 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( + 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/Engines/zPIE.py b/PtyLab/Engines/zPIE.py index 1d03c685..7861cb10 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, @@ -43,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 @@ -56,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 @@ -90,6 +233,89 @@ 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. + + 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() @@ -312,10 +538,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) @@ -329,10 +558,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/PtyLab/Params/Params.py b/PtyLab/Params/Params.py index cc77de05..dc764fa7 100644 --- a/PtyLab/Params/Params.py +++ b/PtyLab/Params/Params.py @@ -153,6 +153,36 @@ 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.purityCalibrationTarget = "z" # aternatively wavelength + + self.purityWavelengthScanRange = 1e-9 + self.purityWavelengthScanPoints = 11 + self.purityWavelengthScanPlot = True + + self.purityZScanRange = 100e-6 + self.purityZScanPoints = 11 + self.purityZScanPlot = True + + # 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-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 # Default values of all OPR parameters 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 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_yourData.py b/example_scripts/example_purity_Zscan_yourData.py new file mode 100644 index 00000000..63d7fc65 --- /dev/null +++ b/example_scripts/example_purity_Zscan_yourData.py @@ -0,0 +1,523 @@ +import logging +from tkinter import Tk, filedialog + +import matplotlib +matplotlib.use("QtAgg") + +import matplotlib.pyplot as plt + +import PtyLab +from PtyLab import Engines + + +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 +# ============================================================================= + +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", "*.*"), + ], + ) + + root.destroy() + + return file_path + + +print() +print("=" * 70) +print("Purity-based parameter 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" +) + +if not filePath: + raise RuntimeError( + "No ptychography dataset was selected." + ) + +print() +print("Selected dataset:") +print(filePath) + + +# ============================================================================= +# Optional seed reconstruction +# ============================================================================= + +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 +# ============================================================================= + +experimentalData, reconstruction, params, monitor, _ = ( + PtyLab.easyInitialize( + filePath, + operationMode="CPM", + ) +) + +# Change only if required by your dataset. +experimentalData.setOrientation(4) + +print() +print("Loaded file:", experimentalData.filename) +print("Ptychogram shape:", experimentalData.ptychogram.shape) +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 + +# 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 + +# ============================================================================= +# Initial calibration guess +# ============================================================================= + +if CALIBRATION_TARGET == "z": + + # 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 + +reconstruction.nosm = 1 +reconstruction.nlambda = 1 +reconstruction.nslice = 1 + + +# ============================================================================= +# Object and probe initialization +# ============================================================================= + +reconstruction.initialObject = "ones" + +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" + ) + + 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.initializeObjectProbe() + +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" + +# 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.absorbingProbeBoundary = False +params.objectContrastSwitch = False +params.absObjectSwitch = False +params.backgroundModeSwitch = False + +params.couplingSwitch = True +params.couplingAleph = 1 + + +monitor.describe_parameters(params) + + +# ============================================================================= +# Run purity-based calibration +# ============================================================================= + +purity_engine = Engines.purityPIE( + reconstruction, + experimentalData, + params, + monitor, +) + +# Number of mPIE iterations used to evaluate EACH candidate. +purity_engine.numIterations = 5 + +best_value, best_purity = ( + purity_engine.reconstruct() +) + + +# ============================================================================= +# Print results +# ============================================================================= + +print() +print("=" * 70) +print("Purity-based calibration finished") +print("=" * 70) + + +if CALIBRATION_TARGET == "z": + + print( + 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" + ) + + 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 + +if continue_with_mPIE: + + mPIE_engine = Engines.mPIE( + reconstruction, + experimentalData, + params, + monitor, + ) + + mPIE_engine.numIterations = 50 + + mPIE_engine.reconstruct() + + +# ============================================================================= +# Save / display +# ============================================================================= + +# reconstruction.saveResults( +# "purity_calibrated_recon.hdf5", +# squeeze=False, +# ) + +plt.show(block=True) \ 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) 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