Reconstruction

PtyLab.Reconstruction

CalibrationFPM

IlluminationCalibration

Source code in PtyLab/Reconstruction/CalibrationFPM.py

class IlluminationCalibration:
    def __init__(
        self, reconstructor: Reconstruction, experimentalData: ExperimentalData
    ):
        # These statements don't copy any data, they just keep a reference to the object
        self.reconstructor = reconstructor
        self.experimentalData = experimentalData

        # initialize search variables
        self.searchGridSize = round(int(reconstructor.Np * 0.1))
        self.gaussSigma = 3
        self.plot = False
        self.calibrateRadius = False
        self.brightfieldIndices = None
        self.fit_mode = "SimilarityTransform"

        # compute system parameters to be calibrated from the loaded reconstruction
        self.dxp = reconstructor.dxp
        self.No = reconstructor.No
        self.Np = reconstructor.Np
        self.wavelength = reconstructor.wavelength
        self.img_size = reconstructor.Np
        # # inverse calculation of the NA since we don't have the values provided by the user at the moment
        # self.reconstructor.NA = reconstructor.entrancePupilDiameter/2 * reconstructor.wavelength / (reconstructor.dxp**2 * reconstructor.Np)
        self.apertRadiusPixel = (
            self.dxp * self.img_size * self.reconstructor.NA / self.wavelength
        )
        self.apertRadiusPixel_init = self.apertRadiusPixel

        lens_range = np.linspace(-self.img_size / 2, self.img_size / 2, self.img_size)
        xlens, ylens = np.meshgrid(lens_range, lens_range)
        self.CTF = np.abs(np.sqrt(xlens**2 + ylens**2) <= self.apertRadiusPixel)
        self.OTF = np.abs(np.sqrt(xlens**2 + ylens**2) <= 2 * self.apertRadiusPixel)

    def initialize_error_search_space(self):
        self.angleRange_x_init = np.sin(np.mgrid[0:360:10] / 180.0 * np.pi)
        self.angleRange_y_init = np.cos(np.mgrid[0:360:10] / 180.0 * np.pi)
        self.gridSearch_x_init = np.mgrid[
            -self.searchGridSize : self.searchGridSize + 1
        ]
        self.gridSearch_y_init = np.mgrid[
            -self.searchGridSize : self.searchGridSize + 1
        ]
        self.x_range = len(self.gridSearch_x_init)
        self.y_range = len(self.gridSearch_y_init)

    def findBrightfielIndices(self, ptychogram):
        """
        Use a threshold to separate brightfield images from darkfield images.
        Threshold is obtained by using a K-means clustering algorithm

        Parameters
        ----------
        ptychogram : 3D array
            image array containing experimental images.

        Returns
        -------
        brightfieldIndices : 1D array
            bool index array where 1s represent brightfield images.

        """
        # sum all the images together into a single 2D image
        intensities = np.sum(ptychogram, (1, 2))

        # kmeans clustering method to find 2 clusters
        cluster = np.array([np.arange(ptychogram.shape[0]), intensities]).T
        kmeans = KMeans(n_clusters=2).fit(cluster)
        clusters = kmeans.labels_
        # one cluster is brightfield other darkfield
        cluster1 = clusters == 0
        cluster2 = clusters == 1

        # if len(cluster1[cluster1]) < len(cluster2[cluster2]):
        if np.mean(ptychogram[cluster1]) > np.mean(ptychogram[cluster2]):
            brightfieldIndices = cluster1
        else:
            brightfieldIndices = cluster2
        return brightfieldIndices

    def convertToFourierSpace(self, ptychogram):
        """
        Convert brightfield image stack to the Fourier domain and also
        increase contrast based on all data analysis.

        Parameters
        ----------
        ptychogram : 3D array
            image array containing experimental images.

        Returns
        -------
        FT_ptychogram : 3D array
            contrast enhanced FFT(ptychogram).

        """
        # windowing in real space will remove ringing artefacts in Fourier domain
        # improves Fourier space contrast for calibration
        ptychogram *= window("hann", ptychogram[0].shape)[None, ...]
        # convert to Fourier domain
        FT_ptychogram = np.abs(ifft2c(ptychogram))
        # get the mean for normalization
        normalization = np.mean(FT_ptychogram, 0)
        # get noise values outside 2x NA
        mean_outside_support = np.mean(normalization[self.OTF == 0])
        # replace noise with a constant term
        # normalization[normalization<(1*mean_outside_support)] = 1*mean_outside_support
        normalization = np.maximum(normalization, 2 * mean_outside_support)

        FT_ptychogram = np.abs(FT_ptychogram / normalization)

        for i in range(FT_ptychogram.shape[0]):
            FT_ptychogram[i] = gaussian(FT_ptychogram[i], sigma=self.gaussSigma)

        return FT_ptychogram

    def generateCircularArcsVectorized(self, initialPositions):
        """
        Generate the circular arcs representing the coherent transfer function
        boundaries at various XY locations. This is required for a grid-search
        method to minimize the risk of gradient descent optimization getting
        stuck in local minima.

        Rather than returning circular perimeters for each XY grid location,
        they will be filtered such that each XY grid locations has a circular
        arc containing the same number of array entries. For this reason the
        arrays will be filtered. The non-equal number of points is caused by
        the circular perimeter going out of bounds and resulting in a cricular
        arc.

        The circular arc array will be a 5D array such that vectorized arrays
        can be used. This reduces readability, but also greatly improves
        performance.

        The returned circular arc array will be a 5D array with entries:
        circularArcs = np.zeros([
        1D  multiple radius locations
        2D  meshgrid for each search grid locations
        1D  intensity values for each pixel along the circular radius
        1D  above value for each X and Y coordinates

        Parameters
        ----------
        initialPositions : 2D array
            X,Y position values for the current image being analysed.

        Returns
        -------
        circularArcs : 5D array
            compact search grid array in incomprehensible format to use
            numpy vectorization for increased speed
        """

        self.gridSearch_x = self.gridSearch_x_init.copy()
        self.gridSearch_y = self.gridSearch_y_init.copy()
        self.angleRange_x = self.angleRange_x_init.copy()
        self.angleRange_y = self.angleRange_y_init.copy()
        self.radiusScanRange = self.radiusScanRangeInit.copy()

        point_number = self.angleRange_x.shape[0]
        R_number = len(self.radiusScanRange)

        # convert radius search array into a 4D array
        self.radiusScanRange = np.reshape(self.radiusScanRange, [-1, 1, 1, 1])
        # convert angular search array into a 4D array
        self.angleRange_x = np.reshape(self.angleRange_x, [1, 1, 1, -1])
        self.angleRange_y = np.reshape(self.angleRange_y, [1, 1, 1, -1])
        # convert grid search mesh grid into a 4D array
        self.gridSearch_y, self.gridSearch_x = np.meshgrid(
            self.gridSearch_x, self.gridSearch_y
        )
        self.gridSearch_x = np.reshape(
            self.gridSearch_x, [1, self.x_range, self.y_range, 1]
        )
        self.gridSearch_y = np.reshape(
            self.gridSearch_y, [1, self.x_range, self.y_range, 1]
        )

        # compute the radial arc values for each spatial frequency position
        # this is given in cartesian coordinates by:
        # x = Rsin(angle) + circle_center_x
        # y = Rcos(angle) + circle_center_y
        xx_circle_arc = np.single(
            self.radiusScanRange * self.angleRange_x
            + self.img_size / 2.0
            + initialPositions[0]
            + self.gridSearch_x
        )
        yy_circle_arc = np.single(
            self.radiusScanRange * self.angleRange_y
            + self.img_size / 2.0
            + initialPositions[1]
            + self.gridSearch_y
        )

        # convert back to a 1D array for filtering whether the circular arcs
        # are within the image boundaries
        xx_circle_arc = np.reshape(xx_circle_arc, [-1, point_number])
        yy_circle_arc = np.reshape(yy_circle_arc, [-1, point_number])

        # remove values such that each arc has the same number of elements
        outliers = (
            (yy_circle_arc > 1)
            * (xx_circle_arc > 1)
            * (xx_circle_arc < (self.img_size - 1))
            * (yy_circle_arc < (self.img_size - 1))
        ) == True
        outliers = np.all(outliers, 0)

        # the final circular arc array contains:
        # 1.   multiple radius locations (R_number)
        # 2-3. meshgrids for each search grid locations
        #      defined by (self.x_range, self.y_range)
        # 4.   intensity values for each pixel along the circular radius
        #      np.count_nonzero(outliers)
        # 5.   split the whole array into XY
        circularArcs = np.zeros(
            [R_number, self.x_range, self.y_range, np.count_nonzero(outliers), 2]
        )

        # generate circle arcs
        self.gridSearch_x = self.gridSearch_x_init.copy()
        self.gridSearch_y = self.gridSearch_y_init.copy()
        self.angleRange_x = self.angleRange_x_init.copy()
        self.angleRange_y = self.angleRange_y_init.copy()
        self.radiusScanRange = self.radiusScanRangeInit.copy()

        # convert radius search array into a 4D array
        self.radiusScanRange = np.reshape(self.radiusScanRange, [-1, 1, 1, 1])
        # convert angular search array into a 4D array
        self.angleRange_x = np.reshape(self.angleRange_x[outliers], [1, 1, 1, -1])
        self.angleRange_y = np.reshape(self.angleRange_y[outliers], [1, 1, 1, -1])
        # convert grid search mesh grid into a 4D array
        self.gridSearch_y, self.gridSearch_x = np.meshgrid(
            self.gridSearch_x, self.gridSearch_y
        )
        self.gridSearch_x = np.reshape(
            self.gridSearch_x, [1, self.x_range, self.y_range, 1]
        )
        self.gridSearch_y = np.reshape(
            self.gridSearch_y, [1, self.x_range, self.y_range, 1]
        )

        circularArcs[:, :, :, :, 0] = np.single(
            self.radiusScanRange * self.angleRange_x
            + self.img_size / 2.0
            + initialPositions[0]
            + self.gridSearch_x
        )
        circularArcs[:, :, :, :, 1] = np.single(
            self.radiusScanRange * self.angleRange_y
            + self.img_size / 2.0
            + initialPositions[1]
            + self.gridSearch_y
        )
        return circularArcs

    def calculateRMSE(self, image, error_candidates, position):
        """
        Calculate the RMSE between the experimental data and low-pass-filtered image,
        The low-pass filter is shifted based on calibrated coordinates.
        Smallest RMSE gives the best k-space shift values

        Parameters
        ----------
        image : 2D array
            raw brightfield experimental image .
        error_candidates : 2D array
            grid search matrix containing the position error candidates.
        position : 2D array
            X,Y position values for the current image being analyses.

        Returns
        -------
        positions_error_x : float
            position error value x.
        positions_error_y : float
            position error value y.
        edge_case : bool
            this shows whether the error was at the edge of the search grid. If
            True then another search will be done using the updated positions.

        """
        xx_img, yy_img = np.mgrid[0 : self.img_size, 0 : self.img_size]

        fft_raw = ifft2c(image)
        low_pass_img = image.copy()
        rms_array = []
        x_arr = []
        y_arr = []
        for error_candidate in error_candidates:
            positions_error_x = self.gridSearch_x_init[int(error_candidate[0])]
            positions_error_y = self.gridSearch_y_init[int(error_candidate[1])]

            # generate a shifted low-pass filter
            pupil = np.zeros(fft_raw.shape)
            pupil_coords = (
                (xx_img - self.img_size / 2.0 + (position[0] + positions_error_x)) ** 2
                + (yy_img - self.img_size / 2.0 + (position[1] + positions_error_y))
                ** 2
            ) <= self.apertRadiusPixel**2
            pupil[pupil_coords] = 1
            # generate a low-pass filtered image with a shifted filter
            estimated_img = np.abs(ifft2c(pupil * fft_raw))

            # see if the shifted positions minimize the error
            error = low_pass_img - estimated_img
            rms = np.sqrt(np.mean(error**2))
            rms_array.append(rms)
            x_arr.append(positions_error_x)
            y_arr.append(positions_error_y)

        err_idx = np.argmin(rms_array)
        positions_error_x = x_arr[err_idx]
        positions_error_y = y_arr[err_idx]

        # check if the value was on the edge of the error map
        b1 = 0
        b2 = len(self.gridSearch_x_init) - 1
        if (error_candidates[err_idx, 0] in [b1, b2]) or (
            error_candidates[err_idx, 1] in [b1, b2]
        ):
            edge_case = True
        else:
            edge_case = False
        return positions_error_x, positions_error_y, edge_case

    def computePositionErrorCandidates(self, circleSearchArray):
        """
        Use the search grid error maps for each radial position and
        compute the gradients along the radial values. Highest gradient
        will indicate the best circular boundary location and these
        grid search maps will be used as initial error candidates.
        1st and 2nd derivatives are used.


        Parameters
        ----------
        circleSearchArray : 3D array
            search grid error values for each radial position

        Returns
        -------
        errorCandidates : 2D array
            binary map indicating position error candidates.

        """
        # error map using first derivative
        # get the 1st derivative along the radial direction
        arc_dx1 = np.diff(circleSearchArray, n=1, axis=0)
        # get the largest radial gradient
        rad_dx_idx = np.argwhere(np.max(arc_dx1) == arc_dx1)[0][0]
        # get the search grid for the largest gradient
        search_grid_dx = np.squeeze(arc_dx1[rad_dx_idx, :, :])
        search_grid_dx = search_grid_dx >= np.max(search_grid_dx) - 0.1 * np.std(
            search_grid_dx
        )  # Select region around max

        # error map using second derivative
        # get the 2nd derivative along the radial direction
        arc_dx2 = np.diff(circleSearchArray, n=2, axis=0)
        # get the largest radial gradient
        rad_dx_dx_idx2 = np.argwhere(np.max(arc_dx2) == arc_dx2)[0][0]
        # get the search grid for the largest gradient
        search_grid_dx_dx = np.squeeze(arc_dx2[rad_dx_dx_idx2, :, :])
        search_grid_dx_dx = search_grid_dx_dx >= np.max(
            search_grid_dx_dx
        ) - 0.25 * np.std(
            search_grid_dx_dx
        )  # Select region around max

        # combine the selected candidates
        # errorCandidates = (search_grid_dx + search_grid_dx_dx)>0
        errorCandidates = (search_grid_dx) > 0

        # plt.subplot(221)
        # plt.imshow(np.mean(arc_dx1,0))
        # plt.subplot(222)
        # plt.imshow(search_grid_dx)
        # plt.subplot(223)
        # plt.imshow(search_grid_dx_dx)
        # plt.subplot(224)
        # plt.imshow(errorCandidates)
        # plt.pause(1)

        errorCandidates = np.argwhere(errorCandidates == True)
        return errorCandidates

    def findCalibratedRadius(self, ptychogram, FT_ptychogram, initialPositions):
        """
        Find the best aperture radius which corresponds to a wrong NA value.
        The actual radius is updated internally and is used for position
        calibration.

        Parameters
        ----------
        ptychogram : 3D array
            image stack.
        FT_ptychogram : 3D array
            FFT(ptychogram).
        initialPositions : 2D array
            initial position vectors.

        Returns
        -------
        oldRadius : float
            old radius pre-calibration.
        newRadius : float
            new radius post-calibration.

        """
        self.initialize_error_search_space()
        self.radiuSearchBounds = 5
        self.radiusSearchStep = 0.5
        self.radiusSearchRange = np.mgrid[
            self.radiuSearchBounds : -self.radiuSearchBounds
            - self.radiusSearchStep : -self.radiusSearchStep
        ]
        # radius tolerance array
        tolerance = [self.radiusSearchStep * 1.1, self.radiusSearchStep * 0.9]
        self.apertRadiusPixel_init = self.apertRadiusPixel

        converged = False
        check_done = False
        radii_array = [self.apertRadiusPixel]
        itr = 0
        while not converged:
            print("Radius {}px, iteration {}".format(self.apertRadiusPixel, itr))
            itr += 1
            # find the optimal radius for each image
            radiiIdxForEachImg = np.zeros(FT_ptychogram.shape[0])
            # radius used for testing
            self.radiusScanRangeInit = self.radiusSearchRange + self.apertRadiusPixel

            # use at most most 20 images
            idx_range = np.arange(FT_ptychogram.shape[0])
            if FT_ptychogram.shape[0] > 20:
                np.random.shuffle(idx_range)
                idx_range = idx_range[:20]

            for idx in idx_range:
                image = FT_ptychogram[idx, :, :]

                # compute circular arc used as candidates and make an error
                # storage array. This returns the radial arcs array
                radial_arcs = self.generateCircularArcsVectorized(
                    initialPositions[idx, :]
                )

                # flatten the coordinate array
                coords = np.array(
                    [
                        radial_arcs[:, :, :, :, 0].flatten(),
                        radial_arcs[:, :, :, :, 1].flatten(),
                    ]
                )
                # interpolate to avoid quantization errors
                image = gaussian(
                    FT_ptychogram[idx, :, :], self.gaussSigma, preserve_range=True
                )
                gridSearchForEachR = map_coordinates(
                    image, coords, order=1, prefilter=False
                )

                # create the mean interpolated intensities along the perimeter for each
                # circular radius and grid search position
                circleSearchArray = np.mean(
                    np.reshape(gridSearchForEachR, radial_arcs.shape[:-1]), 3
                )

                # compute the gradient along the radial values
                arc_dx1 = np.gradient(circleSearchArray, axis=0)
                rad_dx_idx = np.argwhere(np.max(arc_dx1) == arc_dx1)[0][0]
                # best radius index for each image
                radiiIdxForEachImg[idx] = rad_dx_idx

            # find the median radius index
            radius_index = int(np.median(radiiIdxForEachImg))
            # select that radius corresponding to median
            self.apertRadiusPixel = self.radiusScanRangeInit[radius_index]
            radii_array.append(self.apertRadiusPixel)

            if not check_done:
                # check if radius change was small, then we are done
                if abs(radii_array[-1] - radii_array[-2]) <= tolerance[0]:
                    check_done = True
            else:
                # if the check was done
                mean_rad = np.mean(radii_array[-2:])
                if abs(mean_rad - radii_array[-1]) <= tolerance[1]:
                    # if within tolerance update the radius
                    self.apertRadiusPixel = mean_rad
                    converged = True
                else:
                    # if deviation is too large then continue iterating
                    check_done = False

            if (itr >= 20 and not check_done) or self.apertRadiusPixel <= 0:
                self.apertRadiusPixel = self.apertRadiusPixel_init
                converged = True

        oldRadius = self.apertRadiusPixel_init
        newRadius = self.apertRadiusPixel
        return oldRadius, newRadius

    def findPositionCalibrationMatrix(
        self, ptychogram, FT_ptychogram, initialPositions
    ):
        """
        Find the best aperture radius which corresponds to a wrong NA value.
        The actual radius is updated internally and is used for position
        calibration.

        Parameters
        ----------
        ptychogram : 3D array
            image stack.
        FT_ptychogram : 3D array
            FFT(ptychogram).
        initialPositions : 2D array
            initial position vectors.

        Returns
        -------
        calibMatrix : skimage calibration matrix
            best transformation matrix that transforms initial coordinate grid
            to the updated positions.
        updatedPositions : 2D array
            updated positions, without the rigid transformation applied.

        """
        # from cupyx.scipy.ndimage import map_coordinates as map_coordinates_gpu

        self.initialize_error_search_space()
        calibrated_positions = initialPositions.copy()
        failed_position_calib = []
        self.radiusSearchStep = 0.5
        self.radiusSearchRange = np.mgrid[
            max(
                self.gaussSigma + 2 * self.radiusSearchStep, 3 * self.radiusSearchStep
            ) : -self.radiusSearchStep : -self.radiusSearchStep
        ]
        self.radiusScanRangeInit = self.radiusSearchRange + self.apertRadiusPixel

        for idx in range(FT_ptychogram.shape[0]):
            image = FT_ptychogram[idx, :, :]

            #######################################################################
            # repeat the fitting if the solution has not converged
            #######################################################################
            converged = False
            repeat_counter = 0
            try:
                # iterate to account for an error too far away from the initial
                # search grid
                while not converged:
                    ##########################################################
                    # vectorized interpolation and mean value generation
                    ##########################################################
                    # compute circular arc used as candidates and make an error
                    # storage array. This returns the radial arcs array
                    radial_arcs = self.generateCircularArcsVectorized(
                        calibrated_positions[idx]
                    )

                    # flatten the coordinate array
                    coords = np.array(
                        [
                            radial_arcs[:, :, :, :, 0].flatten(),
                            radial_arcs[:, :, :, :, 1].flatten(),
                        ]
                    )

                    # interpolate to avoid quantization errors
                    gridSearchForEachR = map_coordinates(
                        image, coords, order=1, prefilter=False
                    )

                    # gridSearchForEachR = map_coordinates_gpu(cp.array(image),\
                    #                                           cp.array(coords),\
                    #                                           order=1,\
                    #                                           prefilter=False).get()
                    # define
                    circleSearchArray = np.mean(
                        np.reshape(gridSearchForEachR, radial_arcs.shape[:-1]), 3
                    )

                    # compute the position error candidates
                    positionErrorCandidates = self.computePositionErrorCandidates(
                        circleSearchArray
                    )

                    # go through all the possible coordinates and find the best one
                    (
                        positions_error_x,
                        positions_error_y,
                        values_on_edge,
                    ) = self.calculateRMSE(
                        ptychogram[idx, :, :],
                        positionErrorCandidates,
                        calibrated_positions[idx],
                    )

                    # update coordinate grid
                    calibrated_positions[idx, 0] = (
                        calibrated_positions[idx, 0] + positions_error_x
                    )
                    calibrated_positions[idx, 1] = (
                        calibrated_positions[idx, 1] + positions_error_y
                    )

                    # if the initial guess was too far and did not fall within
                    # the initial search grid, try again
                    if values_on_edge and repeat_counter < 5:
                        repeat_counter += 1
                    else:
                        converged = True
            except:  # Exception:
                # traceback.print_exc()
                failed_position_calib.append(idx)

        # once finished fit the whole data
        if len(failed_position_calib) != FT_ptychogram.shape[0]:
            src = np.delete(initialPositions, failed_position_calib, axis=0)
            dst = np.delete(calibrated_positions, failed_position_calib, axis=0)

            calibMatrix, updatedPositions = self.fitCoordinates(src, dst, self.fit_mode)

            updatedPositions = matrix_transform(initialPositions, calibMatrix.params)
        else:
            calibMatrix = SimilarityTransform(np.eye(3))
            updatedPositions = initialPositions
        return calibMatrix, updatedPositions

    def plotCalibration(self, FT_ptychogram, initialPositions, calibrated_positions):
        """
        Plot the fitted circles and the calibrated positions pre and post
        calibration

        Parameters
        ----------
        FT_ptychogram : 2D array
            FFT(ptychogram)
        initialPositions : 2D array
            initial positions pre calibration.
        calibrated_positions : 2D array
            positions post calibration.
        """
        self.initialize_error_search_space()
        try:
            plt.figure(3)

            while True:
                angles_x = np.sin(np.mgrid[0:360:0.1] / 180.0 * np.pi)
                angles_y = np.cos(np.mgrid[0:360:0.1] / 180.0 * np.pi)
                for idx in range(FT_ptychogram.shape[0]):
                    plt.clf()

                    plt.title("position calibration results, image {}".format(idx))
                    image = FT_ptychogram[idx, :, :]

                    initial_row = (
                        self.apertRadiusPixel_init * angles_x
                        + self.img_size / 2
                        + initialPositions[idx, 0]
                    )
                    initial_col = (
                        self.apertRadiusPixel_init * angles_y
                        + self.img_size / 2
                        + initialPositions[idx, 1]
                    )

                    final_row = (
                        self.apertRadiusPixel * angles_x
                        + self.img_size / 2
                        + calibrated_positions[idx, 0]
                    )
                    final_col = (
                        self.apertRadiusPixel * angles_y
                        + self.img_size / 2
                        + calibrated_positions[idx, 1]
                    )

                    plt.imshow(image)
                    plt.scatter(
                        initial_col, initial_row, label="initial guess", c="g", s=1
                    )
                    plt.scatter(final_col, final_row, label="calibrated", c="r", s=1)
                    legend = plt.legend(
                        labels=["initial guess", "calibrated"], loc="upper right"
                    )
                    plt.pause(1)

                    if not plt.fignum_exists(3):
                        break
                if not plt.fignum_exists(3):
                    break
        except:
            pass
        # plot the scattered positions
        plt.figure(4)
        x0 = initialPositions[:, 1]
        y0 = initialPositions[:, 0]
        x1 = calibrated_positions[:, 1]
        y1 = calibrated_positions[:, 0]
        plt.scatter(x0, y0, label="initial guess", c="g")
        plt.scatter(x1, y1, label="calibrated", c="r")
        plt.grid(True)
        plt.legend()
        plt.show()
        return None

    def fitCoordinates(self, src, dst, mode):
        """
        Parameters
        ----------
        src : 2D array
            source coordinates (to be calibrated)
        dst : 2D array
            destination coordinates (reference)
        mode : str
            Transformation mode: Translation, EuclideanTransform, SimilarityTransform, AffineTransform.

        Returns
        -------
        matrix : skimage transformation matrix
            3x3 transformation matrix.
        fitted : 2D array
            transformed source coordinates.

        """
        if mode == "Translation" or mode == "EuclideanTransform":
            tform_mode = EuclideanTransform
        elif mode == "SimilarityTransform":
            tform_mode = SimilarityTransform
        elif mode == "AffineTransform":
            tform_mode = AffineTransform
        else:
            tform_mode = EuclideanTransform
            print("Required fit mode not found, using EuclideanTransform")

        # residual threshold based on a median value
        resThresh = np.median(np.abs(dst - src))

        # compute the transformation matrix between the data points
        matrix, inliers = ransac(
            (src, dst),
            tform_mode,
            min_samples=2,
            residual_threshold=resThresh,
            max_trials=10000,
        )

        if mode == "Translation":
            matrix = tform_mode(translation=matrix.translation)
        fitted = matrix_transform(src, matrix.params)
        return matrix, fitted

    def updatePositions(self):
        """
        Change the "encocder" positions representing real-space LED coordinates
        depending on the calibration outcomes
        Returns
        -------
        None.

        """
        # convert translation from pixels to SI units and update the encoder
        conv = (
            -(1 / self.reconstructor.wavelength)
            * self.reconstructor.dxo
            * self.reconstructor.Np
        )
        z = self.reconstructor.zled

        # convert caibration matrix values into encoder units
        positionCalibMatrix = self.calibMatrix.params.copy()
        encoderCalibMatrix = self.calibMatrix.params.copy()
        encoderCalibMatrix[0:2, 2] = (
            np.sign(conv)
            * positionCalibMatrix[0:2, 2]
            * z
            / (
                np.sqrt(
                    conv**2
                    - positionCalibMatrix[0, 2] ** 2
                    - positionCalibMatrix[1, 2] ** 2
                )
            )
        )
        self.experimentalData.encoder = matrix_transform(
            self.experimentalData.encoder, encoderCalibMatrix
        )

        # self.experimentalData.encoder = np.sign(conv) *  self.positionsFitted * z / (np.sqrt(conv**2-self.positionsFitted[:,0]**2-self.positionsFitted[:,1]**2))[...,None]
        self.reconstructor.positions0 = self.reconstructor.positions.copy()

    def runCalibration(self):
        """
        Perform the calibration steps.
        1. Convert image array to Fourier space and enhacne contrast.
        2. Find brightfield indictes. Could be provided by the user or
        K-means clustering algorithm is used for automation.
        3. Calibrate NA/radius
        4. Calibrate positions

        Returns
        -------
        positionsFitted : 2D array
            positions after applying the correction matrix.
        entrancePupilDiameter : float
            updated aperture diameter.
        calibMatrix : skimage transformation matrix
            calibration matrix.

        """
        # the following arrays need to be copied in memory for processing
        self.ptychogram = deepcopy(self.experimentalData.ptychogram)
        self.initialPositions = deepcopy(self.reconstructor.positions)
        # UNSHFIT DUE TO EXAMPLEDATA SHIFTING BY No//2 and Np//2
        self.initialPositions = (
            self.initialPositions
            - self.reconstructor.No // 2
            + self.reconstructor.Np // 2
        )

        self.initialize_error_search_space()
        # get the FFT(ptychogram) which will also be filtered to enhance
        # contrast
        FFT_ptychogram = self.convertToFourierSpace(self.ptychogram)

        if not hasattr(self.brightfieldIndices, "__len__"):
            # find the brightfield indices
            self.brightfieldIndices = self.findBrightfielIndices(self.ptychogram)

        # calibrate the aperture radius / NA
        if self.calibrateRadius:
            oldRadius, newRadius = self.findCalibratedRadius(
                self.ptychogram[self.brightfieldIndices],
                FFT_ptychogram[self.brightfieldIndices],
                self.initialPositions[self.brightfieldIndices],
            )
            print("Initial radius was {}px".format(np.round(oldRadius, 2)))
            print("Calibrated radius is {}px".format(np.round(newRadius, 2)))
            oldNA = np.round(oldRadius / self.dxp * self.wavelength / self.img_size, 3)
            self.reconstructor.NA = np.round(
                newRadius / self.dxp * self.wavelength / self.img_size, 3
            )
            print("Initial NA was {}".format(oldNA))
            print("Calibrated NA is {}".format(self.reconstructor.NA))
            self.apertRadiusPixel = newRadius
        else:
            self.reconstructor.NA = (
                self.apertRadiusPixel / self.dxp * self.wavelength / self.img_size
            )

        # find the calibration matrix between the initial positions and the ones
        # found based on circle fitting
        self.calibMatrix, updatedPositions = self.findPositionCalibrationMatrix(
            self.ptychogram[self.brightfieldIndices],
            FFT_ptychogram[self.brightfieldIndices, :, :],
            self.initialPositions[self.brightfieldIndices, :],
        )
        # fit the positions
        positionsFitted = matrix_transform(
            self.initialPositions, self.calibMatrix.params
        )

        # update the entrancePupilDiameter
        self.reconstructor.entrancePupilDiameter = self.apertRadiusPixel * self.dxp * 2

        # lastly pre-process positions into the correct form defined in the experimentalData class
        self.positionsFitted = (positionsFitted).astype(int)
        self.updatePositions()

        # plot the results
        if self.plot:
            # self.plotCalibration(FFT_ptychogram, self.initialPositions, positionsFitted)
            self.plotCalibration(
                FFT_ptychogram[self.brightfieldIndices],
                self.initialPositions[self.brightfieldIndices],
                self.reconstructor.positions[self.brightfieldIndices]
                - self.No // 2
                + self.Np // 2,
            )

        return self.calibMatrix

