xptycho blueprints › Algorithm
Algorithm Blueprint
The PMACE algorithm as xptycho runs it (xptycho/pmace.py and xptycho/model.py), organized by the kinds of quantities it uses. The derivation is in the two TCI papers (2023, known probe; 2025, blind multi-mode). The method was first implemented in Qiuchen Zhai's research code ptycho_pmace, which fills in details the papers leave out; xptycho does the same, and this page says where it differs from the papers. The 2023 case is \(K=1\) with the probe held fixed, so this is one algorithm.
1. The quantities
There are five kinds. They differ in who sets them and how long they live. The two kinds of parameters follow the split used in mbirtorch: the forward-model parameters specify the forward model, and the reconstruction parameters specify the reconstruction.
Measured data
Fixed for the whole run. Never modified.
- \(y_j\), one per scan position \(j = 0,\dots,J-1\): the measured amplitude, an \(N_p \times N_p\) real array. It is the square root of the detector counts. (Dark subtraction is preprocessing; the papers do not discuss it.)
The data file also records the scan positions and the instrument facts. Those are forward-model parameters that happen to be read from the file.
Forward-model parameters
Specify the forward model \(A\). Read from the data file or given. Not estimated by the reconstruction; the positions can be replaced by refined ones (refine_probe_positions, then set_params(probe_positions=...)).
- Wavelength (or energy), object-to-detector distance, detector pixel pitch, frame size \(N_p\).
- The scan positions, in meters.
- \(K\), the number of probe modes the model ends with: not a physical fact, but part of the assumed model. A blind run starts with one mode, or with the modes of a given starting probe, and reaches \(K\) on the mode schedule.
From these the model derives the object pixel pitch, \(\lambda L / (N_p \Delta_{\mathrm{det}})\) with \(L\) the detector distance (standard far-field sampling; not stated in the papers), the positions in pixels, and so, with the object grid (a reconstruction parameter), the patch operators \(P_j\). Far-field propagation is a fixed assumption of the model, not a parameter.
Unknowns
The inputs of the forward model. Given, or estimated by the run.
- \(x\): the object, \(N_1 \times N_2\) complex transmittance. Always an output.
- \(d_k\), \(k = 0,\dots,K-1\): the probe modes, each \(N_p \times N_p\) complex. Given and held fixed in the 2023 paper; estimated, and then also an output, in the 2025 paper. The same kind of thing either way.
In the equations \(D_k = \mathrm{diag}(d_k)\) is multiplication by probe mode \(k\). It is the unknown written as a matrix, not an operator of the instrument.
Reconstruction parameters
Chosen by the user. Specify the reconstruction, not the physics.
- \(\alpha_1\): object data-fit weight. \(\alpha_2\): probe data-fit weight.
- \(\kappa \in [1,2]\): exponent on \(|d_k|\) in the averaging weights.
- \(\rho \in (0,1)\): the Mann step.
- The number of iterations. It is an argument of
recon, default 100, not a parameter set with set_params.
- The mode schedule: the iterations at which a new mode is added, and the energy fraction \(f\) a new mode gets (section 5).
- The Fresnel radius used to initialize the probe and to add a mode (sections 4 and 5).
- The object grid: its shape and the position of its first pixel (
object_shape, object_origin). It must hold every patch. By default it is the smallest grid that does. The counterpart of recon_shape in mbirtorch.
- (\(\varepsilon\) in a stable inverse is a fixed formula, \(10^{-6}\sqrt{\|a\|^2/\dim a}\) for the array \(a\) being inverted, not a choice. For the division by the patches, \(a\) is all \(J\) patches together, so one \(\varepsilon\) serves every position.)
Run state
Exists only during the run. Discarded after. The user never sees it.
- \(\mathbf{v} = [v_0,\dots,v_{J-1}]\): one complex \(N_p \times N_p\) patch per scan position, each position's own copy of its part of the object. The same size as the data.
- When the probe is estimated: \(\mathbf{s}_k = [d_{0,k},\dots,d_{J-1,k}]\), each position's own copy of mode \(k\). Also the size of the data, once per mode.
- The other Mann variables: \(\mathbf{w}, \mathbf{z}\) (patches) and \(\mathbf{r}_k, \mathbf{u}_k\) (probes). They are temporary. xptycho never stores \(\mathbf{w}\), \(\mathbf{z}\), or \(\mathbf{r}_k\) for all positions at once: each is formed for one batch of positions and used at once (hardware mapping blueprint, section 2).
PMACE is a consensus method: every scan position holds its own copy of what it can see, and the unknowns are the consensus of the copies. Neither unknown is state. At the end, \(\hat x\) is the weighted average of the \(v_j\) and \(\hat d_k\) is the mean of the \(d_{j,k}\). Inside the loop the steps use consensus values. The object step uses \(\mathbf{z}\) (patches of the average of \(2\mathbf{w}-\mathbf{v}\)) and the modes \(d_k\). The probe step uses the patches \(p_j = P_j \bar x\) of the reported image and the modes \(d_k\). The modes \(d_k\) are the mean of the copies \(\mathbf{s}_k\) after their Mann update.
Not on this page: the device settings. These are the devices, chosen with configure_devices, and the batch size, the parameter batch_size set with set_params. They change where the arithmetic runs, and change the result only by rounding.
2. The forward model and the operators
The forward model is a function of the unknowns, specified by the forward-model parameters:
\[ y_j = \sqrt{\mathrm{Pois}\Big(\sum_{k=0}^{K-1} \big|\mathcal{F} D_k P_j x\big|^2\Big)} \]
Here \(P_j\) cuts patch \(j\) out of the object (from the positions, the pixel pitch, and the object grid), \(D_k = \mathrm{diag}(d_k)\) multiplies by probe mode \(k\), and \(\mathcal{F}\) is the orthonormal 2D DFT, the far-field propagation to the detector. The modes add in intensity, and the detector counts \(I_j = y_j^2\) are Poisson. Everything the algorithm does is built from these pieces and the four steps below.
Building blocks
\(P_j^T\) adds a patch back into the image at position \(j\); \(\mathcal{F}^*\) is the inverse DFT.
A stable inverse of any array \(a\) is \(a^* / (|a|^2 + \varepsilon)\) with \(\varepsilon = 10^{-6}\sqrt{\|a\|^2/\dim a}\); it is written \(D_{k,\varepsilon}^{-1}\), \(X_{j,\varepsilon}^{-1}\), \(\Lambda_{k,\varepsilon}^{-1}\) below. For \(X_{j,\varepsilon}^{-1}\), \(a\) is all \(J\) patches together, so one \(\varepsilon\) serves every position.
\(\Lambda_k = \sum_j P_j^T |D_k|^\kappa P_j\): the coverage map, an image-sized diagonal.
\(w_k = \|d_k\|^2 / \sum_m \|d_m\|^2\): the energy weight of mode \(k\).
In a blind run \(\Lambda_k\) and \(w_k\) are computed from the current consensus modes and change every iteration.
Patch data-fit step \(F^I_j\)
one patch in, one patch out; uses only \(y_j\), \(v_j\), and the consensus modes \(d_k\)
\[\tilde v_{j,k} = D_{k,\varepsilon}^{-1}\,\mathcal{F}^*\!\left( y_j \circ \frac{\mathcal{F} D_k v_j}{\sqrt{\sum_m |\mathcal{F} D_m v_j|^2}} \right), \qquad
F^I_j(v_j) = (1-\alpha_1)\, v_j + \alpha_1 \sum_k w_k\, \tilde v_{j,k}\]
Patch averaging step \(\mathbf{G}^I\)
all patches in, all patches out, through the image. The division by \(\Lambda_k\) is the stable inverse, so a pixel no probe reaches comes out zero and a dimly lit pixel is damped. The object grid is a given shape and may include pixels no patch reaches; those stay zero.
\[\bar x(\mathbf{v}) = \sum_k w_k\, \Lambda_{k,\varepsilon}^{-1} \sum_j P_j^T |D_k|^\kappa v_j, \qquad (\mathbf{G}^I \mathbf{v})_j = P_j\, \bar x(\mathbf{v})\]
Probe data-fit step \(F^P_{j,k}\)
uses the patch \(p_j = P_j \bar x\) of the current average image, \(y_j\), this position's own copy \(d_{j,k}\) for the phase, and the consensus modes \(d_m\) for the amplitude. (The papers write the denominator with the per-position copies and pass \(\mathbf{z}\) for the patch; xptycho does as written here.)
\[\tilde d_{j,k} = X_{j,\varepsilon}^{-1}\,\mathcal{F}^*\!\left( y_j \circ \frac{|\mathcal{F} X_j d_k|}{\sqrt{\sum_m |\mathcal{F} X_j d_m|^2}} \circ \frac{\mathcal{F} X_j d_{j,k}}{|\mathcal{F} X_j d_{j,k}|} \right), \quad X_j = \mathrm{diag}(p_j), \qquad
F^P_{j,k}(d_{j,k}) = (1-\alpha_2)\, d_{j,k} + \alpha_2\, \tilde d_{j,k}\]
Probe averaging step \(\mathbf{G}^P\)
the plain mean over positions, returned as \(J\) identical copies so the Mann update can subtract per position
\[\bar d_k = \frac{1}{J} \sum_j d_{j,k}\]
Outlier replacement
after the Mann update of the copies, per pixel and per mode: copies far from the mean are replaced by it; a mean deviation of zero is replaced by \(10^{-6}\) (in Run.update_probe; not in the papers)
\[ d_k \leftarrow \tfrac{1}{J}\textstyle\sum_j s_{j,k}, \qquad s_{j,k} \leftarrow d_k \ \text{ wherever } \ 0.6745\,\frac{|s_{j,k} - d_k|}{\mathrm{mean}_j |s_{j,k} - d_k|} > 1.5 \]
3. One iteration
As implemented in xptycho (Run.update_object, Run.add_mode, Run.update_probe, called by PtychoModel.recon). xptycho implements the algorithms of the two papers. Each line is tagged local when every position is handled on its own, or global when it needs a sum over all positions.
localw ← FI(v; d)each position fits its own data, using the consensus modes dk
globalz ← GI(2w − v)one scatter-add over all j, divide by Λ, one gather
localv ← v + 2ρ (z − w)Mann update
globalx̄ ← x̄(v); pj ← Pj x̄the image this iteration reports; its patches feed the probe step
xif this iteration is in the mode schedule: add a mode, then make the modes orthogonal if orthogonalize_modessection 5; blind case only
xfor all modes k = 0 … K−1 at once, using the modes d as they stand before this pass:blind case only
localrk ← FPk(sk; p, d)each position refines its own copy of every mode
globaluk ← GP(2rk − sk)one mean over all j, for every mode
localsk ← sk + 2ρ (uk − rk)Mann update
globaldk ← mean of sk; replace outlying copies by dkthe mean and the deviation statistic are sums over j; the replacement is per position
globaldata error of x̄ and dfor the progress line and the convergence record
xoutput: x̂ = x̄; d̂k = dk
Every pass has the same shape:
localeach position: its own \(y_j\), its own state row, the current consensus
→
globalone sum over all positions, then broadcast back
→
localeach position updates its own state row
Once per iteration for the patches, plus one averaging of the updated patches, and once more for the probes, all modes together, when the probe is estimated. That is the entire communication pattern.
4. Initialization
As xptycho does it (PtychoModel.initial_probe, PtychoModel.initial_object). recon uses these formulas only for what is not given. A known probe (probe=) or a starting probe (init_probe=) is used in place of \(d^{(0)}\). A Sample given as init= supplies the starting object, and the starting probe when none is given otherwise.
The probe first, from the mean back-transformed data, Fresnel propagated by \(U_r\), then smoothed by a Gaussian low-pass filter \(S\) with \(\sigma = 1\) pixel on the real and imaginary parts. \(U_r\) is Fresnel propagation with radius \(r\) = probe_fresnel_radius_pixels \(= \sqrt{\lambda z}/\Delta_x\): the radius, in pixels, over which propagation by a distance \(z\) spreads each point of the field, with \(\lambda\) the wavelength and \(\Delta_x\) the object pixel pitch. \(r = 0\) means no propagation, and is the default. Fresnel propagation appears only here and in adding a mode (section 5); the forward model is far-field throughout.
\[ d^{(0)} = S\left\{ U_r\left\{ \frac{1}{J}\sum_j \mathcal{F}^* y_j \right\} \right\} \]
Then the object, real-valued, so that each patch has the right scale: a constant patch of value \(\sqrt{\|y_j\| / \|d^{(0)}\|}\) per position, averaged over the coverage, pixels no patch reaches set to the median, then the same filter on the magnitude. Here \(d^{(0)}\) is the probe the run starts with, all its modes: the formula above, the known probe, or the given starting probe.
\[ x^{(0)} = S\left\{ \left| \Lambda_{0,\varepsilon}^{-1} \sum_j P_j^T \left( \sqrt{\frac{\|y_j\|}{\|d^{(0)}\|}}\; \mathbf{1} \right) \right| \right\}, \qquad \Lambda_0 = \sum_j P_j^T P_j \]
The run state starts as copies: \(v_j^{(0)} = P_j x^{(0)}\), and, when the probe is estimated, \(s_{j,k} = d^{(0)}_k\) for every position and mode. A blind run starts with one mode, or with the modes of the given starting probe.
5. Adding probe modes
A blind run reaches \(K\) modes by adding one at each iteration of the mode schedule (mode_schedule; demo 2 adds the second mode at iteration 20, demo 3 at iteration 10). The addition happens after the image pass of that iteration and before its probe pass (Run.add_mode). With \(E = \sum_k \|d_k\|^2\) the total probe energy before the addition and \(f\) the energy fraction of a new mode (mode_energy_fraction, default 0.05):
- Residual intensity, per position, from the consensus modes and the patches of the current average image, clipped at zero:
\[ I^{\mathrm{res}}_j = \max\!\Big(0,\; I_j - \sum_{k<K} |\mathcal{F} D_k p_j|^2\Big), \qquad p_j = P_j \bar x \]
- The new mode: back-transform the residual amplitude, divide by the patch, average over positions, Fresnel propagate as in the initialization:
\[ d_K = U_r\left\{ \frac{1}{J} \sum_j X_{j,\varepsilon}^{-1}\, \mathcal{F}^* \sqrt{I^{\mathrm{res}}_j} \right\}, \qquad X_j = \mathrm{diag}(p_j) \]
- Scale it to the fraction \(f\) of the total energy: \(d_K \leftarrow d_K \cdot \sqrt{f E} / \|d_K\|\).
- Its per-position copies \(\mathbf{s}_K\) are \(J\) identical copies of \(d_K\).
- Rescale everything: every mode \(d_k\) and every copy \(s_{j,k}\), \(k = 0 \ldots K\), is divided by \(\sqrt{1+f}\), so the total energy is \(E\) again.
- \(K \leftarrow K+1\). The energy weights \(w_k\) follow from the new modes.
Orthogonal modes, optional (orthogonalize_modes, off by default; Run.orthogonalize_modes). Right after the addition, the modes are replaced by an orthogonal set. With the modes as the columns of a matrix \(D = [d_0 \cdots d_{K-1}]\) and its singular value decomposition \(D = U S V^H\), the new modes are the columns of \(U S\), in order of decreasing energy, and every per-position copy \(s_{j,k}\) is set to the new \(d_k\). Since \(U S = D V\) with \(V\) unitary, the new modes predict the same intensities \(\sum_k |\mathcal{F} D_k p_j|^2\) as the old ones. Demo 3, the two-mode reconstruction of the gold-ball data, uses this step; demo 2 does not.
6. The facts that matter for code
- Every pass is local, global, local. The local steps need only their own \(y_j\), their own state row, and the current consensus. The global step is one sum over \(j\). This is what lets the data stream through in batches and across GPUs without changing the answer, provided the order within an iteration (the image pass, the averaging, then the probe pass) is kept.
- The state is as large as the data, times \((1+K)\). xptycho stores \(\mathbf{v}\) and, when the probe is estimated, \(\mathbf{s}_k\) for every position. It does not store \(\mathbf{w}\), \(\mathbf{z}\), or \(\mathbf{r}_k\) for all positions: each is formed for one batch of positions and used at once. Where the per-position state is kept, and how large a scan fits, is the subject of the hardware mapping blueprint.
- The averaging step is scatter-add, divide by coverage, gather. Nothing else touches the whole image. \(\Lambda_k\) is itself a scatter-add of \(|d_k|^\kappa\), recomputed every iteration in a blind run. The iteration does this averaging twice: once on \(2\mathbf{w}-\mathbf{v}\) for the Mann step, once on the updated \(\mathbf{v}\) for the reported image and the probe step.
- The probe pass has four global reductions, each over all modes at once: the \(\varepsilon\) of the division by the patches, the mean of \(2\mathbf{r}_k - \mathbf{s}_k\), the mean of the updated copies, and the per-pixel deviation statistic for outlier replacement, which needs the mean first. Mode addition has two: the \(\varepsilon\) of the division by the patches and the mean of the back-transformed residuals.
- The forward model is fixed by the forward-model parameters and the object grid. \(P_j\) comes from the positions, the pixel pitch, and the object grid (
object_shape, object_origin); \(\mathcal{F}\) is the far-field assumption. The unknowns \(x\) and \(d_k\) are its inputs. Everything else in the algorithm is arithmetic on patches.
7. Parameter names in the API
| API name | symbol | default | meaning |
|---|
object_data_fit | \(\alpha_1\) | 0.6 | weight of the data-fitting step in the object update |
probe_data_fit | \(\alpha_2\) | 0.6 | weight of the data-fitting step in the probe update |
probe_weight_exponent | \(\kappa\) | 1.25 | exponent on the probe magnitude in the averaging weights |
relaxation | \(\rho\) | 0.5 | Mann step |
mode_schedule | | none | the iterations at which a mode is added (section 5) |
mode_energy_fraction | \(f\) | 0.05 | fraction of the probe energy given to a newly added mode (section 5) |
orthogonalize_modes | | off | make the modes orthogonal each time a mode is added (section 5) |
probe_fresnel_radius_pixels | \(r = \sqrt{\lambda z}/\Delta_x\) | 0 | The radius, in pixels, over which the Fresnel propagation of distance \(z\) spreads each point of the field, in the probe initialization and in adding a mode (sections 4 and 5); 0 means no propagation |
object_shape | \(N_1 \times N_2\) | from the positions | rows and columns of the object grid; by default the smallest grid that holds every patch |
object_origin | | from the positions | meters, the position of the center of the first pixel of the object grid |
batch_size | | from the free memory | positions processed together on a device; a device setting that changes the result only by rounding |
num_probe_modes | \(K\) | 1 | number of probe modes; a forward-model parameter, given to the constructor or to from_scan. The rest of this table are reconstruction parameters, set with set_params. |
The defaults are the 2025 paper's values for one mode. With two modes the paper uses \(\alpha_1 = 0.5\).
8. Summary
The five kinds above, with \(K\) a forward-model parameter. The measured data stands alone and is never modified. The run state is internal to the reconstruction and the user never sees it. The probe is an unknown, not part of the physics. The forward model is far-field only; Fresnel propagation appears only in the probe initialization and in adding a mode. Where the papers are silent or differ from xptycho, sections 2 to 5 state what xptycho does.
Which objects hold the parameters and the unknowns is the software architecture page.