Engines

PtyLab.Engines

BaseEngine

BaseEngine

Bases: object

Common properties that are common for all reconstruction Engines are defined here.

Unless you are testing the code, there's hardly any need to create this object. For your own implementation, inherit from this object

Source code in PtyLab/Engines/BaseEngine.py

class BaseEngine(object):
    """
    Common properties that are common for all reconstruction Engines are defined here.

    Unless you are testing the code, there's hardly any need to create this object. For your own implementation,
    inherit from this object

    """

    def __init__(
        self,
        reconstruction: Reconstruction,
        experimentalData: ExperimentalData,
        params: Params,
        monitor: Monitor,
    ):
        # These statements don't copy any data, they just keep a reference to the object
        self.betaObject = 0.25
        self.reconstruction: Reconstruction = reconstruction
        self.experimentalData = experimentalData
        self.params = params
        self.monitor = monitor
        self.monitor.reconstruction = reconstruction

        # datalogger
        self.logger = logging.getLogger("BaseEngine")

    def _prepareReconstruction(self):
        """
        Initialize everything that depends on user changeable attributes.
        :return:
        """
        # check miscellaneous quantities specific for certain Engines
        self._checkMISC()
        self._checkFFT()
        # self._initializeQuadraticPhase()
        self._initialProbePowerCorrection()
        self._probeWindow()
        self._initializeErrors()
        self._setObjectProbeROI()
        self._showInitialGuesses()
        self._initializePCParameters()
        self._checkGPU()  # checkGPU needs to be the last

        # self.reconstruction.probe_storage.push(self.reconstruction.probe, 0, self.experimentalData.ptychogram.shape[0])

    def _setCPSC(self):
        """
        set constrained-pixel-sum constraint:
        -save measured diffraction patterns into ptychograpmDownsampled
        -pad the probe (useful when having a pre-calibrated probe)
        -update the coordinates
        """

        # save the measured ptychogram into ptychograpmDownsampled
        self.experimentalData.ptychogramDownsampled = self.experimentalData.ptychogram

        # pad the probe
        padNum_before = (
            (self.params.CPSCupsamplingFactor - 1) * self.reconstruction.Np // 2
        )
        padNum_after = (
            self.params.CPSCupsamplingFactor - 1
        ) * self.reconstruction.Np - padNum_before
        self.reconstruction.probe = np.pad(
            self.reconstruction.probe,
            (
                (0, 0),
                (0, 0),
                (0, 0),
                (0, 0),
                (padNum_before, padNum_after),
                (padNum_before, padNum_after),
            ),
        )

        # pad the momentums, buffers
        if hasattr(self.reconstruction, "probeBuffer"):
            self.reconstruction.probeBuffer = self.reconstruction.probe.copy()
        if hasattr(self.reconstruction, "probeMomentum"):
            self.reconstruction.probeMomentum = np.pad(
                self.reconstruction.probeMomentum,
                (
                    (0, 0),
                    (0, 0),
                    (0, 0),
                    (0, 0),
                    (padNum_before, padNum_after),
                    (padNum_before, padNum_after),
                ),
            )

        # update coordinates (only need to update the Nd and dxd, the rest updates automatically)
        self.reconstruction.Nd = (
            self.experimentalData.ptychogramDownsampled.shape[-1]
            * self.params.CPSCupsamplingFactor
        )
        self.reconstruction.dxd = (
            self.reconstruction.dxd / self.params.CPSCupsamplingFactor
        )

        self.logger.info("CPSCswitch is on, coordinates(dxd,dxp,dxo) have been updated")

    def update_data(self, experimentalData, reconstruction=None):
        """Update the experimentalData if necessary"""
        self.experimentalData = experimentalData
        if reconstruction is not None:
            self.reconstruction = reconstruction

    def _initializePCParameters(self):
        if self.params.positionCorrectionSwitch:
            # additional pcPIE parameters as they appear in Matlab
            self.daleth = 0.5  # feedback
            self.beth = 0.9  # friction
            self.adaptStep = 1  # adaptive step size
            self.D = np.zeros(
                (self.experimentalData.numFrames, 2)
            )  # position search direction
            # predefine shifts
            rmax = 2
            dy, dx = np.mgrid[-rmax : rmax + 1, -rmax : rmax + 1]

            # self.rowShifts = dy.flatten()#np.array([-1, -1, -1, 0, 0, 0, 1, 1, 1])
            self.rowShifts = np.array([-1, -1, -1, 0, 0, 0, 1, 1, 1])
            # self.colShifts = dx.flatten()#np.array([-1, 0, 1, -1, 0, 1, -1, 0, 1])
            self.colShifts = np.array([-1, 0, 1, -1, 0, 1, -1, 0, 1])
            self.startAtIteration = 1
            self.meanEncoder00 = np.mean(self.experimentalData.encoder[:, 0]).copy()
            self.meanEncoder01 = np.mean(self.experimentalData.encoder[:, 1]).copy()

    def _initializeErrors(self):
        """
        initialize all kinds of errors:
        detectorError is a matrix calculated at each iteration (numFrames,Nd,Nd);
        errorAtPos sums over detectorError at each iteration, (numFrames,1);
        reconstruction.error sums over errorAtPos, one number at each iteration;
        """
        # initialize detector error matrices
        if self.params.saveMemory:
            self.reconstruction.detectorError = 0
        else:
            if not hasattr(self.reconstruction, "detectorError"):
                self.reconstruction.detectorError = np.zeros(
                    (
                        self.experimentalData.numFrames,
                        self.reconstruction.Nd,
                        self.reconstruction.Nd,
                    )
                )
        # initialize energy at each scan position
        if not hasattr(self.reconstruction, "errorAtPos"):
            self.reconstruction.errorAtPos = np.zeros(
                (self.experimentalData.numFrames, 1), dtype=np.float32
            )
        # initialize final error
        if not hasattr(self.reconstruction, "error"):
            self.reconstruction.error = []

    def _initialProbePowerCorrection(self):
        if self.params.probePowerCorrectionSwitch:
            self.reconstruction.probe = (
                self.reconstruction.probe
                / np.sqrt(
                    np.sum(self.reconstruction.probe * self.reconstruction.probe.conj())
                )
                * self.experimentalData.maxProbePower
            )

    def _probeWindow(self):
        # absorbing probe boundary: filter probe with super-gaussian window function
        if not self.params.saveMemory or self.params.absorbingProbeBoundary:
            self.probeWindow = np.exp(
                -(
                    (
                        (self.reconstruction.Xp**2 + self.reconstruction.Yp**2)
                        / (
                            2
                            * (
                                3
                                / 4
                                * self.reconstruction.Np
                                * self.reconstruction.dxp
                                / 2.355
                            )
                            ** 2
                        )
                    )
                    ** 10
                )
            )

        if self.params.probeBoundary:
            self.probeWindow = circ(
                self.reconstruction.Xp,
                self.reconstruction.Yp,
                self.experimentalData.entrancePupilDiameter
                + self.experimentalData.entrancePupilDiameter * 0.2,
            )

    def _setObjectProbeROI(self, update=False):
        """
        Set object/probe ROI for monitoring
        """
        if not hasattr(self.monitor, "objectROI") or update:
            if self.monitor.objectZoom == "full" or self.monitor.objectZoom is None:
                self.monitor.objectROI = [
                    slice(None, None, None),
                    slice(None, None, None),
                ]
            else:
                rx, ry = (
                    (
                        np.max(self.reconstruction.positions, axis=0)
                        - np.min(self.reconstruction.positions, axis=0)
                        + self.reconstruction.Np
                    )
                    / self.monitor.objectZoom
                ).astype(int)
                xc, yc = (
                    (
                        np.max(self.reconstruction.positions, axis=0)
                        + np.min(self.reconstruction.positions, axis=0)
                        + self.reconstruction.Np
                    )
                    / 2
                ).astype(int)

                # self.monitor.objectROI = [
                #     slice(
                #         max(0, yc - ry // 2), min(self.reconstruction.No, yc + ry // 2)
                #     ),
                #     slice(
                #         max(0, xc - rx // 2), min(self.reconstruction.No, xc + rx // 2)
                #     ),
                # ]
                self.monitor.objectROI = [
                    slice(
                        max(0, xc - rx // 2), min(self.reconstruction.No, xc + rx // 2)
                    ),
                    slice(
                        max(0, yc - ry // 2), min(self.reconstruction.No, yc + ry // 2)
                    ),
                ]

        if not hasattr(self.monitor, "probeROI") or update:
            if self.monitor.probeZoom == "full" or self.monitor.probeZoom is None:
                self.monitor.probeROI = [slice(None, None), slice(None, None)]
            else:
                r = int(
                    self.experimentalData.entrancePupilDiameter
                    / self.reconstruction.dxp
                    / self.monitor.probeZoom
                )
                self.monitor.probeROI = [
                    slice(
                        max(0, self.reconstruction.Np // 2 - r),
                        min(self.reconstruction.Np, self.reconstruction.Np // 2 + r),
                    ),
                    slice(
                        max(0, self.reconstruction.Np // 2 - r),
                        min(self.reconstruction.Np, self.reconstruction.Np // 2 + r),
                    ),
                ]

    def _showInitialGuesses(self):
        self.monitor.initializeMonitors()
        objectEstimate = np.squeeze(
            self.reconstruction.object[
                ..., self.monitor.objectROI[0], self.monitor.objectROI[1]
            ]
        )
        probeEstimate = np.squeeze(
            self.reconstruction.probe[
                ..., self.monitor.probeROI[0], self.monitor.probeROI[1]
            ]
        )

        self.monitor.updateObjectProbeErrorMonitor(
            error=self.reconstruction.error,
            object_estimate=objectEstimate,
            probe_estimate=probeEstimate,
            zo=self.reconstruction.zo,
            purity_probe=self.reconstruction.purityProbe,
            purity_object=self.reconstruction.purityObject,
            encoder_positions=self.reconstruction.positions,
        )

        # self.monitor.updateObjectProbeErrorMonitor()

    def _checkMISC(self):
        """
        checks miscellaneous quantities specific certain Engines
        """
        if self.params.backgroundModeSwitch:
            self.reconstruction.background = 1e-1 * np.ones(
                (self.reconstruction.Np, self.reconstruction.Np)
            )

        # preallocate intensity scaling vector
        if self.params.intensityConstraint == "fluctuation":
            self.intensityScaling = np.ones(self.experimentalData.numFrames)

        if self.params.intensityConstraint == "interferometric":
            self.reconstruction.reference = np.ones(
                self.reconstruction.probe[0, 0, 0, 0, ...].shape
            )

        # check if both probePoprobePowerCorrectionSwitch and modulusEnforcedProbeSwitch are on.
        # Since this can cause a contradiction, it raises an error
        if (
            self.params.probePowerCorrectionSwitch
            and self.params.modulusEnforcedProbeSwitch
        ):
            raise ValueError(
                "probePowerCorrectionSwitch and modulusEnforcedProbeSwitch "
                "can not simultaneously be switched on!"
            )

        if self.params.propagatorType == "ASP" and self.params.fftshiftSwitch:
            raise ValueError(
                "ASP propagatorType works only with fftshiftSwitch = False"
            )
        if self.params.propagatorType == "scaledASP" and self.params.fftshiftSwitch:
            raise ValueError(
                "scaledASP propagatorType works only with fftshiftSwitch = False"
            )

        if self.params.CPSCswitch:
            if not hasattr(self.experimentalData, "ptychogramDownsampled"):
                if self.params.CPSCupsamplingFactor == None:
                    raise ValueError(
                        "CPSCswitch is on, CPSCupsamplingFactor need to be set"
                    )
                else:
                    self._setCPSC()

    def _checkFFT(self):
        """
        shift arrays to accelerate fft
        """
        if self.params.fftshiftSwitch:
            if self.params.fftshiftFlag == 0:
                print("check fftshift...")
                print("fftshift data for fast far-field update")
                # shift detector quantities
                self.experimentalData.ptychogram = np.fft.ifftshift(
                    self.experimentalData.ptychogram, axes=(-1, -2)
                )
                if hasattr(self.experimentalData, "ptychogramDownsampled"):
                    self.experimentalData.ptychogramDownsampled = np.fft.ifftshift(
                        self.experimentalData.ptychogramDownsampled, axes=(-1, -2)
                    )
                if hasattr(self.experimentalData, "w"):
                    if self.experimentalData.W is not None:
                        self.experimentalData.W = np.fft.ifftshift(
                            self.experimentalData.W, axes=(-1, -2)
                        )
                if self.experimentalData.emptyBeam is not None:
                    self.experimentalData.emptyBeam = np.fft.ifftshift(
                        self.experimentalData.emptyBeam, axes=(-1, -2)
                    )
                if hasattr(self.experimentalData, "PSD"):
                    if self.experimentalData.PSD is not None:
                        self.experimentalData.PSD = np.fft.ifftshift(
                            self.experimentalData.PSD, axes=(-1, -2)
                        )
                self.params.fftshiftFlag = 1
        else:
            if self.params.fftshiftFlag == 1:
                print("check fftshift...")
                print("ifftshift data")
                self.experimentalData.ptychogram = np.fft.fftshift(
                    self.experimentalData.ptychogram, axes=(-1, -2)
                )
                if hasattr(self.experimentalData, "ptychogramDownsampled"):
                    self.experimentalData.ptychogramDownsampled = np.fft.fftshift(
                        self.experimentalData.ptychogramDownsampled, axes=(-1, -2)
                    )
                if self.experimentalData.W != None:
                    self.experimentalData.W = np.fft.fftshift(
                        self.experimentalData.W, axes=(-1, -2)
                    )
                if self.experimentalData.emptyBeam != None:
                    self.experimentalData.emptyBeam = np.fft.fftshift(
                        self.experimentalData.emptyBeam, axes=(-1, -2)
                    )
                self.params.fftshiftFlag = 0

    def _move_data_to_gpu(self):
        """
        Move the data to the GPU, called when the gpuSwitch is on.
        :return:
        """

        self.reconstruction._move_data_to_gpu()
        self.experimentalData._move_data_to_gpu()

        transfer_fields_to_gpu(
            self,
            [
                "probeWindow",
            ],
            self.logger,
        )  # '.probeWindow = cp.array(self.probeWindow)

        # reconstruction parameters
        # self.reconstruction.probe = cp.array(self.reconstruction.probe, cp.complex64)
        # self.reconstruction.object = cp.array(self.reconstruction.object, cp.complex64)
        # self.reconstruction.detectorError = cp.array(self.reconstruction.detectorError, cp.float32)

        # if self.params.momentumAcceleration:
        #     self.reconstruction.probeBuffer = cp.array(self.reconstruction.probeBuffer, cp.complex64)
        #     self.reconstruction.objectBuffer = cp.array(self.reconstruction.objectBuffer, cp.complex64)
        #     self.reconstruction.probeMomentum = cp.array(self.reconstruction.probeMomentum, cp.complex64)
        #     self.reconstruction.objectMomentum = cp.array(self.reconstruction.objectMomentum, cp.complex64)

        # for doing the coordinate transform and especially the otherwise slow interpolation of aPIE on the gpu
        if hasattr(self.params, "aPIEflag"):
            if self.params.aPIEflag == True:
                fields_to_transfer = [
                    "ptychogramUntransformed",
                    "Uq",
                    "Vq",
                    "theta",
                    "wavelength",
                    "Xd",
                    "Yd",
                    "dxd",
                    "zo",
                ]
                self.theta = self.reconstruction.theta
                self.wavelength = self.reconstruction.wavelength

                transfer_fields_to_gpu(self, fields_to_transfer, self.logger)
                # self.ptychogramUntransformed = cp.array(self.ptychogramUntransformed)
                # self.Uq = cp.array(self.Uq)
                # self.Vq = cp.array(self.Vq)
                # self.theta = cp.array(self.reconstruction.theta)
                # self.wavelength = cp.array(self.reconstruction.wavelength)
                # self.Xd = cp.array(self.Xd)
                # self.Yd = cp.array(self.Yd)
                # self.dxd = cp.array(self.dxd)
                # self.zo = cp.array(self.zo)
                # self.experimentalData.W = cp.array(self.experimentalData.W)

        # non-reconstruction parameters
        # if hasattr(self.experimentalData, 'ptychogramDownsampled'):
        #     self.experimentalData.ptychogramDownsampled = cp.array(self.experimentalData.ptychogramDownsampled,
        #                                                            cp.float32)
        # else:
        #     self.experimentalData.ptychogram = cp.array(self.experimentalData.ptychogram, cp.float32)

        # propagators
        # if self.params.propagatorType == 'Fresnel':
        #     self.reconstruction.quadraticPhase = cp.array(self.reconstruction.quadraticPhase)
        # elif self.params.propagatorType == 'ASP' or self.params.propagatorType == 'polychromeASP':
        #     self.reconstruction.transferFunction = cp.array(self.reconstruction.transferFunction)
        # elif self.params.propagatorType == 'scaledASP' or self.params.propagatorType == 'scaledPolychromeASP':
        #     self.reconstruction.Q1 = cp.array(self.reconstruction.Q1)
        #     self.reconstruction.Q2 = cp.array(self.reconstruction.Q2)
        # elif self.params.propagatorType == 'twoStepPolychrome':
        #     self.reconstruction.quadraticPhase = cp.array(self.reconstruction.quadraticPhase)
        #     self.reconstruction.transferFunction = cp.array(self.reconstruction.transferFunction)

        # other parameters
        # if self.params.backgroundModeSwitch:
        #     self.reconstruction.background = cp.array(self.reconstruction.background)
        # if self.params.absorbingProbeBoundary or self.params.probeBoundary:

        # if self.params.modulusEnforcedProbeSwitch:
        #     self.experimentalData.emptyBeam = cp.array(self.experimentalData.emptyBeam)
        # if self.params.intensityConstraint == 'interferometric':
        #     self.reconstruction.reference = cp.array(self.reconstruction.reference)

    def _move_data_to_cpu(self):
        """
        Move the data to the CPU, called when the gpuSwitch is off.
        :return:
        """
        # reconstruction parameters

        self.reconstruction._move_data_to_cpu()
        self.experimentalData._move_data_to_cpu()
        transfer_fields_to_cpu(
            self,
            [
                "probeWindow",
            ],
            self.logger,
        )

        # self.reconstruction.move_to_CPU()
        # self.params.move_to_CPU()

        # self.reconstruction.probe = asNumpyArray(self.reconstruction.probe)
        # self.reconstruction.object = asNumpyArray(self.reconstruction.object)

        # if self.params.momentumAcceleration:
        # reconstruction_fields_to_transfer = ['probeBuffer', 'objectBuffer', 'probeMomentum',' objectMomentum']
        # for field in reconstruction_fields_to_transfer:
        #
        #     setattr(self.reconstruction, field,)
        # self.reconstruction.probeBuffer = self.reconstruction.probeBuffer.get()
        # self.reconstruction.objectBuffer = self.reconstruction.objectBuffer.get()
        # self.reconstruction.probeMomentum = self.reconstruction.probeMomentum.get()
        # self.reconstruction.objectMomentum = self.reconstruction.objectMomentum.get()

        # for doing the coordinate transform and especially the otherwise slow interpolation of aPIE on the gpu
        # if hasattr(self.params, 'aPIEflag'):
        #     if self.params.aPIEflag:
        #         self.theta = self.theta.get()
        #
        fields_to_transfer = ["theta", "probeWindow"]
        # self.probeWindow = self.probeWindow.get()

        # non-reconstruction parameters
        # if hasattr(self.experimentalData, 'ptychogramDownsampled'):
        #     self.experimentalData.ptychogramDownsampled = self.experimentalData.ptychogramDownsampled.get()
        # else:
        #     self.experimentalData.ptychogram = self.experimentalData.ptychogram.get()
        # self.reconstruction.detectorError = self.reconstruction.detectorError.get()

        # propagators
        # if self.params.propagatorType == 'Fresnel':
        # self.reconstruction.quadraticPhase = self.reconstruction.quadraticPhase.get()
        # elif self.params.propagatorType == 'ASP' or self.params.propagatorType == 'polychromeASP':
        #     self.reconstruction.transferFunction = self.reconstruction.transferFunction.get()
        # elif self.params.propagatorType == 'scaledASP' or self.params.propagatorType == 'scaledPolychromeASP':
        #     self.reconstruction.Q1 = self.reconstruction.Q1.get()
        #     self.reconstruction.Q2 = self.reconstruction.Q2.get()
        # elif self.params.propagatorType == 'twoStepPolychrome':
        #     self.reconstruction.quadraticPhase = self.reconstruction.quadraticPhase.get()
        #     self.reconstruction.transferFunction = self.reconstruction.transferFunction.get()

        # other parameters
        # if self.params.backgroundModeSwitch:
        #     # self.reconstruction.background = self.reconstruction.background.get()
        # if self.params.absorbingProbeBoundary or self.params.probeBoundary:

        # if self.params.modulusEnforcedProbeSwitch:
        #     self.experimentalData.emptyBeam = self.experimentalData.emptyBeam.get()
        # # if self.params.intensityConstraint == 'interferometric':
        #     self.reconstruction.reference = self.reconstruction.reference.get()

    def _checkGPU(self):
        if not hasattr(self.params, "gpuFlag"):
            self.params.gpuFlag = 0

        if self.params._gpuSwitch:
            if cp is None:
                raise ImportError(
                    "Could not import cupy, therefore no GPU reconstruction is possible. To reconstruct, set the params.gpuSwitch to False."
                )
            if not self.params.gpuFlag:
                self.logger.info("switch to gpu")

                # load data to gpu
                self._move_data_to_gpu()
                self.params.gpuFlag = 1
            # always do this as it gets away with hard to debug errors
            self._move_data_to_gpu()
        else:
            self._move_data_to_cpu()
            if self.params.gpuFlag:
                self.logger.info("switch to cpu")
                self._move_data_to_cpu()
                self.params.gpuFlag = 0

    def setPositionOrder(self):
        if self.params.positionOrder == "sequential":
            self.positionIndices = np.arange(self.experimentalData.numFrames)

        elif self.params.positionOrder == "random":
            if len(self.reconstruction.error) == 0:
                self.positionIndices = np.arange(self.experimentalData.numFrames)
            else:
                if len(self.reconstruction.error) < 2:
                    self.positionIndices = np.arange(self.experimentalData.numFrames)
                else:
                    self.positionIndices = np.arange(self.experimentalData.numFrames)
                    np.random.shuffle(self.positionIndices)

        # order by illumiantion angles. Use smallest angles first
        # (i.e. start with brightfield data first, then add the low SNR
        # darkfield)
        # todo check this with Antonios
        elif self.params.positionOrder == "NA":
            rows = self.reconstruction.positions[:, 0] - np.mean(
                self.reconstruction.positions[:, 0]
            )
            cols = self.reconstruction.positions[:, 1] - np.mean(
                self.reconstruction.positions[:, 1]
            )
            dist = np.sqrt(rows**2 + cols**2)
            self.positionIndices = np.argsort(dist)
        else:
            raise ValueError("position order not properly set")

    def changeExperimentalData(self, experimentalData: ExperimentalData):

        if experimentalData is not None:
            if not isinstance(experimentalData, ExperimentalData):
                raise TypeError("Experimental data should be of class ExperimentalData")
            self.experimentalData = experimentalData

    def changeOptimizable(self, optimizable: Reconstruction):

        if optimizable is not None:
            if not isinstance(optimizable, Reconstruction):
                raise TypeError(
                    f"Argument should be an subclass of Reconstruction, but it is {type(optimizable)}"
                )
            self.reconstruction = optimizable

    def convert2single(self):
        """
        Convert the datasets to single precision. Matches: convert2single.m
        :return:
        """
        self.dtype_complex = np.complex64
        self.dtype_real = np.float32
        self._match_dtypes_complex()
        self._match_dtypes_real()

    def _match_dtypes_complex(self):
        raise NotImplementedError()

    def _match_dtypes_real(self):
        raise NotImplementedError()

    def object2detector(self, esw=None):
        """
        Implements object2detector.m. Modifies esw in-place
        :return:
        """
        if esw is None:
            # todo: check this, it seems weird to store it in self.esw
            esw = self.reconstruction.esw
        self.esw, self.reconstruction.ESW = Operators.Operators.object2detector(
            esw, self.params, self.reconstruction
        )

    def detector2object(self, ESW=None):
        """
        Propagate the ESW to the object plane (in-place).

        Matches: detector2object.m
        :return:
        """
        if ESW is None:
            ESW = self.reconstruction.ESW
        esw, eswUpdate = Operators.Operators.detector2object(
            ESW, self.params, self.reconstruction
        )
        # Dirk is not sure why this has to be changed at all but it sometimes is changed for some reason
        self.reconstruction.esw = esw
        # this is the new estimate which will be processed later
        self.reconstruction.eswUpdate = eswUpdate

    def fft2s(self):
        """
        Computes the fourier transform of the exit surface wave.
        :return:
        """
        self.reconstruction.ESW = FT2(
            self.reconstruction.esw, self.params.fftshiftSwitch
        )

    def ifft2s(self):
        """Inverse FFT"""
        # find out if this should be performed on the GPU
        self.reconstruction.eswUpdate = IFT(
            self.reconstruction.ESW, self.params.fftshiftSwitch
        )

    def getBeamWidth(self):
        """
        Calculate probe beam width (Full width half maximum)
        :return:
        """
        xp = getArrayModule(self.reconstruction.probe)
        P = xp.sum(
            abs((self.reconstruction.probe[..., -1, :, :])) ** 2,
            axis=(0, 1, 2),
        )
        P = P / xp.sum(P, axis=(-1, -2))
        P = asNumpyArray(P)
        xMean = np.sum(self.reconstruction.Xp * P, axis=(-1, -2))
        yMean = np.sum(self.reconstruction.Yp * P, axis=(-1, -2))
        xVariance = np.sum((self.reconstruction.Xp - xMean) ** 2 * P, axis=(-1, -2))
        yVariance = np.sum((self.reconstruction.Yp - yMean) ** 2 * P, axis=(-1, -2))

        c = (
            2 * xp.sqrt(2 * xp.log(2))
        )  # constant for converting variance to FWHM (see e.g. https://en.wikipedia.org/wiki/Full_width_at_half_maximum)

        self.reconstruction.beamWidthX = asNumpyArray(c * np.sqrt(xVariance))
        self.reconstruction.beamWidthY = asNumpyArray(c * np.sqrt(yVariance))

        return self.reconstruction.beamWidthY, self.reconstruction.beamWidthX

    def getOverlap(self, ind1, ind2):
        """
        Calculate linear and area overlap between two scan positions indexed ind1 and ind2
        """
        sy = (
            abs(
                self.reconstruction.positions[ind2, 0]
                - self.reconstruction.positions[ind1, 0]
            )
            * self.reconstruction.dxp
        )
        sx = (
            abs(
                self.reconstruction.positions[ind2, 1]
                - self.reconstruction.positions[ind1, 1]
            )
            * self.reconstruction.dxp
        )

        # task 1: get linear overlap
        self.getBeamWidth()
        self.reconstruction.linearOverlap = 1 - np.sqrt(sx**2 + sy**2) / np.minimum(
            self.reconstruction.beamWidthX, self.reconstruction.beamWidthY
        )
        self.reconstruction.linearOverlap = np.maximum(
            self.reconstruction.linearOverlap, 0
        )

        # task 2: get area overlap
        # spatial frequency pixel size
        df = 1 / (self.reconstruction.Np * self.reconstruction.dxp)
        # spatial frequency meshgrid
        fx = np.arange(-self.reconstruction.Np // 2, self.reconstruction.Np // 2) * df
        Fx, Fy = np.meshgrid(fx, fx)
        # absolute value of probe and 2D fft
        P = abs(asNumpyArray(self.reconstruction.probe[:, 0, 0, -1, ...]))
        Q = fft2c(P)
        # calculate overlap between positions
        self.reconstruction.areaOverlap = np.mean(
            abs(
                np.sum(
                    Q**2 * np.exp(-1.0j * 2 * np.pi * (Fx * sx + Fy * sy)),
                    axis=(-1, -2),
                )
            )
            / np.sum(abs(Q) ** 2, axis=(-1, -2)),
            axis=0,
        )

    def getErrorMetrics(self):
        """
        matches getErrorMetrics.m
        :return:
        """
        if not self.params.saveMemory:
            # Calculate mean error for all positions (make separate function for all of that)
            if self.params.FourierMaskSwitch:
                self.reconstruction.errorAtPos = np.sum(
                    np.abs(self.reconstruction.detectorError) * self.experimentalData.W,
                    axis=(-1, -2),
                )
            else:
                self.reconstruction.errorAtPos = np.sum(
                    np.abs(self.reconstruction.detectorError), axis=(-1, -2)
                )
        self.reconstruction.errorAtPos = asNumpyArray(
            self.reconstruction.errorAtPos
        ) / asNumpyArray(self.experimentalData.energyAtPos + 1e-20)
        eAverage = np.sum(self.reconstruction.errorAtPos)

        # append to error vector (for plotting error as function of iteration)
        self.reconstruction.error = np.append(self.reconstruction.error, eAverage)

    def getRMSD(self, positionIndex):
        """
        Root mean square deviation between ptychogram and intensity estimate
        :param positionIndex:
        :return:
        """
        # find out wether or not to use the GPU
        xp = getArrayModule(self.reconstruction.Iestimated)
        self.currentDetectorError = abs(
            self.reconstruction.Imeasured - self.reconstruction.Iestimated
        )

        # todo saveMemory implementation
        if self.params.saveMemory:
            if self.params.FourierMaskSwitch and not self.params.CPSCswitch:
                self.reconstruction.errorAtPos[positionIndex] = xp.sum(
                    self.currentDetectorError * self.experimentalData.W
                )
            elif self.params.FourierMaskSwitch and self.params.CPSCswitch:
                raise NotImplementedError
            else:
                self.reconstruction.errorAtPos[positionIndex] = asNumpyArray(
                    xp.sum(self.currentDetectorError)
                )
        else:
            self.reconstruction.detectorError[positionIndex] = self.currentDetectorError

    def intensityProjection(self, positionIndex):
        """Compute the projected intensity.
        Barebones, need to implement other methods
        """
        # figure out whether or not to use the GPU
        xp = getArrayModule(self.reconstruction.esw)
        # zero division mitigator
        gimmel = 1e-10

        # propagate to detector
        self.object2detector()

        # get estimated intensity (2D array, in the case of multislice, only take the last slice)
        if self.params.intensityConstraint == "interferometric":
            self.reconstruction.Iestimated = xp.sum(
                xp.abs(self.reconstruction.ESW + self.reconstruction.reference) ** 2,
                axis=(0, 1, 2),
            )[-1]
        else:
            self.reconstruction.Iestimated = xp.sum(
                xp.abs(self.reconstruction.ESW) ** 2, axis=(0, 1, 2)
            )[-1]
            self.logger.debug(
                f"Estimated intensity: {self.reconstruction.Iestimated.sum()}, Measured: {self.experimentalData.ptychogram[positionIndex].sum()}"
            )
        if self.params.backgroundModeSwitch:
            self.reconstruction.Iestimated += self.reconstruction.background

        # get measured intensity todo implement kPIE
        if self.params.CPSCswitch:
            self.decompressionProjection(positionIndex)
        else:
            self.reconstruction.Imeasured = self.experimentalData.ptychogram[
                positionIndex
            ]

        self.getRMSD(positionIndex)

        # adaptive denoising
        if self.params.adaptiveDenoisingSwitch:
            self.adaptiveDenoising()

        # intensity projection constraints
        if self.params.intensityConstraint == "fluctuation":
            # scaling
            if self.params.FourierMaskSwitch:
                aleph = xp.sum(
                    self.reconstruction.Imeasured
                    * self.reconstruction.Iestimated
                    * self.experimentalData.W
                ) / xp.sum(
                    self.reconstruction.Imeasured
                    * self.reconstruction.Imeasured
                    * self.experimentalData.W
                )
            else:
                aleph = xp.sum(
                    self.reconstruction.Imeasured * self.reconstruction.Iestimated
                ) / xp.sum(
                    self.reconstruction.Imeasured * self.reconstruction.Imeasured
                )
            self.params.intensityScaling[positionIndex] = aleph
            # scaled projection
            frac = (
                (1 + aleph)
                / 2
                * self.reconstruction.Imeasured
                / (self.reconstruction.Iestimated + gimmel)
            )

        elif self.params.intensityConstraint == "exponential":
            x = self.currentDetectorError / (self.reconstruction.Iestimated + gimmel)
            W = xp.exp(-0.05 * x)
            frac = xp.sqrt(
                self.reconstruction.Imeasured
                / (self.reconstruction.Iestimated + gimmel)
            )
            frac = W * frac + (1 - W)

        elif self.params.intensityConstraint == "poission":
            frac = self.reconstruction.Imeasured / (
                self.reconstruction.Iestimated + gimmel
            )

        elif (
            self.params.intensityConstraint == "standard"
            or self.params.intensityConstraint == "interferometric"
        ):
            frac = xp.sqrt(
                self.reconstruction.Imeasured
                / (self.reconstruction.Iestimated + gimmel)
            )

        else:
            raise ValueError("intensity constraint not properly specified!")

        # apply mask
        if (
            self.params.FourierMaskSwitch
            and self.params.CPSCswitch
            and len(self.reconstruction.error) > 5
        ):
            frac = self.experimentalData.W * frac + (1 - self.experimentalData.W)

        # update ESW
        if self.params.intensityConstraint == "interferometric":
            temp = (
                self.reconstruction.ESW + self.reconstruction.reference
            ) * frac - self.reconstruction.ESW
            self.reconstruction.ESW = (
                self.reconstruction.ESW + self.reconstruction.reference
            ) * frac - self.reconstruction.reference
            self.reconstruction.reference = temp
        else:
            if hasattr(self.params, "intensityMask"):
                if self.params.intensityMask:
                    self.reconstruction.ESW = self.reconstruction.ESW * (
                        frac * (self.reconstruction.intensity_mask)
                        + (self.reconstruction.intensity_mask - 1)
                    )
                else:
                    self.reconstruction.ESW = self.reconstruction.ESW * frac
            else:
                self.reconstruction.ESW = self.reconstruction.ESW * frac

        # update background (see PhD thsis by Peng Li)
        if self.params.backgroundModeSwitch:
            if self.params.FourierMaskSwitch:
                self.reconstruction.background = (
                    self.reconstruction.background
                    * (1 + 1 / self.experimentalData.numFrames * (xp.sqrt(frac) - 1))
                    ** 2
                    * self.experimentalData.W
                )
            else:
                self.reconstruction.background = (
                    self.reconstruction.background
                    * (1 + 1 / self.experimentalData.numFrames * (xp.sqrt(frac) - 1))
                    ** 2
                )

        # back propagate to object plane
        self.detector2object()

    def decompressionProjection(self, positionIndex):
        """
        calculate the upsampled Imeasured from downsampled Imeasured that is actually measured.
        :param positionIndex: index for scan positions
        :return:
        """
        # overwrite the measured intensity (just to have same dimensions as Iestimated)
        xp = getArrayModule(self.reconstruction.Iestimated)

        # determine downsampled fraction (Sl)
        frac = self.experimentalData.ptychogramDownsampled[positionIndex] / (
            xp.sum(
                self.reconstruction.Iestimated.reshape(
                    self.reconstruction.Nd // self.params.CPSCupsamplingFactor,
                    self.params.CPSCupsamplingFactor,
                    self.reconstruction.Nd // self.params.CPSCupsamplingFactor,
                    self.params.CPSCupsamplingFactor,
                ),
                axis=(1, 3),
            )
            + np.finfo(np.float32).eps
        )
        if self.params.FourierMaskSwitch and len(self.reconstruction.error) > 5:
            frac = self.experimentalData.W * frac + (1 - self.experimentalData.W)
        # overwrite up-sampled measured intensity
        self.reconstruction.Imeasured = self.reconstruction.Iestimated * xp.repeat(
            xp.repeat(frac, self.params.CPSCupsamplingFactor, axis=-1),
            self.params.CPSCupsamplingFactor,
            axis=-2,
        )

    def showReconstruction(self, loop):
        """
        Show the reconstruction process.
        :param loop: the iteration number
        :return:
        """
        if np.mod(loop, self.monitor.figureUpdateFrequency) == 0:
            if self.experimentalData.operationMode == "FPM":
                object_estimate = np.squeeze(
                    asNumpyArray(
                        fft2c(self.reconstruction.object)[
                            ..., self.monitor.objectROI[0], self.monitor.objectROI[1]
                        ]
                    )
                )
                probe_estimate = np.squeeze(
                    asNumpyArray(
                        self.reconstruction.probe[
                            ..., self.monitor.probeROI[0], self.monitor.probeROI[1]
                        ]
                    )
                )
            else:
                object_estimate = np.squeeze(
                    asNumpyArray(
                        self.reconstruction.object[
                            ..., self.monitor.objectROI[0], self.monitor.objectROI[1]
                        ]
                    )
                )
                probe_estimate = np.squeeze(
                    asNumpyArray(
                        self.reconstruction.probe[
                            ..., self.monitor.probeROI[0], self.monitor.probeROI[1]
                        ]
                    )
                )
            self.monitor.updateObjectProbeErrorMonitor(
                error=self.reconstruction.error,
                object_estimate=object_estimate,
                probe_estimate=probe_estimate,
                zo=self.reconstruction.zo,
                purity_probe=self.reconstruction.purityProbe,
                purity_object=self.reconstruction.purityObject,
                encoder_positions=self.reconstruction.positions,
            )

            self.monitor.writeEngineName(repr(type(self)))

            self.monitor.update_encoder(
                corrected_positions=self.reconstruction.encoder_corrected,
                original_positions=self.experimentalData.encoder,
            )

            self.monitor.updateBeamWidth(*self.getBeamWidth())

            # self.monitor.visualize_probe_engine(self.reconstruction.probe_storage)

            if self.monitor.verboseLevel == "high":
                if self.params.fftshiftSwitch:
                    Iestimated = np.fft.fftshift(
                        asNumpyArray(self.reconstruction.Iestimated)
                    )
                    Imeasured = np.fft.fftshift(
                        asNumpyArray(self.reconstruction.Imeasured)
                    )
                else:
                    Iestimated = asNumpyArray(self.reconstruction.Iestimated)
                    Imeasured = asNumpyArray(self.reconstruction.Imeasured)

                self.monitor.updateDiffractionDataMonitor(
                    Iestimated=Iestimated, Imeasured=Imeasured
                )

                self.getOverlap(0, 1)

                self.pbar.write("")
                self.pbar.write("iteration: %i" % loop)
                self.pbar.write("error: %.1f" % self.reconstruction.error[-1])
                self.pbar.write(
                    "estimated linear overlap: %.1f %%"
                    % (100 * self.reconstruction.linearOverlap)
                )
                self.pbar.write(
                    "estimated area overlap: %.1f %%"
                    % (100 * self.reconstruction.areaOverlap)
                )

                self.monitor.update_overlap(
                    self.reconstruction.areaOverlap, self.reconstruction.linearOverlap
                )
                # self.pbar.write('coherence structure:')

            if self.params.positionCorrectionSwitch:
                # show reconstruction
                return
                if (
                    len(self.reconstruction.error) > self.startAtIteration
                ):  # & (np.mod(loop,
                    # self.monitor.figureUpdateFrequency) == 0):
                    figure, ax = plt.subplots(
                        1, 1, num=102, squeeze=True, clear=True, figsize=(5, 5)
                    )
                    ax.set_title("Estimated scan grid positions")
                    ax.set_xlabel("(um)")
                    ax.set_ylabel("(um)")
                    # ax.set_xscale('symlog')
                    (line1,) = plt.plot(
                        (
                            self.reconstruction.positions0[:, 1]
                            - self.reconstruction.No // 2
                            + self.reconstruction.Np // 2
                        )
                        * self.reconstruction.dxo
                        * 1e6,
                        (
                            self.reconstruction.positions0[:, 0]
                            - self.reconstruction.No // 2
                            + self.reconstruction.Np // 2
                        )
                        * self.reconstruction.dxo
                        * 1e6,
                        "bo",
                        label="before correction",
                    )
                    (line2,) = plt.plot(
                        (
                            self.reconstruction.positions[:, 1]
                            - self.reconstruction.No // 2
                            + self.reconstruction.Np // 2
                        )
                        * self.reconstruction.dxo
                        * 1e6,
                        (
                            self.reconstruction.positions[:, 0]
                            - self.reconstruction.No // 2
                            + self.reconstruction.Np // 2
                        )
                        * self.reconstruction.dxo
                        * 1e6,
                        "yo",
                        label="after correction",
                    )
                    # plt.xlabel('(um))')
                    # plt.ylabel('(um))')
                    # plt.show()
                    plt.legend(handles=[line1, line2])
                    plt.tight_layout()
                    # plt.show(block=False)

                    figure2, ax2 = plt.subplots(
                        1, 1, num=103, squeeze=True, clear=True, figsize=(5, 5)
                    )
                    ax2.set_title("Displacement")
                    ax2.set_xlabel("(um)")
                    ax2.set_ylabel("(um)")
                    plt.plot(
                        self.D[:, 1] * self.reconstruction.dxo * 1e6,
                        self.D[:, 0] * self.reconstruction.dxo * 1e6,
                        "o",
                    )
                    # ax.set_xscale('symlog')
                    plt.tight_layout()
                    # plt.show(block=False)

                    # elif np.mod(loop, self.monitor.figureUpdateFrequency) == 0:
                    figure.show()
                    figure2.show()
                    figure.canvas.draw()
                    figure.canvas.flush_events()
                    figure2.canvas.draw()
                    figure2.canvas.flush_events()
                    # self.showReconstruction(loop)
            # print('iteration:%i' %len(self.reconstruction.error))
            # print('runtime:')
            # print('error:')

        # Dump each iteration the current object
        if self.params.dump_obj:
            folder_path = "dumps"
            if loop == 0:
                if not os.path.exists(folder_path):
                    # Create the folder
                    os.makedirs(folder_path)
                    print(f"Folder '{folder_path}' created.")
                else:
                    print(f"Folder '{folder_path}' already exists.")

            filename = "obj_dump_" + str(loop) + ".h5py"
            import h5py

            file_path = os.path.join(folder_path, filename)
            with h5py.File(file_path, "w") as hdf:
                obj = self.reconstruction.object.get()
                hdf.create_dataset("Object", data=obj)

    def positionCorrection(self, objectPatch, positionIndex, sy, sx):
        """
        Modified from pcPIE. Position correction is done by using positionCorrection and positionCorrectionUpdate
        :param objectPatch:
        :param positionIndex:
        :param sy:
        :param sx:
        :return:
        """

        xp = getArrayModule(objectPatch)
        if len(self.reconstruction.error) > self.startAtIteration:
            self.logger.debug("Calculating position correction")
            # position gradients
            # shiftedImages = xp.zeros((self.rowShifts.shape + objectPatch.shape))
            cc = xp.zeros((len(self.rowShifts), 1))

            # use the real-space object (FFT for FPM)
            O = self.reconstruction.object
            Opatch = objectPatch
            if self.experimentalData.operationMode == "FPM":
                O = fft2c(self.reconstruction.object)
                Opatch = fft2c(objectPatch)

            if self.params.positionCorrectionSwitch_radius < 2:
                # do the direct one as it's a bit faster

                for shifts in range(len(self.rowShifts)):
                    tempShift = xp.roll(Opatch, self.rowShifts[shifts], axis=-2)
                    # shiftedImages[shifts, ...] = xp.roll(tempShift, self.colShifts[shifts], axis=-1)
                    shiftedImages = xp.roll(tempShift, self.colShifts[shifts], axis=-1)
                    cc[shifts] = xp.squeeze(
                        xp.sum(shiftedImages.conj() * O[..., sy, sx], axis=(-2, -1))
                    )
                    del tempShift, shiftedImages
                    betaGrad = 1000
                    r = 3
            else:
                # print('doing FT position correction')
                ss = slice(
                    -self.params.positionCorrectionSwitch_radius,
                    self.params.positionCorrectionSwitch_radius + 1,
                )
                rowShifts, colShifts = xp.mgrid[ss, ss]
                self.rowShifts = rowShifts.flatten()
                self.colShifts = colShifts.flatten()
                FT_O = xp.fft.fft2(O[..., sy, sx] - O[..., sy, sx].mean())
                FT_Op = xp.fft.fft2(Opatch - O.mean())
                xcor = xp.fft.ifft2(FT_O * FT_Op.conj())
                xcor = abs(xp.fft.fftshift(xcor))
                N = xcor.shape[-1]
                sy = slice(
                    N // 2 - self.params.positionCorrectionSwitch_radius,
                    N // 2 + self.params.positionCorrectionSwitch_radius + 1,
                )
                xcor = xcor[..., sy, sy]
                cc = xcor.flatten()
                betaGrad = 5
                r = 10
                # dy, dx = xp.unravel_index(xp.argmax(xcor), xcor.shape)
                # dx = dx.get()
            # truncated cross - correlation
            # cc = xp.squeeze(xp.sum(shiftedImages.conj() * self.reconstruction.object[..., sy, sx], axis=(-2, -1)))
            cc = abs(cc)

            normFactor = xp.sum(Opatch.conj() * Opatch, axis=(-2, -1)).real
            grad_x = betaGrad * xp.sum(
                (cc.T - xp.mean(cc)) / normFactor * xp.array(self.colShifts)
            )
            grad_y = betaGrad * xp.sum(
                (cc.T - xp.mean(cc)) / normFactor * xp.array(self.rowShifts)
            )
            # r = np.clip(self.params.positionCorrectionSwitch_radius//5, 3, self.reconstruction.Np//10) # maximum shift in pixels?

            if abs(grad_x) > r:
                grad_x = r * grad_x / abs(grad_x)
            if abs(grad_y) > r:
                grad_y = r * grad_y / abs(grad_y)
            grad_y = asNumpyArray(grad_y)
            grad_x = asNumpyArray(grad_x)
            delta_p = self.daleth * np.array([grad_y, grad_x])
            self.D[positionIndex, :] = delta_p + self.beth * self.D[positionIndex, :]
            return delta_p
        return np.zeros(2)

    def position_update_to_change_in_z(self, loop):
        """
        Update the z based on the position updates.
        """
        import jax
        from jax.experimental import optimizers

        if not hasattr(self, "optlib"):
            self.i_z_optimizer = 0
            # from itertools import count
            # count
            op_init, op_update, op_get = optimizers.adam(3e-3)
            state = op_init(self.reconstruction.zo)
            self.optlib = {"op_update": op_update, "op_get": op_get, "state": state}
        else:
            state = self.optlib["state"]
            op_get = self.optlib["op_get"]
            op_update = self.optlib["op_update"]

        X0 = self.reconstruction.encoder_corrected
        Y0 = self.experimentalData.encoder
        msqdisplacement = np.linalg.norm(1e6 * X0 - 1e6 * Y0)

        # center both
        X0 = X0 - X0.mean(axis=0, keepdims=True)
        Y0 = Y0 - Y0.mean(axis=0, keepdims=True)

        # now, find the scaling with respect to the original one
        factor = np.std(X0) / np.std(Y0)

        # update z
        new_z = self.reconstruction.zo / factor
        step = new_z - self.reconstruction.zo
        self.logger.info(f"Naive estimate of new z: {new_z:.3f}, stepsize {step:.3f}")
        step = 5 * step
        # check if the thing should be updated.
        if abs(step) < 1e-4:  # if it's too small, just truncate it,
            # it may be that the distance changed due to some other update.
            # Take that into account as if we don't the steps will be super large.
            self.i_z_optimizer += 1
            step = self.reconstruction.zo - op_get(state)
            self.optlib["state"] = op_update(self.i_z_optimizer, -step, state)

            self.logger.info("Skipping update as step is too small")
            # as we're only updating it for sake of good measure, we don't have to update anything else.
            return
        # now, as we're actually updating, we can increase the step
        self.i_z_optimizer += 1
        self.optlib["state"] = op_update(self.i_z_optimizer, -step, state)
        # get the new value
        z_new = float(jax.device_get(op_get(self.optlib["state"])))

        self.logger.info(f"Loop: {loop} step: {step}")
        self.logger.info(
            f"old z: {self.reconstruction.zo:.3f}\n new z calculated: {z_new:.3f}\n diff: {self.reconstruction.zo - z_new}\n"
        )
        # scale the coordinates accordingly
        factor = self.reconstruction.zo / z_new
        self.reconstruction.zo = z_new
        new_encoder = self.reconstruction.encoder_corrected.copy()
        new_encoder -= new_encoder.mean(axis=0, keepdims=True)
        new_encoder /= factor  # this should be the correct one!

        new_encoder += self.experimentalData.encoder.mean(axis=0, keepdims=True)

        msqdisplacement_a = np.linalg.norm(
            1e6 * new_encoder - self.experimentalData.encoder * 1e6
        )

        self.reconstruction.encoder_corrected = new_encoder
        self.logger.info(
            f"Mean square displacement: before: {msqdisplacement:.3f} after: {msqdisplacement_a:.3f}"
        )

    def positionCorrectionUpdate(self):
        # fit the scaling out, to put in the z
        if len(self.reconstruction.error) > self.startAtIteration:
            self.logger.info("Updating positions")

            # update positions
            if self.experimentalData.operationMode == "FPM":
                conv = (
                    -(1 / self.reconstruction.wavelength)
                    * self.reconstruction.dxo
                    * self.reconstruction.Np
                )
                z = self.reconstruction.zled
                k = (
                    self.reconstruction.positions
                    - self.adaptStep * self.D
                    - self.reconstruction.No // 2
                    + self.reconstruction.Np // 2
                )
                self.reconstruction.encoder_corrected = (
                    np.sign(conv)
                    * k
                    * z
                    / (np.sqrt(conv**2 - k[:, 0] ** 2 - k[:, 1] ** 2))[..., None]
                )
            else:
                new_encoder = (
                    self.reconstruction.encoder_corrected
                    - self.adaptStep * self.D * self.reconstruction.dxo
                )
                new_encoder = new_encoder - new_encoder.mean(axis=0, keepdims=True)
                new_encoder = new_encoder + self.experimentalData.encoder.mean(
                    axis=0, keepdims=True
                )

                self.reconstruction.encoder_corrected = new_encoder
                self.logger.info(
                    f"Average update size: {abs(self.D).mean():.2f} pixels"
                )

    def applyConstraints(self, loop):
        """
        Apply constraints.
        :param loop: loop number
        :return:
        """
        # dirks additions, untested
        if self.params.l2reg:
            #     turns down areas that are not updated. Similar to an
            # l2 regularizer
            self.reconstruction.object *= 1 - self.params.l2reg_object_aleph
            self.reconstruction.probe *= 1 - self.params.l2reg_probe_aleph

        # enforce empty beam constraint
        if self.params.modulusEnforcedProbeSwitch:
            self.modulusEnforcedProbe()

        if self.params.orthogonalizationSwitch:
            if np.mod(loop, self.params.orthogonalizationFrequency) == 0:
                self.orthogonalization()

        # probe normalization to measured PSD todo: check for multiwave and multi object states
        if self.params.probePowerCorrectionSwitch:
            self.reconstruction.probe = (
                self.reconstruction.probe
                / np.sqrt(
                    np.sum(self.reconstruction.probe * self.reconstruction.probe.conj())
                )
                * self.experimentalData.maxProbePower
            )
        if self.params.probeSpectralPowerCorrectionSwitch:
            for wl in range(self.reconstruction.probe.shape[0]):
                self.reconstruction.probe[wl, ...] *= (
                    self.experimentalData.maxProbePower
                    * self.experimentalData.spectralPower[wl]
                    / np.sqrt(
                        np.sum(
                            self.reconstruction.probe[wl, ...]
                            * self.reconstruction.probe[wl, ...].conj()
                        )
                    )
                )

        if (
            self.params.comStabilizationSwitch is not None
            and self.params.comStabilizationSwitch is not False
        ):
            if loop % int(self.params.comStabilizationSwitch) == 0:
                self.comStabilization()

        if self.params.PSDestimationSwitch:
            raise NotImplementedError()

        if self.params.probeBoundary:
            self.reconstruction.probe *= self.probeWindow

        if self.params.absorbingProbeBoundary:
            if self.experimentalData.operationMode == "FPM":
                self.absorbingProbeBoundaryAleph = 1

            self.reconstruction.probe = (
                (1 - self.params.absorbingProbeBoundaryAleph)
                * self.reconstruction.probe
                + self.params.absorbingProbeBoundaryAleph
                * self.reconstruction.probe
                * self.probeWindow
            )

            # experimental: also apply in fourier space
            # self.reconstruction.probe = ifft2c(fft2c(self.reconstruction.probe)*self.probeWindow)

        # Todo: objectSmoothenessSwitch,probeSmoothenessSwitch,
        if self.params.probeSmoothenessSwitch:
            self.reconstruction.probe = smooth_amplitude(
                self.reconstruction.probe,
                self.params.probeSmoothenessWidth,
                self.params.probeSmoothnessAleph,
            )

        if self.params.objectSmoothenessSwitch:
            self.reconstruction.object = smooth_amplitude(
                self.reconstruction.object,
                self.params.objectSmoothenessWidth,
                self.params.objectSmoothnessAleph,
            )

        if self.params.absObjectSwitch:
            self.reconstruction.object = (
                1 - self.params.absObjectBeta
            ) * self.reconstruction.object + self.params.absObjectBeta * abs(
                self.reconstruction.object
            )

        if self.params.absProbeSwitch:
            self.reconstruction.probe = (
                1 - self.params.absProbeBeta
            ) * self.reconstruction.probe + self.params.absProbeBeta * abs(
                self.reconstruction.probe
            )

        # this is intended to slowly push non-measured object region to abs value lower than
        # the max abs inside object ROI allowing for good contrast when monitoring object
        if self.params.objectContrastSwitch:
            self.reconstruction.object = (
                0.995 * self.reconstruction.object
                + 0.005
                * np.mean(
                    abs(
                        self.reconstruction.object[
                            ..., self.monitor.objectROI[0], self.monitor.objectROI[1]
                        ]
                    )
                )
            )
        if self.params.couplingSwitch and self.reconstruction.nlambda > 1:
            self.reconstruction.probe[0] = (
                1 - self.params.couplingAleph
            ) * self.reconstruction.probe[
                0
            ] + self.params.couplingAleph * self.reconstruction.probe[1]
            for lambdaLoop in np.arange(1, self.reconstruction.nlambda - 1):
                self.reconstruction.probe[lambdaLoop] = (
                    1 - self.params.couplingAleph
                ) * self.reconstruction.probe[
                    lambdaLoop
                ] + self.params.couplingAleph * (
                    self.reconstruction.probe[lambdaLoop + 1]
                    + self.reconstruction.probe[lambdaLoop - 1]
                ) / 2

            self.reconstruction.probe[-1] = (
                1 - self.params.couplingAleph
            ) * self.reconstruction.probe[
                -1
            ] + self.params.couplingAleph * self.reconstruction.probe[-2]
        if self.params.binaryProbeSwitch:
            probePeakAmplitude = np.max(abs(self.reconstruction.probe))
            probeThresholded = self.reconstruction.probe.copy()
            probeThresholded[
                (
                    abs(probeThresholded)
                    < self.params.binaryProbeThreshold * probePeakAmplitude
                )
            ] = 0

            self.reconstruction.probe = (
                (1 - self.params.binaryProbeAleph) * self.reconstruction.probe
                + self.params.binaryProbeAleph * probeThresholded
            )

        if self.params.positionCorrectionSwitch:
            self.positionCorrectionUpdate()

        if (
            self.params.map_position_to_z_change
            and (loop % 5 == 1)
            and self.params.positionCorrectionSwitch
        ):
            self.position_update_to_change_in_z(loop)

        if self.params.TV_autofocus:
            merit, AOI_image, allmerits = self.reconstruction.TV_autofocus(
                self.params, loop=loop
            )
            self.monitor.update_focusing_metric(
                merit,
                AOI_image,
                metric_name=self.params.TV_autofocus_metric,
                allmerits=allmerits,
            )

        # if self.params.OPRP and loop % self.params.OPRP_tsvd_interval == 0:
        #     self.reconstruction.probe_storage.tsvd()

    def orthogonalization(self):
        """
        Perform orthogonalization
        :return:
        """
        xp = getArrayModule(self.reconstruction.probe)
        if self.reconstruction.npsm > 1:
            # orthogonalize the probe for each wavelength and each slice
            for id_l in range(self.reconstruction.nlambda):
                for id_s in range(self.reconstruction.nslice):
                    (
                        self.reconstruction.probe[id_l, 0, :, id_s, :, :],
                        self.normalizedEigenvaluesProbe,
                        self.MSPVprobe,
                    ) = orthogonalizeModes(
                        self.reconstruction.probe[id_l, 0, :, id_s, :, :],
                        method="snapShots",
                    )
                    # normalizedEigenvalues can live on either device, and is only
                    # npsm values long -- reducing it on the host is both cheaper
                    # than launching kernels for it and device-independent.
                    eigenvalues = asNumpyArray(self.normalizedEigenvaluesProbe)
                    self.reconstruction.purityProbe = float(
                        np.sqrt(np.sum(eigenvalues**2))
                    )
                    self.reconstruction.purityProbeHist.append(
                        self.reconstruction.purityProbe
                    )
                    # orthogonolize momentum operator
                    if self.params.momentumAcceleration:
                        # orthogonalize probe Buffer
                        p = self.reconstruction.probeBuffer[
                            id_l, 0, :, id_s, :, :
                        ].reshape((self.reconstruction.npsm, self.reconstruction.Np**2))
                        self.reconstruction.probeBuffer[id_l, 0, :, id_s, :, :] = (
                            xp.array(self.MSPVprobe) @ p
                        ).reshape(
                            (
                                self.reconstruction.npsm,
                                self.reconstruction.Np,
                                self.reconstruction.Np,
                            )
                        )
                        # orthogonalize probe momentum
                        p = self.reconstruction.probeMomentum[
                            id_l, 0, :, id_s, :, :
                        ].reshape((self.reconstruction.npsm, self.reconstruction.Np**2))
                        self.reconstruction.probeMomentum[id_l, 0, :, id_s, :, :] = (
                            xp.array(self.MSPVprobe) @ p
                        ).reshape(
                            (
                                self.reconstruction.npsm,
                                self.reconstruction.Np,
                                self.reconstruction.Np,
                            )
                        )

                        # if self.comStabilizationSwitch:
                        #     self.comStabilization()
            # self.reconstruction.probe_storage.push(self.reconstruction.probe, None, len(self.experimentalData.ptychogram), force=True)

        elif self.reconstruction.nosm > 1:
            # orthogonalize the object for each wavelength and each slice
            for id_l in range(self.reconstruction.nlambda):
                for id_s in range(self.reconstruction.nslice):
                    (
                        self.reconstruction.object[id_l, :, 0, id_s, :, :],
                        self.normalizedEigenvaluesObject,
                        self.MSPVobject,
                    ) = orthogonalizeModes(
                        self.reconstruction.object[id_l, :, 0, id_s, :, :],
                        method="snapShots",
                    )
                    eigenvalues = asNumpyArray(self.normalizedEigenvaluesObject)
                    self.reconstruction.purityObject = float(
                        np.sqrt(np.sum(eigenvalues**2))
                    )

                    # orthogonolize momentum operator
                    if self.params.momentumAcceleration:
                        # orthogonalize object Buffer
                        p = self.reconstruction.objectBuffer[
                            id_l, :, 0, id_s, :, :
                        ].reshape((self.reconstruction.nosm, self.reconstruction.No**2))
                        self.reconstruction.objectBuffer[id_l, :, 0, id_s, :, :] = (
                            xp.array(self.MSPVobject) @ p
                        ).reshape(
                            (
                                self.reconstruction.nosm,
                                self.reconstruction.No,
                                self.reconstruction.No,
                            )
                        )
                        # orthogonalize object momentum
                        p = self.reconstruction.objectMomentum[
                            id_l, :, 0, id_s, :, :
                        ].reshape((self.reconstruction.nosm, self.reconstruction.No**2))
                        self.reconstruction.objectMomentum[id_l, :, 0, id_s, :, :] = (
                            xp.array(self.MSPVobject) @ p
                        ).reshape(
                            (
                                self.reconstruction.nosm,
                                self.reconstruction.No,
                                self.reconstruction.No,
                            )
                        )

        else:
            pass

    def comStabilization(self):
        """
        Perform center of mass stabilization (center the probe)
        :return:
        """
        self.logger.info("Doing probe com stabilization")
        xp = getArrayModule(self.reconstruction.probe)
        # calculate center of mass of the probe (for multislice cases, the probe for the last slice is used)
        P2 = xp.sum(
            abs(self.reconstruction.probe[:, :, :, -1, ...]) ** 2, axis=(0, 1, 2)
        )
        P2 = abs(self.reconstruction.probe[0, 0, 0, -1])
        demon = xp.sum(P2) * self.reconstruction.dxp
        xc = int(
            xp.around(xp.sum(xp.array(self.reconstruction.Xp, xp.float32) * P2) / demon)
        )
        yc = int(
            xp.around(xp.sum(xp.array(self.reconstruction.Yp, xp.float32) * P2) / demon)
        )
        # print('Center of mass:', yc, xc)
        # shift only if necessary
        if xc**2 + yc**2 > 1:
            # self.reconstruction.probe_storage._push_hard(self.reconstruction.probe, 100)
            # self.reconstruction.probe_storage.roll(-yc, -xc)

            # shift probe
            self.reconstruction.probe = xp.roll(
                self.reconstruction.probe, (-yc, -xc), axis=(-2, -1)
            )
            # for k in xp.arange(self.reconstruction.npsm):
            #     self.reconstruction.probe[:, :, k, -1, ...] = \
            #         xp.roll(self.reconstruction.probe[:, :, k, -1, ...], (-yc, -xc), axis=(-2, -1))
            #     # for mPIE
            if self.params.momentumAcceleration:
                self.reconstruction.probeMomentum = xp.roll(
                    self.reconstruction.probeMomentum, (-yc, -xc), axis=(-2, -1)
                )
                self.reconstruction.probeBuffer = xp.roll(
                    self.reconstruction.probeBuffer, (-yc, -xc), axis=(-2, -1)
                )

            # shift object
            self.reconstruction.object = xp.roll(
                self.reconstruction.object, (-yc, -xc), axis=(-2, -1)
            )
            # for mPIE
            if self.params.momentumAcceleration:
                self.reconstruction.objectMomentum = xp.roll(
                    self.reconstruction.objectMomentum, (-yc, -xc), axis=(-2, -1)
                )
                self.reconstruction.objectBuffer = xp.roll(
                    self.reconstruction.objectBuffer, (-yc, -xc), axis=(-2, -1)
                )

    def modulusEnforcedProbe(self):
        # propagate probe to detector
        xp = getArrayModule(self.reconstruction.esw)
        self.reconstruction.esw = self.reconstruction.probe
        self.object2detector()

        if self.params.FourierMaskSwitch:
            self.reconstruction.ESW = self.reconstruction.ESW * xp.sqrt(
                self.experimentalData.emptyBeam / 1e-10
                + xp.sum(xp.abs(self.reconstruction.ESW) ** 2, axis=(0, 1, 2, 3))
            ) * self.experimentalData.W + self.reconstruction.ESW * (
                1 - self.experimentalData.W
            )
        else:
            self.reconstruction.ESW = self.reconstruction.ESW * np.sqrt(
                self.experimentalData.emptyBeam
                / (1e-10 + xp.sum(abs(self.reconstruction.ESW) ** 2, axis=(0, 1, 2, 3)))
            )

        self.detector2object()

        if self.params.OPRP:
            pass
            # self.probes.append(self.reconstruction.esw.reshape(-))

        self.reconstruction.probe = self.reconstruction.esw

    def adaptiveDenoising(self):
        """
        Use the difference of mean intensities between the low-resolution
        object estimate and the low-resolution raw data to estimate the
        noise floor to be clipped.
        :return:
        """
        # figure out wether or not to use the GPU
        xp = getArrayModule(self.reconstruction.esw)

        Ameasured = self.reconstruction.Imeasured**0.5
        Aestimated = xp.abs(self.reconstruction.Iestimated) ** 0.5

        noise = xp.abs(xp.mean(Ameasured - Aestimated))

        Ameasured = Ameasured - noise
        Ameasured[Ameasured < 0] = 0
        self.reconstruction.Imeasured = Ameasured**2

    def z_update(self, stepsize=0.01, roi_bounds=[0.3, 0.7], d=10):
        """
        Update Z based on TV
        :param stepsize:
        :return:
        """
        self.reconstruction.TV_autofocus()

    def objectPatchUpdate_TV(self, objectPatch: np.ndarray, DELTA: np.ndarray):
        """
        Update the object patch with a TV regularization.

        :param objectPatch:
        :param DELTA:
        :return:
        """

        xp = getArrayModule(objectPatch)
        frac = self.reconstruction.probe.conj() / xp.max(
            xp.sum(xp.abs(self.reconstruction.probe) ** 2, axis=(0, 1, 2, 3))
        )

        # gradient = xp.gradient(objectPatch, axis=(4, 5))
        #
        # # norm = xp.abs(gradient[0] + gradient[1]) ** 2
        # norm = (gradient[0] + gradient[1]) ** 2
        # temp = [gradient[0] / xp.sqrt(norm + epsilon), gradient[1] / xp.sqrt(norm + epsilon)]
        # TV_update = divergence(temp)
        TV_update = grad_TV(objectPatch, epsilon=1e-2)
        lam = self.params.objectTVregStepSize
        return (
            objectPatch
            + self.betaObject * xp.sum(frac * DELTA, axis=(0, 2, 3), keepdims=True)
            + lam * self.betaObject * TV_update
        )

update_data(experimentalData, reconstruction=None)

Update the experimentalData if necessary

Source code in PtyLab/Engines/BaseEngine.py

def update_data(self, experimentalData, reconstruction=None):
    """Update the experimentalData if necessary"""
    self.experimentalData = experimentalData
    if reconstruction is not None:
        self.reconstruction = reconstruction

convert2single()

Convert the datasets to single precision. Matches: convert2single.m :return:

Source code in PtyLab/Engines/BaseEngine.py

def convert2single(self):
    """
    Convert the datasets to single precision. Matches: convert2single.m
    :return:
    """
    self.dtype_complex = np.complex64
    self.dtype_real = np.float32
    self._match_dtypes_complex()
    self._match_dtypes_real()

object2detector(esw=None)

Implements object2detector.m. Modifies esw in-place :return:

Source code in PtyLab/Engines/BaseEngine.py

def object2detector(self, esw=None):
    """
    Implements object2detector.m. Modifies esw in-place
    :return:
    """
    if esw is None:
        # todo: check this, it seems weird to store it in self.esw
        esw = self.reconstruction.esw
    self.esw, self.reconstruction.ESW = Operators.Operators.object2detector(
        esw, self.params, self.reconstruction
    )

detector2object(ESW=None)

Propagate the ESW to the object plane (in-place).

Matches: detector2object.m :return:

Source code in PtyLab/Engines/BaseEngine.py

def detector2object(self, ESW=None):
    """
    Propagate the ESW to the object plane (in-place).

    Matches: detector2object.m
    :return:
    """
    if ESW is None:
        ESW = self.reconstruction.ESW
    esw, eswUpdate = Operators.Operators.detector2object(
        ESW, self.params, self.reconstruction
    )
    # Dirk is not sure why this has to be changed at all but it sometimes is changed for some reason
    self.reconstruction.esw = esw
    # this is the new estimate which will be processed later
    self.reconstruction.eswUpdate = eswUpdate

fft2s()

Computes the fourier transform of the exit surface wave. :return:

Source code in PtyLab/Engines/BaseEngine.py

def fft2s(self):
    """
    Computes the fourier transform of the exit surface wave.
    :return:
    """
    self.reconstruction.ESW = FT2(
        self.reconstruction.esw, self.params.fftshiftSwitch
    )

ifft2s()

Inverse FFT

Source code in PtyLab/Engines/BaseEngine.py

def ifft2s(self):
    """Inverse FFT"""
    # find out if this should be performed on the GPU
    self.reconstruction.eswUpdate = IFT(
        self.reconstruction.ESW, self.params.fftshiftSwitch
    )

getBeamWidth()

Calculate probe beam width (Full width half maximum) :return:

Source code in PtyLab/Engines/BaseEngine.py

def getBeamWidth(self):
    """
    Calculate probe beam width (Full width half maximum)
    :return:
    """
    xp = getArrayModule(self.reconstruction.probe)
    P = xp.sum(
        abs((self.reconstruction.probe[..., -1, :, :])) ** 2,
        axis=(0, 1, 2),
    )
    P = P / xp.sum(P, axis=(-1, -2))
    P = asNumpyArray(P)
    xMean = np.sum(self.reconstruction.Xp * P, axis=(-1, -2))
    yMean = np.sum(self.reconstruction.Yp * P, axis=(-1, -2))
    xVariance = np.sum((self.reconstruction.Xp - xMean) ** 2 * P, axis=(-1, -2))
    yVariance = np.sum((self.reconstruction.Yp - yMean) ** 2 * P, axis=(-1, -2))

    c = (
        2 * xp.sqrt(2 * xp.log(2))
    )  # constant for converting variance to FWHM (see e.g. https://en.wikipedia.org/wiki/Full_width_at_half_maximum)

    self.reconstruction.beamWidthX = asNumpyArray(c * np.sqrt(xVariance))
    self.reconstruction.beamWidthY = asNumpyArray(c * np.sqrt(yVariance))

    return self.reconstruction.beamWidthY, self.reconstruction.beamWidthX

getOverlap(ind1, ind2)

Calculate linear and area overlap between two scan positions indexed ind1 and ind2

Source code in PtyLab/Engines/BaseEngine.py

def getOverlap(self, ind1, ind2):
    """
    Calculate linear and area overlap between two scan positions indexed ind1 and ind2
    """
    sy = (
        abs(
            self.reconstruction.positions[ind2, 0]
            - self.reconstruction.positions[ind1, 0]
        )
        * self.reconstruction.dxp
    )
    sx = (
        abs(
            self.reconstruction.positions[ind2, 1]
            - self.reconstruction.positions[ind1, 1]
        )
        * self.reconstruction.dxp
    )

    # task 1: get linear overlap
    self.getBeamWidth()
    self.reconstruction.linearOverlap = 1 - np.sqrt(sx**2 + sy**2) / np.minimum(
        self.reconstruction.beamWidthX, self.reconstruction.beamWidthY
    )
    self.reconstruction.linearOverlap = np.maximum(
        self.reconstruction.linearOverlap, 0
    )

    # task 2: get area overlap
    # spatial frequency pixel size
    df = 1 / (self.reconstruction.Np * self.reconstruction.dxp)
    # spatial frequency meshgrid
    fx = np.arange(-self.reconstruction.Np // 2, self.reconstruction.Np // 2) * df
    Fx, Fy = np.meshgrid(fx, fx)
    # absolute value of probe and 2D fft
    P = abs(asNumpyArray(self.reconstruction.probe[:, 0, 0, -1, ...]))
    Q = fft2c(P)
    # calculate overlap between positions
    self.reconstruction.areaOverlap = np.mean(
        abs(
            np.sum(
                Q**2 * np.exp(-1.0j * 2 * np.pi * (Fx * sx + Fy * sy)),
                axis=(-1, -2),
            )
        )
        / np.sum(abs(Q) ** 2, axis=(-1, -2)),
        axis=0,
    )

getErrorMetrics()

matches getErrorMetrics.m :return:

Source code in PtyLab/Engines/BaseEngine.py

def getErrorMetrics(self):
    """
    matches getErrorMetrics.m
    :return:
    """
    if not self.params.saveMemory:
        # Calculate mean error for all positions (make separate function for all of that)
        if self.params.FourierMaskSwitch:
            self.reconstruction.errorAtPos = np.sum(
                np.abs(self.reconstruction.detectorError) * self.experimentalData.W,
                axis=(-1, -2),
            )
        else:
            self.reconstruction.errorAtPos = np.sum(
                np.abs(self.reconstruction.detectorError), axis=(-1, -2)
            )
    self.reconstruction.errorAtPos = asNumpyArray(
        self.reconstruction.errorAtPos
    ) / asNumpyArray(self.experimentalData.energyAtPos + 1e-20)
    eAverage = np.sum(self.reconstruction.errorAtPos)

    # append to error vector (for plotting error as function of iteration)
    self.reconstruction.error = np.append(self.reconstruction.error, eAverage)

getRMSD(positionIndex)

Root mean square deviation between ptychogram and intensity estimate :param positionIndex: :return:

Source code in PtyLab/Engines/BaseEngine.py

def getRMSD(self, positionIndex):
    """
    Root mean square deviation between ptychogram and intensity estimate
    :param positionIndex:
    :return:
    """
    # find out wether or not to use the GPU
    xp = getArrayModule(self.reconstruction.Iestimated)
    self.currentDetectorError = abs(
        self.reconstruction.Imeasured - self.reconstruction.Iestimated
    )

    # todo saveMemory implementation
    if self.params.saveMemory:
        if self.params.FourierMaskSwitch and not self.params.CPSCswitch:
            self.reconstruction.errorAtPos[positionIndex] = xp.sum(
                self.currentDetectorError * self.experimentalData.W
            )
        elif self.params.FourierMaskSwitch and self.params.CPSCswitch:
            raise NotImplementedError
        else:
            self.reconstruction.errorAtPos[positionIndex] = asNumpyArray(
                xp.sum(self.currentDetectorError)
            )
    else:
        self.reconstruction.detectorError[positionIndex] = self.currentDetectorError

intensityProjection(positionIndex)

Compute the projected intensity. Barebones, need to implement other methods

Source code in PtyLab/Engines/BaseEngine.py

def intensityProjection(self, positionIndex):
    """Compute the projected intensity.
    Barebones, need to implement other methods
    """
    # figure out whether or not to use the GPU
    xp = getArrayModule(self.reconstruction.esw)
    # zero division mitigator
    gimmel = 1e-10

    # propagate to detector
    self.object2detector()

    # get estimated intensity (2D array, in the case of multislice, only take the last slice)
    if self.params.intensityConstraint == "interferometric":
        self.reconstruction.Iestimated = xp.sum(
            xp.abs(self.reconstruction.ESW + self.reconstruction.reference) ** 2,
            axis=(0, 1, 2),
        )[-1]
    else:
        self.reconstruction.Iestimated = xp.sum(
            xp.abs(self.reconstruction.ESW) ** 2, axis=(0, 1, 2)
        )[-1]
        self.logger.debug(
            f"Estimated intensity: {self.reconstruction.Iestimated.sum()}, Measured: {self.experimentalData.ptychogram[positionIndex].sum()}"
        )
    if self.params.backgroundModeSwitch:
        self.reconstruction.Iestimated += self.reconstruction.background

    # get measured intensity todo implement kPIE
    if self.params.CPSCswitch:
        self.decompressionProjection(positionIndex)
    else:
        self.reconstruction.Imeasured = self.experimentalData.ptychogram[
            positionIndex
        ]

    self.getRMSD(positionIndex)

    # adaptive denoising
    if self.params.adaptiveDenoisingSwitch:
        self.adaptiveDenoising()

    # intensity projection constraints
    if self.params.intensityConstraint == "fluctuation":
        # scaling
        if self.params.FourierMaskSwitch:
            aleph = xp.sum(
                self.reconstruction.Imeasured
                * self.reconstruction.Iestimated
                * self.experimentalData.W
            ) / xp.sum(
                self.reconstruction.Imeasured
                * self.reconstruction.Imeasured
                * self.experimentalData.W
            )
        else:
            aleph = xp.sum(
                self.reconstruction.Imeasured * self.reconstruction.Iestimated
            ) / xp.sum(
                self.reconstruction.Imeasured * self.reconstruction.Imeasured
            )
        self.params.intensityScaling[positionIndex] = aleph
        # scaled projection
        frac = (
            (1 + aleph)
            / 2
            * self.reconstruction.Imeasured
            / (self.reconstruction.Iestimated + gimmel)
        )

    elif self.params.intensityConstraint == "exponential":
        x = self.currentDetectorError / (self.reconstruction.Iestimated + gimmel)
        W = xp.exp(-0.05 * x)
        frac = xp.sqrt(
            self.reconstruction.Imeasured
            / (self.reconstruction.Iestimated + gimmel)
        )
        frac = W * frac + (1 - W)

    elif self.params.intensityConstraint == "poission":
        frac = self.reconstruction.Imeasured / (
            self.reconstruction.Iestimated + gimmel
        )

    elif (
        self.params.intensityConstraint == "standard"
        or self.params.intensityConstraint == "interferometric"
    ):
        frac = xp.sqrt(
            self.reconstruction.Imeasured
            / (self.reconstruction.Iestimated + gimmel)
        )

    else:
        raise ValueError("intensity constraint not properly specified!")

    # apply mask
    if (
        self.params.FourierMaskSwitch
        and self.params.CPSCswitch
        and len(self.reconstruction.error) > 5
    ):
        frac = self.experimentalData.W * frac + (1 - self.experimentalData.W)

    # update ESW
    if self.params.intensityConstraint == "interferometric":
        temp = (
            self.reconstruction.ESW + self.reconstruction.reference
        ) * frac - self.reconstruction.ESW
        self.reconstruction.ESW = (
            self.reconstruction.ESW + self.reconstruction.reference
        ) * frac - self.reconstruction.reference
        self.reconstruction.reference = temp
    else:
        if hasattr(self.params, "intensityMask"):
            if self.params.intensityMask:
                self.reconstruction.ESW = self.reconstruction.ESW * (
                    frac * (self.reconstruction.intensity_mask)
                    + (self.reconstruction.intensity_mask - 1)
                )
            else:
                self.reconstruction.ESW = self.reconstruction.ESW * frac
        else:
            self.reconstruction.ESW = self.reconstruction.ESW * frac

    # update background (see PhD thsis by Peng Li)
    if self.params.backgroundModeSwitch:
        if self.params.FourierMaskSwitch:
            self.reconstruction.background = (
                self.reconstruction.background
                * (1 + 1 / self.experimentalData.numFrames * (xp.sqrt(frac) - 1))
                ** 2
                * self.experimentalData.W
            )
        else:
            self.reconstruction.background = (
                self.reconstruction.background
                * (1 + 1 / self.experimentalData.numFrames * (xp.sqrt(frac) - 1))
                ** 2
            )

    # back propagate to object plane
    self.detector2object()

decompressionProjection(positionIndex)

calculate the upsampled Imeasured from downsampled Imeasured that is actually measured. :param positionIndex: index for scan positions :return:

Source code in PtyLab/Engines/BaseEngine.py

def decompressionProjection(self, positionIndex):
    """
    calculate the upsampled Imeasured from downsampled Imeasured that is actually measured.
    :param positionIndex: index for scan positions
    :return:
    """
    # overwrite the measured intensity (just to have same dimensions as Iestimated)
    xp = getArrayModule(self.reconstruction.Iestimated)

    # determine downsampled fraction (Sl)
    frac = self.experimentalData.ptychogramDownsampled[positionIndex] / (
        xp.sum(
            self.reconstruction.Iestimated.reshape(
                self.reconstruction.Nd // self.params.CPSCupsamplingFactor,
                self.params.CPSCupsamplingFactor,
                self.reconstruction.Nd // self.params.CPSCupsamplingFactor,
                self.params.CPSCupsamplingFactor,
            ),
            axis=(1, 3),
        )
        + np.finfo(np.float32).eps
    )
    if self.params.FourierMaskSwitch and len(self.reconstruction.error) > 5:
        frac = self.experimentalData.W * frac + (1 - self.experimentalData.W)
    # overwrite up-sampled measured intensity
    self.reconstruction.Imeasured = self.reconstruction.Iestimated * xp.repeat(
        xp.repeat(frac, self.params.CPSCupsamplingFactor, axis=-1),
        self.params.CPSCupsamplingFactor,
        axis=-2,
    )

showReconstruction(loop)

Show the reconstruction process. :param loop: the iteration number :return:

Source code in PtyLab/Engines/BaseEngine.py

def showReconstruction(self, loop):
    """
    Show the reconstruction process.
    :param loop: the iteration number
    :return:
    """
    if np.mod(loop, self.monitor.figureUpdateFrequency) == 0:
        if self.experimentalData.operationMode == "FPM":
            object_estimate = np.squeeze(
                asNumpyArray(
                    fft2c(self.reconstruction.object)[
                        ..., self.monitor.objectROI[0], self.monitor.objectROI[1]
                    ]
                )
            )
            probe_estimate = np.squeeze(
                asNumpyArray(
                    self.reconstruction.probe[
                        ..., self.monitor.probeROI[0], self.monitor.probeROI[1]
                    ]
                )
            )
        else:
            object_estimate = np.squeeze(
                asNumpyArray(
                    self.reconstruction.object[
                        ..., self.monitor.objectROI[0], self.monitor.objectROI[1]
                    ]
                )
            )
            probe_estimate = np.squeeze(
                asNumpyArray(
                    self.reconstruction.probe[
                        ..., self.monitor.probeROI[0], self.monitor.probeROI[1]
                    ]
                )
            )
        self.monitor.updateObjectProbeErrorMonitor(
            error=self.reconstruction.error,
            object_estimate=object_estimate,
            probe_estimate=probe_estimate,
            zo=self.reconstruction.zo,
            purity_probe=self.reconstruction.purityProbe,
            purity_object=self.reconstruction.purityObject,
            encoder_positions=self.reconstruction.positions,
        )

        self.monitor.writeEngineName(repr(type(self)))

        self.monitor.update_encoder(
            corrected_positions=self.reconstruction.encoder_corrected,
            original_positions=self.experimentalData.encoder,
        )

        self.monitor.updateBeamWidth(*self.getBeamWidth())

        # self.monitor.visualize_probe_engine(self.reconstruction.probe_storage)

        if self.monitor.verboseLevel == "high":
            if self.params.fftshiftSwitch:
                Iestimated = np.fft.fftshift(
                    asNumpyArray(self.reconstruction.Iestimated)
                )
                Imeasured = np.fft.fftshift(
                    asNumpyArray(self.reconstruction.Imeasured)
                )
            else:
                Iestimated = asNumpyArray(self.reconstruction.Iestimated)
                Imeasured = asNumpyArray(self.reconstruction.Imeasured)

            self.monitor.updateDiffractionDataMonitor(
                Iestimated=Iestimated, Imeasured=Imeasured
            )

            self.getOverlap(0, 1)

            self.pbar.write("")
            self.pbar.write("iteration: %i" % loop)
            self.pbar.write("error: %.1f" % self.reconstruction.error[-1])
            self.pbar.write(
                "estimated linear overlap: %.1f %%"
                % (100 * self.reconstruction.linearOverlap)
            )
            self.pbar.write(
                "estimated area overlap: %.1f %%"
                % (100 * self.reconstruction.areaOverlap)
            )

            self.monitor.update_overlap(
                self.reconstruction.areaOverlap, self.reconstruction.linearOverlap
            )
            # self.pbar.write('coherence structure:')

        if self.params.positionCorrectionSwitch:
            # show reconstruction
            return
            if (
                len(self.reconstruction.error) > self.startAtIteration
            ):  # & (np.mod(loop,
                # self.monitor.figureUpdateFrequency) == 0):
                figure, ax = plt.subplots(
                    1, 1, num=102, squeeze=True, clear=True, figsize=(5, 5)
                )
                ax.set_title("Estimated scan grid positions")
                ax.set_xlabel("(um)")
                ax.set_ylabel("(um)")
                # ax.set_xscale('symlog')
                (line1,) = plt.plot(
                    (
                        self.reconstruction.positions0[:, 1]
                        - self.reconstruction.No // 2
                        + self.reconstruction.Np // 2
                    )
                    * self.reconstruction.dxo
                    * 1e6,
                    (
                        self.reconstruction.positions0[:, 0]
                        - self.reconstruction.No // 2
                        + self.reconstruction.Np // 2
                    )
                    * self.reconstruction.dxo
                    * 1e6,
                    "bo",
                    label="before correction",
                )
                (line2,) = plt.plot(
                    (
                        self.reconstruction.positions[:, 1]
                        - self.reconstruction.No // 2
                        + self.reconstruction.Np // 2
                    )
                    * self.reconstruction.dxo
                    * 1e6,
                    (
                        self.reconstruction.positions[:, 0]
                        - self.reconstruction.No // 2
                        + self.reconstruction.Np // 2
                    )
                    * self.reconstruction.dxo
                    * 1e6,
                    "yo",
                    label="after correction",
                )
                # plt.xlabel('(um))')
                # plt.ylabel('(um))')
                # plt.show()
                plt.legend(handles=[line1, line2])
                plt.tight_layout()
                # plt.show(block=False)

                figure2, ax2 = plt.subplots(
                    1, 1, num=103, squeeze=True, clear=True, figsize=(5, 5)
                )
                ax2.set_title("Displacement")
                ax2.set_xlabel("(um)")
                ax2.set_ylabel("(um)")
                plt.plot(
                    self.D[:, 1] * self.reconstruction.dxo * 1e6,
                    self.D[:, 0] * self.reconstruction.dxo * 1e6,
                    "o",
                )
                # ax.set_xscale('symlog')
                plt.tight_layout()
                # plt.show(block=False)

                # elif np.mod(loop, self.monitor.figureUpdateFrequency) == 0:
                figure.show()
                figure2.show()
                figure.canvas.draw()
                figure.canvas.flush_events()
                figure2.canvas.draw()
                figure2.canvas.flush_events()
                # self.showReconstruction(loop)
        # print('iteration:%i' %len(self.reconstruction.error))
        # print('runtime:')
        # print('error:')

    # Dump each iteration the current object
    if self.params.dump_obj:
        folder_path = "dumps"
        if loop == 0:
            if not os.path.exists(folder_path):
                # Create the folder
                os.makedirs(folder_path)
                print(f"Folder '{folder_path}' created.")
            else:
                print(f"Folder '{folder_path}' already exists.")

        filename = "obj_dump_" + str(loop) + ".h5py"
        import h5py

        file_path = os.path.join(folder_path, filename)
        with h5py.File(file_path, "w") as hdf:
            obj = self.reconstruction.object.get()
            hdf.create_dataset("Object", data=obj)

positionCorrection(objectPatch, positionIndex, sy, sx)

Modified from pcPIE. Position correction is done by using positionCorrection and positionCorrectionUpdate :param objectPatch: :param positionIndex: :param sy: :param sx: :return:

Source code in PtyLab/Engines/BaseEngine.py

def positionCorrection(self, objectPatch, positionIndex, sy, sx):
    """
    Modified from pcPIE. Position correction is done by using positionCorrection and positionCorrectionUpdate
    :param objectPatch:
    :param positionIndex:
    :param sy:
    :param sx:
    :return:
    """

    xp = getArrayModule(objectPatch)
    if len(self.reconstruction.error) > self.startAtIteration:
        self.logger.debug("Calculating position correction")
        # position gradients
        # shiftedImages = xp.zeros((self.rowShifts.shape + objectPatch.shape))
        cc = xp.zeros((len(self.rowShifts), 1))

        # use the real-space object (FFT for FPM)
        O = self.reconstruction.object
        Opatch = objectPatch
        if self.experimentalData.operationMode == "FPM":
            O = fft2c(self.reconstruction.object)
            Opatch = fft2c(objectPatch)

        if self.params.positionCorrectionSwitch_radius < 2:
            # do the direct one as it's a bit faster

            for shifts in range(len(self.rowShifts)):
                tempShift = xp.roll(Opatch, self.rowShifts[shifts], axis=-2)
                # shiftedImages[shifts, ...] = xp.roll(tempShift, self.colShifts[shifts], axis=-1)
                shiftedImages = xp.roll(tempShift, self.colShifts[shifts], axis=-1)
                cc[shifts] = xp.squeeze(
                    xp.sum(shiftedImages.conj() * O[..., sy, sx], axis=(-2, -1))
                )
                del tempShift, shiftedImages
                betaGrad = 1000
                r = 3
        else:
            # print('doing FT position correction')
            ss = slice(
                -self.params.positionCorrectionSwitch_radius,
                self.params.positionCorrectionSwitch_radius + 1,
            )
            rowShifts, colShifts = xp.mgrid[ss, ss]
            self.rowShifts = rowShifts.flatten()
            self.colShifts = colShifts.flatten()
            FT_O = xp.fft.fft2(O[..., sy, sx] - O[..., sy, sx].mean())
            FT_Op = xp.fft.fft2(Opatch - O.mean())
            xcor = xp.fft.ifft2(FT_O * FT_Op.conj())
            xcor = abs(xp.fft.fftshift(xcor))
            N = xcor.shape[-1]
            sy = slice(
                N // 2 - self.params.positionCorrectionSwitch_radius,
                N // 2 + self.params.positionCorrectionSwitch_radius + 1,
            )
            xcor = xcor[..., sy, sy]
            cc = xcor.flatten()
            betaGrad = 5
            r = 10
            # dy, dx = xp.unravel_index(xp.argmax(xcor), xcor.shape)
            # dx = dx.get()
        # truncated cross - correlation
        # cc = xp.squeeze(xp.sum(shiftedImages.conj() * self.reconstruction.object[..., sy, sx], axis=(-2, -1)))
        cc = abs(cc)

        normFactor = xp.sum(Opatch.conj() * Opatch, axis=(-2, -1)).real
        grad_x = betaGrad * xp.sum(
            (cc.T - xp.mean(cc)) / normFactor * xp.array(self.colShifts)
        )
        grad_y = betaGrad * xp.sum(
            (cc.T - xp.mean(cc)) / normFactor * xp.array(self.rowShifts)
        )
        # r = np.clip(self.params.positionCorrectionSwitch_radius//5, 3, self.reconstruction.Np//10) # maximum shift in pixels?

        if abs(grad_x) > r:
            grad_x = r * grad_x / abs(grad_x)
        if abs(grad_y) > r:
            grad_y = r * grad_y / abs(grad_y)
        grad_y = asNumpyArray(grad_y)
        grad_x = asNumpyArray(grad_x)
        delta_p = self.daleth * np.array([grad_y, grad_x])
        self.D[positionIndex, :] = delta_p + self.beth * self.D[positionIndex, :]
        return delta_p
    return np.zeros(2)

position_update_to_change_in_z(loop)

Update the z based on the position updates.

Source code in PtyLab/Engines/BaseEngine.py

def position_update_to_change_in_z(self, loop):
    """
    Update the z based on the position updates.
    """
    import jax
    from jax.experimental import optimizers

    if not hasattr(self, "optlib"):
        self.i_z_optimizer = 0
        # from itertools import count
        # count
        op_init, op_update, op_get = optimizers.adam(3e-3)
        state = op_init(self.reconstruction.zo)
        self.optlib = {"op_update": op_update, "op_get": op_get, "state": state}
    else:
        state = self.optlib["state"]
        op_get = self.optlib["op_get"]
        op_update = self.optlib["op_update"]

    X0 = self.reconstruction.encoder_corrected
    Y0 = self.experimentalData.encoder
    msqdisplacement = np.linalg.norm(1e6 * X0 - 1e6 * Y0)

    # center both
    X0 = X0 - X0.mean(axis=0, keepdims=True)
    Y0 = Y0 - Y0.mean(axis=0, keepdims=True)

    # now, find the scaling with respect to the original one
    factor = np.std(X0) / np.std(Y0)

    # update z
    new_z = self.reconstruction.zo / factor
    step = new_z - self.reconstruction.zo
    self.logger.info(f"Naive estimate of new z: {new_z:.3f}, stepsize {step:.3f}")
    step = 5 * step
    # check if the thing should be updated.
    if abs(step) < 1e-4:  # if it's too small, just truncate it,
        # it may be that the distance changed due to some other update.
        # Take that into account as if we don't the steps will be super large.
        self.i_z_optimizer += 1
        step = self.reconstruction.zo - op_get(state)
        self.optlib["state"] = op_update(self.i_z_optimizer, -step, state)

        self.logger.info("Skipping update as step is too small")
        # as we're only updating it for sake of good measure, we don't have to update anything else.
        return
    # now, as we're actually updating, we can increase the step
    self.i_z_optimizer += 1
    self.optlib["state"] = op_update(self.i_z_optimizer, -step, state)
    # get the new value
    z_new = float(jax.device_get(op_get(self.optlib["state"])))

    self.logger.info(f"Loop: {loop} step: {step}")
    self.logger.info(
        f"old z: {self.reconstruction.zo:.3f}\n new z calculated: {z_new:.3f}\n diff: {self.reconstruction.zo - z_new}\n"
    )
    # scale the coordinates accordingly
    factor = self.reconstruction.zo / z_new
    self.reconstruction.zo = z_new
    new_encoder = self.reconstruction.encoder_corrected.copy()
    new_encoder -= new_encoder.mean(axis=0, keepdims=True)
    new_encoder /= factor  # this should be the correct one!

    new_encoder += self.experimentalData.encoder.mean(axis=0, keepdims=True)

    msqdisplacement_a = np.linalg.norm(
        1e6 * new_encoder - self.experimentalData.encoder * 1e6
    )

    self.reconstruction.encoder_corrected = new_encoder
    self.logger.info(
        f"Mean square displacement: before: {msqdisplacement:.3f} after: {msqdisplacement_a:.3f}"
    )

applyConstraints(loop)

Apply constraints. :param loop: loop number :return:

Source code in PtyLab/Engines/BaseEngine.py

def applyConstraints(self, loop):
    """
    Apply constraints.
    :param loop: loop number
    :return:
    """
    # dirks additions, untested
    if self.params.l2reg:
        #     turns down areas that are not updated. Similar to an
        # l2 regularizer
        self.reconstruction.object *= 1 - self.params.l2reg_object_aleph
        self.reconstruction.probe *= 1 - self.params.l2reg_probe_aleph

    # enforce empty beam constraint
    if self.params.modulusEnforcedProbeSwitch:
        self.modulusEnforcedProbe()

    if self.params.orthogonalizationSwitch:
        if np.mod(loop, self.params.orthogonalizationFrequency) == 0:
            self.orthogonalization()

    # probe normalization to measured PSD todo: check for multiwave and multi object states
    if self.params.probePowerCorrectionSwitch:
        self.reconstruction.probe = (
            self.reconstruction.probe
            / np.sqrt(
                np.sum(self.reconstruction.probe * self.reconstruction.probe.conj())
            )
            * self.experimentalData.maxProbePower
        )
    if self.params.probeSpectralPowerCorrectionSwitch:
        for wl in range(self.reconstruction.probe.shape[0]):
            self.reconstruction.probe[wl, ...] *= (
                self.experimentalData.maxProbePower
                * self.experimentalData.spectralPower[wl]
                / np.sqrt(
                    np.sum(
                        self.reconstruction.probe[wl, ...]
                        * self.reconstruction.probe[wl, ...].conj()
                    )
                )
            )

    if (
        self.params.comStabilizationSwitch is not None
        and self.params.comStabilizationSwitch is not False
    ):
        if loop % int(self.params.comStabilizationSwitch) == 0:
            self.comStabilization()

    if self.params.PSDestimationSwitch:
        raise NotImplementedError()

    if self.params.probeBoundary:
        self.reconstruction.probe *= self.probeWindow

    if self.params.absorbingProbeBoundary:
        if self.experimentalData.operationMode == "FPM":
            self.absorbingProbeBoundaryAleph = 1

        self.reconstruction.probe = (
            (1 - self.params.absorbingProbeBoundaryAleph)
            * self.reconstruction.probe
            + self.params.absorbingProbeBoundaryAleph
            * self.reconstruction.probe
            * self.probeWindow
        )

        # experimental: also apply in fourier space
        # self.reconstruction.probe = ifft2c(fft2c(self.reconstruction.probe)*self.probeWindow)

    # Todo: objectSmoothenessSwitch,probeSmoothenessSwitch,
    if self.params.probeSmoothenessSwitch:
        self.reconstruction.probe = smooth_amplitude(
            self.reconstruction.probe,
            self.params.probeSmoothenessWidth,
            self.params.probeSmoothnessAleph,
        )

    if self.params.objectSmoothenessSwitch:
        self.reconstruction.object = smooth_amplitude(
            self.reconstruction.object,
            self.params.objectSmoothenessWidth,
            self.params.objectSmoothnessAleph,
        )

    if self.params.absObjectSwitch:
        self.reconstruction.object = (
            1 - self.params.absObjectBeta
        ) * self.reconstruction.object + self.params.absObjectBeta * abs(
            self.reconstruction.object
        )

    if self.params.absProbeSwitch:
        self.reconstruction.probe = (
            1 - self.params.absProbeBeta
        ) * self.reconstruction.probe + self.params.absProbeBeta * abs(
            self.reconstruction.probe
        )

    # this is intended to slowly push non-measured object region to abs value lower than
    # the max abs inside object ROI allowing for good contrast when monitoring object
    if self.params.objectContrastSwitch:
        self.reconstruction.object = (
            0.995 * self.reconstruction.object
            + 0.005
            * np.mean(
                abs(
                    self.reconstruction.object[
                        ..., self.monitor.objectROI[0], self.monitor.objectROI[1]
                    ]
                )
            )
        )
    if self.params.couplingSwitch and self.reconstruction.nlambda > 1:
        self.reconstruction.probe[0] = (
            1 - self.params.couplingAleph
        ) * self.reconstruction.probe[
            0
        ] + self.params.couplingAleph * self.reconstruction.probe[1]
        for lambdaLoop in np.arange(1, self.reconstruction.nlambda - 1):
            self.reconstruction.probe[lambdaLoop] = (
                1 - self.params.couplingAleph
            ) * self.reconstruction.probe[
                lambdaLoop
            ] + self.params.couplingAleph * (
                self.reconstruction.probe[lambdaLoop + 1]
                + self.reconstruction.probe[lambdaLoop - 1]
            ) / 2

        self.reconstruction.probe[-1] = (
            1 - self.params.couplingAleph
        ) * self.reconstruction.probe[
            -1
        ] + self.params.couplingAleph * self.reconstruction.probe[-2]
    if self.params.binaryProbeSwitch:
        probePeakAmplitude = np.max(abs(self.reconstruction.probe))
        probeThresholded = self.reconstruction.probe.copy()
        probeThresholded[
            (
                abs(probeThresholded)
                < self.params.binaryProbeThreshold * probePeakAmplitude
            )
        ] = 0

        self.reconstruction.probe = (
            (1 - self.params.binaryProbeAleph) * self.reconstruction.probe
            + self.params.binaryProbeAleph * probeThresholded
        )

    if self.params.positionCorrectionSwitch:
        self.positionCorrectionUpdate()

    if (
        self.params.map_position_to_z_change
        and (loop % 5 == 1)
        and self.params.positionCorrectionSwitch
    ):
        self.position_update_to_change_in_z(loop)

    if self.params.TV_autofocus:
        merit, AOI_image, allmerits = self.reconstruction.TV_autofocus(
            self.params, loop=loop
        )
        self.monitor.update_focusing_metric(
            merit,
            AOI_image,
            metric_name=self.params.TV_autofocus_metric,
            allmerits=allmerits,
        )

orthogonalization()

Perform orthogonalization :return:

Source code in PtyLab/Engines/BaseEngine.py

def orthogonalization(self):
    """
    Perform orthogonalization
    :return:
    """
    xp = getArrayModule(self.reconstruction.probe)
    if self.reconstruction.npsm > 1:
        # orthogonalize the probe for each wavelength and each slice
        for id_l in range(self.reconstruction.nlambda):
            for id_s in range(self.reconstruction.nslice):
                (
                    self.reconstruction.probe[id_l, 0, :, id_s, :, :],
                    self.normalizedEigenvaluesProbe,
                    self.MSPVprobe,
                ) = orthogonalizeModes(
                    self.reconstruction.probe[id_l, 0, :, id_s, :, :],
                    method="snapShots",
                )
                # normalizedEigenvalues can live on either device, and is only
                # npsm values long -- reducing it on the host is both cheaper
                # than launching kernels for it and device-independent.
                eigenvalues = asNumpyArray(self.normalizedEigenvaluesProbe)
                self.reconstruction.purityProbe = float(
                    np.sqrt(np.sum(eigenvalues**2))
                )
                self.reconstruction.purityProbeHist.append(
                    self.reconstruction.purityProbe
                )
                # orthogonolize momentum operator
                if self.params.momentumAcceleration:
                    # orthogonalize probe Buffer
                    p = self.reconstruction.probeBuffer[
                        id_l, 0, :, id_s, :, :
                    ].reshape((self.reconstruction.npsm, self.reconstruction.Np**2))
                    self.reconstruction.probeBuffer[id_l, 0, :, id_s, :, :] = (
                        xp.array(self.MSPVprobe) @ p
                    ).reshape(
                        (
                            self.reconstruction.npsm,
                            self.reconstruction.Np,
                            self.reconstruction.Np,
                        )
                    )
                    # orthogonalize probe momentum
                    p = self.reconstruction.probeMomentum[
                        id_l, 0, :, id_s, :, :
                    ].reshape((self.reconstruction.npsm, self.reconstruction.Np**2))
                    self.reconstruction.probeMomentum[id_l, 0, :, id_s, :, :] = (
                        xp.array(self.MSPVprobe) @ p
                    ).reshape(
                        (
                            self.reconstruction.npsm,
                            self.reconstruction.Np,
                            self.reconstruction.Np,
                        )
                    )

                    # if self.comStabilizationSwitch:
                    #     self.comStabilization()
        # self.reconstruction.probe_storage.push(self.reconstruction.probe, None, len(self.experimentalData.ptychogram), force=True)

    elif self.reconstruction.nosm > 1:
        # orthogonalize the object for each wavelength and each slice
        for id_l in range(self.reconstruction.nlambda):
            for id_s in range(self.reconstruction.nslice):
                (
                    self.reconstruction.object[id_l, :, 0, id_s, :, :],
                    self.normalizedEigenvaluesObject,
                    self.MSPVobject,
                ) = orthogonalizeModes(
                    self.reconstruction.object[id_l, :, 0, id_s, :, :],
                    method="snapShots",
                )
                eigenvalues = asNumpyArray(self.normalizedEigenvaluesObject)
                self.reconstruction.purityObject = float(
                    np.sqrt(np.sum(eigenvalues**2))
                )

                # orthogonolize momentum operator
                if self.params.momentumAcceleration:
                    # orthogonalize object Buffer
                    p = self.reconstruction.objectBuffer[
                        id_l, :, 0, id_s, :, :
                    ].reshape((self.reconstruction.nosm, self.reconstruction.No**2))
                    self.reconstruction.objectBuffer[id_l, :, 0, id_s, :, :] = (
                        xp.array(self.MSPVobject) @ p
                    ).reshape(
                        (
                            self.reconstruction.nosm,
                            self.reconstruction.No,
                            self.reconstruction.No,
                        )
                    )
                    # orthogonalize object momentum
                    p = self.reconstruction.objectMomentum[
                        id_l, :, 0, id_s, :, :
                    ].reshape((self.reconstruction.nosm, self.reconstruction.No**2))
                    self.reconstruction.objectMomentum[id_l, :, 0, id_s, :, :] = (
                        xp.array(self.MSPVobject) @ p
                    ).reshape(
                        (
                            self.reconstruction.nosm,
                            self.reconstruction.No,
                            self.reconstruction.No,
                        )
                    )

    else:
        pass

comStabilization()

Perform center of mass stabilization (center the probe) :return:

Source code in PtyLab/Engines/BaseEngine.py

def comStabilization(self):
    """
    Perform center of mass stabilization (center the probe)
    :return:
    """
    self.logger.info("Doing probe com stabilization")
    xp = getArrayModule(self.reconstruction.probe)
    # calculate center of mass of the probe (for multislice cases, the probe for the last slice is used)
    P2 = xp.sum(
        abs(self.reconstruction.probe[:, :, :, -1, ...]) ** 2, axis=(0, 1, 2)
    )
    P2 = abs(self.reconstruction.probe[0, 0, 0, -1])
    demon = xp.sum(P2) * self.reconstruction.dxp
    xc = int(
        xp.around(xp.sum(xp.array(self.reconstruction.Xp, xp.float32) * P2) / demon)
    )
    yc = int(
        xp.around(xp.sum(xp.array(self.reconstruction.Yp, xp.float32) * P2) / demon)
    )
    # print('Center of mass:', yc, xc)
    # shift only if necessary
    if xc**2 + yc**2 > 1:
        # self.reconstruction.probe_storage._push_hard(self.reconstruction.probe, 100)
        # self.reconstruction.probe_storage.roll(-yc, -xc)

        # shift probe
        self.reconstruction.probe = xp.roll(
            self.reconstruction.probe, (-yc, -xc), axis=(-2, -1)
        )
        # for k in xp.arange(self.reconstruction.npsm):
        #     self.reconstruction.probe[:, :, k, -1, ...] = \
        #         xp.roll(self.reconstruction.probe[:, :, k, -1, ...], (-yc, -xc), axis=(-2, -1))
        #     # for mPIE
        if self.params.momentumAcceleration:
            self.reconstruction.probeMomentum = xp.roll(
                self.reconstruction.probeMomentum, (-yc, -xc), axis=(-2, -1)
            )
            self.reconstruction.probeBuffer = xp.roll(
                self.reconstruction.probeBuffer, (-yc, -xc), axis=(-2, -1)
            )

        # shift object
        self.reconstruction.object = xp.roll(
            self.reconstruction.object, (-yc, -xc), axis=(-2, -1)
        )
        # for mPIE
        if self.params.momentumAcceleration:
            self.reconstruction.objectMomentum = xp.roll(
                self.reconstruction.objectMomentum, (-yc, -xc), axis=(-2, -1)
            )
            self.reconstruction.objectBuffer = xp.roll(
                self.reconstruction.objectBuffer, (-yc, -xc), axis=(-2, -1)
            )

adaptiveDenoising()

Use the difference of mean intensities between the low-resolution object estimate and the low-resolution raw data to estimate the noise floor to be clipped. :return:

Source code in PtyLab/Engines/BaseEngine.py

def adaptiveDenoising(self):
    """
    Use the difference of mean intensities between the low-resolution
    object estimate and the low-resolution raw data to estimate the
    noise floor to be clipped.
    :return:
    """
    # figure out wether or not to use the GPU
    xp = getArrayModule(self.reconstruction.esw)

    Ameasured = self.reconstruction.Imeasured**0.5
    Aestimated = xp.abs(self.reconstruction.Iestimated) ** 0.5

    noise = xp.abs(xp.mean(Ameasured - Aestimated))

    Ameasured = Ameasured - noise
    Ameasured[Ameasured < 0] = 0
    self.reconstruction.Imeasured = Ameasured**2

z_update(stepsize=0.01, roi_bounds=[0.3, 0.7], d=10)

Update Z based on TV :param stepsize: :return:

Source code in PtyLab/Engines/BaseEngine.py

def z_update(self, stepsize=0.01, roi_bounds=[0.3, 0.7], d=10):
    """
    Update Z based on TV
    :param stepsize:
    :return:
    """
    self.reconstruction.TV_autofocus()

objectPatchUpdate_TV(objectPatch, DELTA)

Update the object patch with a TV regularization.

:param objectPatch: :param DELTA: :return:

Source code in PtyLab/Engines/BaseEngine.py

def objectPatchUpdate_TV(self, objectPatch: np.ndarray, DELTA: np.ndarray):
    """
    Update the object patch with a TV regularization.

    :param objectPatch:
    :param DELTA:
    :return:
    """

    xp = getArrayModule(objectPatch)
    frac = self.reconstruction.probe.conj() / xp.max(
        xp.sum(xp.abs(self.reconstruction.probe) ** 2, axis=(0, 1, 2, 3))
    )

    # gradient = xp.gradient(objectPatch, axis=(4, 5))
    #
    # # norm = xp.abs(gradient[0] + gradient[1]) ** 2
    # norm = (gradient[0] + gradient[1]) ** 2
    # temp = [gradient[0] / xp.sqrt(norm + epsilon), gradient[1] / xp.sqrt(norm + epsilon)]
    # TV_update = divergence(temp)
    TV_update = grad_TV(objectPatch, epsilon=1e-2)
    lam = self.params.objectTVregStepSize
    return (
        objectPatch
        + self.betaObject * xp.sum(frac * DELTA, axis=(0, 2, 3), keepdims=True)
        + lam * self.betaObject * TV_update
    )

smooth_amplitude(field, width, aleph, amplitude_only=True)

Smooth the amplitude of a field. Optional phase can be smoothed as well. Parameters

---

field width aleph amplitude_only

Returns

Source code in PtyLab/Engines/BaseEngine.py

def smooth_amplitude(
    field: np.ndarray, width: float, aleph: float, amplitude_only: bool = True
):
    """
    Smooth the amplitude of a field. Optional phase can be smoothed as well.
    Parameters
    ----------
    field
    width
    aleph
    amplitude_only

    Returns
    -------

    """
    xp = getArrayModule(field)
    smooth_fun = isGpuArray(field) and fourier_gaussian_gpu or fourier_gaussian_cpu
    gimmel = 1e-5
    if amplitude_only:
        ph_field = field / (xp.abs(field) + gimmel)
        A_field = abs(field)
    else:
        ph_field = 1
        A_field = field
    F_field = xp.fft.fft2(A_field)
    for ax in [-2, -1]:
        F_field = smooth_fun(F_field, width, axis=ax)
    field_smooth = xp.fft.ifft2(F_field)

    if amplitude_only:
        field_smooth = abs(field_smooth) * ph_field
    return aleph * field_smooth + (1 - aleph) * field

OPR

OPR

Bases: BaseEngine

Source code in PtyLab/Engines/OPR.py

class OPR(BaseEngine):

    def __init__(
        self,
        reconstruction: Reconstruction,
        experimentalData: ExperimentalData,
        params: Params,
        monitor: Monitor,
    ):
        # This contains reconstruction parameters that are specific to the reconstruction
        # but not necessarily to ePIE reconstruction
        super().__init__(reconstruction, experimentalData, params, monitor)
        self.logger = logging.getLogger("ePIE")
        self.logger.info("Sucesfully created ePIE ePIE_engine")
        self.logger.info("Wavelength attribute: %s", self.reconstruction.wavelength)
        self.initializeReconstructionParams()

    def initializeReconstructionParams(self):
        """
        Set parameters that are specific to the ePIE/OPR engine
        """
        self.alpha = self.params.OPR_alpha
        self.betaProbe = 0.25
        self.betaObject = 0.25
        self.numIterations = 50
        self.OPR_modes = self.params.OPR_modes
        self.n_subspace = self.params.OPR_subspace
        # Working-set budget for the chunked batched orthogonalization. Keeps
        # the transpose scratch bounded independently of the frame count.
        self._orthogonalization_chunk_bytes = 128 * 2**20

    def reconstruct(self):
        self._prepareReconstruction()

        # OPR parameters
        Nmodes = self.OPR_modes.shape[0]
        Np = self.reconstruction.Np
        Nframes = self.experimentalData.numFrames
        mode_slice = self.OPR_modes
        n_subspace = self.n_subspace

        self.reconstruction.probe_stack = cp.zeros(
            (1, 1, Nmodes, 1, Np, Np, Nframes), dtype=cp.complex64
        )

        for i, mode in enumerate(self.OPR_modes):
            # fill the probe-stack with the inital guess of the probes
            self.reconstruction.probe_stack[0, 0, i, 0, :, :, :] = cp.repeat(
                self.reconstruction.probe[0, 0, mode, 0, :, :, cp.newaxis],
                Nframes,
                axis=2,
            )

        # actual reconstruction ePIE_engine
        self.pbar = tqdm.trange(
            self.numIterations, desc="OPR", file=sys.stdout, leave=True
        )
        for loop in self.pbar:
            self.it = loop
            # set position order
            self.setPositionOrder()
            for positionLoop, positionIndex in enumerate(self.positionIndices):
                # get object patch
                row, col = self.reconstruction.positions[positionIndex]
                sy = slice(row, row + self.reconstruction.Np)
                sx = slice(col, col + self.reconstruction.Np)
                # note that object patch has size of probe array
                objectPatch = self.reconstruction.object[..., sy, sx].copy()

                # Get dim reduced probe
                self.reconstruction.probe[:, :, mode_slice, :, :, :] = (
                    self.reconstruction.probe_stack[..., positionIndex]
                )

                # make exit surface wave
                self.reconstruction.esw = objectPatch * self.reconstruction.probe

                # propagate to camera, intensityProjection, propagate back to object
                self.intensityProjection(positionIndex)

                # difference term
                DELTA = self.reconstruction.eswUpdate - self.reconstruction.esw

                if loop % self.params.OPR_tv_freq == 0 and self.params.OPR_tv:
                    self.reconstruction.object[..., sy, sx] = self.objectPatchUpdate_TV(
                        objectPatch, DELTA
                    )
                else:
                    # object update
                    self.reconstruction.object[..., sy, sx] = self.objectPatchUpdate(
                        objectPatch, DELTA
                    )

                # probe update
                self.reconstruction.probe = self.probeUpdate(
                    objectPatch, DELTA, weight=1
                )

                # save first, dominant probe mode
                self.reconstruction.probe_stack[..., positionIndex] = cp.copy(
                    self.reconstruction.probe[:, :, mode_slice, :, :, :]
                )

            # get error metric
            self.getErrorMetrics()

            if self.params.OPR_orthogonalize_modes:
                self.orthogonalizeIncoherentModes()

            self.reconstruction.probe_stack = self.orthogonalizeProbeStack(
                self.reconstruction.probe_stack, n_subspace
            )

            # apply Constraints
            self.applyConstraints(loop)

            # show reconstruction
            self.showReconstruction(loop)

        if self.params.gpuFlag:
            self.logger.info("switch to cpu")
            self._move_data_to_cpu()
            self.params.gpuFlag = 0

    def orthogonalizeIncoherentModes(self):
        """
        Function which cycles through the probe stack and orthogonalizes
        all incoherent modes of all postions
        """
        if self.params.OPR_fast_orthogonalization:
            return self._orthogonalizeIncoherentModes_batched()

        nFrames = self.experimentalData.numFrames
        n = self.reconstruction.Np
        nModes = self.reconstruction.probe_stack.shape[2]
        for pos in range(nFrames):
            probe = self.reconstruction.probe_stack[0, 0, :, 0, :, :, pos]
            probe = probe.reshape(nModes, n * n)

            U, s, Vh = self.svd(probe)

            modes = (s[:, None] * Vh).reshape(nModes, n, n)
            self.reconstruction.probe_stack[0, 0, :, 0, :, :, pos] = modes

    def _orthogonalizeIncoherentModes_batched(self):
        """Batched Gram-matrix equivalent of :meth:`orthogonalizeIncoherentModes`.

        For each frame the loop above computes ``s[:, None] * Vh`` from the SVD
        of a ``(nModes, Np**2)`` matrix P. Since ``P = U S Vh``, that product is
        just ``U^H P``, and U is the left singular matrix of the *tiny*
        ``(nModes, nModes)`` Gram matrix ``P P^H``. So the whole thing reduces to
        a batched factorization of a handful of 4x4 matrices plus one batched
        matmul -- no large SVD, and one kernel launch per chunk instead of one
        per frame.

        Measured 4.1x faster than the loop at 364 px / 202 frames / 4 modes, and
        1.8x at 512 px / 890 frames / 6 modes. The advantage *shrinks* with size:
        the loop's cost is dominated by per-frame launch overhead at small sizes,
        which is exactly what batching removes, while at large sizes the
        factorization itself dominates and batching has less to hide.

        Caveat: eigenvectors are only defined up to a per-mode phase, and when
        two modes carry near-equal power the vectors within that subspace are
        not determined at all. The mode *powers* (singular values) and the
        spanned subspace are reproduced exactly; individual mode vectors may
        differ from LAPACK's arbitrary choice. Guarded by
        ``params.OPR_fast_orthogonalization``.
        """
        stack = self.reconstruction.probe_stack
        xp = getArrayModule(stack)
        n = self.reconstruction.Np
        nModes = stack.shape[2]
        nFrames = stack.shape[-1]

        # Transposing the whole stack at once would allocate a second (and
        # third) copy of it -- 2.4 GB for a 364 px / 202 frame / 4 mode run, on
        # top of the stack itself. Work in frame chunks so the extra allocation
        # stays bounded regardless of frame count; the batched call is already
        # wide enough at a few dozen frames to hide launch overhead.
        elements_per_frame = nModes * n * n
        chunk = int(max(1, self._orthogonalization_chunk_bytes //
                        (elements_per_frame * stack.dtype.itemsize)))

        flat = stack[0, 0, :, 0, :, :, :].reshape(nModes, n * n, nFrames)
        for start in range(0, nFrames, chunk):
            stop = min(start + chunk, nFrames)
            # (nModes, Np**2, chunk) -> (chunk, nModes, Np**2)
            P = xp.ascontiguousarray(xp.moveaxis(flat[:, :, start:stop], 2, 0))
            G = P @ P.conj().transpose(0, 2, 1)
            # batched SVD of the tiny Hermitian Gram matrices; see gram_tsvd for
            # why this is used in preference to eigh. Already ordered by
            # descending mode power.
            U, _w, _Vh = xp.linalg.svd(G)
            modes = U.conj().transpose(0, 2, 1) @ P
            flat[:, :, start:stop] = xp.moveaxis(modes, 0, 2)
            del P, G, U, modes

    def average(self, arr):
        """
        Calculates the average from neighboring values of a numpy array
        :param arr: 1-dimensional input array, which is used to
        calculate the average
        :return: 1-dimensionl array with the same shape as the input array
        """
        arr_start = arr[:-1]
        arr_end = arr[1:]
        arr_end = cp.append(arr_end, 0)
        arr_start = cp.append(0, arr_start)
        divider = cp.ones_like(arr) * 3
        divider[0] = 2
        divider[-1] = 2
        return (arr + arr_end + arr_start) / divider

    def svd(self, P):
        if isGpuArray(P):
            try:
                return cp.linalg.svd(P, full_matrices=False)
            except:
                print(
                    "Something is wrong with SVD on cuda. Probably an installation error"
                )
                raise
        A, v, At = np.linalg.svd(asNumpyArray(P), full_matrices=False)
        if isGpuArray(P):
            A = cp.array(A)
            v = cp.array(v)
            At = cp.array(At)
        return A, v, At

    def rsvd(self, P, n_dim):
        return rsvd(P, n_dim)
        # A, v, At = self.svd(P)
        # v[n_dim:] = 0
        # return A, v, At

    @staticmethod
    def gram_tsvd(A, n_dim):
        """Rank-``n_dim`` truncated SVD of a tall matrix via its Gram matrix.

        ``A`` is ``(Np**2, nFrames)`` -- very tall and thin. A full SVD of it
        costs O(Np**2 * nFrames**2) and allocates a ``(Np**2, nFrames)`` U. The
        right singular vectors are the eigenvectors of the much smaller
        ``(nFrames, nFrames)`` Gram matrix ``A^H A``, so::

            V, s**2 = eigh(A^H A)      ->      U = A V / s

        Measured against the full SVD of the same matrix: 1.8x faster on 5.3x
        less peak memory at 364 px / 202 frames, and 1.9x on 4.3x less at
        512 px / 890 frames. The memory saving is the point -- it is what keeps
        a large OPR run inside a 32 GB card.

        The Gram matrix squares the condition number, so it is formed in double
        precision -- it is only ``nFrames x nFrames``, which is negligible next
        to the probe stack.

        Returns ``(U, s, Vh)`` truncated to ``n_dim`` components, matching the
        layout of ``xp.linalg.svd(..., full_matrices=False)`` after zeroing the
        tail of ``s``.
        """
        xp = getArrayModule(A)
        n_dim = int(min(n_dim, A.shape[1]))

        G = (A.conj().T @ A).astype(xp.complex128)
        # G is Hermitian positive semi-definite, so its SVD and its
        # eigendecomposition coincide: the left singular vectors are the
        # eigenvectors and the singular values are the eigenvalues, already in
        # descending order. We use svd rather than eigh because eigh routes
        # through cupyx.cusolver, which is not importable in every CuPy/CUDA
        # installation (it needs libcusolver at a version cupy-cuda12x does not
        # always ship), whereas svd works through cupy's own bindings.
        V, w, _Vh = xp.linalg.svd(G)
        w = w[:n_dim]
        V = V[:, :n_dim]

        s = xp.sqrt(xp.clip(w.real, 0.0, None))
        V = V.astype(A.dtype)
        # guard the division for numerically-zero singular values
        s_safe = xp.where(s > 0, s, 1.0)
        U = (A @ V) / s_safe.astype(A.real.dtype)[None, :]
        return U, s.astype(A.real.dtype), V.conj().T

    def orthogonalizeProbeStack(self, probe_stack, n_dim):
        """
        Takes the probe stack maps it by a truncated singular value decomposition in to
        a lower dimensional (n_dim) space.
        :param probe_stack: Probes of all positions
        :param n_dim: Dimension of the lower dimensional sub space
        :return: reduced probe stack
        """
        xp = getArrayModule(probe_stack)
        n = self.reconstruction.Np
        nFrames = self.experimentalData.numFrames

        for i, mode in enumerate(self.OPR_modes):
            A = probe_stack[:, :, i, :, :, :].reshape(n * n, nFrames)

            if self.params.OPR_tsvd_type == "randomized":
                U, s, Vh = self.rsvd(A, n_dim)
            elif self.params.OPR_tsvd_type == "gram":
                U, s, Vh = self.gram_tsvd(A, n_dim)
            elif self.params.OPR_tsvd_type == "numpy":
                U, s, Vh = xp.linalg.svd(A, full_matrices=False)
                s = s.copy()
                s[n_dim:] = 0
            else:
                raise ValueError(
                    f"unknown OPR_tsvd_type {self.params.OPR_tsvd_type!r}; "
                    f"expected 'numpy', 'gram' or 'randomized'"
                )

            if self.params.OPR_neighbor_constraint:
                # Calculate the average of neigboring singular values
                content = s[:, None] * Vh
                for j in range(min(n_dim, content.shape[0])):
                    content[j] = self.average(content[j])

                probe_stack[:, :, i, :, :, :] = self.alpha * probe_stack[
                    :, :, i, :, :, :
                ] + (1 - self.alpha) * (U @ content).reshape(n, n, nFrames)
            else:
                update = (U @ (s[:, None] * Vh)).reshape(n, n, nFrames)
                probe_stack[:, :, i, :, :, :] *= self.alpha
                probe_stack[:, :, i, :, :, :] += (1 - self.alpha) * update

        return probe_stack

    def objectPatchUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray):
        """
        ePIE object update function
        :param objectPatch: Slice of the object array
        :param DELTA:
        :return: updated object patch
        """
        # find out which array module to use, numpy or cupy (or other...)
        xp = getArrayModule(objectPatch)

        frac = self.reconstruction.probe.conj() / xp.max(
            xp.sum(xp.abs(self.reconstruction.probe) ** 2, axis=(0, 1, 2, 3))
        )
        return objectPatch + self.betaObject * xp.sum(
            frac * DELTA, axis=(0, 2, 3), keepdims=True
        )

    def probeUpdate(
        self, objectPatch: np.ndarray, DELTA: np.ndarray, weight: float, gimmel=0.1
    ):
        """
        Update the probe
        :param objectPatch: Slice of the object array
        :param DELTA:
        :return: updated probe
        """
        # find out which array module to use, numpy or cupy (or other...)
        xp = getArrayModule(objectPatch)
        frac = objectPatch.conj() / (
            xp.max(xp.sum(xp.abs(objectPatch) ** 2, axis=(0, 1, 2, 3))) + gimmel
        )
        frac = frac * weight
        r = self.reconstruction.probe + self.betaProbe * xp.sum(
            frac * DELTA, axis=(0, 1, 3), keepdims=True
        )
        return r

initializeReconstructionParams()

Set parameters that are specific to the ePIE/OPR engine

Source code in PtyLab/Engines/OPR.py

def initializeReconstructionParams(self):
    """
    Set parameters that are specific to the ePIE/OPR engine
    """
    self.alpha = self.params.OPR_alpha
    self.betaProbe = 0.25
    self.betaObject = 0.25
    self.numIterations = 50
    self.OPR_modes = self.params.OPR_modes
    self.n_subspace = self.params.OPR_subspace
    # Working-set budget for the chunked batched orthogonalization. Keeps
    # the transpose scratch bounded independently of the frame count.
    self._orthogonalization_chunk_bytes = 128 * 2**20

orthogonalizeIncoherentModes()

Function which cycles through the probe stack and orthogonalizes all incoherent modes of all postions

Source code in PtyLab/Engines/OPR.py

def orthogonalizeIncoherentModes(self):
    """
    Function which cycles through the probe stack and orthogonalizes
    all incoherent modes of all postions
    """
    if self.params.OPR_fast_orthogonalization:
        return self._orthogonalizeIncoherentModes_batched()

    nFrames = self.experimentalData.numFrames
    n = self.reconstruction.Np
    nModes = self.reconstruction.probe_stack.shape[2]
    for pos in range(nFrames):
        probe = self.reconstruction.probe_stack[0, 0, :, 0, :, :, pos]
        probe = probe.reshape(nModes, n * n)

        U, s, Vh = self.svd(probe)

        modes = (s[:, None] * Vh).reshape(nModes, n, n)
        self.reconstruction.probe_stack[0, 0, :, 0, :, :, pos] = modes

average(arr)

Calculates the average from neighboring values of a numpy array :param arr: 1-dimensional input array, which is used to calculate the average :return: 1-dimensionl array with the same shape as the input array

Source code in PtyLab/Engines/OPR.py

def average(self, arr):
    """
    Calculates the average from neighboring values of a numpy array
    :param arr: 1-dimensional input array, which is used to
    calculate the average
    :return: 1-dimensionl array with the same shape as the input array
    """
    arr_start = arr[:-1]
    arr_end = arr[1:]
    arr_end = cp.append(arr_end, 0)
    arr_start = cp.append(0, arr_start)
    divider = cp.ones_like(arr) * 3
    divider[0] = 2
    divider[-1] = 2
    return (arr + arr_end + arr_start) / divider

gram_tsvd(A, n_dim) staticmethod

Rank-n_dim truncated SVD of a tall matrix via its Gram matrix.

A is (Np**2, nFrames) -- very tall and thin. A full SVD of it costs O(Np2 * nFrames2) and allocates a (Np**2, nFrames) U. The right singular vectors are the eigenvectors of the much smaller (nFrames, nFrames) Gram matrix A^H A, so::

V, s**2 = eigh(A^H A)      ->      U = A V / s

Measured against the full SVD of the same matrix: 1.8x faster on 5.3x less peak memory at 364 px / 202 frames, and 1.9x on 4.3x less at 512 px / 890 frames. The memory saving is the point -- it is what keeps a large OPR run inside a 32 GB card.

The Gram matrix squares the condition number, so it is formed in double precision -- it is only nFrames x nFrames, which is negligible next to the probe stack.

Returns (U, s, Vh) truncated to n_dim components, matching the layout of xp.linalg.svd(..., full_matrices=False) after zeroing the tail of s.

Source code in PtyLab/Engines/OPR.py

@staticmethod
def gram_tsvd(A, n_dim):
    """Rank-``n_dim`` truncated SVD of a tall matrix via its Gram matrix.

    ``A`` is ``(Np**2, nFrames)`` -- very tall and thin. A full SVD of it
    costs O(Np**2 * nFrames**2) and allocates a ``(Np**2, nFrames)`` U. The
    right singular vectors are the eigenvectors of the much smaller
    ``(nFrames, nFrames)`` Gram matrix ``A^H A``, so::

        V, s**2 = eigh(A^H A)      ->      U = A V / s

    Measured against the full SVD of the same matrix: 1.8x faster on 5.3x
    less peak memory at 364 px / 202 frames, and 1.9x on 4.3x less at
    512 px / 890 frames. The memory saving is the point -- it is what keeps
    a large OPR run inside a 32 GB card.

    The Gram matrix squares the condition number, so it is formed in double
    precision -- it is only ``nFrames x nFrames``, which is negligible next
    to the probe stack.

    Returns ``(U, s, Vh)`` truncated to ``n_dim`` components, matching the
    layout of ``xp.linalg.svd(..., full_matrices=False)`` after zeroing the
    tail of ``s``.
    """
    xp = getArrayModule(A)
    n_dim = int(min(n_dim, A.shape[1]))

    G = (A.conj().T @ A).astype(xp.complex128)
    # G is Hermitian positive semi-definite, so its SVD and its
    # eigendecomposition coincide: the left singular vectors are the
    # eigenvectors and the singular values are the eigenvalues, already in
    # descending order. We use svd rather than eigh because eigh routes
    # through cupyx.cusolver, which is not importable in every CuPy/CUDA
    # installation (it needs libcusolver at a version cupy-cuda12x does not
    # always ship), whereas svd works through cupy's own bindings.
    V, w, _Vh = xp.linalg.svd(G)
    w = w[:n_dim]
    V = V[:, :n_dim]

    s = xp.sqrt(xp.clip(w.real, 0.0, None))
    V = V.astype(A.dtype)
    # guard the division for numerically-zero singular values
    s_safe = xp.where(s > 0, s, 1.0)
    U = (A @ V) / s_safe.astype(A.real.dtype)[None, :]
    return U, s.astype(A.real.dtype), V.conj().T

orthogonalizeProbeStack(probe_stack, n_dim)

Takes the probe stack maps it by a truncated singular value decomposition in to a lower dimensional (n_dim) space. :param probe_stack: Probes of all positions :param n_dim: Dimension of the lower dimensional sub space :return: reduced probe stack

Source code in PtyLab/Engines/OPR.py

def orthogonalizeProbeStack(self, probe_stack, n_dim):
    """
    Takes the probe stack maps it by a truncated singular value decomposition in to
    a lower dimensional (n_dim) space.
    :param probe_stack: Probes of all positions
    :param n_dim: Dimension of the lower dimensional sub space
    :return: reduced probe stack
    """
    xp = getArrayModule(probe_stack)
    n = self.reconstruction.Np
    nFrames = self.experimentalData.numFrames

    for i, mode in enumerate(self.OPR_modes):
        A = probe_stack[:, :, i, :, :, :].reshape(n * n, nFrames)

        if self.params.OPR_tsvd_type == "randomized":
            U, s, Vh = self.rsvd(A, n_dim)
        elif self.params.OPR_tsvd_type == "gram":
            U, s, Vh = self.gram_tsvd(A, n_dim)
        elif self.params.OPR_tsvd_type == "numpy":
            U, s, Vh = xp.linalg.svd(A, full_matrices=False)
            s = s.copy()
            s[n_dim:] = 0
        else:
            raise ValueError(
                f"unknown OPR_tsvd_type {self.params.OPR_tsvd_type!r}; "
                f"expected 'numpy', 'gram' or 'randomized'"
            )

        if self.params.OPR_neighbor_constraint:
            # Calculate the average of neigboring singular values
            content = s[:, None] * Vh
            for j in range(min(n_dim, content.shape[0])):
                content[j] = self.average(content[j])

            probe_stack[:, :, i, :, :, :] = self.alpha * probe_stack[
                :, :, i, :, :, :
            ] + (1 - self.alpha) * (U @ content).reshape(n, n, nFrames)
        else:
            update = (U @ (s[:, None] * Vh)).reshape(n, n, nFrames)
            probe_stack[:, :, i, :, :, :] *= self.alpha
            probe_stack[:, :, i, :, :, :] += (1 - self.alpha) * update

    return probe_stack

objectPatchUpdate(objectPatch, DELTA)

ePIE object update function :param objectPatch: Slice of the object array :param DELTA: :return: updated object patch

Source code in PtyLab/Engines/OPR.py

def objectPatchUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray):
    """
    ePIE object update function
    :param objectPatch: Slice of the object array
    :param DELTA:
    :return: updated object patch
    """
    # find out which array module to use, numpy or cupy (or other...)
    xp = getArrayModule(objectPatch)

    frac = self.reconstruction.probe.conj() / xp.max(
        xp.sum(xp.abs(self.reconstruction.probe) ** 2, axis=(0, 1, 2, 3))
    )
    return objectPatch + self.betaObject * xp.sum(
        frac * DELTA, axis=(0, 2, 3), keepdims=True
    )

probeUpdate(objectPatch, DELTA, weight, gimmel=0.1)

Update the probe :param objectPatch: Slice of the object array :param DELTA: :return: updated probe

Source code in PtyLab/Engines/OPR.py

def probeUpdate(
    self, objectPatch: np.ndarray, DELTA: np.ndarray, weight: float, gimmel=0.1
):
    """
    Update the probe
    :param objectPatch: Slice of the object array
    :param DELTA:
    :return: updated probe
    """
    # find out which array module to use, numpy or cupy (or other...)
    xp = getArrayModule(objectPatch)
    frac = objectPatch.conj() / (
        xp.max(xp.sum(xp.abs(objectPatch) ** 2, axis=(0, 1, 2, 3))) + gimmel
    )
    frac = frac * weight
    r = self.reconstruction.probe + self.betaProbe * xp.sum(
        frac * DELTA, axis=(0, 1, 3), keepdims=True
    )
    return r

aPIE

aPIE

Bases: BaseEngine

aPIE: angle correction PIE: ePIE combined with Luus-Jaakola algorithm (the latter for angle correction) + momentum

Source code in PtyLab/Engines/aPIE.py

class aPIE(BaseEngine):
    """
    aPIE: angle correction PIE: ePIE combined with Luus-Jaakola algorithm (the latter for angle correction) + momentum
    """

    def __init__(
        self,
        reconstruction: Reconstruction,
        experimentalData: ExperimentalData,
        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")
        self.logger.info("Wavelength attribute: %s", self.reconstruction.wavelength)
        self.initializeReconstructionParams()

    def initializeReconstructionParams(self):
        """
        Set parameters that are specific to the ePIE settings.
        :return:
        """
        self.betaProbe = 0.25
        self.betaObject = 0.25
        self.aPIEfriction = 0.7
        self.feedback = 0.5
        self.numIterations = 50

        if not hasattr(self.reconstruction, "thetaMomentum"):
            self.reconstruction.thetaMomentum = 0
        if not hasattr(self.reconstruction, "thetaHistory"):
            self.reconstruction.thetaHistory = np.array([])

        self.thetaSearchRadiusMin = 0.01
        self.thetaSearchRadiusMax = 0.1
        self.ptychogramUntransformed = self.experimentalData.ptychogram.copy()
        self.experimentalData.W = np.ones_like(self.reconstruction.Xd)

        if self.reconstruction.theta == None:
            raise ValueError("theta value is not given")

    def doReconstruction(self):
        self._prepareReconstruction()

        xp = getArrayModule(self.reconstruction.object)

        # linear search
        thetaSearchRadiusList = np.linspace(
            self.thetaSearchRadiusMax, self.thetaSearchRadiusMin, self.numIterations
        )

        self.pbar = tqdm.trange(
            self.numIterations, desc="aPIE", file=sys.stdout, leave=True
        )
        for loop in self.pbar:
            # save theta search history
            self.reconstruction.thetaHistory = np.append(
                self.reconstruction.thetaHistory,
                asNumpyArray(self.reconstruction.theta),
            )

            # select two angles (todo check if three angles behave better)
            theta = (
                np.array(
                    [
                        self.reconstruction.theta,
                        self.reconstruction.theta
                        + thetaSearchRadiusList[loop] * (-1 + 2 * np.random.rand()),
                    ]
                )
                + self.reconstruction.thetaMomentum
            )

            # save object and probe
            probeTemp = self.reconstruction.probe.copy()
            objectTemp = self.reconstruction.object.copy()

            # probe and object buffer (todo maybe there's more elegant way )
            probeBuffer = xp.zeros_like(
                probeTemp
            )  # shape=(np.array([probeTemp, probeTemp])).shape)
            probeBuffer = [probeBuffer, probeBuffer]
            objectBuffer = xp.zeros_like(
                objectTemp
            )  # , shape=(np.array([objectTemp, objectTemp])).shape)  # for polychromatic case this will need to be multimode
            objectBuffer = [objectBuffer, objectBuffer]
            # initialize error
            errorTemp = np.zeros((2, 1))

            for k in range(2):
                self.reconstruction.probe = probeTemp
                self.reconstruction.object = objectTemp
                # reset ptychogram (transform into estimate coordinates)
                Xq = T_inv(
                    self.reconstruction.Xd,
                    self.reconstruction.Yd,
                    self.reconstruction.zo,
                    theta[k],
                )  # todo check if 1D is enough to save time
                for l in range(self.experimentalData.numFrames):
                    temp = self.ptychogramUntransformed[l]
                    f = interp2d(
                        self.reconstruction.xd,
                        self.reconstruction.xd,
                        temp,
                        kind="linear",
                        fill_value=0,
                    )
                    temp2 = abs(f(Xq[0], self.reconstruction.xd))
                    temp2 = np.nan_to_num(temp2)
                    temp2[temp2 < 0] = 0
                    self.experimentalData.ptychogram[l] = xp.array(temp2)

                # renormalization(for energy conservation) # todo not layer by layer?
                self.experimentalData.ptychogram = (
                    self.experimentalData.ptychogram
                    / np.linalg.norm(self.experimentalData.ptychogram)
                    * np.linalg.norm(self.ptychogramUntransformed)
                )

                self.experimentalData.W = np.ones_like(self.reconstruction.Xd)
                fw = interp2d(
                    self.reconstruction.xd,
                    self.reconstruction.xd,
                    self.experimentalData.W,
                    kind="linear",
                    fill_value=0,
                )
                self.experimentalData.W = abs(fw(Xq[0], self.reconstruction.xd))
                self.experimentalData.W = np.nan_to_num(self.experimentalData.W)
                self.experimentalData.W[self.experimentalData.W == 0] = 1e-3
                self.experimentalData.W = xp.array(self.experimentalData.W)

                # todo check if it is right
                if self.params.fftshiftSwitch:
                    self.experimentalData.ptychogram = xp.fft.ifftshift(
                        self.experimentalData.ptychogram, axes=(-1, -2)
                    )
                    self.experimentalData.W = xp.fft.ifftshift(
                        self.experimentalData.W, axes=(-1, -2)
                    )

                # set position order
                self.setPositionOrder()

                for positionLoop, positionIndex in enumerate(self.positionIndices):
                    ### patch1 ###
                    # get object patch1
                    row1, col1 = self.reconstruction.positions[positionIndex]
                    sy = slice(row1, row1 + self.reconstruction.Np)
                    sx = slice(col1, col1 + self.reconstruction.Np)
                    # note that object patch has size of probe array
                    objectPatch = self.reconstruction.object[..., sy, sx].copy()

                    # make exit surface wave
                    self.reconstruction.esw = objectPatch * self.reconstruction.probe

                    # propagate to camera, intensityProjection, propagate back to object
                    self.intensityProjection(positionIndex)

                    # difference term1
                    DELTA = self.reconstruction.eswUpdate - self.reconstruction.esw

                    # object update
                    self.reconstruction.object[..., sy, sx] = self.objectPatchUpdate(
                        objectPatch, DELTA
                    )

                    # probe update
                    self.reconstruction.probe = self.probeUpdate(objectPatch, DELTA)

                # get error metric
                self.getErrorMetrics()
                # remove error from error history
                errorTemp[k] = self.reconstruction.error[-1]
                self.reconstruction.error = np.delete(self.reconstruction.error, -1)

                # apply Constraints
                self.applyConstraints(loop)
                # update buffer
                probeBuffer[k] = self.reconstruction.probe
                objectBuffer[k] = self.reconstruction.object

            if errorTemp[1] < errorTemp[0]:
                dtheta = theta[1] - theta[0]
                self.reconstruction.theta = theta[1]
                self.reconstruction.probe = probeBuffer[1]
                self.reconstruction.object = objectBuffer[1]
                self.reconstruction.error = np.append(
                    self.reconstruction.error, errorTemp[1]
                )
            else:
                dtheta = 0
                self.reconstruction.theta = theta[0]
                self.reconstruction.probe = probeBuffer[0]
                self.reconstruction.object = objectBuffer[0]
                self.reconstruction.error = np.append(
                    self.reconstruction.error, errorTemp[0]
                )

            self.reconstruction.thetaMomentum = (
                self.feedback * dtheta
                + self.aPIEfriction * self.reconstruction.thetaMomentum
            )
            # print updated theta
            self.pbar.set_description(
                "aPIE: update a=%.3f deg (search radius=%.3f deg, thetaMomentum=%.3f deg)"
                % (
                    self.reconstruction.theta,
                    thetaSearchRadiusList[loop],
                    self.reconstruction.thetaMomentum,
                )
            )

            # show reconstruction
            if loop == 0:
                figure, ax = plt.subplots(
                    1, 1, num=777, squeeze=True, clear=True, figsize=(5, 5)
                )
                ax.set_title("Estimated angle")
                ax.set_xlabel("iteration")
                ax.set_ylabel("estimated theta [deg]")
                ax.set_xscale("symlog")
                line = plt.plot(0, self.reconstruction.theta, "o-")[0]
                plt.tight_layout()
                plt.show(block=False)

            elif np.mod(loop, self.monitor.figureUpdateFrequency) == 0:
                idx = np.linspace(
                    0,
                    np.log10(len(self.reconstruction.thetaHistory) - 1),
                    np.minimum(len(self.reconstruction.thetaHistory), 100),
                )
                idx = np.rint(10**idx).astype("int")

                line.set_xdata(idx)
                line.set_ydata(np.array(self.reconstruction.thetaHistory)[idx])
                ax.set_xlim(0, np.max(idx))
                ax.set_ylim(
                    min(self.reconstruction.thetaHistory),
                    max(self.reconstruction.thetaHistory),
                )

                figure.canvas.draw()
                figure.canvas.flush_events()

            self.showReconstruction(loop)

        if self.params.gpuFlag:
            self.logger.info("switch to cpu")
            self._move_data_to_cpu()
            self.params.gpuFlag = 0

        # self.thetaSearchRadiusMax = thetaSearchRadiusList[loop]

    def objectPatchUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray):
        """
        Todo add docstring
        :param objectPatch:
        :param DELTA:
        :return:
        """
        # find out which array module to use, numpy or cupy (or other...)
        xp = getArrayModule(objectPatch)

        frac = self.reconstruction.probe.conj() / xp.max(
            xp.sum(xp.abs(self.reconstruction.probe) ** 2, axis=(0, 1, 2, 3))
        )
        return objectPatch + self.betaObject * xp.sum(
            frac * DELTA, axis=(0, 2, 3), keepdims=True
        )

    def probeUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray):
        """
        Todo add docstring
        :param objectPatch:
        :param DELTA:
        :return:
        """
        # find out which array module to use, numpy or cupy (or other...)
        xp = getArrayModule(objectPatch)
        frac = objectPatch.conj() / xp.max(
            xp.sum(xp.abs(objectPatch) ** 2, axis=(0, 1, 2, 3))
        )
        r = self.reconstruction.probe + self.betaProbe * xp.sum(
            frac * DELTA, axis=(0, 1, 3), keepdims=True
        )
        return r

initializeReconstructionParams()

Set parameters that are specific to the ePIE settings. :return:

Source code in PtyLab/Engines/aPIE.py

def initializeReconstructionParams(self):
    """
    Set parameters that are specific to the ePIE settings.
    :return:
    """
    self.betaProbe = 0.25
    self.betaObject = 0.25
    self.aPIEfriction = 0.7
    self.feedback = 0.5
    self.numIterations = 50

    if not hasattr(self.reconstruction, "thetaMomentum"):
        self.reconstruction.thetaMomentum = 0
    if not hasattr(self.reconstruction, "thetaHistory"):
        self.reconstruction.thetaHistory = np.array([])

    self.thetaSearchRadiusMin = 0.01
    self.thetaSearchRadiusMax = 0.1
    self.ptychogramUntransformed = self.experimentalData.ptychogram.copy()
    self.experimentalData.W = np.ones_like(self.reconstruction.Xd)

    if self.reconstruction.theta == None:
        raise ValueError("theta value is not given")

objectPatchUpdate(objectPatch, DELTA)

Todo add docstring :param objectPatch: :param DELTA: :return:

Source code in PtyLab/Engines/aPIE.py

def objectPatchUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray):
    """
    Todo add docstring
    :param objectPatch:
    :param DELTA:
    :return:
    """
    # find out which array module to use, numpy or cupy (or other...)
    xp = getArrayModule(objectPatch)

    frac = self.reconstruction.probe.conj() / xp.max(
        xp.sum(xp.abs(self.reconstruction.probe) ** 2, axis=(0, 1, 2, 3))
    )
    return objectPatch + self.betaObject * xp.sum(
        frac * DELTA, axis=(0, 2, 3), keepdims=True
    )

probeUpdate(objectPatch, DELTA)

Todo add docstring :param objectPatch: :param DELTA: :return:

Source code in PtyLab/Engines/aPIE.py

def probeUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray):
    """
    Todo add docstring
    :param objectPatch:
    :param DELTA:
    :return:
    """
    # find out which array module to use, numpy or cupy (or other...)
    xp = getArrayModule(objectPatch)
    frac = objectPatch.conj() / xp.max(
        xp.sum(xp.abs(objectPatch) ** 2, axis=(0, 1, 2, 3))
    )
    r = self.reconstruction.probe + self.betaProbe * xp.sum(
        frac * DELTA, axis=(0, 1, 3), keepdims=True
    )
    return r

T(x, y, z, theta)

Coordinate transformation

Source code in PtyLab/Engines/aPIE.py

def T(x, y, z, theta):
    """
    Coordinate transformation
    """
    r0 = np.sqrt(x**2 + y**2 + z**2)
    yd = y
    xd = x * np.cos(toDegree(theta)) - np.sin(toDegree(theta)) * (r0 - z)
    return xd, yd

T_inv(xd, yd, z, theta)

inverse coordinate transformation

Source code in PtyLab/Engines/aPIE.py

def T_inv(xd, yd, z, theta):
    """
    inverse coordinate transformation
    """
    if theta != 45:
        rootTerm = np.sqrt(
            (z * np.cos(toDegree(theta))) ** 2
            + xd**2
            + yd**2 * np.cos(toDegree(2 * theta))
            - 2 * xd * z * np.sin(toDegree(theta))
        )
        x = (
            xd * np.cos(toDegree(theta))
            - z * np.sin(toDegree(theta)) * np.cos(toDegree(theta))
            + np.sin(toDegree(theta)) * rootTerm
        ) / np.cos(toDegree(2 * theta))
    else:
        x = (xd**2 - (yd**2) / 2 - xd * np.sqrt(2) * z) / (xd * np.sqrt(2) - z)
    return x

e3PIE

e3PIE

Bases: BaseEngine

Source code in PtyLab/Engines/e3PIE.py

class e3PIE(BaseEngine):

    def __init__(
        self,
        reconstruction: Reconstruction,
        experimentalData: ExperimentalData,
        params: Params,
        monitor: Monitor,
    ):
        # This contains reconstruction parameters that are specific to the reconstruction
        # but not necessarily to e3PIE reconstruction
        super().__init__(reconstruction, experimentalData, params, monitor)
        self.logger = logging.getLogger("e3PIE")
        self.logger.info("Sucesfully created e3PIE e3PIE_engine")

        self.logger.info("Wavelength attribute: %s", self.reconstruction.wavelength)

        self.initializeReconstructionParams()

    def initializeReconstructionParams(self):
        """
        Set parameters that are specific to the e3PIE settings.
        :return:
        """
        # these are read back as self.betaProbe / self.betaObject in reconstruct()
        # and objectPatchUpdate(), matching every other engine (cf. ePIE.py)
        self.betaProbe = 0.25
        self.betaObject = 0.25
        self.numIterations = 50

        # preallocate transfer function. This runs from __init__, before
        # _checkGPU has moved anything, so the probe is still on the host here;
        # H is listed in Reconstruction.possible_GPU_fields and travels with the
        # rest of the state when the engine switches to the GPU.
        xp = getArrayModule(self.reconstruction.probe)
        self.reconstruction.H = aspw(
            xp.squeeze(self.reconstruction.probe[0, 0, 0, 0, ...]),
            self.reconstruction.dz,
            self.reconstruction.wavelength / self.reconstruction.refrIndex,
            self.reconstruction.Lp,
        )[1]
        # shift transfer function to avoid fftshifts for FFTS
        self.reconstruction.H = xp.fft.ifftshift(self.reconstruction.H)

    def reconstruct(self):
        self._prepareReconstruction()

        # initialize esw
        self.reconstruction.esw = self.reconstruction.probe.copy()
        # get module
        xp = getArrayModule(self.reconstruction.object)

        self.pbar = tqdm.trange(
            self.numIterations, desc="e3PIE", file=sys.stdout, leave=True
        )

        # self.pbar = (1, 2)

        for loop in self.pbar:
            if loop == self.numIterations - 1:
                noreason = True
            self.setPositionOrder()
            for positionLoop, positionIndex in enumerate(self.positionIndices):
                # get object patch
                row, col = self.reconstruction.positions[positionIndex]
                sy = slice(row, row + self.reconstruction.Np)
                sx = slice(col, col + self.reconstruction.Np)
                # note that object patch has size of probe array
                objectPatch = self.reconstruction.object[..., sy, sx].copy()
                # objectPatch2 = self.reconstruction.object[..., :, :].copy()

                # form first slice esw (exit surface wave)
                self.reconstruction.esw[:, :, :, 0, ...] = (
                    objectPatch[:, :, :, 0, ...]
                    * self.reconstruction.probe[:, :, :, 0, ...]
                )

                # propagate through object
                for sliceLoop in range(1, self.reconstruction.nslice):
                    self.reconstruction.probe[:, :, :, sliceLoop, ...] = xp.fft.ifft2(
                        xp.fft.fft2(
                            self.reconstruction.esw[:, :, :, sliceLoop - 1, ...]
                        )
                        * self.reconstruction.H
                    )
                    self.reconstruction.esw[:, :, :, sliceLoop, ...] = (
                        self.reconstruction.probe[:, :, :, sliceLoop, ...]
                        * objectPatch[:, :, :, sliceLoop, ...]
                    )

                # propagate to camera, intensityProjection, propagate back to object
                self.intensityProjection(positionIndex)

                # difference term
                DELTA = (self.reconstruction.eswUpdate - self.reconstruction.esw)[
                    :, :, :, -1, ...
                ]
                # update object slice
                for loopTemp in range(self.reconstruction.nslice - 1):

                    sliceLoop = self.reconstruction.nslice - 1 - loopTemp

                    # temp_delta = self.reconstruction.esw[..., sliceLoop, sy, sx]

                    # compute and update current object slice
                    self.reconstruction.object[..., sliceLoop, sy, sx] = (
                        self.objectPatchUpdate(
                            objectPatch[:, :, :, sliceLoop, ...],
                            DELTA,
                            self.reconstruction.probe[:, :, :, sliceLoop, ...],
                        )
                    )
                    # eswTemp update (here probe incident on last slice)
                    beth = 0.9 # todo, why need beth, not betaProbe, changable?
                    self.reconstruction.probe[:, :, :, sliceLoop, ...] = (
                        self.probeUpdate(
                            objectPatch[:, :, :, sliceLoop, ...],
                            DELTA,
                            self.reconstruction.probe[:, :, :, sliceLoop, ...],
                            beth,
                        )
                    )

                    # back-propagate and calculate gradient term
                    DELTA = (
                        xp.fft.ifft2(
                            xp.fft.fft2(
                                self.reconstruction.probe[:, :, :, sliceLoop, ...]
                            )
                            * self.reconstruction.H.conj()
                        )
                        - self.reconstruction.esw[:, :, :, sliceLoop - 1, ...]
                    )

                # update last object slice
                self.reconstruction.object[..., 0, sy, sx] = self.objectPatchUpdate(
                    objectPatch[:, :, :, 0, ...],
                    DELTA,
                    self.reconstruction.probe[:, :, :, 0, ...],
                )
                # update probe
                self.reconstruction.probe[:, :, :, 0, ...] = self.probeUpdate(
                    objectPatch[:, :, :, 0, ...],
                    DELTA,
                    self.reconstruction.probe[:, :, :, 0, ...],
                    self.betaProbe,
                )

            # set porduct of all object slices
            self.reconstruction.objectProd = np.prod(self.reconstruction.object, 3)

            # get error metric
            self.getErrorMetrics()

            # apply Constraints todo uncomment orthogonalization? check object smootheness regularization
            self.applyConstraints(loop)

            # show reconstruction
            self.showReconstruction(loop)

        if self.params.gpuFlag:
            self.logger.info("switch to cpu")
            self._move_data_to_cpu()
            self.params.gpuFlag = 0

    def objectPatchUpdate(
        self, objectPatch: np.ndarray, DELTA: np.ndarray, localProbe: np.ndarray
    ):
        """
        Todo add docstring
        :param objectPatch:
        :param DELTA:
        :return:
        """
        # find out which array module to use, numpy or cupy (or other...)
        xp = getArrayModule(objectPatch)
        frac = localProbe.conj() / xp.max(
            xp.sum(xp.abs(localProbe) ** 2, axis=(0, 1, 2))
        )
        return objectPatch + self.betaObject * xp.sum(
            frac * DELTA, axis=(0, 2), keepdims=True
        )

    def probeUpdate(
        self, objectPatch: np.ndarray, DELTA: np.ndarray, localProbe: np.ndarray, beth
    ):
        """
        Todo add docstring
        :param objectPatch:
        :param DELTA:
        :return:
        """
        # find out which array module to use, numpy or cupy (or other...)
        xp = getArrayModule(objectPatch)
        frac = objectPatch.conj() / xp.max(
            xp.sum(xp.abs(objectPatch) ** 2, axis=(0, 1, 2))
        )
        r = localProbe + beth * xp.sum(frac * DELTA, axis=(0, 1), keepdims=True)
        return r

initializeReconstructionParams()

Set parameters that are specific to the e3PIE settings. :return:

Source code in PtyLab/Engines/e3PIE.py

def initializeReconstructionParams(self):
    """
    Set parameters that are specific to the e3PIE settings.
    :return:
    """
    # these are read back as self.betaProbe / self.betaObject in reconstruct()
    # and objectPatchUpdate(), matching every other engine (cf. ePIE.py)
    self.betaProbe = 0.25
    self.betaObject = 0.25
    self.numIterations = 50

    # preallocate transfer function. This runs from __init__, before
    # _checkGPU has moved anything, so the probe is still on the host here;
    # H is listed in Reconstruction.possible_GPU_fields and travels with the
    # rest of the state when the engine switches to the GPU.
    xp = getArrayModule(self.reconstruction.probe)
    self.reconstruction.H = aspw(
        xp.squeeze(self.reconstruction.probe[0, 0, 0, 0, ...]),
        self.reconstruction.dz,
        self.reconstruction.wavelength / self.reconstruction.refrIndex,
        self.reconstruction.Lp,
    )[1]
    # shift transfer function to avoid fftshifts for FFTS
    self.reconstruction.H = xp.fft.ifftshift(self.reconstruction.H)

objectPatchUpdate(objectPatch, DELTA, localProbe)

Todo add docstring :param objectPatch: :param DELTA: :return:

Source code in PtyLab/Engines/e3PIE.py

def objectPatchUpdate(
    self, objectPatch: np.ndarray, DELTA: np.ndarray, localProbe: np.ndarray
):
    """
    Todo add docstring
    :param objectPatch:
    :param DELTA:
    :return:
    """
    # find out which array module to use, numpy or cupy (or other...)
    xp = getArrayModule(objectPatch)
    frac = localProbe.conj() / xp.max(
        xp.sum(xp.abs(localProbe) ** 2, axis=(0, 1, 2))
    )
    return objectPatch + self.betaObject * xp.sum(
        frac * DELTA, axis=(0, 2), keepdims=True
    )

probeUpdate(objectPatch, DELTA, localProbe, beth)

Todo add docstring :param objectPatch: :param DELTA: :return:

Source code in PtyLab/Engines/e3PIE.py

def probeUpdate(
    self, objectPatch: np.ndarray, DELTA: np.ndarray, localProbe: np.ndarray, beth
):
    """
    Todo add docstring
    :param objectPatch:
    :param DELTA:
    :return:
    """
    # find out which array module to use, numpy or cupy (or other...)
    xp = getArrayModule(objectPatch)
    frac = objectPatch.conj() / xp.max(
        xp.sum(xp.abs(objectPatch) ** 2, axis=(0, 1, 2))
    )
    r = localProbe + beth * xp.sum(frac * DELTA, axis=(0, 1), keepdims=True)
    return r

ePIE

ePIE

Bases: BaseEngine

Source code in PtyLab/Engines/ePIE.py

class ePIE(BaseEngine):
    def __init__(
        self,
        reconstruction: Reconstruction,
        experimentalData: ExperimentalData,
        params: Params,
        monitor: Monitor,
    ):
        # This contains reconstruction parameters that are specific to the reconstruction
        # but not necessarily to ePIE reconstruction
        super().__init__(reconstruction, experimentalData, params, monitor)
        self.logger = logging.getLogger("ePIE")
        self.logger.info("Sucesfully created ePIE ePIE_engine")
        self.logger.info("Wavelength attribute: %s", self.reconstruction.wavelength)
        self.initializeReconstructionParams()

    def initializeReconstructionParams(self):
        """
        Set parameters that are specific to the ePIE settings.
        :return:
        """
        self.betaProbe = 0.25
        self.betaObject = 0.25
        self.numIterations = 50

    def reconstruct(self, experimentalData: ExperimentalData = None):
        """Run the reconstruction to completion.

        Use :meth:`reconstruct_stepwise` instead if you want to interleave your
        own work between scan positions.
        """
        for _ in self.reconstruct_stepwise(experimentalData):
            pass

    def reconstruct_stepwise(self, experimentalData: ExperimentalData = None):
        """Generator variant of :meth:`reconstruct`.

        Yields ``(iteration, positionLoop)`` after every scan position. Nothing
        happens until the generator is consumed.
        """
        if experimentalData is not None:
            self.reconstruction.data = experimentalData
            self.experimentalData = experimentalData
        self._prepareReconstruction()

        # actual reconstruction ePIE_engine
        self.pbar = tqdm.trange(
            self.numIterations, desc="ePIE", file=sys.stdout, leave=True
        )
        for loop in self.pbar:
            # set position order
            self.setPositionOrder()
            if self.params.OPRP:
                # make the initial guess the default storage
                self.reconstruction.probe_storage.push(
                    self.reconstruction.probe,
                    0,
                    self.experimentalData.ptychogram.shape[0],
                )
            for positionLoop, positionIndex in enumerate(self.positionIndices):
                # get object patch
                if self.params.OPRP:
                    self.reconstruction.probe = self.reconstruction.probe_storage.get(
                        positionIndex
                    )
                row, col = self.reconstruction.positions[positionIndex]
                sy = slice(row, row + self.reconstruction.Np)
                sx = slice(col, col + self.reconstruction.Np)
                # note that object patch has size of probe array
                objectPatch = self.reconstruction.object[..., sy, sx].copy()

                # make exit surface wave
                self.reconstruction.esw = objectPatch * self.reconstruction.probe

                # propagate to camera, intensityProjection, propagate back to object
                self.intensityProjection(positionIndex)

                # difference term
                DELTA = self.reconstruction.eswUpdate - self.reconstruction.esw

                # object update
                self.reconstruction.object[..., sy, sx] = self.objectPatchUpdate(
                    objectPatch, DELTA
                )

                # probe update
                self.reconstruction.probe = self.probeUpdate(objectPatch, DELTA)
                if self.params.OPRP:
                    self.reconstruction.probe_storage.push(
                        self.reconstruction.probe,
                        positionIndex,
                        self.experimentalData.ptychogram.shape[0],
                    )
                yield loop, positionLoop

            # get error metric
            self.getErrorMetrics()

            # apply Constraints
            self.applyConstraints(loop)

            # show reconstruction
            # self.showReconstruction(loop)

        if self.params.gpuFlag:
            self.logger.info("switch to cpu")
            self._move_data_to_cpu()
            self.params.gpuFlag = 0

    def objectPatchUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray):
        """
        Todo add docstring
        :param objectPatch:
        :param DELTA:
        :return:
        """
        # find out which array module to use, numpy or cupy (or other...)
        xp = getArrayModule(objectPatch)

        frac = self.reconstruction.probe.conj() / xp.max(
            xp.sum(xp.abs(self.reconstruction.probe) ** 2, axis=(0, 1, 2, 3))
        )
        return objectPatch + self.betaObject * xp.sum(
            frac * DELTA, axis=(0, 2, 3), keepdims=True
        )

    def probeUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray):
        """
        Todo add docstring
        :param objectPatch:
        :param DELTA:
        :return:
        """
        # find out which array module to use, numpy or cupy (or other...)
        xp = getArrayModule(objectPatch)
        frac = objectPatch.conj() / xp.max(
            xp.sum(xp.abs(objectPatch) ** 2, axis=(0, 1, 2, 3))
        )
        r = self.reconstruction.probe + self.betaProbe * xp.sum(
            frac * DELTA, axis=(0, 1, 3), keepdims=True
        )
        return r

initializeReconstructionParams()

Set parameters that are specific to the ePIE settings. :return:

Source code in PtyLab/Engines/ePIE.py

def initializeReconstructionParams(self):
    """
    Set parameters that are specific to the ePIE settings.
    :return:
    """
    self.betaProbe = 0.25
    self.betaObject = 0.25
    self.numIterations = 50

reconstruct(experimentalData=None)

Run the reconstruction to completion.

Use :meth:reconstruct_stepwise instead if you want to interleave your own work between scan positions.

Source code in PtyLab/Engines/ePIE.py

def reconstruct(self, experimentalData: ExperimentalData = None):
    """Run the reconstruction to completion.

    Use :meth:`reconstruct_stepwise` instead if you want to interleave your
    own work between scan positions.
    """
    for _ in self.reconstruct_stepwise(experimentalData):
        pass

reconstruct_stepwise(experimentalData=None)

Generator variant of :meth:reconstruct.

Yields (iteration, positionLoop) after every scan position. Nothing happens until the generator is consumed.

Source code in PtyLab/Engines/ePIE.py

def reconstruct_stepwise(self, experimentalData: ExperimentalData = None):
    """Generator variant of :meth:`reconstruct`.

    Yields ``(iteration, positionLoop)`` after every scan position. Nothing
    happens until the generator is consumed.
    """
    if experimentalData is not None:
        self.reconstruction.data = experimentalData
        self.experimentalData = experimentalData
    self._prepareReconstruction()

    # actual reconstruction ePIE_engine
    self.pbar = tqdm.trange(
        self.numIterations, desc="ePIE", file=sys.stdout, leave=True
    )
    for loop in self.pbar:
        # set position order
        self.setPositionOrder()
        if self.params.OPRP:
            # make the initial guess the default storage
            self.reconstruction.probe_storage.push(
                self.reconstruction.probe,
                0,
                self.experimentalData.ptychogram.shape[0],
            )
        for positionLoop, positionIndex in enumerate(self.positionIndices):
            # get object patch
            if self.params.OPRP:
                self.reconstruction.probe = self.reconstruction.probe_storage.get(
                    positionIndex
                )
            row, col = self.reconstruction.positions[positionIndex]
            sy = slice(row, row + self.reconstruction.Np)
            sx = slice(col, col + self.reconstruction.Np)
            # note that object patch has size of probe array
            objectPatch = self.reconstruction.object[..., sy, sx].copy()

            # make exit surface wave
            self.reconstruction.esw = objectPatch * self.reconstruction.probe

            # propagate to camera, intensityProjection, propagate back to object
            self.intensityProjection(positionIndex)

            # difference term
            DELTA = self.reconstruction.eswUpdate - self.reconstruction.esw

            # object update
            self.reconstruction.object[..., sy, sx] = self.objectPatchUpdate(
                objectPatch, DELTA
            )

            # probe update
            self.reconstruction.probe = self.probeUpdate(objectPatch, DELTA)
            if self.params.OPRP:
                self.reconstruction.probe_storage.push(
                    self.reconstruction.probe,
                    positionIndex,
                    self.experimentalData.ptychogram.shape[0],
                )
            yield loop, positionLoop

        # get error metric
        self.getErrorMetrics()

        # apply Constraints
        self.applyConstraints(loop)

        # show reconstruction
        # self.showReconstruction(loop)

    if self.params.gpuFlag:
        self.logger.info("switch to cpu")
        self._move_data_to_cpu()
        self.params.gpuFlag = 0

objectPatchUpdate(objectPatch, DELTA)

Todo add docstring :param objectPatch: :param DELTA: :return:

Source code in PtyLab/Engines/ePIE.py

def objectPatchUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray):
    """
    Todo add docstring
    :param objectPatch:
    :param DELTA:
    :return:
    """
    # find out which array module to use, numpy or cupy (or other...)
    xp = getArrayModule(objectPatch)

    frac = self.reconstruction.probe.conj() / xp.max(
        xp.sum(xp.abs(self.reconstruction.probe) ** 2, axis=(0, 1, 2, 3))
    )
    return objectPatch + self.betaObject * xp.sum(
        frac * DELTA, axis=(0, 2, 3), keepdims=True
    )