findBrightfielIndices(ptychogram)

Use a threshold to separate brightfield images from darkfield images. Threshold is obtained by using a K-means clustering algorithm

Parameters

ptychogram : 3D array image array containing experimental images.

Returns

brightfieldIndices : 1D array bool index array where 1s represent brightfield images.

Source code in PtyLab/Reconstruction/CalibrationFPM.py

def findBrightfielIndices(self, ptychogram):
    """
    Use a threshold to separate brightfield images from darkfield images.
    Threshold is obtained by using a K-means clustering algorithm

    Parameters
    ----------
    ptychogram : 3D array
        image array containing experimental images.

    Returns
    -------
    brightfieldIndices : 1D array
        bool index array where 1s represent brightfield images.

    """
    # sum all the images together into a single 2D image
    intensities = np.sum(ptychogram, (1, 2))

    # kmeans clustering method to find 2 clusters
    cluster = np.array([np.arange(ptychogram.shape[0]), intensities]).T
    kmeans = KMeans(n_clusters=2).fit(cluster)
    clusters = kmeans.labels_
    # one cluster is brightfield other darkfield
    cluster1 = clusters == 0
    cluster2 = clusters == 1

    # if len(cluster1[cluster1]) < len(cluster2[cluster2]):
    if np.mean(ptychogram[cluster1]) > np.mean(ptychogram[cluster2]):
        brightfieldIndices = cluster1
    else:
        brightfieldIndices = cluster2
    return brightfieldIndices

