04 – CPM Position Correction¶
This tutorial demonstrates lateral scan position correction using the pcPIE engine. Position errors — caused by mechanical stage imperfections or thermal drift — are estimated and corrected jointly with the object and probe during reconstruction.
What you'll learn:
- How to enable and configure position correction with
pcPIE - How to visualize corrected vs. original scan positions
- How position correction improves reconstruction quality
!!! note "Dataset"
The dataset (simu.hdf5) is generated automatically into the example_data/ folder at the project root the first time you run the notebook. If the file already exists it is not re-generated.
!!! tip "Prerequisites" Complete 01 – CPM Simulation before this tutorial.
Position correction in real-space ptychography¶
This tutorial explains how to use position correction in ptyLab
Last updated: 22nd June 2021
Position correction algorithm¶
The position correction code implemented in ptyLab is a modified version of the position correction algorithm described in
- Zhang, Fucai, et al. "Translation position determination in ptychographic coherent diffraction imaging." Optics express 21.11 (2013): 13592-13606.
We recommend citing that paper when using position correction for your results.
The main idea of the algorithm is based on the observation that the object update inside ePIE is slightly shifted towards the true position of each respective object patch. Therefore, the authors use registration to determine the relative shift between each object patch estimate before and after ePIE. Let $\textbf{d}_k=(\Delta_{x,k}, \,\Delta_{y,k})^T$ be the relative shift (in units of pixels) between the object before and after ePIE (at the $k^{th}$ position). This signal may be used to provide a feedback on the position estimate $\textbf{t}_k$ at the respective position. The updated estimated position at position $k$ and iteration $n$ is then given by $\textbf{t}_{n,k} = \textbf{t}_{n−1,k}+\gamma \textbf{d}_{k,n}.$
Modification in ptyLab¶
In ptyLab we use an adam (*) optimizer to speed up the convergence rate of Zhang's position correction algorithm. At the point of this writing, we have not carried out exhaustive hyperparameter tests, but observe a convergence speed up. Please e-mail us if you have suggestions for optimal parameters.
The adam update rule for our position correction is given recursively by the equations $\textbf{d}_k= \delta \cdot (\Delta_{x,k} ,\,\Delta_{y,k})^T + (1−\delta)\cdot \textbf{d}_{k−1}$,
$s^2_k = \beta s^2_{k−1} + (1−\beta)∥\textbf{d}_k∥^2$,
where we observed a good convergence rate for $\gamma=0.5$ , $\beta=0.999$, and $s_0=1$ .
For more information on adam, see:
(*) Kingma, Diederik P., and Jimmy Ba. "Adam: A method for stochastic optimization." arXiv preprint arXiv:1412.6980 (2014).
Setup¶
The dataset (simu.hdf5) is generated automatically into the example_data/ folder the first time the next cell is run. The file contains a synthetic CPM dataset produced by the same simulation as Tutorial 01.
First, we import the libraries and generate (or load) the dataset.
import matplotlib
import PtyLab
from PtyLab.io import getExampleDataFolder, generate_simu_hdf5
import logging
logging.basicConfig(level=logging.INFO)
import numpy as np
from PtyLab import Params
from PtyLab import Monitor
from PtyLab import Reconstruction
from PtyLab import ExperimentalData
from PtyLab import Engines
fileName = "simu.hdf5"
filePath = getExampleDataFolder() / fileName
if not filePath.exists():
print("Generating simu.hdf5...")
generate_simu_hdf5(filePath)
print("Done.")
else:
print(f"{fileName} found, skipping generation.")
simu.hdf5 found, skipping generation.
# easy initialization
exampleData, reconstruction, params, monitor, ePIE_engine = PtyLab.easyInitialize(
filePath
)
INFO:GPU:cupy and CUDA available, switching to GPU INFO:Reconstruction:Copying attribute wavelength INFO:Reconstruction:Copying attribute dxd INFO:Reconstruction:Copying attribute theta INFO:Reconstruction:Copying attribute spectralDensity INFO:Reconstruction:Copying attribute entrancePupilDiameter INFO:Reconstruction:Initial object set to ones INFO:Reconstruction:Initial probe set to circ INFO:ePIE:Sucesfully created ePIE ePIE_engine
INFO:ePIE:Wavelength attribute: 6.327999813038332e-07
Found encoder with shape (100, 2)
Before the reconstruction is started, we perturb the scan grid, using the following commands. These lines produce the following plot. The blue points are the correct scan grid. The yellow points are the perturbed scan grid.
# perturb encoder positions
maxPosError = 10
encoder0 = exampleData.encoder.copy()
exampleData.encoder = encoder0 + maxPosError * reconstruction.dxo * (
2 * np.random.rand(encoder0.shape[0], encoder0.shape[1]) - 1
)
import matplotlib.pyplot as plt
figure, ax = plt.subplots(1, 1, num=112, squeeze=True, clear=True, figsize=(5, 5))
ax.set_title("Original and perturbed scan grid positions")
ax.set_xlabel("(um)")
ax.set_ylabel("(um)")
(line1,) = plt.plot(encoder0[:, 1] * 1e6, encoder0[:, 0] * 1e6, "bo", label="correct")
(line2,) = plt.plot(
exampleData.encoder[:, 1] * 1e6,
exampleData.encoder[:, 0] * 1e6,
"yo",
label="false",
)
plt.legend(handles=[line1, line2])
plt.tight_layout()
plt.show(block=False)
We use the incorrect scan grid as an input for the reconstruction. After altering the encoder data, we need to reset the reconstruction with the perturbed data as original data.
# initialize again reconstruction with the wrong encoder data
reconstruction = Reconstruction(exampleData, params)
INFO:Reconstruction:Copying attribute wavelength INFO:Reconstruction:Copying attribute dxd INFO:Reconstruction:Copying attribute theta INFO:Reconstruction:Copying attribute spectralDensity INFO:Reconstruction:Copying attribute entrancePupilDiameter
Now we can set initial values for variables included in the reconstruction, monitor and in params and run ePIE, exactly like in Tutorial #2.
exampleData.showPtychogram()
Min max ptychogram: 2.462452171423518e-14, 32774.45703125 Min max ptychogram: 0.0, 4.5155487060546875 0.0 4.5155487
VBox(children=(IntSlider(value=0, description='Frame', max=99), Output()))
# now, all our experimental data is loaded into experimental_data and we don't have to worry about it anymore.
# now create an object to hold everything we're eventually interested in
reconstruction.npsm = 1 # Number of probe modes to reconstruct
reconstruction.nosm = 1 # Number of object modes to reconstruct
reconstruction.nlambda = 1 # len(exampleData.spectralDensity) # Number of wavelength
reconstruction.nslice = 1 # Number of object slice
reconstruction.initialProbe = "circ"
exampleData.entrancePupilDiameter = (
reconstruction.Np / 3 * reconstruction.dxp
) # initial estimate of beam diameter
reconstruction.initialObject = "ones"
# initialize probe and object and related params
reconstruction.initializeObjectProbe()
# customize initial probe quadratic phase
reconstruction.probe = reconstruction.probe * np.exp(
1.0j
* 2
* np.pi
/ reconstruction.wavelength
* (reconstruction.Xp**2 + reconstruction.Yp**2)
/ (2 * 6e-3)
)
# this will copy any attributes from experimental data that we might care to optimize
# # Set monitor properties
monitor = Monitor()
monitor.figureUpdateFrequency = 1
monitor.objectPlot = "complex" # complex abs angle
monitor.verboseLevel = "high" # high: plot two figures, low: plot only one figure
monitor.objectZoom = 1.5 # control object plot FoV
monitor.probeZoom = 0.5 # control probe plot FoV
# Run the reconstruction
params = Params()
## main parameters
params.positionOrder = "random" # 'sequential' or 'random'
params.propagator = (
"Fresnel" # Fraunhofer Fresnel ASP scaledASP polychromeASP scaledPolychromeASP
)
## how do we want to reconstruct?
params.gpuSwitch = True
params.probePowerCorrectionSwitch = True
params.modulusEnforcedProbeSwitch = False
params.comStabilizationSwitch = True
params.orthogonalizationSwitch = False
params.orthogonalizationFrequency = 10
params.fftshiftSwitch = False
params.intensityConstraint = "standard" # standard fluctuation exponential poission
params.absorbingProbeBoundary = False
params.objectContrastSwitch = False
params.absObjectSwitch = False
params.backgroundModeSwitch = False
params.couplingSwitch = True
params.couplingAleph = 1
params.positionCorrectionSwitch = False
INFO:Reconstruction:Initial object set to ones INFO:Reconstruction:Initial probe set to circ /mnt/home/shantanu/projects/PtyLab.py/PtyLab/Monitor/Monitor.py:178: UserWarning: For faster update of the reconstruction plot, set `monitor.figureUpdateFrequency = 5` or higher. warnings.warn( /mnt/home/shantanu/projects/PtyLab.py/PtyLab/Monitor/Monitor.py:191: UserWarning: For diffraction data plot, preferably use an interactive matplotlib backend or set `monitor.verboseLevel = "low"`. warnings.warn( INFO:GPU:cupy and CUDA available, switching to GPU
If you use ePIE, this is what we get:
ePIE_engine = Engines.ePIE(reconstruction, exampleData, params, monitor)
## choose engine
# ePIE
engine_ePIE = Engines.ePIE(reconstruction, exampleData, params, monitor)
engine_ePIE.numIterations = 50
engine_ePIE.betaProbe = 0.25
engine_ePIE.betaObject = 0.25
engine_ePIE.reconstruct()
INFO:ePIE:Sucesfully created ePIE ePIE_engine INFO:ePIE:Wavelength attribute: 6.327999813038332e-07 INFO:ePIE:Sucesfully created ePIE ePIE_engine INFO:ePIE:Wavelength attribute: 6.327999813038332e-07
<generator object ePIE.reconstruct at 0x7ba65fdbbb50>
To improve this, select the position correction engine. First, we reset the guesses for the probe and the object to the initial ones ('circ' for the probe and 'ones' for the object). The pcPIE engine has inherently the feature for momentum acceleration in the reconstruction of the object and the probe, therefore, if you want to compare it with ePIE, you need to turn off the momentum update. Also, position correction actually starts after $startAtIteration$ iterations. The default value is 20. With pcPIE and momentum accelaration, this is what we get:
# reset object and probe to initial guesses
reconstruction.initialProbe = "circ"
exampleData.entrancePupilDiameter = (
reconstruction.Np / 3 * reconstruction.dxp
) # initial estimate of beam diameter
reconstruction.initialObject = "ones"
# initialize probe and object and related params
reconstruction.initializeObjectProbe()
# customize initial probe quadratic phase
reconstruction.probe = reconstruction.probe * np.exp(
1.0j
* 2
* np.pi
/ reconstruction.wavelength
* (reconstruction.Xp**2 + reconstruction.Yp**2)
/ (2 * 6e-3)
)
reconstruction.error = []
params.positionCorrectionSwitch = True
engine_pcPIE = Engines.pcPIE(reconstruction, exampleData, params, monitor)
engine_pcPIE.numIterations = 70
engine_pcPIE.reconstruct()
iteration: 69 error: 44.6 estimated linear overlap: 0.0 % estimated area overlap: 52.0 % pcPIE: 100%|██████████| 70/70 [00:28<00:00, 2.49it/s]
INFO:pcPIE:switch to cpu
The estimated scan grid (obtained from pcPIE) is shown by the yellow points in the following plot.
# compare original encoder and result from pcPIE
figure, ax = plt.subplots(1, 1, num=113, squeeze=True, clear=True, figsize=(5, 5))
ax.set_title("Original and corrected after perturbation scan grid positions")
ax.set_xlabel("(um)")
ax.set_ylabel("(um)")
(line1,) = plt.plot(encoder0[:, 1] * 1e6, encoder0[:, 0] * 1e6, "bo", label="correct")
(line2,) = plt.plot(
(reconstruction.positions[:, 1] - reconstruction.No // 2 + reconstruction.Np // 2)
* reconstruction.dxo
* 1e6,
(reconstruction.positions[:, 0] - reconstruction.No // 2 + reconstruction.Np // 2)
* reconstruction.dxo
* 1e6,
"yo",
label="estimated",
)
plt.legend(handles=[line1, line2])
plt.tight_layout()
plt.show(block=False)
We see that the positions estimated by pcPIE match "well" with the ground truth. There is some deviation left, but we can't always win. Finally, we can save the data.
# now save the data
# reconstruction.saveResults('simulation_ePIE_comparison_pcPIE.hdf5')