probeUpdate(objectPatch, DELTA)

Todo add docstring :param objectPatch: :param DELTA: :return:

Source code in PtyLab/Engines/ePIE.py

def probeUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray):
    """
    Todo add docstring
    :param objectPatch:
    :param DELTA:
    :return:
    """
    # find out which array module to use, numpy or cupy (or other...)
    xp = getArrayModule(objectPatch)
    frac = objectPatch.conj() / xp.max(
        xp.sum(xp.abs(objectPatch) ** 2, axis=(0, 1, 2, 3))
    )
    r = self.reconstruction.probe + self.betaProbe * xp.sum(
        frac * DELTA, axis=(0, 1, 3), keepdims=True
    )
    return r

ePIE_TV

ePIE_TV

Bases: BaseEngine

Source code in PtyLab/Engines/ePIE_TV.py

class ePIE_TV(BaseEngine):

    def __init__(
        self,
        reconstruction: Reconstruction,
        experimentalData: ExperimentalData,
        params: Params,
        monitor: Monitor,
    ):
        # This contains reconstruction parameters that are specific to the reconstruction
        # but not necessarily to ePIE reconstruction
        super().__init__(reconstruction, experimentalData, params, monitor)
        self.logger = logging.getLogger("ePIE")
        self.logger.info("Sucesfully created ePIE ePIE_engine")
        self.logger.info("Wavelength attribute: %s", self.reconstruction.wavelength)
        self.initializeReconstructionParams()

    def initializeReconstructionParams(self):
        """
        Set parameters that are specific to the ePIE settings.
        :return:
        """
        self.betaProbe = 0.25
        self.betaObject = 0.25
        self.numIterations = 50

    def reconstruct(self):
        self._prepareReconstruction()

        # actual reconstruction ePIE_engine
        self.pbar = tqdm.trange(
            self.numIterations, desc="ePIE", file=sys.stdout, leave=True
        )
        for loop in self.pbar:
            # set position order
            self.setPositionOrder()
            for positionLoop, positionIndex in enumerate(self.positionIndices):
                # get object patch
                row, col = self.reconstruction.positions[positionIndex]
                sy = slice(row, row + self.reconstruction.Np)
                sx = slice(col, col + self.reconstruction.Np)
                # note that object patch has size of probe array
                objectPatch = self.reconstruction.object[..., sy, sx].copy()

                # make exit surface wave
                self.reconstruction.esw = objectPatch * self.reconstruction.probe

                # propagate to camera, intensityProjection, propagate back to object
                self.intensityProjection(positionIndex)

                # difference term
                DELTA = self.reconstruction.eswUpdate - self.reconstruction.esw
                if loop % 5 == 0:
                    # object update
                    self.reconstruction.object[..., sy, sx] = self.objectPatchUpdate_TV(
                        objectPatch, DELTA
                    )
                else:
                    self.reconstruction.object[..., sy, sx] = self.objectPatchUpdate(
                        objectPatch, DELTA
                    )

                # probe update
                # self.reconstruction.probe = self.probeUpdate(objectPatch, DELTA)
                self.probeUpdate_new(objectPatch, DELTA)

            # get error metric
            self.getErrorMetrics()

            # apply Constraints
            self.applyConstraints(loop)

            # show reconstruction
            self.showReconstruction(loop)

        if self.params.gpuFlag:
            self.logger.info("switch to cpu")
            self._move_data_to_cpu()
            self.params.gpuFlag = 0

    def objectPatchUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray):
        """
        Todo add docstring
        :param objectPatch:
        :param DELTA:
        :return:
        """
        # find out which array module to use, numpy or cupy (or other...)
        xp = getArrayModule(objectPatch)

        frac = self.reconstruction.probe.conj() / xp.max(
            xp.sum(xp.abs(self.reconstruction.probe) ** 2, axis=(0, 1, 2, 3))
        )
        return objectPatch + self.betaObject * xp.sum(
            frac * DELTA, axis=(0, 2, 3), keepdims=True
        )

    def objectPatchUpdate_TV(self, objectPatch: np.ndarray, DELTA: np.ndarray):
        """
        Todo add docstring
        :param objectPatch:
        :param DELTA:
        :return:
        """

        def divergence(f):
            xp = getArrayModule(f[0])
            return xp.gradient(f[0], axis=(4, 5))[0] + xp.gradient(f[1], axis=(4, 5))[1]

        xp = getArrayModule(objectPatch)
        frac = self.reconstruction.probe.conj() / xp.max(
            xp.sum(xp.abs(self.reconstruction.probe) ** 2, axis=(0, 1, 2, 3))
        )

        epsilon = 1e-2
        gradient = xp.gradient(objectPatch, axis=(4, 5))
        # norm = xp.abs(gradient[0] + gradient[1]) ** 2
        norm = (gradient[0] + gradient[1]) ** 2
        temp = [
            gradient[0] / xp.sqrt(norm + epsilon),
            gradient[1] / xp.sqrt(norm + epsilon),
        ]
        TV_update = divergence(temp)
        lam = 5e-4
        return (
            objectPatch
            + self.betaObject * xp.sum(frac * DELTA, axis=(0, 2, 3), keepdims=True)
            + lam * self.betaObject * TV_update
        )

    def probeUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray):
        """
        Todo add docstring
        :param objectPatch:
        :param DELTA:
        :return:
        """
        # find out which array module to use, numpy or cupy (or other...)
        xp = getArrayModule(objectPatch)
        frac = objectPatch.conj() / xp.max(
            xp.sum(xp.abs(objectPatch) ** 2, axis=(0, 1, 2, 3))
        )
        r = self.reconstruction.probe + self.betaProbe * xp.sum(
            frac * DELTA, axis=(0, 1, 3), keepdims=True
        )
        return r

    def probeUpdate_new(self, objectPatch: np.ndarray, DELTA: np.ndarray):
        """
        Todo add docstring
        :param objectPatch:
        :param DELTA:
        :return:
        """
        # find out which array module to use, numpy or cupy (or other...)
        xp = getArrayModule(objectPatch)
        # frac = objectPatch.conj() / xp.max(xp.sum(xp.abs(objectPatch) ** 2, axis=(0,1,2,3)))
        self.reconstruction.probe += self.betaProbe * xp.sum(
            objectPatch.conj()
            / xp.max(xp.sum(xp.abs(objectPatch) ** 2, axis=(0, 1, 2, 3)))
            * DELTA,
            axis=(0, 1, 3),
            keepdims=True,
        )

