Ptycho: multibeam support & more generic mode description
Enable multibeam support, i.e. when:
- the illumination consists of multiple, incoherent beams
- the illuminated parts of the sample are disjoint
See: https://arxiv.org/abs/2402.12082
The disjoint nature of the reconstruction can be handled in different ways: either reconstruct independently and merge at the end (similarly to asynchronous MPI ptycho), or assume the scan is large enough to overlap the reconstructed patches from different parts.
This requires a change in the ptycho modes, as each 'multibeam' mode corresponds to one probe and one object:
- currently the number of modes is
nb_probe x nb_obj
and the shape of psi is[nb_obj, nb_probe, stack_size, ny, nx]
- it could also be possible to use coherent probe modes, where
nb_probe
only leads to 1 incoherent mode, but different for each position (a linear -coherent- combination of the probe modes), though that is only implemented in holotomo so far - by adding a
nb_beam
property, the number of object, probes, and incoherent modes would all be multiplied bynb_beam
Possible implementation:
- switch to using
@property
- Add
nb_beam
property, and associated functions - All
obj,probe<->Psi
need to be updated in operators- including the way normalisation is made (with multiple beams, the normalisation is not the same anymore for all modes)
- need to update:
- export
- illumination calculations
- ...
Instead of having multiple loops for obj/probe/beam/frames, could a more generic approach be used ? Just use:
- a psi shape
[nb_mode, stack_size, ny, nx]
- for incoherent modes: for each mode keep track of which object and probe mode it corresponds to, using two index arrays of shape
[nb_frame, nb_mode]
:- for a standard configuration (1 object, 2 probe modes) the index arrays would be e.g.
idxobj=[0,0,...,0]
andidxprobe=[0,1,0,1,0,1...0,1]
- for a multibeam (2 beams, 1 obj, 1 probe) configuration the index arrays would be
idxobj=[0,1,0,1,...,0,1]
andidxprobe=[0,1,0,1,0,1...0,1]
, and for 2 probe modes (2 beams, 1 obj, 2 probes) would beidxobj=[0,0,1,1,0,0,1,1,...,0,0,1,1]
andidxprobe=[0,1,2,3,0,1,2,3...0,1,2,3]
(with probe modes 0 and 1 for the first object mode, and probe modes 2 and 3 for the second object mode). - when looping over object and probe modes, it should be possible to use a static array to store up to N object and probe values to avoid multiple read/writes (probably not essential for reads, more for writes).
- for a standard configuration (1 object, 2 probe modes) the index arrays would be e.g.
- for coherent (probe) modes:
- each frame would use a linear combination of probe modes, and so an array of probe coefficients is required instead of a single index.
- the size of the object index array would remain
[nb_frame, nb_mode]
- the size of the probe coefficient array would be
[nb_frame, nb_mode, nb_probe]
- this can also be used for incoherent probe modes (and also allows to take into account a scale factor for each illumination position), but with more overhead (extra loop for probe modes, breaks a lot of GPU parallelism, though on modern GPU should not be a problem ?).
Edited by Vincent Favre-Nicolin