convertToFourierSpace(ptychogram)

Convert brightfield image stack to the Fourier domain and also increase contrast based on all data analysis.

Parameters

ptychogram : 3D array image array containing experimental images.

Returns

FT_ptychogram : 3D array contrast enhanced FFT(ptychogram).

Source code in PtyLab/Reconstruction/CalibrationFPM.py

def convertToFourierSpace(self, ptychogram):
    """
    Convert brightfield image stack to the Fourier domain and also
    increase contrast based on all data analysis.

    Parameters
    ----------
    ptychogram : 3D array
        image array containing experimental images.

    Returns
    -------
    FT_ptychogram : 3D array
        contrast enhanced FFT(ptychogram).

    """
    # windowing in real space will remove ringing artefacts in Fourier domain
    # improves Fourier space contrast for calibration
    ptychogram *= window("hann", ptychogram[0].shape)[None, ...]
    # convert to Fourier domain
    FT_ptychogram = np.abs(ifft2c(ptychogram))
    # get the mean for normalization
    normalization = np.mean(FT_ptychogram, 0)
    # get noise values outside 2x NA
    mean_outside_support = np.mean(normalization[self.OTF == 0])
    # replace noise with a constant term
    # normalization[normalization<(1*mean_outside_support)] = 1*mean_outside_support
    normalization = np.maximum(normalization, 2 * mean_outside_support)

    FT_ptychogram = np.abs(FT_ptychogram / normalization)

    for i in range(FT_ptychogram.shape[0]):
        FT_ptychogram[i] = gaussian(FT_ptychogram[i], sigma=self.gaussSigma)

    return FT_ptychogram

generateCircularArcsVectorized(initialPositions)

Generate the circular arcs representing the coherent transfer function boundaries at various XY locations. This is required for a grid-search method to minimize the risk of gradient descent optimization getting stuck in local minima.

Rather than returning circular perimeters for each XY grid location, they will be filtered such that each XY grid locations has a circular arc containing the same number of array entries. For this reason the arrays will be filtered. The non-equal number of points is caused by the circular perimeter going out of bounds and resulting in a cricular arc.

The circular arc array will be a 5D array such that vectorized arrays can be used. This reduces readability, but also greatly improves performance.