initializeReconstructionParams()

Set parameters that are specific to the ePIE settings. :return:

Source code in PtyLab/Engines/ePIE_TV.py

def initializeReconstructionParams(self):
    """
    Set parameters that are specific to the ePIE settings.
    :return:
    """
    self.betaProbe = 0.25
    self.betaObject = 0.25
    self.numIterations = 50

objectPatchUpdate(objectPatch, DELTA)

Todo add docstring :param objectPatch: :param DELTA: :return:

Source code in PtyLab/Engines/ePIE_TV.py

def objectPatchUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray):
    """
    Todo add docstring
    :param objectPatch:
    :param DELTA:
    :return:
    """
    # find out which array module to use, numpy or cupy (or other...)
    xp = getArrayModule(objectPatch)

    frac = self.reconstruction.probe.conj() / xp.max(
        xp.sum(xp.abs(self.reconstruction.probe) ** 2, axis=(0, 1, 2, 3))
    )
    return objectPatch + self.betaObject * xp.sum(
        frac * DELTA, axis=(0, 2, 3), keepdims=True
    )

objectPatchUpdate_TV(objectPatch, DELTA)

Todo add docstring :param objectPatch: :param DELTA: :return:

Source code in PtyLab/Engines/ePIE_TV.py

def objectPatchUpdate_TV(self, objectPatch: np.ndarray, DELTA: np.ndarray):
    """
    Todo add docstring
    :param objectPatch:
    :param DELTA:
    :return:
    """

    def divergence(f):
        xp = getArrayModule(f[0])
        return xp.gradient(f[0], axis=(4, 5))[0] + xp.gradient(f[1], axis=(4, 5))[1]

    xp = getArrayModule(objectPatch)
    frac = self.reconstruction.probe.conj() / xp.max(
        xp.sum(xp.abs(self.reconstruction.probe) ** 2, axis=(0, 1, 2, 3))
    )

    epsilon = 1e-2
    gradient = xp.gradient(objectPatch, axis=(4, 5))
    # norm = xp.abs(gradient[0] + gradient[1]) ** 2
    norm = (gradient[0] + gradient[1]) ** 2
    temp = [
        gradient[0] / xp.sqrt(norm + epsilon),
        gradient[1] / xp.sqrt(norm + epsilon),
    ]
    TV_update = divergence(temp)
    lam = 5e-4
    return (
        objectPatch
        + self.betaObject * xp.sum(frac * DELTA, axis=(0, 2, 3), keepdims=True)
        + lam * self.betaObject * TV_update
    )

probeUpdate(objectPatch, DELTA)

Todo add docstring :param objectPatch: :param DELTA: :return:

Source code in PtyLab/Engines/ePIE_TV.py

def probeUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray):
    """
    Todo add docstring
    :param objectPatch:
    :param DELTA:
    :return:
    """
    # find out which array module to use, numpy or cupy (or other...)
    xp = getArrayModule(objectPatch)
    frac = objectPatch.conj() / xp.max(
        xp.sum(xp.abs(objectPatch) ** 2, axis=(0, 1, 2, 3))
    )
    r = self.reconstruction.probe + self.betaProbe * xp.sum(
        frac * DELTA, axis=(0, 1, 3), keepdims=True
    )
    return r

probeUpdate_new(objectPatch, DELTA)

Todo add docstring :param objectPatch: :param DELTA: :return:

Source code in PtyLab/Engines/ePIE_TV.py

def probeUpdate_new(self, objectPatch: np.ndarray, DELTA: np.ndarray):
    """
    Todo add docstring
    :param objectPatch:
    :param DELTA:
    :return:
    """
    # find out which array module to use, numpy or cupy (or other...)
    xp = getArrayModule(objectPatch)
    # frac = objectPatch.conj() / xp.max(xp.sum(xp.abs(objectPatch) ** 2, axis=(0,1,2,3)))
    self.reconstruction.probe += self.betaProbe * xp.sum(
        objectPatch.conj()
        / xp.max(xp.sum(xp.abs(objectPatch) ** 2, axis=(0, 1, 2, 3)))
        * DELTA,
        axis=(0, 1, 3),
        keepdims=True,
    )