The returned circular arc array will be a 5D array with entries: circularArcs = np.zeros([ 1D multiple radius locations 2D meshgrid for each search grid locations 1D intensity values for each pixel along the circular radius 1D above value for each X and Y coordinates

Parameters

initialPositions : 2D array X,Y position values for the current image being analysed.

Returns

circularArcs : 5D array compact search grid array in incomprehensible format to use numpy vectorization for increased speed

Source code in PtyLab/Reconstruction/CalibrationFPM.py

def generateCircularArcsVectorized(self, initialPositions):
    """
    Generate the circular arcs representing the coherent transfer function
    boundaries at various XY locations. This is required for a grid-search
    method to minimize the risk of gradient descent optimization getting
    stuck in local minima.

    Rather than returning circular perimeters for each XY grid location,
    they will be filtered such that each XY grid locations has a circular
    arc containing the same number of array entries. For this reason the
    arrays will be filtered. The non-equal number of points is caused by
    the circular perimeter going out of bounds and resulting in a cricular
    arc.

    The circular arc array will be a 5D array such that vectorized arrays
    can be used. This reduces readability, but also greatly improves
    performance.

    The returned circular arc array will be a 5D array with entries:
    circularArcs = np.zeros([
    1D  multiple radius locations
    2D  meshgrid for each search grid locations
    1D  intensity values for each pixel along the circular radius
    1D  above value for each X and Y coordinates

    Parameters
    ----------
    initialPositions : 2D array
        X,Y position values for the current image being analysed.

    Returns
    -------
    circularArcs : 5D array
        compact search grid array in incomprehensible format to use
        numpy vectorization for increased speed
    """

    self.gridSearch_x = self.gridSearch_x_init.copy()
    self.gridSearch_y = self.gridSearch_y_init.copy()
    self.angleRange_x = self.angleRange_x_init.copy()
    self.angleRange_y = self.angleRange_y_init.copy()
    self.radiusScanRange = self.radiusScanRangeInit.copy()

    point_number = self.angleRange_x.shape[0]
    R_number = len(self.radiusScanRange)

    # convert radius search array into a 4D array
    self.radiusScanRange = np.reshape(self.radiusScanRange, [-1, 1, 1, 1])
    # convert angular search array into a 4D array
    self.angleRange_x = np.reshape(self.angleRange_x, [1, 1, 1, -1])
    self.angleRange_y = np.reshape(self.angleRange_y, [1, 1, 1, -1])
    # convert grid search mesh grid into a 4D array
    self.gridSearch_y, self.gridSearch_x = np.meshgrid(
        self.gridSearch_x, self.gridSearch_y
    )
    self.gridSearch_x = np.reshape(
        self.gridSearch_x, [1, self.x_range, self.y_range, 1]
    )
    self.gridSearch_y = np.reshape(
        self.gridSearch_y, [1, self.x_range, self.y_range, 1]
    )

    # compute the radial arc values for each spatial frequency position
    # this is given in cartesian coordinates by:
    # x = Rsin(angle) + circle_center_x
    # y = Rcos(angle) + circle_center_y
    xx_circle_arc = np.single(
        self.radiusScanRange * self.angleRange_x
        + self.img_size / 2.0
        + initialPositions[0]
        + self.gridSearch_x
    )
    yy_circle_arc = np.single(
        self.radiusScanRange * self.angleRange_y
        + self.img_size / 2.0
        + initialPositions[1]
        + self.gridSearch_y
    )

    # convert back to a 1D array for filtering whether the circular arcs
    # are within the image boundaries
    xx_circle_arc = np.reshape(xx_circle_arc, [-1, point_number])
    yy_circle_arc = np.reshape(yy_circle_arc, [-1, point_number])

    # remove values such that each arc has the same number of elements
    outliers = (
        (yy_circle_arc > 1)
        * (xx_circle_arc > 1)
        * (xx_circle_arc < (self.img_size - 1))
        * (yy_circle_arc < (self.img_size - 1))
    ) == True
    outliers = np.all(outliers, 0)

    # the final circular arc array contains:
    # 1.   multiple radius locations (R_number)
    # 2-3. meshgrids for each search grid locations
    #      defined by (self.x_range, self.y_range)
    # 4.   intensity values for each pixel along the circular radius
    #      np.count_nonzero(outliers)
    # 5.   split the whole array into XY
    circularArcs = np.zeros(
        [R_number, self.x_range, self.y_range, np.count_nonzero(outliers), 2]
    )

    # generate circle arcs
    self.gridSearch_x = self.gridSearch_x_init.copy()
    self.gridSearch_y = self.gridSearch_y_init.copy()
    self.angleRange_x = self.angleRange_x_init.copy()
    self.angleRange_y = self.angleRange_y_init.copy()
    self.radiusScanRange = self.radiusScanRangeInit.copy()

    # convert radius search array into a 4D array
    self.radiusScanRange = np.reshape(self.radiusScanRange, [-1, 1, 1, 1])
    # convert angular search array into a 4D array
    self.angleRange_x = np.reshape(self.angleRange_x[outliers], [1, 1, 1, -1])
    self.angleRange_y = np.reshape(self.angleRange_y[outliers], [1, 1, 1, -1])
    # convert grid search mesh grid into a 4D array
    self.gridSearch_y, self.gridSearch_x = np.meshgrid(
        self.gridSearch_x, self.gridSearch_y
    )
    self.gridSearch_x = np.reshape(
        self.gridSearch_x, [1, self.x_range, self.y_range, 1]
    )
    self.gridSearch_y = np.reshape(
        self.gridSearch_y, [1, self.x_range, self.y_range, 1]
    )

    circularArcs[:, :, :, :, 0] = np.single(
        self.radiusScanRange * self.angleRange_x
        + self.img_size / 2.0
        + initialPositions[0]
        + self.gridSearch_x
    )
    circularArcs[:, :, :, :, 1] = np.single(
        self.radiusScanRange * self.angleRange_y
        + self.img_size / 2.0
        + initialPositions[1]
        + self.gridSearch_y
    )
    return circularArcs

calculateRMSE(image, error_candidates, position)

Calculate the RMSE between the experimental data and low-pass-filtered image, The low-pass filter is shifted based on calibrated coordinates. Smallest RMSE gives the best k-space shift values

Parameters

image : 2D array raw brightfield experimental image . error_candidates : 2D array grid search matrix containing the position error candidates. position : 2D array X,Y position values for the current image being analyses.

Returns

positions_error_x : float position error value x. positions_error_y : float position error value y. edge_case : bool this shows whether the error was at the edge of the search grid. If True then another search will be done using the updated positions.

Source code in PtyLab/Reconstruction/CalibrationFPM.py

def calculateRMSE(self, image, error_candidates, position):
    """
    Calculate the RMSE between the experimental data and low-pass-filtered image,
    The low-pass filter is shifted based on calibrated coordinates.
    Smallest RMSE gives the best k-space shift values

    Parameters
    ----------
    image : 2D array
        raw brightfield experimental image .
    error_candidates : 2D array
        grid search matrix containing the position error candidates.
    position : 2D array
        X,Y position values for the current image being analyses.

    Returns
    -------
    positions_error_x : float
        position error value x.
    positions_error_y : float
        position error value y.
    edge_case : bool
        this shows whether the error was at the edge of the search grid. If
        True then another search will be done using the updated positions.

    """
    xx_img, yy_img = np.mgrid[0 : self.img_size, 0 : self.img_size]

    fft_raw = ifft2c(image)
    low_pass_img = image.copy()
    rms_array = []
    x_arr = []
    y_arr = []
    for error_candidate in error_candidates:
        positions_error_x = self.gridSearch_x_init[int(error_candidate[0])]
        positions_error_y = self.gridSearch_y_init[int(error_candidate[1])]

        # generate a shifted low-pass filter
        pupil = np.zeros(fft_raw.shape)
        pupil_coords = (
            (xx_img - self.img_size / 2.0 + (position[0] + positions_error_x)) ** 2
            + (yy_img - self.img_size / 2.0 + (position[1] + positions_error_y))
            ** 2
        ) <= self.apertRadiusPixel**2
        pupil[pupil_coords] = 1
        # generate a low-pass filtered image with a shifted filter
        estimated_img = np.abs(ifft2c(pupil * fft_raw))

        # see if the shifted positions minimize the error
        error = low_pass_img - estimated_img
        rms = np.sqrt(np.mean(error**2))
        rms_array.append(rms)
        x_arr.append(positions_error_x)
        y_arr.append(positions_error_y)

    err_idx = np.argmin(rms_array)
    positions_error_x = x_arr[err_idx]
    positions_error_y = y_arr[err_idx]

    # check if the value was on the edge of the error map
    b1 = 0
    b2 = len(self.gridSearch_x_init) - 1
    if (error_candidates[err_idx, 0] in [b1, b2]) or (
        error_candidates[err_idx, 1] in [b1, b2]
    ):
        edge_case = True
    else:
        edge_case = False
    return positions_error_x, positions_error_y, edge_case

computePositionErrorCandidates(circleSearchArray)

Use the search grid error maps for each radial position and compute the gradients along the radial values. Highest gradient will indicate the best circular boundary location and these grid search maps will be used as initial error candidates. 1st and 2nd derivatives are used.

Parameters

circleSearchArray : 3D array search grid error values for each radial position

Returns

errorCandidates : 2D array binary map indicating position error candidates.

Source code in PtyLab/Reconstruction/CalibrationFPM.py

def computePositionErrorCandidates(self, circleSearchArray):
    """
    Use the search grid error maps for each radial position and
    compute the gradients along the radial values. Highest gradient
    will indicate the best circular boundary location and these
    grid search maps will be used as initial error candidates.
    1st and 2nd derivatives are used.


    Parameters
    ----------
    circleSearchArray : 3D array
        search grid error values for each radial position

    Returns
    -------
    errorCandidates : 2D array
        binary map indicating position error candidates.

    """
    # error map using first derivative
    # get the 1st derivative along the radial direction
    arc_dx1 = np.diff(circleSearchArray, n=1, axis=0)
    # get the largest radial gradient
    rad_dx_idx = np.argwhere(np.max(arc_dx1) == arc_dx1)[0][0]
    # get the search grid for the largest gradient
    search_grid_dx = np.squeeze(arc_dx1[rad_dx_idx, :, :])
    search_grid_dx = search_grid_dx >= np.max(search_grid_dx) - 0.1 * np.std(
        search_grid_dx
    )  # Select region around max

    # error map using second derivative
    # get the 2nd derivative along the radial direction
    arc_dx2 = np.diff(circleSearchArray, n=2, axis=0)
    # get the largest radial gradient
    rad_dx_dx_idx2 = np.argwhere(np.max(arc_dx2) == arc_dx2)[0][0]
    # get the search grid for the largest gradient
    search_grid_dx_dx = np.squeeze(arc_dx2[rad_dx_dx_idx2, :, :])
    search_grid_dx_dx = search_grid_dx_dx >= np.max(
        search_grid_dx_dx
    ) - 0.25 * np.std(
        search_grid_dx_dx
    )  # Select region around max

    # combine the selected candidates
    # errorCandidates = (search_grid_dx + search_grid_dx_dx)>0
    errorCandidates = (search_grid_dx) > 0

    # plt.subplot(221)
    # plt.imshow(np.mean(arc_dx1,0))
    # plt.subplot(222)
    # plt.imshow(search_grid_dx)
    # plt.subplot(223)
    # plt.imshow(search_grid_dx_dx)
    # plt.subplot(224)
    # plt.imshow(errorCandidates)
    # plt.pause(1)

    errorCandidates = np.argwhere(errorCandidates == True)
    return errorCandidates

findCalibratedRadius(ptychogram, FT_ptychogram, initialPositions)

Find the best aperture radius which corresponds to a wrong NA value. The actual radius is updated internally and is used for position calibration.

Parameters

ptychogram : 3D array image stack. FT_ptychogram : 3D array FFT(ptychogram). initialPositions : 2D array initial position vectors.

Returns

oldRadius : float old radius pre-calibration. newRadius : float new radius post-calibration.

Source code in PtyLab/Reconstruction/CalibrationFPM.py

def findCalibratedRadius(self, ptychogram, FT_ptychogram, initialPositions):
    """
    Find the best aperture radius which corresponds to a wrong NA value.
    The actual radius is updated internally and is used for position
    calibration.

    Parameters
    ----------
    ptychogram : 3D array
        image stack.
    FT_ptychogram : 3D array
        FFT(ptychogram).
    initialPositions : 2D array
        initial position vectors.

    Returns
    -------
    oldRadius : float
        old radius pre-calibration.
    newRadius : float
        new radius post-calibration.

    """
    self.initialize_error_search_space()
    self.radiuSearchBounds = 5
    self.radiusSearchStep = 0.5
    self.radiusSearchRange = np.mgrid[
        self.radiuSearchBounds : -self.radiuSearchBounds
        - self.radiusSearchStep : -self.radiusSearchStep
    ]
    # radius tolerance array
    tolerance = [self.radiusSearchStep * 1.1, self.radiusSearchStep * 0.9]
    self.apertRadiusPixel_init = self.apertRadiusPixel

    converged = False
    check_done = False
    radii_array = [self.apertRadiusPixel]
    itr = 0
    while not converged:
        print("Radius {}px, iteration {}".format(self.apertRadiusPixel, itr))
        itr += 1
        # find the optimal radius for each image
        radiiIdxForEachImg = np.zeros(FT_ptychogram.shape[0])
        # radius used for testing
        self.radiusScanRangeInit = self.radiusSearchRange + self.apertRadiusPixel

        # use at most most 20 images
        idx_range = np.arange(FT_ptychogram.shape[0])
        if FT_ptychogram.shape[0] > 20:
            np.random.shuffle(idx_range)
            idx_range = idx_range[:20]

        for idx in idx_range:
            image = FT_ptychogram[idx, :, :]

            # compute circular arc used as candidates and make an error
            # storage array. This returns the radial arcs array
            radial_arcs = self.generateCircularArcsVectorized(
                initialPositions[idx, :]
            )

            # flatten the coordinate array
            coords = np.array(
                [
                    radial_arcs[:, :, :, :, 0].flatten(),
                    radial_arcs[:, :, :, :, 1].flatten(),
                ]
            )
            # interpolate to avoid quantization errors
            image = gaussian(
                FT_ptychogram[idx, :, :], self.gaussSigma, preserve_range=True
            )
            gridSearchForEachR = map_coordinates(
                image, coords, order=1, prefilter=False
            )

            # create the mean interpolated intensities along the perimeter for each
            # circular radius and grid search position
            circleSearchArray = np.mean(
                np.reshape(gridSearchForEachR, radial_arcs.shape[:-1]), 3
            )

            # compute the gradient along the radial values
            arc_dx1 = np.gradient(circleSearchArray, axis=0)
            rad_dx_idx = np.argwhere(np.max(arc_dx1) == arc_dx1)[0][0]
            # best radius index for each image
            radiiIdxForEachImg[idx] = rad_dx_idx

        # find the median radius index
        radius_index = int(np.median(radiiIdxForEachImg))
        # select that radius corresponding to median
        self.apertRadiusPixel = self.radiusScanRangeInit[radius_index]
        radii_array.append(self.apertRadiusPixel)

        if not check_done:
            # check if radius change was small, then we are done
            if abs(radii_array[-1] - radii_array[-2]) <= tolerance[0]:
                check_done = True
        else:
            # if the check was done
            mean_rad = np.mean(radii_array[-2:])
            if abs(mean_rad - radii_array[-1]) <= tolerance[1]:
                # if within tolerance update the radius
                self.apertRadiusPixel = mean_rad
                converged = True
            else:
                # if deviation is too large then continue iterating
                check_done = False

        if (itr >= 20 and not check_done) or self.apertRadiusPixel <= 0:
            self.apertRadiusPixel = self.apertRadiusPixel_init
            converged = True

    oldRadius = self.apertRadiusPixel_init
    newRadius = self.apertRadiusPixel
    return oldRadius, newRadius

findPositionCalibrationMatrix(ptychogram, FT_ptychogram, initialPositions)

Find the best aperture radius which corresponds to a wrong NA value. The actual radius is updated internally and is used for position calibration.

Parameters

ptychogram : 3D array image stack. FT_ptychogram : 3D array FFT(ptychogram). initialPositions : 2D array initial position vectors.

Returns

calibMatrix : skimage calibration matrix best transformation matrix that transforms initial coordinate grid to the updated positions. updatedPositions : 2D array updated positions, without the rigid transformation applied.

Source code in PtyLab/Reconstruction/CalibrationFPM.py

def findPositionCalibrationMatrix(
    self, ptychogram, FT_ptychogram, initialPositions
):
    """
    Find the best aperture radius which corresponds to a wrong NA value.
    The actual radius is updated internally and is used for position
    calibration.

    Parameters
    ----------
    ptychogram : 3D array
        image stack.
    FT_ptychogram : 3D array
        FFT(ptychogram).
    initialPositions : 2D array
        initial position vectors.

    Returns
    -------
    calibMatrix : skimage calibration matrix
        best transformation matrix that transforms initial coordinate grid
        to the updated positions.
    updatedPositions : 2D array
        updated positions, without the rigid transformation applied.

    """
    # from cupyx.scipy.ndimage import map_coordinates as map_coordinates_gpu

    self.initialize_error_search_space()
    calibrated_positions = initialPositions.copy()
    failed_position_calib = []
    self.radiusSearchStep = 0.5
    self.radiusSearchRange = np.mgrid[
        max(
            self.gaussSigma + 2 * self.radiusSearchStep, 3 * self.radiusSearchStep
        ) : -self.radiusSearchStep : -self.radiusSearchStep
    ]
    self.radiusScanRangeInit = self.radiusSearchRange + self.apertRadiusPixel

    for idx in range(FT_ptychogram.shape[0]):
        image = FT_ptychogram[idx, :, :]

        #######################################################################
        # repeat the fitting if the solution has not converged
        #######################################################################
        converged = False
        repeat_counter = 0
        try:
            # iterate to account for an error too far away from the initial
            # search grid
            while not converged:
                ##########################################################
                # vectorized interpolation and mean value generation
                ##########################################################
                # compute circular arc used as candidates and make an error
                # storage array. This returns the radial arcs array
                radial_arcs = self.generateCircularArcsVectorized(
                    calibrated_positions[idx]
                )

                # flatten the coordinate array
                coords = np.array(
                    [
                        radial_arcs[:, :, :, :, 0].flatten(),
                        radial_arcs[:, :, :, :, 1].flatten(),
                    ]
                )

                # interpolate to avoid quantization errors
                gridSearchForEachR = map_coordinates(
                    image, coords, order=1, prefilter=False
                )

                # gridSearchForEachR = map_coordinates_gpu(cp.array(image),\
                #                                           cp.array(coords),\
                #                                           order=1,\
                #                                           prefilter=False).get()
                # define
                circleSearchArray = np.mean(
                    np.reshape(gridSearchForEachR, radial_arcs.shape[:-1]), 3
                )

                # compute the position error candidates
                positionErrorCandidates = self.computePositionErrorCandidates(
                    circleSearchArray
                )

                # go through all the possible coordinates and find the best one
                (
                    positions_error_x,
                    positions_error_y,
                    values_on_edge,
                ) = self.calculateRMSE(
                    ptychogram[idx, :, :],
                    positionErrorCandidates,
                    calibrated_positions[idx],
                )

                # update coordinate grid
                calibrated_positions[idx, 0] = (
                    calibrated_positions[idx, 0] + positions_error_x
                )
                calibrated_positions[idx, 1] = (
                    calibrated_positions[idx, 1] + positions_error_y
                )

                # if the initial guess was too far and did not fall within
                # the initial search grid, try again
                if values_on_edge and repeat_counter < 5:
                    repeat_counter += 1
                else:
                    converged = True
        except:  # Exception:
            # traceback.print_exc()
            failed_position_calib.append(idx)

    # once finished fit the whole data
    if len(failed_position_calib) != FT_ptychogram.shape[0]:
        src = np.delete(initialPositions, failed_position_calib, axis=0)
        dst = np.delete(calibrated_positions, failed_position_calib, axis=0)

        calibMatrix, updatedPositions = self.fitCoordinates(src, dst, self.fit_mode)

        updatedPositions = matrix_transform(initialPositions, calibMatrix.params)
    else:
        calibMatrix = SimilarityTransform(np.eye(3))
        updatedPositions = initialPositions
    return calibMatrix, updatedPositions

plotCalibration(FT_ptychogram, initialPositions, calibrated_positions)

Plot the fitted circles and the calibrated positions pre and post calibration

Parameters

FT_ptychogram : 2D array FFT(ptychogram) initialPositions : 2D array initial positions pre calibration. calibrated_positions : 2D array positions post calibration.

Source code in PtyLab/Reconstruction/CalibrationFPM.py

def plotCalibration(self, FT_ptychogram, initialPositions, calibrated_positions):
    """
    Plot the fitted circles and the calibrated positions pre and post
    calibration

    Parameters
    ----------
    FT_ptychogram : 2D array
        FFT(ptychogram)
    initialPositions : 2D array
        initial positions pre calibration.
    calibrated_positions : 2D array
        positions post calibration.
    """
    self.initialize_error_search_space()
    try:
        plt.figure(3)

        while True:
            angles_x = np.sin(np.mgrid[0:360:0.1] / 180.0 * np.pi)
            angles_y = np.cos(np.mgrid[0:360:0.1] / 180.0 * np.pi)
            for idx in range(FT_ptychogram.shape[0]):
                plt.clf()

                plt.title("position calibration results, image {}".format(idx))
                image = FT_ptychogram[idx, :, :]

                initial_row = (
                    self.apertRadiusPixel_init * angles_x
                    + self.img_size / 2
                    + initialPositions[idx, 0]
                )
                initial_col = (
                    self.apertRadiusPixel_init * angles_y
                    + self.img_size / 2
                    + initialPositions[idx, 1]
                )

                final_row = (
                    self.apertRadiusPixel * angles_x
                    + self.img_size / 2
                    + calibrated_positions[idx, 0]
                )
                final_col = (
                    self.apertRadiusPixel * angles_y
                    + self.img_size / 2
                    + calibrated_positions[idx, 1]
                )

                plt.imshow(image)
                plt.scatter(
                    initial_col, initial_row, label="initial guess", c="g", s=1
                )
                plt.scatter(final_col, final_row, label="calibrated", c="r", s=1)
                legend = plt.legend(
                    labels=["initial guess", "calibrated"], loc="upper right"
                )
                plt.pause(1)

                if not plt.fignum_exists(3):
                    break
            if not plt.fignum_exists(3):
                break
    except:
        pass
    # plot the scattered positions
    plt.figure(4)
    x0 = initialPositions[:, 1]
    y0 = initialPositions[:, 0]
    x1 = calibrated_positions[:, 1]
    y1 = calibrated_positions[:, 0]
    plt.scatter(x0, y0, label="initial guess", c="g")
    plt.scatter(x1, y1, label="calibrated", c="r")
    plt.grid(True)
    plt.legend()
    plt.show()
    return None

fitCoordinates(src, dst, mode)

Parameters

src : 2D array source coordinates (to be calibrated) dst : 2D array destination coordinates (reference) mode : str Transformation mode: Translation, EuclideanTransform, SimilarityTransform, AffineTransform.

Returns

matrix : skimage transformation matrix 3x3 transformation matrix. fitted : 2D array transformed source coordinates.

Source code in PtyLab/Reconstruction/CalibrationFPM.py

def fitCoordinates(self, src, dst, mode):
    """
    Parameters
    ----------
    src : 2D array
        source coordinates (to be calibrated)
    dst : 2D array
        destination coordinates (reference)
    mode : str
        Transformation mode: Translation, EuclideanTransform, SimilarityTransform, AffineTransform.

    Returns
    -------
    matrix : skimage transformation matrix
        3x3 transformation matrix.
    fitted : 2D array
        transformed source coordinates.

    """
    if mode == "Translation" or mode == "EuclideanTransform":
        tform_mode = EuclideanTransform
    elif mode == "SimilarityTransform":
        tform_mode = SimilarityTransform
    elif mode == "AffineTransform":
        tform_mode = AffineTransform
    else:
        tform_mode = EuclideanTransform
        print("Required fit mode not found, using EuclideanTransform")

    # residual threshold based on a median value
    resThresh = np.median(np.abs(dst - src))

    # compute the transformation matrix between the data points
    matrix, inliers = ransac(
        (src, dst),
        tform_mode,
        min_samples=2,
        residual_threshold=resThresh,
        max_trials=10000,
    )

    if mode == "Translation":
        matrix = tform_mode(translation=matrix.translation)
    fitted = matrix_transform(src, matrix.params)
    return matrix, fitted

updatePositions()

Change the "encocder" positions representing real-space LED coordinates depending on the calibration outcomes Returns

---

None.

Source code in PtyLab/Reconstruction/CalibrationFPM.py

def updatePositions(self):
    """
    Change the "encocder" positions representing real-space LED coordinates
    depending on the calibration outcomes
    Returns
    -------
    None.

    """
    # convert translation from pixels to SI units and update the encoder
    conv = (
        -(1 / self.reconstructor.wavelength)
        * self.reconstructor.dxo
        * self.reconstructor.Np
    )
    z = self.reconstructor.zled

    # convert caibration matrix values into encoder units
    positionCalibMatrix = self.calibMatrix.params.copy()
    encoderCalibMatrix = self.calibMatrix.params.copy()
    encoderCalibMatrix[0:2, 2] = (
        np.sign(conv)
        * positionCalibMatrix[0:2, 2]
        * z
        / (
            np.sqrt(
                conv**2
                - positionCalibMatrix[0, 2] ** 2
                - positionCalibMatrix[1, 2] ** 2
            )
        )
    )
    self.experimentalData.encoder = matrix_transform(
        self.experimentalData.encoder, encoderCalibMatrix
    )

    # self.experimentalData.encoder = np.sign(conv) *  self.positionsFitted * z / (np.sqrt(conv**2-self.positionsFitted[:,0]**2-self.positionsFitted[:,1]**2))[...,None]
    self.reconstructor.positions0 = self.reconstructor.positions.copy()

runCalibration()

Perform the calibration steps. 1. Convert image array to Fourier space and enhacne contrast. 2. Find brightfield indictes. Could be provided by the user or K-means clustering algorithm is used for automation. 3. Calibrate NA/radius 4. Calibrate positions

Returns

positionsFitted : 2D array positions after applying the correction matrix. entrancePupilDiameter : float updated aperture diameter. calibMatrix : skimage transformation matrix calibration matrix.

Source code in PtyLab/Reconstruction/CalibrationFPM.py

def runCalibration(self):
    """
    Perform the calibration steps.
    1. Convert image array to Fourier space and enhacne contrast.
    2. Find brightfield indictes. Could be provided by the user or
    K-means clustering algorithm is used for automation.
    3. Calibrate NA/radius
    4. Calibrate positions

    Returns
    -------
    positionsFitted : 2D array
        positions after applying the correction matrix.
    entrancePupilDiameter : float
        updated aperture diameter.
    calibMatrix : skimage transformation matrix
        calibration matrix.

    """
    # the following arrays need to be copied in memory for processing
    self.ptychogram = deepcopy(self.experimentalData.ptychogram)
    self.initialPositions = deepcopy(self.reconstructor.positions)
    # UNSHFIT DUE TO EXAMPLEDATA SHIFTING BY No//2 and Np//2
    self.initialPositions = (
        self.initialPositions
        - self.reconstructor.No // 2
        + self.reconstructor.Np // 2
    )

    self.initialize_error_search_space()
    # get the FFT(ptychogram) which will also be filtered to enhance
    # contrast
    FFT_ptychogram = self.convertToFourierSpace(self.ptychogram)

    if not hasattr(self.brightfieldIndices, "__len__"):
        # find the brightfield indices
        self.brightfieldIndices = self.findBrightfielIndices(self.ptychogram)

    # calibrate the aperture radius / NA
    if self.calibrateRadius:
        oldRadius, newRadius = self.findCalibratedRadius(
            self.ptychogram[self.brightfieldIndices],
            FFT_ptychogram[self.brightfieldIndices],
            self.initialPositions[self.brightfieldIndices],
        )
        print("Initial radius was {}px".format(np.round(oldRadius, 2)))
        print("Calibrated radius is {}px".format(np.round(newRadius, 2)))
        oldNA = np.round(oldRadius / self.dxp * self.wavelength / self.img_size, 3)
        self.reconstructor.NA = np.round(
            newRadius / self.dxp * self.wavelength / self.img_size, 3
        )
        print("Initial NA was {}".format(oldNA))
        print("Calibrated NA is {}".format(self.reconstructor.NA))
        self.apertRadiusPixel = newRadius
    else:
        self.reconstructor.NA = (
            self.apertRadiusPixel / self.dxp * self.wavelength / self.img_size
        )

    # find the calibration matrix between the initial positions and the ones
    # found based on circle fitting
    self.calibMatrix, updatedPositions = self.findPositionCalibrationMatrix(
        self.ptychogram[self.brightfieldIndices],
        FFT_ptychogram[self.brightfieldIndices, :, :],
        self.initialPositions[self.brightfieldIndices, :],
    )
    # fit the positions
    positionsFitted = matrix_transform(
        self.initialPositions, self.calibMatrix.params
    )

    # update the entrancePupilDiameter
    self.reconstructor.entrancePupilDiameter = self.apertRadiusPixel * self.dxp * 2

    # lastly pre-process positions into the correct form defined in the experimentalData class
    self.positionsFitted = (positionsFitted).astype(int)
    self.updatePositions()

    # plot the results
    if self.plot:
        # self.plotCalibration(FFT_ptychogram, self.initialPositions, positionsFitted)
        self.plotCalibration(
            FFT_ptychogram[self.brightfieldIndices],
            self.initialPositions[self.brightfieldIndices],
            self.reconstructor.positions[self.brightfieldIndices]
            - self.No // 2
            + self.Np // 2,
        )

    return self.calibMatrix

Reconstruction

Reconstruction

Bases: object

This object will contain all the things that can be modified by a reconstruction.

In itself, it's little more than a data holder. It is initialized with an ExperimentalData object.

Some parameters which are "immutable" within the ExperimentalData can be modified (e.g. zo modification by zPIE during the reconstruction routine). All of them are defined in the listOfReconstructionProperties

Source code in PtyLab/Reconstruction/Reconstruction.py

class Reconstruction(object):
    """
    This object will contain all the things that can be modified by a reconstruction.

    In itself, it's little more than a data holder. It is initialized with an ExperimentalData object.

    Some parameters which are "immutable" within the ExperimentalData can be modified
    (e.g. zo modification by zPIE during the reconstruction routine). All of them
    are defined in the listOfReconstructionProperties
    """

    _Nd = None

    # Note: zo, the sample-detector distance, is always read.
    listOfReconstructionPropertiesCPM = [
        "wavelength",
        # 'zo',
        "dxd",
        "theta",
        "spectralDensity",
        "entrancePupilDiameter",
    ]
    listOfReconstructionPropertiesFPM = [
        "wavelength",
        # 'zo',
        "dxd",
        "zled",
        "NA",
    ]

    def __init__(self, data: ExperimentalData, params: Params):

        self.zMomentum = 0
        self.wavelength = None
        self._zo = None
        self.dxd = None
        self.theta = None

        # positions including possible misalignment correction
        self.encoder_corrected = None

        self.logger = logging.getLogger("Reconstruction")
        self.data = data
        self.params = params
        self.copyAttributesFromExperiment(data)
        self.computeParameters()
        self.initializeSettings()

        # list of the fields that have to be transfered back and forth from the GPU
        self.possible_GPU_fields = [
            "probe",
            "object",
            "probeBuffer",
            "objectBuffer",
            "probeMomentum",
            "objectMomentum",
            "detectorError",
            "background",
            "reference",
            "intensity_mask",
            # multislice (e3PIE) transfer function, built on the host in
            # e3PIE.initializeReconstructionParams and used inside the position loop
            "H",
        ]

    # @property
    # def probe(self):
    #     # convenience function. Updates the temporary probe. Nothing in probe is updated
    #     # return self._probe
    #     return self.probe_storage.get_temporary()#_probe_storage.get(None)
    #
    # @probe.setter
    # def probe(self, new_probe):
    #     # ignore this for now
    #     # self._probe = new_probe
    #     # self.probe_storage.set_temporary(new_probe)

    def copyAttributesFromExperiment(self, data: ExperimentalData):
        """
        Copy all the attributes from the experiment that are in listOfReconstructionProperties (CPM or FPM)
        """
        self.logger.debug("Copying attributes from Experimental Data")
        if self.data.operationMode == "CPM":
            listOfReconstructionProperties = self.listOfReconstructionPropertiesCPM
        elif self.data.operationMode == "FPM":
            listOfReconstructionProperties = self.listOfReconstructionPropertiesFPM
        for key in listOfReconstructionProperties:
            self.logger.info("Copying attribute %s", key)
            # setattr(self, key, copy(np.array(getattr(data, key))))
            setattr(self, key, copy(getattr(data, key)))

        # set the distance, this has to be last
        # In FPM the sample to detector distance is irrelevant
        # LED-to-sample distance is the more important factor that affects
        # wave propagation and illumination angle
        if self.data.operationMode == "CPM":
            self.zo = getattr(data, "zo")

        # set the original positions
        if self.encoder_corrected is None:
            self.encoder_corrected = data.encoder.copy()

    def reset_positioncorrection(self):
        """Reset the position corrections."""
        self.encoder_corrected = self.data.encoder.copy()

    @property
    def zo(self):
        """Distance from sample to detector. Also updates all derived qualities."""
        return self._zo

    @zo.setter
    def zo(self, new_value):
        self._zo = new_value
        if self.data.operationMode == "CPM":
            self.logger.debug(f"Changing sample-detector distance to {new_value}")
            self.dxp = self.wavelength * self._zo / self.Ld
        elif self.data.operationMode == "FPM":
             self.logger.debug(f"Changing illumination-to-sample distance to {new_value}")
             self.zled = self._zo

    def computeParameters(self):
        """
        compute parameters that can be altered by the user later.
        """

        if self.data.operationMode == "CPM":
            # CPM dxp (depending on the propagatorType, if none given, assum Fraunhofer/Fresnel)
            # self.dxp = self.wavelength * self._zo / self.Ld
            # if entrancePupilDiameter is not provided in the hdf5 file, set it to be one third of the probe FoV.
            if self.data.entrancePupilDiameter is None:
                self.data.entrancePupilDiameter = self.Lp / 3
            # if spectralDensity is not provided in the hdf5 file, set it to be a 1d array of the wavelength
            if isinstance(self.spectralDensity, type(None)):
                # this is a confusing name, it should be the wavelengths, not the intensity of the different
                # wavelengths
                self.spectralDensity = np.atleast_1d(self.wavelength)

        elif self.data.operationMode == "FPM":
            # FPM dxp (different from CPM due to lens-based systems)
            self.dxp = self.dxd / self.data.magnification
            # the propagation distance that is meaningful in this context is the
            # illumination to sample distance for LED array based microscopes
            self.zo = self.zled
            # if NA is not provided in the hdf5 file, set Fourier pupil entrance diameter it to be half of the Fourier space FoV.
            # then estimate the NA from the pupil diameter in the Fourier plane
            if isinstance(self.NA, type(None)):
                self.data.entrancePupilDiameter = self.Lp / 2
                self.NA = (
                    self.data.entrancePupilDiameter
                    * self.wavelength
                    / (2 * self.dxp**2 * self.Np)
                )
            else:
                # compute the pupil radius in the Fourier plane
                self.data.entrancePupilDiameter = (
                    2 * self.dxp**2 * self.Np * self.NA / self.wavelength
                )

        # set object pixel numbers
        if not hasattr(self, 'No'):
            self.No = (
                self.Np * 2**2
            )  # unimportant but leave it here as it's required for self.positions
            # we need space for the probe as well, on both sides that would be half the probe
            range_pixels = np.max(self.positions, axis=0) - np.min(self.positions, axis=0)
            # print(range_pixels)
            range_pixels = np.max(range_pixels) + self.Np * 2
            if range_pixels % 2 == 1:
                range_pixels += 1
            self.No = np.max([self.Np, range_pixels])

    def make_alignment_plot(self, saveit=False):
        import time

        t0 = time.time()
        p_new = self.positions.T
        p_old = self.positions0.T

        from bokeh.plotting import figure, output_file, save
        from bokeh.layouts import row

        from pathlib import Path

        if saveit:
            output = Path("plots/alignment.html")
            output.parent.mkdir(exist_ok=True)
            # set output to static HTML file

            output_file(filename=output, title="Static HTML file", mode="inline")

        # create a new plot with a specific size
        p = figure(
            sizing_mode="stretch_width",
            max_width=500,
            height=500,
            title=f'alignment (updated {time.strftime("%Y%h%d, %H:%M:%S")})',
        )
        p.match_aspect = True
        square = p.square(
            p_old[0], p_old[1], fill_color="yellow", size=5, legend_label="original"
        )
        # add a circle renderer for the new points
        circle = p.circle(
            p_new[0], p_new[1], fill_color="red", size=5, legend_label="new"
        )

        p.xaxis.axis_label = "Position x [um]"
        p.yaxis.axis_label = "Position y [um]"

        p2 = None
        p3 = None
        p4 = None

        figsize = 500  # px

        if hasattr(self, "zHistory"):  # display the plot of the defocus
            p2 = figure(
                sizing_mode="stretch_width",
                max_width=figsize,
                height=figsize,
                title="focus history",
            )
            p2.circle(np.arange(len(self.zHistory)), np.array(self.zHistory) * 1e3)
            p2.xaxis.axis_label = "Iteration #"
            p2.yaxis.axis_label = "Position [mm]"
            # p = vplot(p, p2)

        if hasattr(self, "merit"):  # display the merit as well for defocii
            p3 = figure(
                sizing_mode="stretch_width",
                max_width=figsize,
                height=figsize,
                title="merit TV",
            )
            p3.circle(self.dz * 1e3, np.array(self.merit), legend_label="original")
            p3.square(
                -self.dz * 1e3,
                np.array(self.merit),
                legend_label="mirrored",
                color="red",
            )
            p3.xaxis.axis_label = "Defocus [mm]"
            p3.yaxis.axis_label = "Score [a.u.]"
            # p = vplot(p, p3)
        if hasattr(self, "TV_history"):
            if len(self.TV_history) >= 1:
                p4 = figure(
                    sizing_mode="stretch_width",
                    max_width=figsize,
                    height=figsize,
                    title="TV history",
                )
                p4.square(np.arange(len(self.TV_history)), self.TV_history)
                p4.xaxis.axis_label = "Iteration"
                p4.yaxis.axis_label = "TV score"
        # only add the plots that are available
        p_list = filter(lambda x: x is not None, [p, p2, p4, p3])
        p = row(*p_list)

        if saveit:
            save(
                p,
            )
        t1 = time.time()
        print(f"Alignment display took {t1-t0} secs")
        return p

    def initializeSettings(self):
        """
        Initialize the attributes that have to do with a reconstruction
        or experimentalData fields which will become "reconstruction"

        This method just sets the settings. It sets the what kind of initial guess should be used for initialObject
        and initialProbe but it does not compute them yet. That will be done by calling initializeObjectProbe()

        :return:
        """
        # create a 6D object where which allows to have:
        # 1. polychromatic = nlambda
        # 2. mixed state object - nosm
        # 3. mixed state probe - npsm
        # 4. multislice object (thick) - nslice
        self.nlambda = 1
        self.nosm = 1
        self.npsm = 1
        self.nslice = 1

        # beam and object purity (# default initial value for plots.)
        self.purityProbe = 1.0
        self.purityObject = 1.0
        self.purityProbeHist = []

        self.positions0 = self.positions.copy()

        if self.data.operationMode == "FPM":
            self.initialObject = "upsampled"
            self.initialProbe = "circ"
        elif self.data.operationMode == "CPM":
            self.initialProbe = "circ"
            self.initialObject = "ones"
        else:
            self.initialProbe = "circ"
            self.initialObject = "ones"

    def prepare_probe(self, i):
        """Replace probe with the i-th TSVD estimate.

        This function is used in OPRP
        """
        raise NotImplementedError()

    def initializeObjectProbe(self, force=True):

        # initialize object and probe
        self.initializeObject(force=force)
        self.initializeProbe(force=force)

        # set object and probe objects
        self.object = self.initialGuessObject.copy()
        self.probe = self.initialGuessProbe.copy()

    def initializeObject(self, type_of_init=None, force=True):
        if not force:
            raise NotImplementedError()
        if type_of_init is not None:
            self.initialObject = type_of_init
        self.logger.info("Initial object set to %s", self.initialObject)
        self.shape_O = (
            self.nlambda,
            self.nosm,
            1,
            self.nslice,
            self.No,
            self.No,
        )
        if self.initialObject == 'recon':
            # Load the object from an existing reconstruction
            self.initialGuessObject = self.loadResults(self.initialProbe_filename, datatype='object')
        else:
            self.initialGuessObject = initialProbeOrObject(self.shape_O, self.initialObject, self, self.logger).astype(np.complex64)

        # self.initialGuessObject *= 1e-2

    @staticmethod
    def loadResults(fileName, datatype='probe'):
        '''
        Loads data from a ptylab reconstruction file.
        '''
        with h5py.File(fileName) as archive:
            data = np.copy(np.array(archive[datatype]))
        return data

    def initializeProbe(self, force=False):
        if self.data.entrancePupilDiameter is None:
            # if it is not set, set it to something reasonable
            self.logger.warning(
                "entrancePupilDiameter not set. Setting to one third of the FoV of the probe."
            )
            self.data.entrancePupilDiameter = self.Lp / 3
        self.logger.info("Initial probe set to %s", self.initialProbe)
        self.shape_P = (
            self.nlambda,
            1,
            self.npsm,
            self.nslice,
            int(self.Np),
            int(self.Np),
        )

        if self.initialProbe == 'recon':
            self.initialGuessProbe = self.loadResults(self.initialProbe_filename, datatype='probe')
        else:
            if force:
                self.initialGuessProbe = None
            # if force:
            #     self.initialProbe = "circ"
            self.initialGuessProbe = initialProbeOrObject(
                self.shape_P, self.initialProbe, self
            ).astype(np.complex64)

    # initialize momentum, called in specific engines with momentum accelaration
    def initializeObjectMomentum(self):
        self.objectMomentum = np.zeros_like(self.initialGuessObject)

    def initializeProbeMomentum(self):
        self.probeMomentum = np.zeros_like(self.initialGuessProbe)

    def load_object(self, filename):
        """
        Load the object from a previous reconstruction

        Parameters
        ----------
        filename: .hdf5 file
            Filenamne of the reconstruction whose object should be loaded.

        Returns
        -------

        """
        with h5py.File(filename, "r") as archive:
            obj = np.array(archive["object"])
            obj = obj[
                : self.shape_O[0],
                : self.shape_O[1],
                : self.shape_O[2],
                : self.shape_O[3],
                : self.shape_O[4],
                : self.shape_O[5],
            ]
            if np.all(np.array(obj.shape) == np.array(self.shape_O)):
                self.object = obj
            else:
                raise RuntimeError(
                    f'Shape of saved probe cannot be extended to shape of required probe. File: {archive["object"].shape}. Need: {self.shape_O}'
                )

    def load_probe(self, filename, expand_npsm=False, center_phase=False):
        """
        Load the probe from a previous reconstruction.

        Parameters
        ----------
        filename: .hdf5 file
            The filename of the reconstruction whose probe should be loaded.

        Returns
        -------

        """
        with h5py.File(filename, "r") as archive:
            probe = np.array(archive["probe"])
            N_probe_read = probe.shape[-1]
            # roughly extract the center
            ss = slice(np.clip(N_probe_read//2-self.Np//2, 0, None), np.clip(N_probe_read//2-self.Np//2+int(self.Np), 0, N_probe_read))
            probe = probe[
                : self.nlambda,
                :1,
                : self.npsm,
                : self.nslice,
                ss,ss
            ]
            if np.all(np.array(probe.shape) == np.array(self.shape_P)):
                self.probe = probe
            else:
                raise RuntimeError(
                    f'Shape of saved probe cannot be extended to shape of required probe. File: {archive["probe"].shape}. Need: {self.shape_P}'
                )
        if center_phase:
            self._center_probe_angle()

    def _center_probe_angle(self):
        """ Center the angle of propagation for the probe. """
        from skimage.registration import phase_cross_correlation
        from scipy.ndimage import fourier_shift
        p0 = np.squeeze(self.probe)[0]
        shift = phase_cross_correlation(p0, 0 * p0 + 1, normalization=None, space='fourier')[0]
        phexp = np.fft.fftshift(fourier_shift(0 * p0 + 1j, -shift / 2))
        self.probe *= phexp

    def load(self, filename):
        """Load the results given by saveResults."""
        with h5py.File(filename, "r") as archive:

            self.probe = np.array(archive["probe"])
            self.object = np.array(archive["object"])
            self.error = np.array(archive["error"])
            self.wavelength = np.array(archive["wavelength"])
            self.dxp = np.array(archive["dxp"])
            self.purityProbe = np.array(archive["purityProbe"])
            self.purityObject = np.array(archive["purityObject"])
            self.zo = np.array(archive["zo"])
            if "theta" in archive.keys():
                self.theta = np.array(archive["theta"])

    def saveResults(self, fileName="recent", type="all", squeeze=False):
        """
        Save reconstruction results.


        Parameters
        ----------
        fileName
        type
        squeeze


        Returns
        -------

        """

        allowed_save_types = ["all", "object", "probe", "probe_stack"]
        if type not in allowed_save_types:
            raise NotImplementedError(
                f"Only {allowed_save_types} are allowed keywords for type"
            )
        if not squeeze:
            squeezefun = lambda x: x
        else:
            squeezefun = np.squeeze
        if type == "all":
            if self.data.operationMode == "CPM":
                with h5py.File(fileName, "w") as hf:
                    hf.create_dataset("probe", data=self.probe, dtype="complex64")
                    hf.create_dataset("object", data=self.object, dtype="complex64")
                    hf.create_dataset("error", data=self.error, dtype="f")
                    hf.create_dataset("zo", data=self._zo, dtype="f")
                    hf.create_dataset("wavelength", data=self.wavelength, dtype="f")
                    hf.create_dataset("dxp", data=self.dxp, dtype="f")
                    hf.create_dataset("purityProbe", data=self.purityProbe, dtype="f")
                    hf.create_dataset("purityObject", data=self.purityObject, dtype="f")
                    hf.create_dataset('I object', data=abs(self.object), dtype='f')
                    hf.create_dataset('I probe', data=abs(self.probe), dtype='f')
                    hf.create_dataset('encoder_corrected', data=self.encoder_corrected)

                    if hasattr(self, "theta"):
                        if self.theta != None:
                            hf.create_dataset("theta", data=self.theta, dtype="f")

            if self.data.operationMode == "FPM":
                hf = h5py.File(fileName, "w")
                hf.create_dataset("probe", data=self.probe, dtype="complex64")
                hf.create_dataset("object", data=self.object, dtype="complex64")
                hf.create_dataset("error", data=self.error, dtype="f")
                hf.create_dataset("zled", data=self.zled, dtype="f")
                hf.create_dataset("wavelength", data=self.wavelength, dtype="f")
                hf.create_dataset("dxp", data=self.dxp, dtype="f")
        elif type == "probe":
            with h5py.File(fileName, "w") as hf:
                hf.create_dataset(
                    "probe", data=squeezefun(self.probe), dtype="complex64"
                )
        elif type == "object":
            with h5py.File(fileName, "w") as hf:
                hf.create_dataset(
                    "object", data=squeezefun(self.object), dtype="complex64"
                )
        elif type == "probe_stack":
            hf = h5py.File(fileName + '_probe_stack.hdf5', 'w')
            hf.create_dataset('probe_stack', data=self.probe_stack.get(), dtype='complex64')
        print("The reconstruction results (%s) have been saved" % type)

    # detector coordinates
    @property
    def Nd(self):
        return self.data.ptychogram.shape[1]

    @property
    def xd(self):
        """Detector coordinates 1D"""
        return np.linspace(-self.Nd / 2, self.Nd / 2, np.int(self.Nd)) * self.dxd

    @property
    def Xd(self):
        """Detector coordinates 2D"""
        Xd, Yd = np.meshgrid(self.xd, self.xd)
        return Xd

    @property
    def Yd(self):
        """Detector coordinates 2D"""
        Xd, Yd = np.meshgrid(self.xd, self.xd)
        return Yd

    @property
    def Ld(self):
        """Detector size in SI units."""
        return self.Nd * self.dxd

    # probe coordinates
    @property
    def Np(self):
        """Probe pixel numbers"""
        Np = self.Nd
        return Np

    @property
    def Lp(self):
        """probe size in SI units"""
        Lp = self.Np * self.dxp
        return Lp

    @property
    def xp(self):
        """Probe coordinates 1D"""
        try:
            return np.linspace(-self.Np / 2, self.Np / 2, int(self.Np)) * self.dxp
        except AttributeError as e:
            raise AttributeError(
                e, 'probe pixel number "Np" and/or probe sampling "dxp" not defined yet'
            )

    @property
    def Xp(self):
        """Probe coordinates 2D"""
        Xp, Yp = np.meshgrid(self.xp, self.xp)
        return Xp

    @property
    def Yp(self):
        """Probe coordinates 2D"""
        Xp, Yp = np.meshgrid(self.xp, self.xp)
        return Yp

    # Object coordinates
    @property
    def dxo(self):
        """object pixel size, always equal to probe pixel size."""
        dxo = self.dxp
        return dxo

    @property
    def Lo(self):
        """Field of view (entrance pupil plane)"""
        return self.No * self.dxo

    @property
    def dxo_fpm(self):
        """Real-space object pixel size for FPM.
        """
        return self.dxp * self.Np / self.No

    @property
    def Lo_fpm(self):
        """Real-space field of view of the FPM object, equal to that of the raw images."""
        return self.No * self.dxo_fpm

    @property
    def dfp(self):
        """Spatial-frequency pixel size of the probe grid, 1 / Lp."""
        return 1 / self.Lp

    @property
    def xo(self):
        """object coordinates 1D"""
        try:
            return np.linspace(-self.No / 2, self.No / 2, np.int(self.No)) * self.dxo
        except AttributeError as e:
            raise AttributeError(
                e, 'object pixel number "No" and/or pixel size "dxo" not defined yet'
            )

    @property
    def Xo(self):
        """Object coordinates 2D"""
        Xo, Yo = np.meshgrid(self.xo, self.xo)
        return Xo

    @property
    def Yo(self):
        """Object coordinates 2D"""
        Xo, Yo = np.meshgrid(self.xo, self.xo)
        return Yo

    # scan positions in pixel
    @property
    def positions(self):
        """estimated positions in pixel numbers(real space for CPM, Fourier space for FPM)
        note: Positions are given in row-column order and refer to the
        pixel in the upper left corner of the respective data matrix;
        -1st example: suppose the 2nd row of positions0 is [3, 4] and the
        operation mode is 'CPM'. That implies that the second intensity
        in the spectrogram updates an object patch that has
        its left uppper corner pixel at the pixel coordinates [3, 4]
        -2nd example: suppose the 2nd row of positions0 is [3, 4] and the
        operation mode is 'FPM'. That implies that the second intensity
        in the spectrogram is updates a patch which has pixel coordinates
        [3,4] in the high-resolution Fourier transform
        """
        if self.data.operationMode == "FPM":
            conv = -(1 / self.wavelength) * self.dxo * self.Np
            positions = np.round(
                conv
                * self.encoder_corrected
                / np.sqrt(
                    self.encoder_corrected[:, 0] ** 2
                    + self.encoder_corrected[:, 1] ** 2
                    + self.zled**2
                )[..., None]
            )

            try:
                positions = positions + self.No // 2 - self.Np // 2
            except:
                pass

            return positions.astype(int)
        else:
            return calculate_pixel_positions(
                self.encoder_corrected, self.dxo, self.No, self.Np, asint=True
            )

    # system property list
    @property
    def NAd(self):
        """Detection NA"""
        NAd = self.Ld / (2 * self.zo)
        return NAd

    @property
    def DoF(self):
        """expected Depth of field"""
        DoF = self.wavelength / self.NAd**2
        # self.Dof2 = 5.2 *self.dxp**2 /self.wavelength
        return DoF

    def _move_data_to_cpu(self):
        """
        Move all the required fields to the CPU
        :return:
        """
        transfer_fields_to_cpu(self, self.possible_GPU_fields, self.logger)

    def _move_data_to_gpu(self):
        transfer_fields_to_gpu(self, self.possible_GPU_fields, self.logger)

    def describe_reconstruction(self):
        minmax_tv = ''
        try:
            minmax_tv = f'(min: {self.params.TV_autofocus_min_z*1e3}, max: {self.params.TV_autofocus_max_z*1e3}.)'
        except TypeError: # one of them is none
            pass
        info = f"""
        Experimental data:
        - Number of ptychograms: {self.data.ptychogram.shape}
        - Number of pixels ptychogram: {self.data.Nd}
        - Ptychogram size: {self.data.Ld*1e3} mm
        - Pixel pitch: {self.data.dxd*1e6} um
        - Scan size: {1e3*(self.data.encoder.max(axis=0) - self.data.encoder.min(axis=0))} mm 

        Reconstruction:
        - number of pixels: {self.No}
        - Pixel pitch: {self.dxo*1e6} um
        - Field of view: {self.Lo*1e3} mm
        - Scan size in pixels: {self.positions.max(axis=0)- self.positions.min(axis=0)}
        - Propagation distance: {self.zo * 1e3} mm {minmax_tv}
        - Probe FoV: {self.Lp*1e3} mm

        Derived parameters:
        - NA detector: {self.NAd}
        - DOF: {self.DoF*1e6} um

        """
        self.logger.info(info)
        return info

    @property
    def quadraticPhase(self):
        """These functions are cached internally in Python and therefore no longer required."""
        raise NotImplementedError("Quadratic phase is no longer cached. ")

    @property
    def transferFunction(self):
        raise NotImplementedError("Quad phase is not longer cached")

    @property
    def Q1(self):
        raise NotImplementedError("Q1 is no longer available")

    @property
    def Q2(self):
        raise NotImplementedError("Q2 is no longer available")

    def TV_autofocus(self, params: Params, loop):

        """Perform an autofocusing step based on optimizing the total variation.

        If not required, returns none. Otherwise, returns the value of the TV at the current z0."""
        start_time = time.time()

        if self.data.operationMode == "FPM":
            raise NotImplementedError(
                f"Not implemented/tested for FPM. Set params.TV_autofocus to False. Got {params.TV_autofocus}"
            )
        if not params.TV_autofocus:
            return None, None, None
        if loop is not None:
            if loop % params.TV_autofocus_run_every != 0:
                return None, None, None

        if params.l2reg:
            self.logger.warning(
                "Both TV_autofocus and L2reg are turned on. This usually leads to poor performance. Consider disabling l2reg if the probe collapses to focal points"
            )

        d = params.TV_autofocus_range_dof
        nplanes = params.TV_autofocus_nplanes
        dz = np.linspace(-1, 1, nplanes) * d * self.DoF

        if params.TV_autofocus_what == "object":
            field = self.object[self.nlambda // 2, 0, 0, self.nslice // 2, :, :]
        elif params.TV_autofocus_what == "probe":
            field = self.probe[self.nlambda // 2, 0, 0, self.nslice // 2, :, :]
        else:
            raise NotImplementedError(
                f"So far, only object and probe are valid options for params.T_autofocus_what. Got {params.TV_autofocus_what}"
            )

        ss = params.TV_autofocus_roi
        if isinstance(ss, list):
            # semi-smart way to set up an AOI.
            # if the coordinates are a list, expand the list for y and x
            ss = np.array(ss)
            if ss.ndim == 1:
                ss = np.repeat(ss[None], axis=0, repeats=2)

            N = field.shape[-1]
            sy, sx = [slice(int(s[0] * N), int(s[1] * N)) for s in ss]
            # make them the same size if they're not
            sy = slice(sy.start, sy.start + sx.stop - sx.start)
        else:
            sy, sx = ss, ss

        merit, OEs = metric_at(
            field,
            dz,
            self.dxo,  # same as dxp
            self.wavelength,
            (sy, sx),
            intensity_only=self.params.TV_autofocus_intensityonly,
            metric=self.params.TV_autofocus_metric,
            return_propagated=True,
        )
        # from here on we are looking at 11 data points, work on CPU
        # as it's much more convenient and faster
        feedback = np.sum(dz * merit) / np.sum(merit)

        scores = np.vstack([self.zo + dz, merit])

        self.zMomentum *= params.TV_autofocus_friction
        self.zMomentum += params.TV_autofocus_stepsize * feedback
        # now, clip it to the bounds
        delta_z = self.zo - np.clip(
            self.zo + self.zMomentum,
            self.params.TV_autofocus_min_z,
            self.params.TV_autofocus_max_z,
        )
        self.zo -= delta_z
        end_time = time.time()
        self.logger.info(
            f"TV autofocus took {end_time-start_time} seconds, and moved focus by {-delta_z*1e6} micron"
        )
        indices = [nplanes//2, np.argmax(merit)]
        OEs = OEs[indices]
        phexp = OEs.sum((-2,-1), keepdims=True).conj()
        phexp = phexp / abs(phexp)
        OEs *= phexp
        return merit[nplanes//2] / asNumpyArray(abs(self.object[..., sy, sx]).mean()), np.hstack(OEs), (scores, self.zo)

    def reset_TV_autofocus(self):
        """Reset the settings of TV autofocus. Can be useful to reset the memory effect if the steps are getting really large."""
        self.zMomentum = 0

    @property
    def TV(self):
        """Return the TV of the object"""
        return TV(self.object, 1e-2)

zo property writable

Distance from sample to detector. Also updates all derived qualities.

xd property

Detector coordinates 1D

Xd property

Detector coordinates 2D

Yd property

Detector coordinates 2D

Ld property

Detector size in SI units.

Np property

Probe pixel numbers

Lp property

probe size in SI units

xp property

Probe coordinates 1D

Xp property

Probe coordinates 2D

Yp property

Probe coordinates 2D

dxo property

object pixel size, always equal to probe pixel size.

Lo property

Field of view (entrance pupil plane)

dxo_fpm property

Real-space object pixel size for FPM.

Lo_fpm property

Real-space field of view of the FPM object, equal to that of the raw images.

dfp property

Spatial-frequency pixel size of the probe grid, 1 / Lp.

xo property

object coordinates 1D

Xo property

Object coordinates 2D

Yo property

Object coordinates 2D

positions property

estimated positions in pixel numbers(real space for CPM, Fourier space for FPM) note: Positions are given in row-column order and refer to the pixel in the upper left corner of the respective data matrix; -1st example: suppose the 2nd row of positions0 is [3, 4] and the operation mode is 'CPM'. That implies that the second intensity in the spectrogram updates an object patch that has its left uppper corner pixel at the pixel coordinates [3, 4] -2nd example: suppose the 2nd row of positions0 is [3, 4] and the operation mode is 'FPM'. That implies that the second intensity in the spectrogram is updates a patch which has pixel coordinates [3,4] in the high-resolution Fourier transform

NAd property

Detection NA

DoF property

expected Depth of field

quadraticPhase property

These functions are cached internally in Python and therefore no longer required.

TV property

Return the TV of the object

copyAttributesFromExperiment(data)

Copy all the attributes from the experiment that are in listOfReconstructionProperties (CPM or FPM)

Source code in PtyLab/Reconstruction/Reconstruction.py

def copyAttributesFromExperiment(self, data: ExperimentalData):
    """
    Copy all the attributes from the experiment that are in listOfReconstructionProperties (CPM or FPM)
    """
    self.logger.debug("Copying attributes from Experimental Data")
    if self.data.operationMode == "CPM":
        listOfReconstructionProperties = self.listOfReconstructionPropertiesCPM
    elif self.data.operationMode == "FPM":
        listOfReconstructionProperties = self.listOfReconstructionPropertiesFPM
    for key in listOfReconstructionProperties:
        self.logger.info("Copying attribute %s", key)
        # setattr(self, key, copy(np.array(getattr(data, key))))
        setattr(self, key, copy(getattr(data, key)))

    # set the distance, this has to be last
    # In FPM the sample to detector distance is irrelevant
    # LED-to-sample distance is the more important factor that affects
    # wave propagation and illumination angle
    if self.data.operationMode == "CPM":
        self.zo = getattr(data, "zo")

    # set the original positions
    if self.encoder_corrected is None:
        self.encoder_corrected = data.encoder.copy()

reset_positioncorrection()

Reset the position corrections.

Source code in PtyLab/Reconstruction/Reconstruction.py

def reset_positioncorrection(self):
    """Reset the position corrections."""
    self.encoder_corrected = self.data.encoder.copy()

computeParameters()

compute parameters that can be altered by the user later.

Source code in PtyLab/Reconstruction/Reconstruction.py

def computeParameters(self):
    """
    compute parameters that can be altered by the user later.
    """

    if self.data.operationMode == "CPM":
        # CPM dxp (depending on the propagatorType, if none given, assum Fraunhofer/Fresnel)
        # self.dxp = self.wavelength * self._zo / self.Ld
        # if entrancePupilDiameter is not provided in the hdf5 file, set it to be one third of the probe FoV.
        if self.data.entrancePupilDiameter is None:
            self.data.entrancePupilDiameter = self.Lp / 3
        # if spectralDensity is not provided in the hdf5 file, set it to be a 1d array of the wavelength
        if isinstance(self.spectralDensity, type(None)):
            # this is a confusing name, it should be the wavelengths, not the intensity of the different
            # wavelengths
            self.spectralDensity = np.atleast_1d(self.wavelength)

    elif self.data.operationMode == "FPM":
        # FPM dxp (different from CPM due to lens-based systems)
        self.dxp = self.dxd / self.data.magnification
        # the propagation distance that is meaningful in this context is the
        # illumination to sample distance for LED array based microscopes
        self.zo = self.zled
        # if NA is not provided in the hdf5 file, set Fourier pupil entrance diameter it to be half of the Fourier space FoV.
        # then estimate the NA from the pupil diameter in the Fourier plane
        if isinstance(self.NA, type(None)):
            self.data.entrancePupilDiameter = self.Lp / 2
            self.NA = (
                self.data.entrancePupilDiameter
                * self.wavelength
                / (2 * self.dxp**2 * self.Np)
            )
        else:
            # compute the pupil radius in the Fourier plane
            self.data.entrancePupilDiameter = (
                2 * self.dxp**2 * self.Np * self.NA / self.wavelength
            )

    # set object pixel numbers
    if not hasattr(self, 'No'):
        self.No = (
            self.Np * 2**2
        )  # unimportant but leave it here as it's required for self.positions
        # we need space for the probe as well, on both sides that would be half the probe
        range_pixels = np.max(self.positions, axis=0) - np.min(self.positions, axis=0)
        # print(range_pixels)
        range_pixels = np.max(range_pixels) + self.Np * 2
        if range_pixels % 2 == 1:
            range_pixels += 1
        self.No = np.max([self.Np, range_pixels])

initializeSettings()

Initialize the attributes that have to do with a reconstruction or experimentalData fields which will become "reconstruction"

This method just sets the settings. It sets the what kind of initial guess should be used for initialObject and initialProbe but it does not compute them yet. That will be done by calling initializeObjectProbe()

:return:

Source code in PtyLab/Reconstruction/Reconstruction.py

def initializeSettings(self):
    """
    Initialize the attributes that have to do with a reconstruction
    or experimentalData fields which will become "reconstruction"

    This method just sets the settings. It sets the what kind of initial guess should be used for initialObject
    and initialProbe but it does not compute them yet. That will be done by calling initializeObjectProbe()

    :return:
    """
    # create a 6D object where which allows to have:
    # 1. polychromatic = nlambda
    # 2. mixed state object - nosm
    # 3. mixed state probe - npsm
    # 4. multislice object (thick) - nslice
    self.nlambda = 1
    self.nosm = 1
    self.npsm = 1
    self.nslice = 1

    # beam and object purity (# default initial value for plots.)
    self.purityProbe = 1.0
    self.purityObject = 1.0
    self.purityProbeHist = []

    self.positions0 = self.positions.copy()

    if self.data.operationMode == "FPM":
        self.initialObject = "upsampled"
        self.initialProbe = "circ"
    elif self.data.operationMode == "CPM":
        self.initialProbe = "circ"
        self.initialObject = "ones"
    else:
        self.initialProbe = "circ"
        self.initialObject = "ones"

prepare_probe(i)

Replace probe with the i-th TSVD estimate.

This function is used in OPRP

Source code in PtyLab/Reconstruction/Reconstruction.py

def prepare_probe(self, i):
    """Replace probe with the i-th TSVD estimate.

    This function is used in OPRP
    """
    raise NotImplementedError()

loadResults(fileName, datatype='probe') staticmethod

Loads data from a ptylab reconstruction file.

Source code in PtyLab/Reconstruction/Reconstruction.py

@staticmethod
def loadResults(fileName, datatype='probe'):
    '''
    Loads data from a ptylab reconstruction file.
    '''
    with h5py.File(fileName) as archive:
        data = np.copy(np.array(archive[datatype]))
    return data

load_object(filename)

Load the object from a previous reconstruction

Parameters

filename: .hdf5 file Filenamne of the reconstruction whose object should be loaded.

Returns

Source code in PtyLab/Reconstruction/Reconstruction.py

def load_object(self, filename):
    """
    Load the object from a previous reconstruction

    Parameters
    ----------
    filename: .hdf5 file
        Filenamne of the reconstruction whose object should be loaded.

    Returns
    -------

    """
    with h5py.File(filename, "r") as archive:
        obj = np.array(archive["object"])
        obj = obj[
            : self.shape_O[0],
            : self.shape_O[1],
            : self.shape_O[2],
            : self.shape_O[3],
            : self.shape_O[4],
            : self.shape_O[5],
        ]
        if np.all(np.array(obj.shape) == np.array(self.shape_O)):
            self.object = obj
        else:
            raise RuntimeError(
                f'Shape of saved probe cannot be extended to shape of required probe. File: {archive["object"].shape}. Need: {self.shape_O}'
            )

load_probe(filename, expand_npsm=False, center_phase=False)

Load the probe from a previous reconstruction.

Parameters

filename: .hdf5 file The filename of the reconstruction whose probe should be loaded.

Returns

Source code in PtyLab/Reconstruction/Reconstruction.py

def load_probe(self, filename, expand_npsm=False, center_phase=False):
    """
    Load the probe from a previous reconstruction.

    Parameters
    ----------
    filename: .hdf5 file
        The filename of the reconstruction whose probe should be loaded.

    Returns
    -------

    """
    with h5py.File(filename, "r") as archive:
        probe = np.array(archive["probe"])
        N_probe_read = probe.shape[-1]
        # roughly extract the center
        ss = slice(np.clip(N_probe_read//2-self.Np//2, 0, None), np.clip(N_probe_read//2-self.Np//2+int(self.Np), 0, N_probe_read))
        probe = probe[
            : self.nlambda,
            :1,
            : self.npsm,
            : self.nslice,
            ss,ss
        ]
        if np.all(np.array(probe.shape) == np.array(self.shape_P)):
            self.probe = probe
        else:
            raise RuntimeError(
                f'Shape of saved probe cannot be extended to shape of required probe. File: {archive["probe"].shape}. Need: {self.shape_P}'
            )
    if center_phase:
        self._center_probe_angle()

load(filename)

Load the results given by saveResults.

Source code in PtyLab/Reconstruction/Reconstruction.py

def load(self, filename):
    """Load the results given by saveResults."""
    with h5py.File(filename, "r") as archive:

        self.probe = np.array(archive["probe"])
        self.object = np.array(archive["object"])
        self.error = np.array(archive["error"])
        self.wavelength = np.array(archive["wavelength"])
        self.dxp = np.array(archive["dxp"])
        self.purityProbe = np.array(archive["purityProbe"])
        self.purityObject = np.array(archive["purityObject"])
        self.zo = np.array(archive["zo"])
        if "theta" in archive.keys():
            self.theta = np.array(archive["theta"])

saveResults(fileName='recent', type='all', squeeze=False)

Save reconstruction results.

Parameters

fileName type squeeze

Returns

Source code in PtyLab/Reconstruction/Reconstruction.py

def saveResults(self, fileName="recent", type="all", squeeze=False):
    """
    Save reconstruction results.


    Parameters
    ----------
    fileName
    type
    squeeze


    Returns
    -------

    """

    allowed_save_types = ["all", "object", "probe", "probe_stack"]
    if type not in allowed_save_types:
        raise NotImplementedError(
            f"Only {allowed_save_types} are allowed keywords for type"
        )
    if not squeeze:
        squeezefun = lambda x: x
    else:
        squeezefun = np.squeeze
    if type == "all":
        if self.data.operationMode == "CPM":
            with h5py.File(fileName, "w") as hf:
                hf.create_dataset("probe", data=self.probe, dtype="complex64")
                hf.create_dataset("object", data=self.object, dtype="complex64")
                hf.create_dataset("error", data=self.error, dtype="f")
                hf.create_dataset("zo", data=self._zo, dtype="f")
                hf.create_dataset("wavelength", data=self.wavelength, dtype="f")
                hf.create_dataset("dxp", data=self.dxp, dtype="f")
                hf.create_dataset("purityProbe", data=self.purityProbe, dtype="f")
                hf.create_dataset("purityObject", data=self.purityObject, dtype="f")
                hf.create_dataset('I object', data=abs(self.object), dtype='f')
                hf.create_dataset('I probe', data=abs(self.probe), dtype='f')
                hf.create_dataset('encoder_corrected', data=self.encoder_corrected)

                if hasattr(self, "theta"):
                    if self.theta != None:
                        hf.create_dataset("theta", data=self.theta, dtype="f")

        if self.data.operationMode == "FPM":
            hf = h5py.File(fileName, "w")
            hf.create_dataset("probe", data=self.probe, dtype="complex64")
            hf.create_dataset("object", data=self.object, dtype="complex64")
            hf.create_dataset("error", data=self.error, dtype="f")
            hf.create_dataset("zled", data=self.zled, dtype="f")
            hf.create_dataset("wavelength", data=self.wavelength, dtype="f")
            hf.create_dataset("dxp", data=self.dxp, dtype="f")
    elif type == "probe":
        with h5py.File(fileName, "w") as hf:
            hf.create_dataset(
                "probe", data=squeezefun(self.probe), dtype="complex64"
            )
    elif type == "object":
        with h5py.File(fileName, "w") as hf:
            hf.create_dataset(
                "object", data=squeezefun(self.object), dtype="complex64"
            )
    elif type == "probe_stack":
        hf = h5py.File(fileName + '_probe_stack.hdf5', 'w')
        hf.create_dataset('probe_stack', data=self.probe_stack.get(), dtype='complex64')
    print("The reconstruction results (%s) have been saved" % type)

TV_autofocus(params, loop)

Perform an autofocusing step based on optimizing the total variation.

If not required, returns none. Otherwise, returns the value of the TV at the current z0.

Source code in PtyLab/Reconstruction/Reconstruction.py

def TV_autofocus(self, params: Params, loop):

    """Perform an autofocusing step based on optimizing the total variation.

    If not required, returns none. Otherwise, returns the value of the TV at the current z0."""
    start_time = time.time()

    if self.data.operationMode == "FPM":
        raise NotImplementedError(
            f"Not implemented/tested for FPM. Set params.TV_autofocus to False. Got {params.TV_autofocus}"
        )
    if not params.TV_autofocus:
        return None, None, None
    if loop is not None:
        if loop % params.TV_autofocus_run_every != 0:
            return None, None, None

    if params.l2reg:
        self.logger.warning(
            "Both TV_autofocus and L2reg are turned on. This usually leads to poor performance. Consider disabling l2reg if the probe collapses to focal points"
        )

    d = params.TV_autofocus_range_dof
    nplanes = params.TV_autofocus_nplanes
    dz = np.linspace(-1, 1, nplanes) * d * self.DoF

    if params.TV_autofocus_what == "object":
        field = self.object[self.nlambda // 2, 0, 0, self.nslice // 2, :, :]
    elif params.TV_autofocus_what == "probe":
        field = self.probe[self.nlambda // 2, 0, 0, self.nslice // 2, :, :]
    else:
        raise NotImplementedError(
            f"So far, only object and probe are valid options for params.T_autofocus_what. Got {params.TV_autofocus_what}"
        )

    ss = params.TV_autofocus_roi
    if isinstance(ss, list):
        # semi-smart way to set up an AOI.
        # if the coordinates are a list, expand the list for y and x
        ss = np.array(ss)
        if ss.ndim == 1:
            ss = np.repeat(ss[None], axis=0, repeats=2)

        N = field.shape[-1]
        sy, sx = [slice(int(s[0] * N), int(s[1] * N)) for s in ss]
        # make them the same size if they're not
        sy = slice(sy.start, sy.start + sx.stop - sx.start)
    else:
        sy, sx = ss, ss

    merit, OEs = metric_at(
        field,
        dz,
        self.dxo,  # same as dxp
        self.wavelength,
        (sy, sx),
        intensity_only=self.params.TV_autofocus_intensityonly,
        metric=self.params.TV_autofocus_metric,
        return_propagated=True,
    )
    # from here on we are looking at 11 data points, work on CPU
    # as it's much more convenient and faster
    feedback = np.sum(dz * merit) / np.sum(merit)

    scores = np.vstack([self.zo + dz, merit])

    self.zMomentum *= params.TV_autofocus_friction
    self.zMomentum += params.TV_autofocus_stepsize * feedback
    # now, clip it to the bounds
    delta_z = self.zo - np.clip(
        self.zo + self.zMomentum,
        self.params.TV_autofocus_min_z,
        self.params.TV_autofocus_max_z,
    )
    self.zo -= delta_z
    end_time = time.time()
    self.logger.info(
        f"TV autofocus took {end_time-start_time} seconds, and moved focus by {-delta_z*1e6} micron"
    )
    indices = [nplanes//2, np.argmax(merit)]
    OEs = OEs[indices]
    phexp = OEs.sum((-2,-1), keepdims=True).conj()
    phexp = phexp / abs(phexp)
    OEs *= phexp
    return merit[nplanes//2] / asNumpyArray(abs(self.object[..., sy, sx]).mean()), np.hstack(OEs), (scores, self.zo)

reset_TV_autofocus()

Reset the settings of TV autofocus. Can be useful to reset the memory effect if the steps are getting really large.

Source code in PtyLab/Reconstruction/Reconstruction.py

def reset_TV_autofocus(self):
    """Reset the settings of TV autofocus. Can be useful to reset the memory effect if the steps are getting really large."""
    self.zMomentum = 0

calculate_pixel_positions(encoder_corrected, dxo, No, Np, asint)

Calculate the pixel positions.

Source code in PtyLab/Reconstruction/Reconstruction.py

def calculate_pixel_positions(encoder_corrected, dxo, No, Np, asint):
    """
    Calculate the pixel positions.
    """
    positions = np.round(
        encoder_corrected / dxo
    )  # encoder is in m, positions0 and positions are in pixels
    positions = positions + No // 2 - Np // 2
    if asint:
        positions = positions.astype(int)
    return positions