mPIE

mPIE

Bases: BaseEngine

Source code in PtyLab/Engines/mPIE.py

class mPIE(BaseEngine):
    def __init__(
        self,
        reconstruction: Reconstruction,
        experimentalData: ExperimentalData,
        params: Params,
        monitor: Monitor,
    ):
        # This contains reconstruction parameters that are specific to the reconstruction
        # but not necessarily to ePIE reconstruction
        super().__init__(reconstruction, experimentalData, params, monitor)
        self.logger = logging.getLogger("mPIE")
        self.logger.info("Sucesfully created mPIE mPIE_engine")
        self.logger.info("Wavelength attribute: %s", self.reconstruction.wavelength)
        # initialize mPIE Params
        self.initializeReconstructionParams()
        self.params.momentumAcceleration = True
        self.name = "mPIE"

    @property
    def keepPatches(self):
        """Wether or not to keep track of the individual object update patches.

        This strongly increases the amount of memory required, only use when absolutely required.

        """
        return hasattr(self, "patches")

    @keepPatches.setter
    def keepPatches(self, keep_them):

        if keep_them:
            self.logger.info("Keeping patches!")
            self.patches = np.zeros(
                (
                    self.experimentalData.ptychogram.shape[0],
                    *self.reconstruction.shape_O,
                ),
                np.complex64,
            )
        else:
            self.logger.info("Not keeping patches")
            if hasattr(self, "patches"):
                del self.patches

    def initializeReconstructionParams(self):
        """
        Set parameters that are specific to the mPIE settings.
        :return:
        """
        # self.eswUpdate = self.reconstruction.esw.copy()
        self.betaProbe = 0.25
        self.betaObject = 0.25
        self.alphaProbe = 0.1  # probe regularization
        self.alphaObject = 0.1  # object regularization
        self.feedbackM = 0.3  # feedback
        self.frictionM = 0.7  # friction
        self.numIterations = 50

        # initialize momentum
        self.reconstruction.initializeObjectMomentum()
        self.reconstruction.initializeProbeMomentum()
        # set object and probe buffers
        self.reconstruction.objectBuffer = self.reconstruction.object.copy()
        self.reconstruction.probeBuffer = self.reconstruction.probe.copy()

        self.reconstruction.probeWindow = np.abs(self.reconstruction.probe)

    def reconstruct(self, experimentalData=None, reconstruction=None, vis_after_each_iteration=None):
        """Reconstruct object. If experimentalData is given, it replaces the current data. Idem for reconstruction."""

        self.changeExperimentalData(experimentalData)
        self.changeOptimizable(reconstruction)

        self._prepareReconstruction()
        # set object and probe buffers, in case object and probe are changed in the _prepareReconstruction() step
        self.reconstruction.objectBuffer = self.reconstruction.object.copy()
        self.reconstruction.probeBuffer = self.reconstruction.probe.copy()
        # actual reconstruction MPIE_engine
        self.pbar = tqdm.trange(
            self.numIterations, desc="mPIE", file=sys.stdout, leave=True
        )
        for loop in self.pbar:
            # set position order
            self.setPositionOrder()
            self.pbar_pos = tqdm.tqdm(
                self.positionIndices, leave=False, desc="ptychogram", file=sys.stdout
            )
            for positionLoop, positionIndex in enumerate(self.pbar_pos):
                # get object patch, stored as self.probe
                # self.reconstruction.make_probe(positionIndex)

                row, col = self.reconstruction.positions[positionIndex]
                sy = slice(row, row + self.reconstruction.Np)
                sx = slice(col, col + self.reconstruction.Np)
                # note that object patch has size of probe array
                objectPatch = self.reconstruction.object[..., sy, sx].copy()

                # make exit surface wave
                self.reconstruction.esw = objectPatch * self.reconstruction.probe

                # propagate to camera, intensityProjection, propagate back to object
                self.intensityProjection(positionIndex)

                # difference term
                DELTA = self.reconstruction.eswUpdate - self.reconstruction.esw
                # self.viewer.layers['update'].data[positionIndex] = abs(DELTA ** 2).get()
                # import pyqtgraph as pg
                # pg.QtGui.QGuiApplication.processEvents()

                # object update
                if (
                    self.params.objectTVregSwitch
                    and loop % self.params.objectTVfreq == 0
                ):
                    object_patch = self.objectPatchUpdate_TV(objectPatch, DELTA)
                else:
                    object_patch = self.objectPatchUpdate(objectPatch, DELTA)

                if self.keepPatches:
                    self.patches[positionIndex, ..., sy, sx] = asNumpyArray(
                        abs(object_patch) ** 2
                    )
                else:
                    self.reconstruction.object[..., sy, sx] = object_patch

                # probe update
                weight = 1
                if self.params.weigh_probe_updates_by_intensity:
                    weight = self.experimentalData.relative_intensity(positionIndex)
                    # print(f'for position {positionIndex}, using weight {weight}')

                self.reconstruction.probe = self.probeUpdate(objectPatch, DELTA, weight)
                # self.reconstruction.push_probe_update(self.reconstruction.probe, positionIndex, self.experimentalData.ptychogram.shape[0])

                if self.params.positionCorrectionSwitch:
                    shifter = self.positionCorrection(
                        objectPatch, positionIndex, sy, sx
                    )
                    # self.pbar_pos.write(f'Corr: {shifter[0]*1e6:.2f} um x {shifter[1]*1e6:.2f} um')

                # momentum updates
                if np.random.rand(1) > 0.95:
                    self.objectMomentumUpdate()
                    self.probeMomentumUpdate()
                # yield positionLoop, positionIndex

            # get error metric
            self.getErrorMetrics()
            # yield 1,1

            # apply Constraints
            self.applyConstraints(loop)
            # yield 1, 1

            # show reconstruction
            self.showReconstruction(loop)

            if callable(vis_after_each_iteration):
                vis_after_each_iteration(loop, self.reconstruction)

        if self.params.gpuFlag:
            self.logger.info("switch to cpu")
            self._move_data_to_cpu()
            self.params.gpuFlag = 0

            # todo clearMemory implementation

    def objectMomentumUpdate(self):
        """
        momentum update object, save updated objectMomentum and objectBuffer.
        :return:
        """
        gradient = self.reconstruction.objectBuffer - self.reconstruction.object
        self.reconstruction.objectMomentum = (
            gradient + self.frictionM * self.reconstruction.objectMomentum
        )
        self.reconstruction.object = (
            self.reconstruction.object
            - self.feedbackM * self.reconstruction.objectMomentum
        )
        self.reconstruction.objectBuffer = self.reconstruction.object.copy()

    def probeMomentumUpdate(self):
        """
        momentum update probe, save updated probeMomentum and probeBuffer.
        :return:
        """
        gradient = self.reconstruction.probeBuffer - self.reconstruction.probe
        self.reconstruction.probeMomentum = (
            gradient + self.frictionM * self.reconstruction.probeMomentum
        )
        self.reconstruction.probe = (
            self.reconstruction.probe
            - self.feedbackM * self.reconstruction.probeMomentum
        )
        self.reconstruction.probeBuffer = self.reconstruction.probe.copy()

    def objectPatchUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray):
        """
        Todo add docstring
        :param objectPatch:
        :param DELTA:
        :return:
        """
        # find out which array module to use, numpy or cupy (or other...)
        xp = getArrayModule(objectPatch)
        absP2 = xp.abs(self.reconstruction.probe) ** 2
        Pmax = xp.max(xp.sum(absP2, axis=(0, 1, 2, 3)), axis=(-1, -2))
        if self.experimentalData.operationMode == "FPM":
            frac = (
                abs(self.reconstruction.probe)
                / Pmax
                * self.reconstruction.probe.conj()
                / (self.alphaObject * Pmax + (1 - self.alphaObject) * absP2)
            )
        else:
            frac = self.reconstruction.probe.conj() / (
                self.alphaObject * Pmax + (1 - self.alphaObject) * absP2
            )

        return objectPatch + self.betaObject * xp.sum(
            frac * DELTA, axis=2, keepdims=True
        )

    def probeUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray, weight: float):
        """
        Todo add docstring
        :param objectPatch:
        :param DELTA:
        :return:
        """
        # find out which array module to use, numpy or cupy (or other...)
        xp = getArrayModule(objectPatch)
        absO2 = xp.abs(objectPatch) ** 2
        Omax = xp.max(xp.sum(absO2, axis=(0, 1, 2, 3)), axis=(-1, -2))
        frac = objectPatch.conj() / (
            self.alphaProbe * Omax + (1 - self.alphaProbe) * absO2
        )
        r = self.reconstruction.probe + weight * self.betaProbe * xp.sum(
            frac * DELTA, axis=1, keepdims=True
        )
        return r

keepPatches property writable

Wether or not to keep track of the individual object update patches.

This strongly increases the amount of memory required, only use when absolutely required.

initializeReconstructionParams()

Set parameters that are specific to the mPIE settings. :return:

Source code in PtyLab/Engines/mPIE.py

def initializeReconstructionParams(self):
    """
    Set parameters that are specific to the mPIE settings.
    :return:
    """
    # self.eswUpdate = self.reconstruction.esw.copy()
    self.betaProbe = 0.25
    self.betaObject = 0.25
    self.alphaProbe = 0.1  # probe regularization
    self.alphaObject = 0.1  # object regularization
    self.feedbackM = 0.3  # feedback
    self.frictionM = 0.7  # friction
    self.numIterations = 50

    # initialize momentum
    self.reconstruction.initializeObjectMomentum()
    self.reconstruction.initializeProbeMomentum()
    # set object and probe buffers
    self.reconstruction.objectBuffer = self.reconstruction.object.copy()
    self.reconstruction.probeBuffer = self.reconstruction.probe.copy()

    self.reconstruction.probeWindow = np.abs(self.reconstruction.probe)

reconstruct(experimentalData=None, reconstruction=None, vis_after_each_iteration=None)

Reconstruct object. If experimentalData is given, it replaces the current data. Idem for reconstruction.

Source code in PtyLab/Engines/mPIE.py

def reconstruct(self, experimentalData=None, reconstruction=None, vis_after_each_iteration=None):
    """Reconstruct object. If experimentalData is given, it replaces the current data. Idem for reconstruction."""

    self.changeExperimentalData(experimentalData)
    self.changeOptimizable(reconstruction)

    self._prepareReconstruction()
    # set object and probe buffers, in case object and probe are changed in the _prepareReconstruction() step
    self.reconstruction.objectBuffer = self.reconstruction.object.copy()
    self.reconstruction.probeBuffer = self.reconstruction.probe.copy()
    # actual reconstruction MPIE_engine
    self.pbar = tqdm.trange(
        self.numIterations, desc="mPIE", file=sys.stdout, leave=True
    )
    for loop in self.pbar:
        # set position order
        self.setPositionOrder()
        self.pbar_pos = tqdm.tqdm(
            self.positionIndices, leave=False, desc="ptychogram", file=sys.stdout
        )
        for positionLoop, positionIndex in enumerate(self.pbar_pos):
            # get object patch, stored as self.probe
            # self.reconstruction.make_probe(positionIndex)

            row, col = self.reconstruction.positions[positionIndex]
            sy = slice(row, row + self.reconstruction.Np)
            sx = slice(col, col + self.reconstruction.Np)
            # note that object patch has size of probe array
            objectPatch = self.reconstruction.object[..., sy, sx].copy()

            # make exit surface wave
            self.reconstruction.esw = objectPatch * self.reconstruction.probe

            # propagate to camera, intensityProjection, propagate back to object
            self.intensityProjection(positionIndex)

            # difference term
            DELTA = self.reconstruction.eswUpdate - self.reconstruction.esw
            # self.viewer.layers['update'].data[positionIndex] = abs(DELTA ** 2).get()
            # import pyqtgraph as pg
            # pg.QtGui.QGuiApplication.processEvents()

            # object update
            if (
                self.params.objectTVregSwitch
                and loop % self.params.objectTVfreq == 0
            ):
                object_patch = self.objectPatchUpdate_TV(objectPatch, DELTA)
            else:
                object_patch = self.objectPatchUpdate(objectPatch, DELTA)

            if self.keepPatches:
                self.patches[positionIndex, ..., sy, sx] = asNumpyArray(
                    abs(object_patch) ** 2
                )
            else:
                self.reconstruction.object[..., sy, sx] = object_patch

            # probe update
            weight = 1
            if self.params.weigh_probe_updates_by_intensity:
                weight = self.experimentalData.relative_intensity(positionIndex)
                # print(f'for position {positionIndex}, using weight {weight}')

            self.reconstruction.probe = self.probeUpdate(objectPatch, DELTA, weight)
            # self.reconstruction.push_probe_update(self.reconstruction.probe, positionIndex, self.experimentalData.ptychogram.shape[0])

            if self.params.positionCorrectionSwitch:
                shifter = self.positionCorrection(
                    objectPatch, positionIndex, sy, sx
                )
                # self.pbar_pos.write(f'Corr: {shifter[0]*1e6:.2f} um x {shifter[1]*1e6:.2f} um')

            # momentum updates
            if np.random.rand(1) > 0.95:
                self.objectMomentumUpdate()
                self.probeMomentumUpdate()
            # yield positionLoop, positionIndex

        # get error metric
        self.getErrorMetrics()
        # yield 1,1

        # apply Constraints
        self.applyConstraints(loop)
        # yield 1, 1

        # show reconstruction
        self.showReconstruction(loop)

        if callable(vis_after_each_iteration):
            vis_after_each_iteration(loop, self.reconstruction)

    if self.params.gpuFlag:
        self.logger.info("switch to cpu")
        self._move_data_to_cpu()
        self.params.gpuFlag = 0

objectMomentumUpdate()

momentum update object, save updated objectMomentum and objectBuffer. :return:

Source code in PtyLab/Engines/mPIE.py

def objectMomentumUpdate(self):
    """
    momentum update object, save updated objectMomentum and objectBuffer.
    :return:
    """
    gradient = self.reconstruction.objectBuffer - self.reconstruction.object
    self.reconstruction.objectMomentum = (
        gradient + self.frictionM * self.reconstruction.objectMomentum
    )
    self.reconstruction.object = (
        self.reconstruction.object
        - self.feedbackM * self.reconstruction.objectMomentum
    )
    self.reconstruction.objectBuffer = self.reconstruction.object.copy()

probeMomentumUpdate()

momentum update probe, save updated probeMomentum and probeBuffer. :return:

Source code in PtyLab/Engines/mPIE.py

def probeMomentumUpdate(self):
    """
    momentum update probe, save updated probeMomentum and probeBuffer.
    :return:
    """
    gradient = self.reconstruction.probeBuffer - self.reconstruction.probe
    self.reconstruction.probeMomentum = (
        gradient + self.frictionM * self.reconstruction.probeMomentum
    )
    self.reconstruction.probe = (
        self.reconstruction.probe
        - self.feedbackM * self.reconstruction.probeMomentum
    )
    self.reconstruction.probeBuffer = self.reconstruction.probe.copy()

objectPatchUpdate(objectPatch, DELTA)

Todo add docstring :param objectPatch: :param DELTA: :return:

Source code in PtyLab/Engines/mPIE.py

def objectPatchUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray):
    """
    Todo add docstring
    :param objectPatch:
    :param DELTA:
    :return:
    """
    # find out which array module to use, numpy or cupy (or other...)
    xp = getArrayModule(objectPatch)
    absP2 = xp.abs(self.reconstruction.probe) ** 2
    Pmax = xp.max(xp.sum(absP2, axis=(0, 1, 2, 3)), axis=(-1, -2))
    if self.experimentalData.operationMode == "FPM":
        frac = (
            abs(self.reconstruction.probe)
            / Pmax
            * self.reconstruction.probe.conj()
            / (self.alphaObject * Pmax + (1 - self.alphaObject) * absP2)
        )
    else:
        frac = self.reconstruction.probe.conj() / (
            self.alphaObject * Pmax + (1 - self.alphaObject) * absP2
        )

    return objectPatch + self.betaObject * xp.sum(
        frac * DELTA, axis=2, keepdims=True
    )

probeUpdate(objectPatch, DELTA, weight)

Todo add docstring :param objectPatch: :param DELTA: :return:

Source code in PtyLab/Engines/mPIE.py

def probeUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray, weight: float):
    """
    Todo add docstring
    :param objectPatch:
    :param DELTA:
    :return:
    """
    # find out which array module to use, numpy or cupy (or other...)
    xp = getArrayModule(objectPatch)
    absO2 = xp.abs(objectPatch) ** 2
    Omax = xp.max(xp.sum(absO2, axis=(0, 1, 2, 3)), axis=(-1, -2))
    frac = objectPatch.conj() / (
        self.alphaProbe * Omax + (1 - self.alphaProbe) * absO2
    )
    r = self.reconstruction.probe + weight * self.betaProbe * xp.sum(
        frac * DELTA, axis=1, keepdims=True
    )
    return r

mPIE_tv

mPIE_tv

Bases: BaseEngine

Source code in PtyLab/Engines/mPIE_tv.py

class mPIE_tv(BaseEngine):

    def __init__(
        self,
        reconstruction: Reconstruction,
        experimentalData: ExperimentalData,
        params: Params,
        monitor: Monitor,
    ):
        # This contains reconstruction parameters that are specific to the reconstruction
        # but not necessarily to ePIE reconstruction
        super().__init__(reconstruction, experimentalData, params, monitor)
        self.logger = logging.getLogger("mPIE")
        self.logger.info("Sucesfully created mPIE mPIE_engine")
        self.logger.info("Wavelength attribute: %s", self.reconstruction.wavelength)
        # initialize mPIE Params
        self.initializeReconstructionParams()
        self.params.momentumAcceleration = True

    def initializeReconstructionParams(self):
        """
        Set parameters that are specific to the mPIE settings.
        :return:
        """
        # self.eswUpdate = self.reconstruction.esw.copy()
        self.betaProbe = 0.25
        self.betaObject = 0.25
        self.alphaProbe = 0.1  # probe regularization
        self.alphaObject = 0.1  # object regularization
        self.feedbackM = 0.3  # feedback
        self.frictionM = 0.7  # friction
        self.numIterations = 50

        # initialize momentum
        self.reconstruction.initializeObjectMomentum()
        self.reconstruction.initializeProbeMomentum()
        # set object and probe buffers
        self.reconstruction.objectBuffer = self.reconstruction.object.copy()
        self.reconstruction.probeBuffer = self.reconstruction.probe.copy()

        self.reconstruction.probeWindow = np.abs(self.reconstruction.probe)

    def reconstruct(self):
        self._prepareReconstruction()

        # actual reconstruction MPIE_engine
        self.pbar = tqdm.trange(
            self.numIterations, desc="mPIE", file=sys.stdout, leave=True
        )
        for loop in self.pbar:
            # set position order
            self.setPositionOrder()

            for positionLoop, positionIndex in enumerate(self.positionIndices):
                # get object patch
                row, col = self.reconstruction.positions[positionIndex]
                sy = slice(row, row + self.reconstruction.Np)
                sx = slice(col, col + self.reconstruction.Np)
                # note that object patch has size of probe array
                objectPatch = self.reconstruction.object[..., sy, sx].copy()

                # make exit surface wave
                self.reconstruction.esw = objectPatch * self.reconstruction.probe

                # propagate to camera, intensityProjection, propagate back to object
                self.intensityProjection(positionIndex)

                # difference term
                DELTA = self.reconstruction.eswUpdate - self.reconstruction.esw

                tv_freq = 1
                if loop % tv_freq == 0:
                    # object update
                    self.reconstruction.object[..., sy, sx] = self.objectPatchUpdate_TV(
                        objectPatch, DELTA
                    )
                else:
                    self.reconstruction.object[..., sy, sx] = self.objectPatchUpdate(
                        objectPatch, DELTA
                    )

                # probe update
                self.reconstruction.probe = self.probeUpdate(objectPatch, DELTA)

                # momentum updates
                if np.random.rand(1) > 0.95:
                    self.objectMomentumUpdate()
                    self.probeMomentumUpdate()

            # get error metric
            self.getErrorMetrics()

            # apply Constraints
            self.applyConstraints(loop)

            # show reconstruction
            self.showReconstruction(loop)

        if self.params.gpuFlag:
            self.logger.info("switch to cpu")
            self._move_data_to_cpu()
            self.params.gpuFlag = 0

            # todo clearMemory implementation

    def objectMomentumUpdate(self):
        """
        momentum update object, save updated objectMomentum and objectBuffer.
        :return:
        """
        gradient = self.reconstruction.objectBuffer - self.reconstruction.object
        self.reconstruction.objectMomentum = (
            gradient + self.frictionM * self.reconstruction.objectMomentum
        )
        self.reconstruction.object = (
            self.reconstruction.object
            - self.feedbackM * self.reconstruction.objectMomentum
        )
        self.reconstruction.objectBuffer = self.reconstruction.object.copy()

    def probeMomentumUpdate(self):
        """
        momentum update probe, save updated probeMomentum and probeBuffer.
        :return:
        """
        gradient = self.reconstruction.probeBuffer - self.reconstruction.probe
        self.reconstruction.probeMomentum = (
            gradient + self.frictionM * self.reconstruction.probeMomentum
        )
        self.reconstruction.probe = (
            self.reconstruction.probe
            - self.feedbackM * self.reconstruction.probeMomentum
        )
        self.reconstruction.probeBuffer = self.reconstruction.probe.copy()

    def objectPatchUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray):
        """
        Todo add docstring
        :param objectPatch:
        :param DELTA:
        :return:
        """
        # find out which array module to use, numpy or cupy (or other...)
        xp = getArrayModule(objectPatch)
        absP2 = xp.abs(self.reconstruction.probe) ** 2
        Pmax = xp.max(xp.sum(absP2, axis=(0, 1, 2, 3)), axis=(-1, -2))
        if self.experimentalData.operationMode == "FPM":
            frac = (
                abs(self.reconstruction.probe)
                / Pmax
                * self.reconstruction.probe.conj()
                / (self.alphaObject * Pmax + (1 - self.alphaObject) * absP2)
            )
        else:
            frac = self.reconstruction.probe.conj() / (
                self.alphaObject * Pmax + (1 - self.alphaObject) * absP2
            )
        return objectPatch + self.betaObject * xp.sum(
            frac * DELTA, axis=2, keepdims=True
        )

    def objectPatchUpdate_TV(self, objectPatch: np.ndarray, DELTA: np.ndarray):
        """
        Todo add docstring
        :param objectPatch:
        :param DELTA:
        :return:
        """

        def divergence(f):
            xp = getArrayModule(f[0])
            return xp.gradient(f[0], axis=(4, 5))[0] + xp.gradient(f[1], axis=(4, 5))[1]

        xp = getArrayModule(objectPatch)
        frac = self.reconstruction.probe.conj() / xp.max(
            xp.sum(xp.abs(self.reconstruction.probe) ** 2, axis=(0, 1, 2, 3))
        )

        epsilon = 1e-2
        gradient = xp.gradient(objectPatch, axis=(4, 5))
        # norm = xp.abs(gradient[0] + gradient[1]) ** 2
        norm = (gradient[0] + gradient[1]) ** 2
        temp = [
            gradient[0] / xp.sqrt(norm + epsilon),
            gradient[1] / xp.sqrt(norm + epsilon),
        ]
        TV_update = divergence(temp)
        """
        plt.figure()
        plt.imshow(np.abs(TV_update.get()[0, 0, 0, 0, :, :]))
        plt.figure()
        plt.imshow(np.angle(TV_update.get()[0, 0, 0, 0, :, :]))
        plt.show()
        """
        lam = self.params.TV_lam
        return (
            objectPatch
            + self.betaObject * xp.sum(frac * DELTA, axis=(0, 2, 3), keepdims=True)
            + lam * self.betaObject * TV_update
        )

    def probeUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray):
        """
        Todo add docstring
        :param objectPatch:
        :param DELTA:
        :return:
        """
        # find out which array module to use, numpy or cupy (or other...)
        xp = getArrayModule(objectPatch)
        absO2 = xp.abs(objectPatch) ** 2
        Omax = xp.max(xp.sum(absO2, axis=(0, 1, 2, 3)), axis=(-1, -2))
        frac = objectPatch.conj() / (
            self.alphaProbe * Omax + (1 - self.alphaProbe) * absO2
        )
        r = self.reconstruction.probe + self.betaProbe * xp.sum(
            frac * DELTA, axis=1, keepdims=True
        )
        return r

initializeReconstructionParams()

Set parameters that are specific to the mPIE settings. :return:

Source code in PtyLab/Engines/mPIE_tv.py

def initializeReconstructionParams(self):
    """
    Set parameters that are specific to the mPIE settings.
    :return:
    """
    # self.eswUpdate = self.reconstruction.esw.copy()
    self.betaProbe = 0.25
    self.betaObject = 0.25
    self.alphaProbe = 0.1  # probe regularization
    self.alphaObject = 0.1  # object regularization
    self.feedbackM = 0.3  # feedback
    self.frictionM = 0.7  # friction
    self.numIterations = 50

    # initialize momentum
    self.reconstruction.initializeObjectMomentum()
    self.reconstruction.initializeProbeMomentum()
    # set object and probe buffers
    self.reconstruction.objectBuffer = self.reconstruction.object.copy()
    self.reconstruction.probeBuffer = self.reconstruction.probe.copy()

    self.reconstruction.probeWindow = np.abs(self.reconstruction.probe)

objectMomentumUpdate()

momentum update object, save updated objectMomentum and objectBuffer. :return:

Source code in PtyLab/Engines/mPIE_tv.py

def objectMomentumUpdate(self):
    """
    momentum update object, save updated objectMomentum and objectBuffer.
    :return:
    """
    gradient = self.reconstruction.objectBuffer - self.reconstruction.object
    self.reconstruction.objectMomentum = (
        gradient + self.frictionM * self.reconstruction.objectMomentum
    )
    self.reconstruction.object = (
        self.reconstruction.object
        - self.feedbackM * self.reconstruction.objectMomentum
    )
    self.reconstruction.objectBuffer = self.reconstruction.object.copy()

probeMomentumUpdate()

momentum update probe, save updated probeMomentum and probeBuffer. :return:

Source code in PtyLab/Engines/mPIE_tv.py

def probeMomentumUpdate(self):
    """
    momentum update probe, save updated probeMomentum and probeBuffer.
    :return:
    """
    gradient = self.reconstruction.probeBuffer - self.reconstruction.probe
    self.reconstruction.probeMomentum = (
        gradient + self.frictionM * self.reconstruction.probeMomentum
    )
    self.reconstruction.probe = (
        self.reconstruction.probe
        - self.feedbackM * self.reconstruction.probeMomentum
    )
    self.reconstruction.probeBuffer = self.reconstruction.probe.copy()

objectPatchUpdate(objectPatch, DELTA)

Todo add docstring :param objectPatch: :param DELTA: :return:

Source code in PtyLab/Engines/mPIE_tv.py

def objectPatchUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray):
    """
    Todo add docstring
    :param objectPatch:
    :param DELTA:
    :return:
    """
    # find out which array module to use, numpy or cupy (or other...)
    xp = getArrayModule(objectPatch)
    absP2 = xp.abs(self.reconstruction.probe) ** 2
    Pmax = xp.max(xp.sum(absP2, axis=(0, 1, 2, 3)), axis=(-1, -2))
    if self.experimentalData.operationMode == "FPM":
        frac = (
            abs(self.reconstruction.probe)
            / Pmax
            * self.reconstruction.probe.conj()
            / (self.alphaObject * Pmax + (1 - self.alphaObject) * absP2)
        )
    else:
        frac = self.reconstruction.probe.conj() / (
            self.alphaObject * Pmax + (1 - self.alphaObject) * absP2
        )
    return objectPatch + self.betaObject * xp.sum(
        frac * DELTA, axis=2, keepdims=True
    )

objectPatchUpdate_TV(objectPatch, DELTA)

Todo add docstring :param objectPatch: :param DELTA: :return:

Source code in PtyLab/Engines/mPIE_tv.py

def objectPatchUpdate_TV(self, objectPatch: np.ndarray, DELTA: np.ndarray):
    """
    Todo add docstring
    :param objectPatch:
    :param DELTA:
    :return:
    """

    def divergence(f):
        xp = getArrayModule(f[0])
        return xp.gradient(f[0], axis=(4, 5))[0] + xp.gradient(f[1], axis=(4, 5))[1]

    xp = getArrayModule(objectPatch)
    frac = self.reconstruction.probe.conj() / xp.max(
        xp.sum(xp.abs(self.reconstruction.probe) ** 2, axis=(0, 1, 2, 3))
    )

    epsilon = 1e-2
    gradient = xp.gradient(objectPatch, axis=(4, 5))
    # norm = xp.abs(gradient[0] + gradient[1]) ** 2
    norm = (gradient[0] + gradient[1]) ** 2
    temp = [
        gradient[0] / xp.sqrt(norm + epsilon),
        gradient[1] / xp.sqrt(norm + epsilon),
    ]
    TV_update = divergence(temp)
    """
    plt.figure()
    plt.imshow(np.abs(TV_update.get()[0, 0, 0, 0, :, :]))
    plt.figure()
    plt.imshow(np.angle(TV_update.get()[0, 0, 0, 0, :, :]))
    plt.show()
    """
    lam = self.params.TV_lam
    return (
        objectPatch
        + self.betaObject * xp.sum(frac * DELTA, axis=(0, 2, 3), keepdims=True)
        + lam * self.betaObject * TV_update
    )

probeUpdate(objectPatch, DELTA)

Todo add docstring :param objectPatch: :param DELTA: :return:

Source code in PtyLab/Engines/mPIE_tv.py

def probeUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray):
    """
    Todo add docstring
    :param objectPatch:
    :param DELTA:
    :return:
    """
    # find out which array module to use, numpy or cupy (or other...)
    xp = getArrayModule(objectPatch)
    absO2 = xp.abs(objectPatch) ** 2
    Omax = xp.max(xp.sum(absO2, axis=(0, 1, 2, 3)), axis=(-1, -2))
    frac = objectPatch.conj() / (
        self.alphaProbe * Omax + (1 - self.alphaProbe) * absO2
    )
    r = self.reconstruction.probe + self.betaProbe * xp.sum(
        frac * DELTA, axis=1, keepdims=True
    )
    return r

mqNewton

mqNewton

Bases: BaseEngine

Source code in PtyLab/Engines/mqNewton.py

class mqNewton(BaseEngine):
    def __init__(
        self,
        reconstruction: Reconstruction,
        experimentalData: ExperimentalData,
        params: Params,
        monitor: Monitor,
    ):
        # This contains reconstruction parameters that are specific to the reconstruction
        # but not necessarily to ePIE reconstruction
        super().__init__(reconstruction, experimentalData, params, monitor)
        self.logger = logging.getLogger("mqNewton")
        self.logger.info("Sucesfully created momentum accelerated qNewton mqNewton")

        self.logger.info("Wavelength attribute: %s", self.reconstruction.wavelength)
        self.initializeReconstructionParams()
        # initialize momentum
        self.reconstruction.initializeObjectMomentum()
        self.reconstruction.initializeProbeMomentum()
        # set object and probe buffers
        self.reconstruction.objectBuffer = self.reconstruction.object.copy()
        self.reconstruction.probeBuffer = self.reconstruction.probe.copy()
        self.params.momentumAcceleration = True

    def initializeReconstructionParams(self):
        """
        Set parameters that are specific to the qNewton settings.
        :return:
        """
        self.betaProbe = 1
        self.betaObject = 1
        self.regObject = 1
        self.regProbe = 1
        self.beta1 = 0.5
        self.beta2 = 0.5
        self.betaProbe_m = 0.25
        self.betaObject_m = 0.25
        self.feedbackM = 0.3  # feedback
        self.frictionM = 0.7  # friction
        self.momentum_method = "ADAM"  # which optimizer to use for momentum updates
        self.numIterations = 50

    def initializeAdaptiveMomentum(self):
        self.momentum_engine = getattr(mqNewton, self.momentum_method)
        print("Momentum Engines implemented: momentum, ADAM, NADAM")
        print("Momentum mqNewton used: {}".format(self.momentum_method))
        if self.momentum_method in ["ADAM", "NADAM"]:
            # 2nd order momentum terms
            self.reconstruction.objectMomentum_v = (
                self.reconstruction.objectMomentum.copy()
            )
            self.reconstruction.probeMomentum_v = (
                self.reconstruction.probeMomentum.copy()
            )

    def reconstruct(self, experimentalData: ExperimentalData = None):
        if experimentalData is not None:
            self.experimentalData = experimentalData
            self.reconstruction.data = experimentalData
        self._prepareReconstruction()
        self.initializeAdaptiveMomentum()

        self.pbar = tqdm.trange(
            self.numIterations, desc="mqNewton", file=sys.stdout, leave=True
        )
        for loop in self.pbar:
            # set position order
            self.setPositionOrder()

            for positionLoop, positionIndex in enumerate(self.positionIndices):
                # get object patch
                row, col = self.reconstruction.positions[positionIndex]
                sy = slice(row, row + self.reconstruction.Np)
                sx = slice(col, col + self.reconstruction.Np)
                # note that object patch has size of probe array
                objectPatch = self.reconstruction.object[..., sy, sx].copy()

                # make exit surface wave
                self.reconstruction.esw = objectPatch * self.reconstruction.probe

                # propagate to camera, intensityProjection, propagate back to object
                self.intensityProjection(positionIndex)

                # difference term
                DELTA = self.reconstruction.eswUpdate - self.reconstruction.esw

                # object update
                self.reconstruction.object[..., sy, sx] = self.objectPatchUpdate(
                    objectPatch, DELTA
                )

                # probe update
                self.reconstruction.probe = self.probeUpdate(objectPatch, DELTA)

                # momentum updates
                self.objectMomentumUpdate(loop)
                self.probeMomentumUpdate(loop)

                if self.params.positionCorrectionSwitch:
                    self.positionCorrection(objectPatch, positionIndex, sy, sx)

            # get error metric
            self.getErrorMetrics()

            # apply Constraints
            self.applyConstraints(loop)

            # show reconstruction
            self.showReconstruction(loop)

        if self.params.gpuFlag:
            self.logger.info("switch to cpu")
            self._move_data_to_cpu()
            self.params.gpuFlag = 0
            # todo clearMemory implementation

    def ADAM(self, grad, mt, vt, itr):
        xp = getArrayModule(grad)
        beta1_scale = 1 - self.beta1**itr
        beta2_scale = 1 - self.beta2**itr
        mt = self.beta1 * mt + (1 - self.beta1) * grad
        vt = (
            self.beta2 * vt
            + (1 - self.beta2) * xp.linalg.norm(grad.flatten().squeeze(), 2) ** 2
        )
        m_hat = mt / beta1_scale
        v_hat = vt / beta2_scale
        return m_hat / (v_hat**0.5 + 1e-8), mt, vt

    def NADAM(self, grad, mt, vt, itr):
        """
        NADAM optimizer uses adaptive momentum updates (ADAM) with Nesterov
        momentum acceleration
        :return:
        """
        xp = getArrayModule(grad)

        beta1_scale = 1 - self.beta1**itr
        beta2_scale = 1 - self.beta2**itr

        norm_sq = xp.linalg.norm(grad.flatten(), 2) ** 2
        mt = self.beta1 * mt + (1 - self.beta1) * grad
        vt = self.beta2 * vt + (1 - self.beta2) * norm_sq
        m_hat = mt / beta1_scale
        v_hat = vt / beta2_scale
        update = (self.beta1 * m_hat + grad * (1 - self.beta1) / beta1_scale) / (
            v_hat**0.5 + 1e-8
        )
        return update, mt, vt

    def momentum(self, grad, mt, vt, itr):
        """
        standard momentum update
        :return:
        """
        mt = grad + self.frictionM * mt
        return mt, mt, vt

    def objectMomentumUpdate(self, loop):
        """
        momentum update object, save updated objectMomentum and objectBuffer.
        :return:
        """
        gradient = self.reconstruction.objectBuffer - self.reconstruction.object
        (
            update,
            self.reconstruction.objectMomentum,
            self.reconstruction.objectMomentum_v,
        ) = self.momentum_engine(
            self,
            gradient,
            self.reconstruction.objectMomentum,
            self.reconstruction.objectMomentum_v,
            loop + 1,
        )

        self.reconstruction.object -= self.betaObject_m * update
        self.reconstruction.objectBuffer = self.reconstruction.object.copy()

    def probeMomentumUpdate(self, loop):
        """
        momentum update probe, save updated probeMomentum and probeBuffer.
        :return:
        """
        gradient = self.reconstruction.probeBuffer - self.reconstruction.probe
        (
            update,
            self.reconstruction.probeMomentum,
            self.reconstruction.probeMomentum_v,
        ) = self.momentum_engine(
            self,
            gradient,
            self.reconstruction.probeMomentum,
            self.reconstruction.probeMomentum_v,
            loop + 1,
        )

        self.reconstruction.probe -= self.betaProbe_m * update
        self.reconstruction.probeBuffer = self.reconstruction.probe.copy()

    def objectPatchUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray):
        xp = getArrayModule(objectPatch)
        Pmax = xp.max(xp.sum(xp.abs(self.reconstruction.probe), axis=(0, 1, 2, 3)))
        frac = (
            xp.abs(self.reconstruction.probe)
            / Pmax
            * self.reconstruction.probe.conj()
            / (xp.abs(self.reconstruction.probe) ** 2 + self.regObject)
        )
        return objectPatch + self.betaObject * xp.sum(
            frac * DELTA, axis=(0, 2, 3), keepdims=True
        )

    def probeUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray):
        xp = getArrayModule(objectPatch)
        Omax = xp.max(xp.sum(xp.abs(self.reconstruction.object), axis=(0, 1, 2, 3)))
        frac = (
            xp.abs(objectPatch)
            / Omax
            * objectPatch.conj()
            / (xp.abs(objectPatch) ** 2 + self.regProbe)
        )
        r = self.reconstruction.probe + self.betaProbe * xp.sum(
            frac * DELTA, axis=(0, 1, 3), keepdims=True
        )
        return r

initializeReconstructionParams()

Set parameters that are specific to the qNewton settings. :return:

Source code in PtyLab/Engines/mqNewton.py

def initializeReconstructionParams(self):
    """
    Set parameters that are specific to the qNewton settings.
    :return:
    """
    self.betaProbe = 1
    self.betaObject = 1
    self.regObject = 1
    self.regProbe = 1
    self.beta1 = 0.5
    self.beta2 = 0.5
    self.betaProbe_m = 0.25
    self.betaObject_m = 0.25
    self.feedbackM = 0.3  # feedback
    self.frictionM = 0.7  # friction
    self.momentum_method = "ADAM"  # which optimizer to use for momentum updates
    self.numIterations = 50

NADAM(grad, mt, vt, itr)

NADAM optimizer uses adaptive momentum updates (ADAM) with Nesterov momentum acceleration :return:

Source code in PtyLab/Engines/mqNewton.py

def NADAM(self, grad, mt, vt, itr):
    """
    NADAM optimizer uses adaptive momentum updates (ADAM) with Nesterov
    momentum acceleration
    :return:
    """
    xp = getArrayModule(grad)

    beta1_scale = 1 - self.beta1**itr
    beta2_scale = 1 - self.beta2**itr

    norm_sq = xp.linalg.norm(grad.flatten(), 2) ** 2
    mt = self.beta1 * mt + (1 - self.beta1) * grad
    vt = self.beta2 * vt + (1 - self.beta2) * norm_sq
    m_hat = mt / beta1_scale
    v_hat = vt / beta2_scale
    update = (self.beta1 * m_hat + grad * (1 - self.beta1) / beta1_scale) / (
        v_hat**0.5 + 1e-8
    )
    return update, mt, vt

momentum(grad, mt, vt, itr)

standard momentum update :return:

Source code in PtyLab/Engines/mqNewton.py

def momentum(self, grad, mt, vt, itr):
    """
    standard momentum update
    :return:
    """
    mt = grad + self.frictionM * mt
    return mt, mt, vt

objectMomentumUpdate(loop)

momentum update object, save updated objectMomentum and objectBuffer. :return:

Source code in PtyLab/Engines/mqNewton.py

def objectMomentumUpdate(self, loop):
    """
    momentum update object, save updated objectMomentum and objectBuffer.
    :return:
    """
    gradient = self.reconstruction.objectBuffer - self.reconstruction.object
    (
        update,
        self.reconstruction.objectMomentum,
        self.reconstruction.objectMomentum_v,
    ) = self.momentum_engine(
        self,
        gradient,
        self.reconstruction.objectMomentum,
        self.reconstruction.objectMomentum_v,
        loop + 1,
    )

    self.reconstruction.object -= self.betaObject_m * update
    self.reconstruction.objectBuffer = self.reconstruction.object.copy()

probeMomentumUpdate(loop)

momentum update probe, save updated probeMomentum and probeBuffer. :return:

Source code in PtyLab/Engines/mqNewton.py

def probeMomentumUpdate(self, loop):
    """
    momentum update probe, save updated probeMomentum and probeBuffer.
    :return:
    """
    gradient = self.reconstruction.probeBuffer - self.reconstruction.probe
    (
        update,
        self.reconstruction.probeMomentum,
        self.reconstruction.probeMomentum_v,
    ) = self.momentum_engine(
        self,
        gradient,
        self.reconstruction.probeMomentum,
        self.reconstruction.probeMomentum_v,
        loop + 1,
    )

    self.reconstruction.probe -= self.betaProbe_m * update
    self.reconstruction.probeBuffer = self.reconstruction.probe.copy()

multiPIE

multiPIE

Bases: BaseEngine

Source code in PtyLab/Engines/multiPIE.py

class multiPIE(BaseEngine):
    def __init__(
        self,
        reconstruction: Reconstruction,
        experimentalData: ExperimentalData,
        params: Params,
        monitor: Monitor,
    ):
        # This contains reconstruction parameters that are specific to the reconstruction
        # but not necessarily to ePIE reconstruction
        super().__init__(reconstruction, experimentalData, params, monitor)
        self.logger = logging.getLogger("multiPIE")
        self.logger.info("Sucesfully created multiPIE multiPIE_engine")
        self.logger.info("Wavelength attribute: %s", self.reconstruction.wavelength)
        # initialize multiPIE Params
        self.initializeReconstructionParams()
        self.params.momentumAcceleration = True

    def initializeReconstructionParams(self):
        """
        Set parameters that are specific to the multiPIE settings.
        :return:
        """
        # self.eswUpdate = self.reconstruction.esw.copy()
        self.betaProbe = 0.25
        self.betaObject = 0.25
        self.alphaProbe = 0.1  # probe regularization
        self.alphaObject = 0.1  # object regularization
        self.betaM = 0.3  # feedback
        self.stepM = 0.7  # friction
        # self.reconstruction.probeWindow = np.abs(self.reconstruction.probe)
        self.numIterations = 50

        # initialize momentum
        self.reconstruction.initializeObjectMomentum()
        self.reconstruction.initializeProbeMomentum()
        # set object and probe buffers
        self.reconstruction.objectBuffer = self.reconstruction.object.copy()
        self.reconstruction.probeBuffer = self.reconstruction.probe.copy()

    def reconstruct(self):
        self._prepareReconstruction()

        self.pbar = tqdm.trange(
            self.numIterations, desc="multiPIE", file=sys.stdout, leave=True
        )

        for loop in self.pbar:
            # set position order
            self.setPositionOrder()

            for positionLoop, positionIndex in enumerate(self.positionIndices):
                # get object patch
                row, col = self.reconstruction.positions[positionIndex]
                sy = slice(row, row + self.reconstruction.Np)
                sx = slice(col, col + self.reconstruction.Np)
                # note that object patch has size of probe array
                objectPatch = self.reconstruction.object[..., sy, sx].copy()

                # make exit surface wave
                self.reconstruction.esw = objectPatch * self.reconstruction.probe

                # propagate to camera, intensityProjection, propagate back to object
                self.intensityProjection(positionIndex)

                # difference term
                DELTA = self.reconstruction.eswUpdate - self.reconstruction.esw

                # object update
                self.reconstruction.object[..., sy, sx] = self.objectPatchUpdate(
                    objectPatch, DELTA
                )

                # probe update
                self.reconstruction.probe = self.probeUpdate(objectPatch, DELTA)

                # momentum updates
                if np.random.rand(1) > 0.95:
                    self.objectMomentumUpdate()
                    self.probeMomentumUpdate()

            # get error metric
            self.getErrorMetrics()

            # apply Constraints
            self.applyConstraints(loop)

            # show reconstruction
            self.showReconstruction(loop)

            # todo clearMemory implementation

        if self.params.gpuFlag:
            self.logger.info("switch to cpu")
            self._move_data_to_cpu()
            self.params.gpuFlag = 0

    def objectMomentumUpdate(self):
        """
        momentum update object, save updated objectMomentum and objectBuffer.
        :return:
        """
        gradient = self.reconstruction.objectBuffer - self.reconstruction.object
        self.reconstruction.objectMomentum = (
            gradient + self.stepM * self.reconstruction.objectMomentum
        )
        self.reconstruction.object = (
            self.reconstruction.object - self.betaM * self.reconstruction.objectMomentum
        )
        self.reconstruction.objectBuffer = self.reconstruction.object.copy()

    def probeMomentumUpdate(self):
        """
        momentum update probe, save updated probeMomentum and probeBuffer.
        :return:
        """
        gradient = self.reconstruction.probeBuffer - self.reconstruction.probe
        self.reconstruction.probeMomentum = (
            gradient + self.stepM * self.reconstruction.probeMomentum
        )
        self.reconstruction.probe = (
            self.reconstruction.probe - self.betaM * self.reconstruction.probeMomentum
        )
        self.reconstruction.probeBuffer = self.reconstruction.probe.copy()

    def objectPatchUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray):
        """
        Todo add docstring
        :param objectPatch:
        :param DELTA:
        :return:
        """
        # find out which array module to use, numpy or cupy (or other...)
        # xp = getArrayModule(objectPatch)
        # absP2 = xp.abs(self.reconstruction.probe[0]) ** 2
        # Pmax = xp.max(xp.sum(absP2, axis=(0, 1, 2)), axis=(-1, -2))
        # if self.experimentalData.operationMode == 'FPM':
        #     frac = abs(self.reconstruction.probe) / Pmax * \
        #            self.reconstruction.probe[0].conj() / (self.alphaObject * Pmax + (1 - self.alphaObject) * absP2)
        # else:
        #     frac = self.reconstruction.probe[0].conj() / (self.alphaObject * Pmax + (1 - self.alphaObject) * absP2)
        # return objectPatch + self.betaObject * frac * DELTA
        xp = getArrayModule(objectPatch)
        absP2 = xp.abs(self.reconstruction.probe) ** 2
        Pmax = xp.max(xp.sum(absP2, axis=(0, 1, 2, 3)), axis=(-1, -2))
        if self.experimentalData.operationMode == "FPM":
            frac = (
                abs(self.reconstruction.probe)
                / Pmax
                * self.reconstruction.probe.conj()
                / (self.alphaObject * Pmax + (1 - self.alphaObject) * absP2)
            )
        else:
            frac = self.reconstruction.probe.conj() / (
                self.alphaObject * Pmax + (1 - self.alphaObject) * absP2
            )
        return objectPatch + self.betaObject * xp.sum(
            frac * DELTA, axis=2, keepdims=True
        )

    def probeUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray):
        """
        Todo add docstring
        :param objectPatch:
        :param DELTA:
        :return:
        """
        # find out which array module to use, numpy or cupy (or other...)
        xp = getArrayModule(objectPatch)
        absO2 = xp.abs(objectPatch) ** 2
        Omax = xp.max(xp.sum(absO2, axis=(0, 1, 2, 3)), axis=(-1, -2))
        frac = objectPatch.conj() / (
            self.alphaProbe * Omax + (1 - self.alphaProbe) * absO2
        )
        r = self.reconstruction.probe + self.betaProbe * xp.sum(
            frac * DELTA, axis=(0, 1), keepdims=True
        )
        return r

initializeReconstructionParams()

Set parameters that are specific to the multiPIE settings. :return:

Source code in PtyLab/Engines/multiPIE.py

def initializeReconstructionParams(self):
    """
    Set parameters that are specific to the multiPIE settings.
    :return:
    """
    # self.eswUpdate = self.reconstruction.esw.copy()
    self.betaProbe = 0.25
    self.betaObject = 0.25
    self.alphaProbe = 0.1  # probe regularization
    self.alphaObject = 0.1  # object regularization
    self.betaM = 0.3  # feedback
    self.stepM = 0.7  # friction
    # self.reconstruction.probeWindow = np.abs(self.reconstruction.probe)
    self.numIterations = 50

    # initialize momentum
    self.reconstruction.initializeObjectMomentum()
    self.reconstruction.initializeProbeMomentum()
    # set object and probe buffers
    self.reconstruction.objectBuffer = self.reconstruction.object.copy()
    self.reconstruction.probeBuffer = self.reconstruction.probe.copy()

objectMomentumUpdate()

momentum update object, save updated objectMomentum and objectBuffer. :return:

Source code in PtyLab/Engines/multiPIE.py

def objectMomentumUpdate(self):
    """
    momentum update object, save updated objectMomentum and objectBuffer.
    :return:
    """
    gradient = self.reconstruction.objectBuffer - self.reconstruction.object
    self.reconstruction.objectMomentum = (
        gradient + self.stepM * self.reconstruction.objectMomentum
    )
    self.reconstruction.object = (
        self.reconstruction.object - self.betaM * self.reconstruction.objectMomentum
    )
    self.reconstruction.objectBuffer = self.reconstruction.object.copy()

probeMomentumUpdate()

momentum update probe, save updated probeMomentum and probeBuffer. :return:

Source code in PtyLab/Engines/multiPIE.py

def probeMomentumUpdate(self):
    """
    momentum update probe, save updated probeMomentum and probeBuffer.
    :return:
    """
    gradient = self.reconstruction.probeBuffer - self.reconstruction.probe
    self.reconstruction.probeMomentum = (
        gradient + self.stepM * self.reconstruction.probeMomentum
    )
    self.reconstruction.probe = (
        self.reconstruction.probe - self.betaM * self.reconstruction.probeMomentum
    )
    self.reconstruction.probeBuffer = self.reconstruction.probe.copy()

objectPatchUpdate(objectPatch, DELTA)

Todo add docstring :param objectPatch: :param DELTA: :return:

Source code in PtyLab/Engines/multiPIE.py

def objectPatchUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray):
    """
    Todo add docstring
    :param objectPatch:
    :param DELTA:
    :return:
    """
    # find out which array module to use, numpy or cupy (or other...)
    # xp = getArrayModule(objectPatch)
    # absP2 = xp.abs(self.reconstruction.probe[0]) ** 2
    # Pmax = xp.max(xp.sum(absP2, axis=(0, 1, 2)), axis=(-1, -2))
    # if self.experimentalData.operationMode == 'FPM':
    #     frac = abs(self.reconstruction.probe) / Pmax * \
    #            self.reconstruction.probe[0].conj() / (self.alphaObject * Pmax + (1 - self.alphaObject) * absP2)
    # else:
    #     frac = self.reconstruction.probe[0].conj() / (self.alphaObject * Pmax + (1 - self.alphaObject) * absP2)
    # return objectPatch + self.betaObject * frac * DELTA
    xp = getArrayModule(objectPatch)
    absP2 = xp.abs(self.reconstruction.probe) ** 2
    Pmax = xp.max(xp.sum(absP2, axis=(0, 1, 2, 3)), axis=(-1, -2))
    if self.experimentalData.operationMode == "FPM":
        frac = (
            abs(self.reconstruction.probe)
            / Pmax
            * self.reconstruction.probe.conj()
            / (self.alphaObject * Pmax + (1 - self.alphaObject) * absP2)
        )
    else:
        frac = self.reconstruction.probe.conj() / (
            self.alphaObject * Pmax + (1 - self.alphaObject) * absP2
        )
    return objectPatch + self.betaObject * xp.sum(
        frac * DELTA, axis=2, keepdims=True
    )

probeUpdate(objectPatch, DELTA)

Todo add docstring :param objectPatch: :param DELTA: :return:

Source code in PtyLab/Engines/multiPIE.py

def probeUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray):
    """
    Todo add docstring
    :param objectPatch:
    :param DELTA:
    :return:
    """
    # find out which array module to use, numpy or cupy (or other...)
    xp = getArrayModule(objectPatch)
    absO2 = xp.abs(objectPatch) ** 2
    Omax = xp.max(xp.sum(absO2, axis=(0, 1, 2, 3)), axis=(-1, -2))
    frac = objectPatch.conj() / (
        self.alphaProbe * Omax + (1 - self.alphaProbe) * absO2
    )
    r = self.reconstruction.probe + self.betaProbe * xp.sum(
        frac * DELTA, axis=(0, 1), keepdims=True
    )
    return r

pcPIE

pcPIE

Bases: BaseEngine

Source code in PtyLab/Engines/pcPIE.py

class pcPIE(BaseEngine):
    def __init__(
        self,
        reconstruction: Reconstruction,
        experimentalData: ExperimentalData,
        params: Params,
        monitor: Monitor,
    ):
        # This contains reconstruction parameters that are specific to the reconstruction
        # but not necessarily to ePIE reconstruction
        super().__init__(reconstruction, experimentalData, params, monitor)
        self.logger = logging.getLogger("pcPIE")
        self.logger.info("Successfully created pcPIE pcPIE_engine")
        self.logger.info("Wavelength attribute: %s", self.reconstruction.wavelength)
        # initialize pcPIE Params
        self.initializeReconstructionParams()
        # initialize momentum
        self.reconstruction.initializeObjectMomentum()
        self.reconstruction.initializeProbeMomentum()
        # set object and probe buffers
        self.reconstruction.objectBuffer = self.reconstruction.object.copy()
        self.reconstruction.probeBuffer = self.reconstruction.probe.copy()

        self.params.momentumAcceleration = True

    def initializeReconstructionParams(self):
        """
        Set parameters that are specific to the pcPIE settings.
        :return:
        """
        # these are same as mPIE
        # self.eswUpdate = self.reconstruction.esw.copy()
        self.betaProbe = 0.25
        self.betaObject = 0.25
        self.alphaProbe = 0.1  # probe regularization
        self.alphaObject = 0.1  # object regularization
        self.betaM = 0.3  # feedback
        self.stepM = 0.7  # friction
        # self.probeWindow = np.abs(self.reconstruction.probe)
        self.numIterations = 50

    def reconstruct(self):
        self._prepareReconstruction()

        # actual reconstruction ePIE_engine

        self.pbar = tqdm.trange(
            self.numIterations, desc="pcPIE", file=sys.stdout, leave=True
        )  # in order to change description to the tqdm progress bar
        for loop in self.pbar:
            # set position order
            self.setPositionOrder()

            for positionLoop, positionIndex in enumerate(self.positionIndices):
                # get object patch
                row, col = self.reconstruction.positions[positionIndex]
                sy = slice(row, row + self.reconstruction.Np)
                sx = slice(col, col + self.reconstruction.Np)
                # note that object patch has size of probe array
                objectPatch = self.reconstruction.object[..., sy, sx].copy()

                # make exit surface wave
                self.reconstruction.esw = objectPatch * self.reconstruction.probe

                # propagate to camera, intensityProjection, propagate back to object
                self.intensityProjection(positionIndex)

                # difference term
                DELTA = self.reconstruction.eswUpdate - self.reconstruction.esw

                # object update
                self.reconstruction.object[..., sy, sx] = self.objectPatchUpdate(
                    objectPatch, DELTA
                )

                # probe update
                self.reconstruction.probe = self.probeUpdate(objectPatch, DELTA)
                if self.params.positionCorrectionSwitch:
                    self.positionCorrection(objectPatch, positionIndex, sy, sx)

                # momentum updates
                if np.random.rand(1) > 0.95:
                    self.objectMomentumUpdate()
                    self.probeMomentumUpdate()

            # get error metric
            self.getErrorMetrics()

            # apply Constraints
            self.applyConstraints(loop)

            # show reconstruction
            self.showReconstruction(loop)

            # todo clearMemory implementation

        if self.params.gpuFlag:
            self.logger.info("switch to cpu")
            self._move_data_to_cpu()
            self.params.gpuFlag = 0

    def objectMomentumUpdate(self):
        """
        momentum update object, save updated objectMomentum and objectBuffer.
        :return:
        """
        gradient = self.reconstruction.objectBuffer - self.reconstruction.object
        self.reconstruction.objectMomentum = (
            gradient + self.stepM * self.reconstruction.objectMomentum
        )
        self.reconstruction.object = (
            self.reconstruction.object - self.betaM * self.reconstruction.objectMomentum
        )
        self.reconstruction.objectBuffer = self.reconstruction.object.copy()

    def probeMomentumUpdate(self):
        """
        momentum update probe, save updated probeMomentum and probeBuffer.
        :return:
        """
        gradient = self.reconstruction.probeBuffer - self.reconstruction.probe
        self.reconstruction.probeMomentum = (
            gradient + self.stepM * self.reconstruction.probeMomentum
        )
        self.reconstruction.probe = (
            self.reconstruction.probe - self.betaM * self.reconstruction.probeMomentum
        )
        self.reconstruction.probeBuffer = self.reconstruction.probe.copy()

    def objectPatchUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray):
        """
        Todo add docstring
        :param objectPatch:
        :param DELTA:
        :return:
        """
        # find out which array module to use, numpy or cupy (or other...)
        xp = getArrayModule(objectPatch)
        absP2 = xp.abs(self.reconstruction.probe) ** 2
        Pmax = xp.max(xp.sum(absP2, axis=(0, 1, 2, 3)), axis=(-1, -2))
        if self.experimentalData.operationMode == "FPM":
            frac = (
                abs(self.reconstruction.probe)
                / Pmax
                * self.reconstruction.probe.conj()
                / (self.alphaObject * Pmax + (1 - self.alphaObject) * absP2)
            )
        else:
            frac = self.reconstruction.probe.conj() / (
                self.alphaObject * Pmax + (1 - self.alphaObject) * absP2
            )
        return objectPatch + self.betaObject * xp.sum(
            frac * DELTA, axis=2, keepdims=True
        )

    def probeUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray):
        """
        Todo add docstring
        :param objectPatch:
        :param DELTA:
        :return:
        """
        # find out which array module to use, numpy or cupy (or other...)
        xp = getArrayModule(objectPatch)
        absO2 = xp.abs(objectPatch) ** 2
        Omax = xp.max(xp.sum(absO2, axis=(0, 1, 2, 3)), axis=(-1, -2))
        frac = objectPatch.conj() / (
            self.alphaProbe * Omax + (1 - self.alphaProbe) * absO2
        )
        r = self.reconstruction.probe + self.betaProbe * xp.sum(
            frac * DELTA, axis=1, keepdims=True
        )
        return r

initializeReconstructionParams()

Set parameters that are specific to the pcPIE settings. :return:

Source code in PtyLab/Engines/pcPIE.py

def initializeReconstructionParams(self):
    """
    Set parameters that are specific to the pcPIE settings.
    :return:
    """
    # these are same as mPIE
    # self.eswUpdate = self.reconstruction.esw.copy()
    self.betaProbe = 0.25
    self.betaObject = 0.25
    self.alphaProbe = 0.1  # probe regularization
    self.alphaObject = 0.1  # object regularization
    self.betaM = 0.3  # feedback
    self.stepM = 0.7  # friction
    # self.probeWindow = np.abs(self.reconstruction.probe)
    self.numIterations = 50

objectMomentumUpdate()

momentum update object, save updated objectMomentum and objectBuffer. :return:

Source code in PtyLab/Engines/pcPIE.py

def objectMomentumUpdate(self):
    """
    momentum update object, save updated objectMomentum and objectBuffer.
    :return:
    """
    gradient = self.reconstruction.objectBuffer - self.reconstruction.object
    self.reconstruction.objectMomentum = (
        gradient + self.stepM * self.reconstruction.objectMomentum
    )
    self.reconstruction.object = (
        self.reconstruction.object - self.betaM * self.reconstruction.objectMomentum
    )
    self.reconstruction.objectBuffer = self.reconstruction.object.copy()

probeMomentumUpdate()

momentum update probe, save updated probeMomentum and probeBuffer. :return:

Source code in PtyLab/Engines/pcPIE.py

def probeMomentumUpdate(self):
    """
    momentum update probe, save updated probeMomentum and probeBuffer.
    :return:
    """
    gradient = self.reconstruction.probeBuffer - self.reconstruction.probe
    self.reconstruction.probeMomentum = (
        gradient + self.stepM * self.reconstruction.probeMomentum
    )
    self.reconstruction.probe = (
        self.reconstruction.probe - self.betaM * self.reconstruction.probeMomentum
    )
    self.reconstruction.probeBuffer = self.reconstruction.probe.copy()

objectPatchUpdate(objectPatch, DELTA)

Todo add docstring :param objectPatch: :param DELTA: :return:

Source code in PtyLab/Engines/pcPIE.py

def objectPatchUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray):
    """
    Todo add docstring
    :param objectPatch:
    :param DELTA:
    :return:
    """
    # find out which array module to use, numpy or cupy (or other...)
    xp = getArrayModule(objectPatch)
    absP2 = xp.abs(self.reconstruction.probe) ** 2
    Pmax = xp.max(xp.sum(absP2, axis=(0, 1, 2, 3)), axis=(-1, -2))
    if self.experimentalData.operationMode == "FPM":
        frac = (
            abs(self.reconstruction.probe)
            / Pmax
            * self.reconstruction.probe.conj()
            / (self.alphaObject * Pmax + (1 - self.alphaObject) * absP2)
        )
    else:
        frac = self.reconstruction.probe.conj() / (
            self.alphaObject * Pmax + (1 - self.alphaObject) * absP2
        )
    return objectPatch + self.betaObject * xp.sum(
        frac * DELTA, axis=2, keepdims=True
    )

probeUpdate(objectPatch, DELTA)

Todo add docstring :param objectPatch: :param DELTA: :return:

Source code in PtyLab/Engines/pcPIE.py

def probeUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray):
    """
    Todo add docstring
    :param objectPatch:
    :param DELTA:
    :return:
    """
    # find out which array module to use, numpy or cupy (or other...)
    xp = getArrayModule(objectPatch)
    absO2 = xp.abs(objectPatch) ** 2
    Omax = xp.max(xp.sum(absO2, axis=(0, 1, 2, 3)), axis=(-1, -2))
    frac = objectPatch.conj() / (
        self.alphaProbe * Omax + (1 - self.alphaProbe) * absO2
    )
    r = self.reconstruction.probe + self.betaProbe * xp.sum(
        frac * DELTA, axis=1, keepdims=True
    )
    return r

qNewton

qNewton

Bases: BaseEngine

Source code in PtyLab/Engines/qNewton.py

class qNewton(BaseEngine):
    def __init__(
        self,
        reconstruction: Reconstruction,
        experimentalData: ExperimentalData,
        params: Params,
        monitor: Monitor,
    ):
        # This contains reconstruction parameters that are specific to the reconstruction
        # but not necessarily to ePIE reconstruction
        super().__init__(reconstruction, experimentalData, params, monitor)
        self.logger = logging.getLogger("qNewton")
        self.logger.info("Sucesfully created qNewton qNewton_engine")

        self.logger.info("Wavelength attribute: %s", self.reconstruction.wavelength)
        self.initializeReconstructionParams()

    def initializeReconstructionParams(self):
        """
        Set parameters that are specific to the qNewton settings.
        :return:
        """
        self.betaProbe = 1
        self.betaObject = 1
        self.regObject = 1
        self.regProbe = 1
        self.numIterations = 50

    def reconstruct(self, experimentalData: ExperimentalData = None):
        if experimentalData is not None:
            self.reconstruction.data = experimentalData
            self.experimentalData = experimentalData
        self._prepareReconstruction()

        self.pbar = tqdm.trange(
            self.numIterations, desc="qNewton", file=sys.stdout, leave=True
        )
        for loop in self.pbar:
            # set position order
            self.setPositionOrder()

            for positionLoop, positionIndex in enumerate(self.positionIndices):
                # get object patch
                row, col = self.reconstruction.positions[positionIndex]
                sy = slice(row, row + self.reconstruction.Np)
                sx = slice(col, col + self.reconstruction.Np)
                # note that object patch has size of probe array
                objectPatch = self.reconstruction.object[..., sy, sx].copy()

                # make exit surface wave
                self.reconstruction.esw = objectPatch * self.reconstruction.probe

                # propagate to camera, intensityProjection, propagate back to object
                self.intensityProjection(positionIndex)

                # difference term
                DELTA = self.reconstruction.eswUpdate - self.reconstruction.esw

                # object update
                self.reconstruction.object[..., sy, sx] = self.objectPatchUpdate(
                    objectPatch, DELTA
                )

                # probe update
                self.reconstruction.probe = self.probeUpdate(objectPatch, DELTA)

            # get error metric
            self.getErrorMetrics()

            # apply Constraints
            self.applyConstraints(loop)

            # show reconstruction
            self.showReconstruction(loop)

        if self.params.gpuFlag:
            self.logger.info("switch to cpu")
            self._move_data_to_cpu()
            self.params.gpuFlag = 0

            # todo clearMemory implementation

    def objectPatchUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray):
        """
        Temporary barebones update
        """
        xp = getArrayModule(objectPatch)
        Pmax = xp.max(xp.sum(xp.abs(self.reconstruction.probe), axis=(0, 1, 2, 3)))
        frac = (
            xp.abs(self.reconstruction.probe)
            / Pmax
            * self.reconstruction.probe.conj()
            / (xp.abs(self.reconstruction.probe) ** 2 + self.regObject)
        )
        return objectPatch + self.betaObject * xp.sum(
            frac * DELTA, axis=(0, 2, 3), keepdims=True
        )

    def probeUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray):
        """
        Temporary barebones update

        """
        xp = getArrayModule(objectPatch)
        Omax = xp.max(xp.sum(xp.abs(self.reconstruction.object), axis=(0, 1, 2, 3)))
        frac = (
            xp.abs(objectPatch)
            / Omax
            * objectPatch.conj()
            / (xp.abs(objectPatch) ** 2 + self.regProbe)
        )
        r = self.reconstruction.probe + self.betaProbe * xp.sum(
            frac * DELTA, axis=(0, 1, 3), keepdims=True
        )
        return r

initializeReconstructionParams()

Set parameters that are specific to the qNewton settings. :return:

Source code in PtyLab/Engines/qNewton.py

def initializeReconstructionParams(self):
    """
    Set parameters that are specific to the qNewton settings.
    :return:
    """
    self.betaProbe = 1
    self.betaObject = 1
    self.regObject = 1
    self.regProbe = 1
    self.numIterations = 50

objectPatchUpdate(objectPatch, DELTA)

Temporary barebones update

Source code in PtyLab/Engines/qNewton.py

def objectPatchUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray):
    """
    Temporary barebones update
    """
    xp = getArrayModule(objectPatch)
    Pmax = xp.max(xp.sum(xp.abs(self.reconstruction.probe), axis=(0, 1, 2, 3)))
    frac = (
        xp.abs(self.reconstruction.probe)
        / Pmax
        * self.reconstruction.probe.conj()
        / (xp.abs(self.reconstruction.probe) ** 2 + self.regObject)
    )
    return objectPatch + self.betaObject * xp.sum(
        frac * DELTA, axis=(0, 2, 3), keepdims=True
    )

probeUpdate(objectPatch, DELTA)

Temporary barebones update

Source code in PtyLab/Engines/qNewton.py

def probeUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray):
    """
    Temporary barebones update

    """
    xp = getArrayModule(objectPatch)
    Omax = xp.max(xp.sum(xp.abs(self.reconstruction.object), axis=(0, 1, 2, 3)))
    frac = (
        xp.abs(objectPatch)
        / Omax
        * objectPatch.conj()
        / (xp.abs(objectPatch) ** 2 + self.regProbe)
    )
    r = self.reconstruction.probe + self.betaProbe * xp.sum(
        frac * DELTA, axis=(0, 1, 3), keepdims=True
    )
    return r

zPIE

zPIE

Bases: BaseEngine

Source code in PtyLab/Engines/zPIE.py

class zPIE(BaseEngine):
    def __init__(
        self,
        reconstruction: Reconstruction,
        experimentalData: ExperimentalData,
        params: Params,
        monitor: Monitor,
    ):
        # This contains reconstruction parameters that are specific to the reconstruction
        # but not necessarily to ePIE reconstruction
        super().__init__(reconstruction, experimentalData, params, monitor)
        self.logger = logging.getLogger("zPIE")
        self.logger.info("Sucesfully created zPIE zPIE_engine")
        self.logger.info("Wavelength attribute: %s", self.reconstruction.wavelength)
        self.initializeReconstructionParams()
        self.name = "zPIE"

    def initializeReconstructionParams(self):
        """
        Set parameters that are specific to the ePIE settings.
        :return:
        """
        self.betaProbe = 0.25
        self.betaObject = 0.25
        self.numIterations = 50
        self.DoF = self.reconstruction.DoF
        self.zPIEgradientStepSize = 100  # gradient step size for axial position correction (typical range [1, 100])
        self.zPIEfriction = 0.7
        self.focusObject = True
        self.zMomentun = 0

    def show_defocus(self, viewer=None, scanrange_times_dof=1000, N_points=10):
        z = np.linspace(-1, 1, N_points) * scanrange_times_dof * self.reconstruction.DoF

        from PtyLab.Operators.Operators import aspw

        reconstruction = self.reconstruction
        defocii = np.abs(
            np.array(
                [
                    aspw(
                        reconstruction.object,
                        dz,
                        reconstruction.wavelength,
                        reconstruction.Lo,
                    )[0]
                    for dz in z
                ]
            )
            ** 2
        )

        if viewer is None:
            # currently a hacky way for this, these napari implementations must
            # later be moved to an optional sub-package.
            try:
                import napari

                viewer = napari.Viewer()
            except ImportError:
                msg = "Install napari to access this `NapariMonitor` implementation"
                raise ImportError(msg)

        viewer.add_image(defocii)

    def reconstruct(self, experimentalData=None, reconstruction=None):
        self.changeExperimentalData(experimentalData)
        self.changeOptimizable(reconstruction)
        self._prepareReconstruction()

        ###################################### actual reconstruction zPIE_engine #######################################

        xp = getArrayModule(self.reconstruction.object)
        if not hasattr(self.reconstruction, "zHistory"):
            self.reconstruction.zHistory = []

        # preallocate grids
        if self.params.propagatorType == "ASP":
            n = self.reconstruction.Np * 1
        else:
            n = 2 * self.reconstruction.Np

        if not self.focusObject:
            n = self.reconstruction.Np

        X, Y = xp.meshgrid(xp.arange(-n // 2, n // 2), xp.arange(-n // 2, n // 2))
        w = xp.exp(-((xp.sqrt(X**2 + Y**2) / self.reconstruction.Np) ** 4))

        self.pbar = tqdm.trange(
            self.numIterations, desc="zPIE", file=sys.stdout, leave=True
        )  # in order to change description to the tqdm progress bar
        for loop in self.pbar:
            # set position order
            self.setPositionOrder()
            imProps = []

            # get positions
            if loop == 1:
                zNew = self.reconstruction.zo.copy()
            else:
                d = 10

                dz = np.linspace(-1, 1, 11) * d * self.DoF
                self.dz = dz

                merit = []
                # todo, mixed states implementation, check if more need to be put on GPU to speed up
                for k in np.arange(len(dz)):
                    imProp = None
                    if self.focusObject:
                        roi = slice(
                            self.reconstruction.No // 2 - n // 2,
                            self.reconstruction.No // 2 + n // 2,
                        )
                        imProp, _ = aspw(
                            u=xp.squeeze(self.reconstruction.object[..., roi, roi]),
                            z=dz[k],
                            wavelength=self.reconstruction.wavelength,
                            L=self.reconstruction.dxo * n,
                            bandlimit=False,
                        )
                    else:
                        if self.reconstruction.nlambda == 1:
                            imProp, _ = aspw(
                                u=xp.squeeze(self.reconstruction.probe[..., :, :]),
                                z=dz[k],
                                wavelength=self.reconstruction.wavelength,
                                L=self.reconstruction.Lp,
                            )
                        else:
                            nlambda = self.reconstruction.nlambda // 2
                            imProp, _ = aspw(
                                xp.squeeze(
                                    self.reconstruction.probe[nlambda, ..., :, :]
                                ),
                                dz[k],
                                self.reconstruction.spectralDensity[nlambda],
                                self.reconstruction.Lp,
                            )
                    imProps.append(imProp.get())
                    # TV approach
                    aleph = 1e-2
                    gradx = xp.roll(imProp, -1, axis=-1) - xp.roll(imProp, 1, axis=-1)
                    grady = xp.roll(imProp, -1, axis=-2) - xp.roll(imProp, 1, axis=-2)
                    merit.append(
                        xp.sum(xp.sqrt(abs(gradx) ** 2 + abs(grady) ** 2 + aleph))
                    )
                    # take a tiny break, we may overask the GPU
                    # yield 0, 0

                merit = xp.array(merit)
                if not hasattr(self.reconstruction, "TV_history"):
                    self.reconstruction.TV_history = []

                self.reconstruction.TV_history.append(
                    float(merit[len(merit) // 2].get())
                )
                if xp is not np:
                    merit = merit.get()
                feedback = np.sum(dz * merit) / np.sum(
                    merit
                )  # at optimal z, feedback term becomes 0

                print("Step size: ", feedback)
                self.zMomentun = (
                    self.zPIEfriction * self.zMomentun
                    + self.zPIEgradientStepSize * feedback
                )
                zNew = self.reconstruction.zo + self.zMomentun

                # asdlkcmasldk

            self.reconstruction.zHistory.append(self.reconstruction.zo)

            # print updated z
            self.pbar.set_description(
                "zPIE: update z = %.3f mm (dz = %.1f um)"
                % (self.reconstruction.zo * 1e3, self.zMomentun * 1e6)
            )

            # reset coordinates
            self.reconstruction.zo = zNew

            # re-sample is automatically done by using @property
            if self.params.propagatorType != "ASP":
                self.reconstruction.dxp = (
                    self.reconstruction.wavelength
                    * self.reconstruction.zo
                    / self.reconstruction.Ld
                )
                # reset propagatorType
                # self.reconstruction.quadraticPhase = xp.array(np.exp(1.j * np.pi / (self.reconstruction.wavelength * self.reconstruction.zo)
                #                                                      * (self.reconstruction.Xp ** 2 + self.reconstruction.Yp ** 2)))
            ##################################################################################################################

            for positionLoop, positionIndex in enumerate(self.positionIndices):
                # print('Starting normal reconstruction loop')
                ### patch1 ###
                # get object patch
                row, col = self.reconstruction.positions[positionIndex]
                sy = slice(row, row + self.reconstruction.Np)
                sx = slice(col, col + self.reconstruction.Np)
                # note that object patch has size of probe array
                objectPatch = self.reconstruction.object[..., sy, sx].copy()

                # make exit surface wave
                self.reconstruction.esw = objectPatch * self.reconstruction.probe

                # propagate to camera, intensityProjection, propagate back to object
                self.intensityProjection(positionIndex)

                # difference term
                DELTA = self.reconstruction.eswUpdate - self.reconstruction.esw

                # object update
                self.reconstruction.object[..., sy, sx] = self.objectPatchUpdate(
                    objectPatch, DELTA
                )

                # probe update
                self.reconstruction.probe = self.probeUpdate(objectPatch, DELTA)
            # yield positionLoop, positionIndex

            # get error metric
            self.getErrorMetrics()

            # apply Constraints
            self.applyConstraints(loop)
            # display it
            # self.showReconstruction(loop)

            self.merit = merit
            self.zNew = zNew
            self.reconstruction.merit = merit
            self.reconstruction.dz = dz

            self.reconstruction.make_alignment_plot(True)
            # show reconstruction
            if False:
                if loop == 0:
                    figure, axes = plt.subplots(
                        1, 3, num=666, squeeze=True, clear=True, figsize=(5, 5)
                    )
                    ax = axes[0]
                    ax_score = axes[1]
                    ax.set_title("Estimated distance (object-camera)")
                    ax.set_xlabel("iteration")
                    ax.set_ylabel("estimated z (mm)")
                    ax.set_xscale("symlog")

                    ax_score.set_title("TV score")
                    ax_score.set_xlabel("Distance [um]")
                    ax_score.set_ylabel("TV")
                    (score_line,) = ax_score.plot(dz * 1e6, merit)
                    (line,) = ax.plot(0, zNew, "o-")
                    plt.tight_layout()
                    plt.show(block=False)

                elif np.mod(loop, self.monitor.figureUpdateFrequency) == 0:
                    idx = np.linspace(
                        0,
                        np.log10(len(self.reconstruction.zHistory) - 1),
                        np.minimum(len(self.reconstruction.zHistory), 100),
                    )
                    idx = np.rint(10**idx).astype("int")

                    line.set_xdata(idx)
                    line.set_ydata(np.array(self.reconstruction.zHistory)[idx] * 1e3)

                    score_line.set_ydata(merit)
                    ax_score.set_ylim(merit.min() - 1, merit.max() + 1)
                    ax.set_xlim(0, np.max(idx))
                    ax.set_ylim(
                        np.min(self.reconstruction.zHistory) * 1e3,
                        np.max(self.reconstruction.zHistory) * 1e3,
                    )

                    figure.canvas.draw()
                    figure.canvas.flush_events()
            self.showReconstruction(loop)

        if self.params.gpuFlag:
            self.logger.info("switch to cpu")
            self._move_data_to_cpu()
            self.params.gpuFlag = 0

    def objectPatchUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray):
        """
        Todo add docstring
        :param objectPatch:
        :param DELTA:
        :return:
        """
        # find out which array module to use, numpy or cupy (or other...)
        xp = getArrayModule(objectPatch)

        frac = self.reconstruction.probe.conj() / xp.max(
            xp.sum(xp.abs(self.reconstruction.probe) ** 2, axis=(0, 1, 2, 3))
        )
        return objectPatch + self.betaObject * xp.sum(
            frac * DELTA, axis=(0, 2, 3), keepdims=True
        )

    def probeUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray):
        """
        Todo add docstring
        :param objectPatch:
        :param DELTA:
        :return:
        """
        # find out which array module to use, numpy or cupy (or other...)
        xp = getArrayModule(objectPatch)
        frac = objectPatch.conj() / xp.max(
            xp.sum(xp.abs(objectPatch) ** 2, axis=(0, 1, 2, 3))
        )
        r = self.reconstruction.probe + self.betaProbe * xp.sum(
            frac * DELTA, axis=(0, 1, 3), keepdims=True
        )
        return r

initializeReconstructionParams()

Set parameters that are specific to the ePIE settings. :return:

Source code in PtyLab/Engines/zPIE.py

def initializeReconstructionParams(self):
    """
    Set parameters that are specific to the ePIE settings.
    :return:
    """
    self.betaProbe = 0.25
    self.betaObject = 0.25
    self.numIterations = 50
    self.DoF = self.reconstruction.DoF
    self.zPIEgradientStepSize = 100  # gradient step size for axial position correction (typical range [1, 100])
    self.zPIEfriction = 0.7
    self.focusObject = True
    self.zMomentun = 0

objectPatchUpdate(objectPatch, DELTA)

Todo add docstring :param objectPatch: :param DELTA: :return:

Source code in PtyLab/Engines/zPIE.py

def objectPatchUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray):
    """
    Todo add docstring
    :param objectPatch:
    :param DELTA:
    :return:
    """
    # find out which array module to use, numpy or cupy (or other...)
    xp = getArrayModule(objectPatch)

    frac = self.reconstruction.probe.conj() / xp.max(
        xp.sum(xp.abs(self.reconstruction.probe) ** 2, axis=(0, 1, 2, 3))
    )
    return objectPatch + self.betaObject * xp.sum(
        frac * DELTA, axis=(0, 2, 3), keepdims=True
    )

probeUpdate(objectPatch, DELTA)

Todo add docstring :param objectPatch: :param DELTA: :return:

Source code in PtyLab/Engines/zPIE.py

def probeUpdate(self, objectPatch: np.ndarray, DELTA: np.ndarray):
    """
    Todo add docstring
    :param objectPatch:
    :param DELTA:
    :return:
    """
    # find out which array module to use, numpy or cupy (or other...)
    xp = getArrayModule(objectPatch)
    frac = objectPatch.conj() / xp.max(
        xp.sum(xp.abs(objectPatch) ** 2, axis=(0, 1, 2, 3))
    )
    r = self.reconstruction.probe + self.betaProbe * xp.sum(
        frac * DELTA, axis=(0, 1, 3), keepdims=True
    )
    return r