Utilities

PtyLab.utils

alignment

show_alignment(reconstruction, data, params, engine)

Show a viewer which gives information about the alignment of the probe and the object

Source code in PtyLab/utils/alignment.py

def show_alignment(
    reconstruction: Reconstruction,
    data: ExperimentalData,
    params: Params,
    engine: Engines.BaseEngine,
):
    """Show a viewer which gives information about the alignment of the probe and the object"""
    # currently a hacky way for this, these napari implementations must 
    # later be moved an optional sub-package. 
    try:
        import napari
        viewer = napari.Viewer()
    except ImportError:
        msg = "Install napari to access this `NapariMonitor` implementation"
        raise ImportError(msg)

    # sort all the images by distance to the center
    mean_pos = reconstruction.positions.mean(0, keepdims=True)
    order = np.argsort(np.linalg.norm(reconstruction.positions - mean_pos, axis=-1))[
        :15
    ]
    ptycho_ordered = data.ptychogram[order]
    # do one iteration of all of them

    from PtyLab.Operators.Operators import detector2object, object2detector

    reconstruction.initializeObjectProbe()
    reconstruction.esw = reconstruction.probe
    # do all the propagators
    from PtyLab.Operators.Operators import fresnelPropagator

    z0_0 = reconstruction.zo
    # esw, updated_esw = detector2object(ptycho_ordered-ptycho_ordered.mean(), params, reconstruction)
    # updated_esw = updated_esw[0,0,0]
    # print(updated_esw.shape)
    # viewer.add_image(abs(updated_esw**2), name='refocused')
    # move them so they 'overlap'
    # for i, position in enumerate(reconstruction.positions[order]-reconstruction.positions[order][0]):
    #     # do one iteration for all of them and keep them in memory
    #     row, col = position
    #     updated_esw[i] = np.roll(np.roll(updated_esw[i], axis=-2, shift=col),
    #                              axis=-1, shift=row)
    # viewer.add_image(abs(updated_esw**2), name='aligned refocused')

    viewer.add_image(ptycho_ordered, name="ptychogram radially ordered")
    # propagate them
    viewer.show()

downloader

download_with_progress(url, filename)

Download a file and show a progress bar while doing so.

Parameters

url: URL of data filename: path to save it

Returns

None

Source code in PtyLab/utils/downloader.py

def download_with_progress(url, filename):
    """
    Download a file and show a progress bar while doing so.

    Parameters
    ----------
    url: URL of data
    filename: path to save it

    Returns
    -------
    None
    """
    pbar = None
    last_position = 0

    def show_progress(block_num, block_size, total_size):
        nonlocal  pbar
        nonlocal last_position
        if pbar is None:
            pbar = tqdm.tqdm(total=total_size, unit='Mb', unit_scale=1e-6 )#True, leave=False )#.ProgressBar(maxval=total_size)
        #pbar.start()


        downloaded = block_num * block_size
        update = downloaded - last_position
        last_position = downloaded
        if downloaded < total_size:
            #pbar.moveto(downloaded/total_size)
            pbar.update(update)

    urllib.request.urlretrieve(url, filename, show_progress)

fsvd

============================================================================= Randomized SVD. See Halko, Martinsson, Tropp's 2011 SIAM paper:

This file has been adopted to fit in PtyLab by making it GPU-aware by Dirk Boonzajer

"Finding structure with randomness: Probabilistic algorithms for constructing approximate matrix decompositions" =============================================================================

rsvd(A, rank, n_oversamples=None, n_subspace_iters=None, return_range=False)

Randomized SVD (p. 227 of Halko et al).

:param A: (m x n) matrix. :param rank: Desired rank approximation. :param n_oversamples: Oversampling parameter for Gaussian random samples. :param n_subspace_iters: Number of power iterations. :param return_range: If True, return basis for approximate range of A. :return: U, S, and Vt as in truncated SVD.

Source code in PtyLab/utils/fsvd.py

def rsvd(A, rank, n_oversamples=None, n_subspace_iters=None,
         return_range=False):
    """Randomized SVD (p. 227 of Halko et al).

    :param A:                (m x n) matrix.
    :param rank:             Desired rank approximation.
    :param n_oversamples:    Oversampling parameter for Gaussian random samples.
    :param n_subspace_iters: Number of power iterations.
    :param return_range:     If `True`, return basis for approximate range of A.
    :return:                 U, S, and Vt as in truncated SVD.
    """
    xp = getArrayModule(A)
    if n_oversamples is None:
        # This is the default used in the paper.
        n_samples = 2 * rank
    else:
        n_samples = rank + n_oversamples

    # Stage A.
    Q = find_range(A, n_samples, n_subspace_iters)

    # Stage B.
    B = Q.T.conj() @ A
    U_tilde, S, Vt = xp.linalg.svd(B, full_matrices=False)
    U = Q @ U_tilde

    # Truncate.
    U, S, Vt = U[:, :rank], S[:rank], Vt[:rank, :]

    # This is useful for computing the actual error of our approximation.
    if return_range:
        return U, S, Vt, Q
    return U, S, Vt

find_range(A, n_samples, n_subspace_iters=None)

Algorithm 4.1: Randomized range finder (p. 240 of Halko et al).

Given a matrix A and a number of samples, computes an orthonormal matrix that approximates the range of A.

:param A: (m x n) matrix. :param n_samples: Number of Gaussian random samples. :param n_subspace_iters: Number of subspace iterations. :return: Orthonormal basis for approximate range of A.

Source code in PtyLab/utils/fsvd.py

def find_range(A, n_samples, n_subspace_iters=None):
    """Algorithm 4.1: Randomized range finder (p. 240 of Halko et al).

    Given a matrix A and a number of samples, computes an orthonormal matrix
    that approximates the range of A.

    :param A:                (m x n) matrix.
    :param n_samples:        Number of Gaussian random samples.
    :param n_subspace_iters: Number of subspace iterations.
    :return:                 Orthonormal basis for approximate range of A.
    """
    xp = getArrayModule(A)
    m, n = A.shape
    O = 1j*xp.random.normal(0,1.0, size=(n, n_samples))
    O +=xp.random.normal(0,1.0, size=(n, n_samples))
    O = O.astype(xp.complex64)
    #O = xp.random.randn(n, n_samples) + 1j * xp.random.randn(n, n_samples)
    Y = A @ O

    if n_subspace_iters:
        return subspace_iter(A, Y, n_subspace_iters)
    else:
        return ortho_basis(Y)

subspace_iter(A, Y0, n_iters)

Algorithm 4.4: Randomized subspace iteration (p. 244 of Halko et al).

Uses a numerically stable subspace iteration algorithm to down-weight smaller singular values.

:param A: (m x n) matrix. :param Y0: Initial approximate range of A. :param n_iters: Number of subspace iterations. :return: Orthonormalized approximate range of A after power iterations.

Source code in PtyLab/utils/fsvd.py

def subspace_iter(A, Y0, n_iters):
    """Algorithm 4.4: Randomized subspace iteration (p. 244 of Halko et al).

    Uses a numerically stable subspace iteration algorithm to down-weight
    smaller singular values.

    :param A:       (m x n) matrix.
    :param Y0:      Initial approximate range of A.
    :param n_iters: Number of subspace iterations.
    :return:        Orthonormalized approximate range of A after power
                    iterations.
    """
    Q = ortho_basis(Y0)
    for _ in range(n_iters):
        Z = ortho_basis(A.T.conj() @ Q)
        Q = ortho_basis(A @ Z)
    return Q

ortho_basis(M)

Computes an orthonormal basis for a matrix.

:param M: (m x n) matrix. :return: An orthonormal basis for M.

Source code in PtyLab/utils/fsvd.py

def ortho_basis(M):
    """Computes an orthonormal basis for a matrix.

    :param M: (m x n) matrix.
    :return:  An orthonormal basis for M.
    """
    xp = getArrayModule(M)
    Q, _ = xp.linalg.qr(M)
    return Q

gpuUtils

getArrayModule(*args, **kwargs)

Return a numerical array processing module based on wether the array lives on the CPU or on the GPU.

See cupy.getArrayModule for details. :param args: :param kwargs: :return:

Source code in PtyLab/utils/gpuUtils.py

def getArrayModule(*args, **kwargs):
    """
    Return a numerical array processing module based on wether the array lives on the CPU or on the GPU.

    See cupy.getArrayModule for details.
    :param args:
    :param kwargs:
    :return:
    """
    if CP_AVAILABLE:
        return cp.get_array_module(*args, **kwargs)
    else:
        return np

asNumpyArray(ary)

Return a numpy.ndarray version of ary.

:param ary: numpy or cupy ndarray :return: cpu-version of ary

Source code in PtyLab/utils/gpuUtils.py

def asNumpyArray(ary) -> np.ndarray:
    """
    Return a numpy.ndarray version of `ary`.

    :param ary: numpy or cupy ndarray
    :return: cpu-version of ary

    """
    if CP_AVAILABLE:
        return cp.asnumpy(ary)
    else:
        return ary

transfer_fields_to_gpu(self, fields, logger, dtype='auto')

Move any fields defined in fields to the CPU. Fields has to be a list of strings with field names :param self: :param fields: :param logger: :param dtype: data type. If 'auto', will be set to np.float32 for real-valued data and np.complex64 for complex :return:

Source code in PtyLab/utils/gpuUtils.py

def transfer_fields_to_gpu(
    self: object, fields: list[str], logger: logging.Logger, dtype="auto"
):
    """
    Move any fields defined in fields to the CPU. Fields has to be a list of strings with field names
    :param self:
    :param fields:
    :param logger:
    :param dtype: data type. If 'auto', will be set to np.float32 for real-valued data and np.complex64 for complex
    :return:
    """
    for field in fields:
        if hasattr(self, field):  # This field is defined
            # move it to the CPU
            attribute = getattr(self, field)
            try:
                setattr(self, field, asCupyArray(attribute, dtype=dtype))
            except AttributeError:
                logger.error(f"Cannot set attribute {field}")
                raise
            logger.debug(f"Moved {field} to GPU")
        else:
            logger.debug(f"Skipped {field} as it is not defined")

transfer_fields_to_cpu(self, fields, logger)

Move any fields defined in fields to the CPU. Fields has to be a list of strings with field names :param self: :param fields: :param logger: :return:

Source code in PtyLab/utils/gpuUtils.py

def transfer_fields_to_cpu(self: object, fields: list[str], logger: logging.Logger):
    """
    Move any fields defined in fields to the CPU. Fields has to be a list of strings with field names
    :param self:
    :param fields:
    :param logger:
    :return:
    """
    for field in fields:
        if hasattr(self, field):  # This field is defined
            # move it to the CPU
            attribute = getattr(self, field)
            setattr(self, field, asNumpyArray(attribute))
            logger.debug(f"Moved {field} to CPU")
        else:
            logger.debug(f"Skipped {field} as it is not defined")

initializationFunctions

initialProbeOrObject(shape, type_of_init, data, logger=None)

Initialization objects are created for the reconstruction. Currently implemented: ones - every element is set to 1 + random noise circ - same as 'ones' but with a circular boundary constraint upsampled - upsampled low-resolution estimate (used for FPM)

Random noise is added to the arrays to enforce linear independence required for orthogonalization of modes

:return:

Source code in PtyLab/utils/initializationFunctions.py

def initialProbeOrObject(shape, type_of_init, data, logger: logging.Logger = None):
    """
    Initialization objects are created for the reconstruction. Currently
    implemented:
        ones - every element is set to 1 + random noise
        circ - same as 'ones' but with a circular boundary constraint
        upsampled - upsampled low-resolution estimate (used for FPM)

    Random noise is added to the arrays to enforce linear independence required
    for orthogonalization of modes

    :return:
    """
    if type(type_of_init) is np.ndarray:  # it has already been run
        if logger is not None:
            logger.warning(
                "initialObjectOrProbe was called but the object has already "
                "been initialized. Skipping."
            )
        return type_of_init
    supported_shapes = ["circ", 'circ_smooth', "rand", "gaussian", "ones", "upsampled"]
    if type_of_init not in supported_shapes:
        raise NotImplementedError(f'Got {type_of_init} for shape. Supported shapes are: {supported_shapes}')

    if type_of_init == "ones":
        return np.ones(shape) + 0.001 * np.random.rand(*shape)

    if type_of_init in ["circ", 'circ_smooth']:
        try:
            # BUG: This only works for the probe, not for the object
            pupil = circ(data.Xp, data.Yp, data.data.entrancePupilDiameter)
            initial_field = np.ones(shape, dtype=np.complex64) + 0.001 * np.random.rand(*shape)

            if 'smooth' in type_of_init:
                dia_pixel = data.data.entrancePupilDiameter / data.dxo
                pupil = ndimage.gaussian_filter(pupil.astype(np.float64), 0.1 * dia_pixel)

            initial_field *= pupil
            return initial_field

        except AttributeError as e:
            raise AttributeError(
                e, "probe/aperture/entrancePupilDiameter was not defined"
            )


    if type_of_init == "upsampled":
        low_res = ifft2c(np.sqrt(np.mean(data.data.ptychogram, 0)))
        pad_size = (int((data.No - data.Np) / 2), int((data.No - data.Np) / 2))
        upsampled = np.pad(
            low_res, pad_size, mode="constant", constant_values=0
        )  # * data.No / data.Np
        return np.ones(shape) * upsampled

scanGrids

tsp_ga

genetic algorithm for traveling salesman problem.

Source code in PtyLab/utils/scanGrids.py

class tsp_ga:
    """
    genetic algorithm for traveling salesman problem.
    """

    def __init__(
        self, R, C, start="0", population_size=5, iterations=100, plotUpdateFrequency=20
    ):
        self.xy = np.vstack((R, C)).T
        self.start = start
        self.population_size = population_size
        self.iterations = int(iterations)
        self.plotUpdateFrequency = plotUpdateFrequency

        self.figure = plt.figure(num=999, clear=True, figsize=(5, 5))
        self.ax_scanGridOpt = self.figure.add_subplot(111)
        self.ax_scanGridOpt.set_title("optimized scan grid")
        self.ax_scanGridOpt.set_xlabel("um")
        self.ax_scanGridOpt.set_ylabel("um")
        (self.ax_scanGridOpt_plot,) = self.ax_scanGridOpt.plot(
            self.xy[:, 0], self.xy[:, 1], "o-"
        )
        self.figure.show()

        n = self.xy.shape[0]
        numEl = np.sum(x for x in range(n))
        cities_ = np.linspace(0, n - 1, n).astype(int)
        cities = ["" for i in range(cities_.size)]
        for i in range(cities_.size):
            cities[i] = np.array2string(cities_[i])

        edges = []
        totalDist = 0
        dist_dict = {c: {} for c in cities}
        for idx_1 in range(0, len(cities) - 1):
            for idx_2 in range(idx_1 + 1, len(cities)):
                city_a = cities[idx_1]  # xy[idx_1,0].astype(int)
                city_b = cities[idx_2]  # xy[idx_2,0].astype(int)
                # dist = get_distance(xy[int(city_a), 0], xy[int(city_a), 1], xy[int(city_b), 0], xy[int(city_b), 1], )
                dist = np.sqrt(
                    (self.xy[int(city_a), 0] - self.xy[int(city_b), 0]) ** 2
                    + (self.xy[int(city_a), 1] - self.xy[int(city_b), 1]) ** 2
                )
                totalDist += dist
                dist_dict[city_a][city_b] = dist
                edges.append((city_a, city_b, dist))

        self.meanDist = totalDist / numEl
        self.hash_map = dist_dict
        self.cities = [k for k in self.hash_map.keys()]
        self.cities.remove(start)
        self.genes = []
        self.generate_genes = vectorize(self.generate_genes)
        self.generate_genes()

    def generate_genes(self):
        for i in range(self.population_size):
            gene = [self.start]
            city_a = "0"
            options = [k for k in self.cities]
            loop = 0
            threshold = len(self.cities)
            while len(gene) < len(self.cities) + 1:
                city_b = random.choice(options)
                try:
                    dist = self.hash_map[city_a][city_b]
                except:
                    dist = self.hash_map[city_b][city_a]
                if (dist > self.meanDist / 4) & (loop < threshold):
                    loop += 1
                    continue
                if (dist > self.meanDist / 3) & (loop < 2 * threshold):
                    loop += 1
                    continue
                if (dist > self.meanDist / 2) & (loop < 3 * threshold):
                    loop += 1
                    continue
                if (dist > self.meanDist) & (loop < 4 * threshold):
                    loop += 1
                    continue
                loc = options.index(city_b)
                loop = 0
                gene.append(city_b)
                del options[loc]
                city_a = city_b
            gene.append(self.start)
            self.genes.append(gene)
        # return self.genes

    def GA_Matlab(self):
        globalMin = np.Inf
        distHistory = np.zeros((self.iterations, 1))
        tmpPop = []
        for i in range(4):
            tmpPop.append(self.genes[i])
        newPop = []

        counter = 0
        for iter in range(self.iterations):
            counter = counter + 1
            totalDist = []
            # Evaluate Each Population Member(Calculate Total Distance)
            for geneP in self.genes:  # range(self.population_size):
                d = 0  # self.hash_map[geneP[-1]][geneP[0]] # Closed Path
                for k in range(1, len(geneP)):
                    try:
                        city_a = geneP[k - 1]
                        city_b = geneP[k]
                        dist = self.hash_map[city_a][city_b]
                    except:
                        dist = self.hash_map[city_b][city_a]
                    d += dist
                totalDist.append(d)

            #  Find the Best Route in the Population
            if iter == 0:
                minDist = totalDist[0]
                index = 0
                distHistory[iter] = minDist
            else:
                minDist = np.amin(totalDist)
                index = np.argmin(totalDist)
                distHistory[iter] = minDist
            if minDist < globalMin:
                globalMin = minDist
                optRoute = self.genes[index]
            del newPop
            newPop = []
            # Genetic Algorithm Operators
            randomOrder = np.random.permutation(self.population_size)
            for p in range(4, self.population_size + 1, 4):
                randomOrderP = randomOrder[p - 4 : p]
                rtes = []
                dists = []
                for ii in range(4):
                    randomOrderP2 = randomOrderP[ii]
                    rtes.append(self.genes[randomOrderP2])
                    dists.append(totalDist[randomOrderP2])
                idx = np.argmin(dists)
                bestOf4Route = rtes[idx]
                routeInsertionPoints = np.transpose(
                    np.sort(np.ceil((len(geneP) - 2) * np.random.rand(1, 2)))
                )
                I = int(routeInsertionPoints[0])
                J = int(routeInsertionPoints[1])
                for k in range(4):  # Mutate the Best to get Three New Routes
                    tmpPop[k] = bestOf4Route.copy()
                    if k == 1:  # Flip
                        index_a = []
                        for ii in range(I, J + 1):
                            index_a.append(tmpPop[k][J + I - ii])
                        jj = 0
                        for ii in range(I, J + 1):
                            tmpPop[k][ii] = index_a[jj]
                            jj += 1
                    elif k == 2:  # Swap
                        index_a = tmpPop[k][I]
                        index_b = tmpPop[k][J]
                        tmpPop[k][I] = index_b
                        tmpPop[k][J] = index_a
                    elif k == 3:  # Slide
                        index_a = []
                        for ii in range(I + 1, J + 1):
                            index_a.append(tmpPop[k][ii])
                        index_b = tmpPop[k][I]
                        jj = 0
                        for ii in range(I, J):
                            tmpPop[k][ii] = index_a[jj]
                            jj += 1
                        tmpPop[k][J] = index_b
                    newPop.append(tmpPop[k])
            if iter % (self.iterations // self.plotUpdateFrequency) == 0:
                print("distance: %i um, numIteration: %i, " % (globalMin, iter))
                current_best_gene_r = np.array(optRoute).astype(int)

                self.ax_scanGridOpt_plot.set_data(
                    self.xy[current_best_gene_r, 0], self.xy[current_best_gene_r, 1]
                )
                self.figure.canvas.draw()
                self.figure.canvas.flush_events()

            del self.genes
            self.genes = newPop.copy()
        return globalMin, optRoute

    def converge(self):
        values = self.GA_Matlab()
        current_score = values[0]
        current_best_gene = values[1]
        return np.array(current_best_gene).astype(int)

GenerateNonUniformFermat(n, radius=1000, power=1)

generate spiral patterns :param n: number of points generated :param radius: radius in micrometer :param power = 1 is standard Fermat, power>1 yields more points towards the center of grid :return: R: row C: column

Source code in PtyLab/utils/scanGrids.py

def GenerateNonUniformFermat(n, radius=1000, power=1):
    """
    generate spiral patterns
    :param n: number of points generated
    :param radius: radius in micrometer
    :param power = 1 is standard Fermat, power>1 yields more points towards the center of grid
    :return:
    R: row
    C: column
    """
    # golden ratio
    r = np.sqrt(np.arange(0, n) / n)
    theta0 = 137.508 / 180 * np.pi
    theta = np.arange(0, n) * theta0
    C = radius * r**power * np.cos(theta)
    R = radius * r**power * np.sin(theta)
    return R, C

GenerateFermatSpiral(n, c)

generate Fermat Spiral :param n: number of points generated :param c: optional argument that controls scaling of spiral :return: R: row C: column

Source code in PtyLab/utils/scanGrids.py

def GenerateFermatSpiral(n, c):
    """
    generate Fermat Spiral
    :param n: number of points generated
    :param c: optional argument that controls scaling of spiral
    :return:
    R: row
    C: column
    """
    # golden ratio
    r = np.sqrt(np.arange(0, n)) * c
    theta0 = 137.508 / 180 * np.pi
    theta = np.arange(0, n) * theta0
    C = r * np.cos(theta)
    R = r * np.sin(theta)
    return R, C

GenerateConcentricGrid(Nr, s, rend)

generate concentric circles :param Nr: number of circles (or shells) :param s: number of pixels between points on each circle, roughly calculated as rend/Nr :param rend: end radius size (in pixel units) :return: R: row C: column

Source code in PtyLab/utils/scanGrids.py

def GenerateConcentricGrid(Nr, s, rend):
    """
    generate concentric circles
    :param Nr: number of circles (or shells)
    :param s: number of pixels between points on each circle, roughly calculated as rend/Nr
    :param rend: end radius size (in pixel units)
    :return:
    R: row
    C: column
    """
    dx = 1  # max Resolution (Schritt von einem zum anderen Pixel)
    rstart = dx
    r = np.linspace(rstart, rend, Nr)
    # determine number of positions on k'th shell
    nop = np.zeros(Nr, dtype=int)
    for k in np.arange(Nr):
        nop[k] = int(np.floor(2 * np.pi * r[k] / s))
    positions = np.zeros((sum(nop) + 1, 2))
    ind = 1
    for k in np.arange(1, Nr):
        dtheta = 2 * np.pi / nop[k]
        theta = np.arange(1, nop[k] + 1) * dtheta + 2 * np.pi / (k + 1)
        for l in np.arange(nop[k]):
            positions[ind, :] = r[k] * np.array([np.cos(theta[l]), np.sin(theta[l])])
            ind += 1
    positions = (np.floor(positions / dx)).astype(int)

    R = positions[:, 0]
    C = positions[:, 1]
    return R, C

GenerateRasterGrid(n, ds, randomOffset=False, amplitude=1)

generate a raster grid containing n*n points with a period of ds in pixelsize, with the option of adding randomOffsets. :param n: number of points per dimension :param ds: period (# of pixels) per dimension :param randomOffset: optional to add random offsets, default: False :param amplitude: amplitude for the random offsets, default: 1 :return: R: row C: column

Source code in PtyLab/utils/scanGrids.py

def GenerateRasterGrid(n, ds, randomOffset=False, amplitude=1):
    """
    generate a raster grid containing n*n points with a period of ds in pixelsize,
    with the option of adding randomOffsets.
    :param n: number of points per dimension
    :param ds: period (# of pixels) per dimension
    :param randomOffset: optional to add random offsets, default: False
    :param amplitude: amplitude for the random offsets, default: 1
    :return:
    R: row
    C: column
    """
    I, J = np.meshgrid(np.arange(n), np.arange(n))
    C = I.reshape(n**2) * ds
    R = J.reshape(n**2) * ds

    # center the scan grid at [0,0]
    # for even numbers
    if np.mod(n, 2) == 0:
        C = C - n * ds / 2
        R = R - n * ds / 2
    # for odd numbers
    else:
        C = C - (n - 1) * ds / 2
        R = R - (n - 1) * ds / 2

    if randomOffset:
        C = C + np.round(amplitude * (-1 + 2 * np.random.rand(C.shape)))
        R = R + np.round(amplitude * (-1 + 2 * np.random.rand(R.shape)))

    R = np.round(R - np.mean(R)).astype(int)
    C = np.round(C - np.mean(C)).astype(int)

    return R, C

utils

fft2c(field, fftshiftSwitch=False, *args, **kwargs)

performs 2 - dimensional unitary Fourier transformation, where energy is preserved sum( abs(g)2 ) == sum( abs(fft2c(g))2 ) if g is two - dimensional, fft2c(g) yields the 2D DFT of g if g is multi - dimensional, fft2c(g) yields the 2D DFT of g along the last two axes :param array: :return:

Source code in PtyLab/utils/utils.py

def fft2c(field, fftshiftSwitch=False, *args, **kwargs):
    """
    performs 2 - dimensional unitary Fourier transformation, where energy is preserved sum( abs(g)**2 ) == sum( abs(fft2c(g))**2 )
    if g is two - dimensional, fft2c(g) yields the 2D DFT of g
    if g is multi - dimensional, fft2c(g) yields the 2D DFT of g along the last two axes
    :param array:
    :return:
    """
    xp = getArrayModule(field)

    if fftshiftSwitch:
        return xp.fft.fft2(field, norm="ortho")
    else:
        axes = (-2, -1)
        return xp.fft.fftshift(
            xp.fft.fft2(xp.fft.ifftshift(field, axes=axes), norm="ortho"), axes=axes
        )

ifft2c(field, fftshiftSwitch=False)

performs 2 - dimensional inverse Fourier transformation, where energy is preserved sum( abs(G)2 ) == sum( abs(fft2c(g))2 ) if G is two - dimensional, fft2c(G) yields the 2D iDFT of G if G is multi - dimensional, fft2c(G) yields the 2D iDFT of G along the last two axes :param array: :return:

Source code in PtyLab/utils/utils.py

def ifft2c(field, fftshiftSwitch=False):
    """
    performs 2 - dimensional inverse Fourier transformation, where energy is preserved sum( abs(G)**2 ) == sum( abs(fft2c(g))**2 ) 
    if G is two - dimensional, fft2c(G) yields the 2D iDFT of G
    if G is multi - dimensional, fft2c(G) yields the 2D iDFT of G along the last two axes
    :param array:
    :return:
    """
    xp = getArrayModule(field)

    if fftshiftSwitch:
        return xp.fft.ifft2(field, norm="ortho")
    else:
        axes = (-2, -1)
        return xp.fft.fftshift(
            xp.fft.ifft2(xp.fft.ifftshift(field, axes=axes), norm="ortho"), axes=axes
        )

circ(x, y, D)

generate a binary array containing a circle on a 2D grid :param x: 2D x coordinate, normally calculated from meshgrid: x,y = np.meshgird((,)) :param y: 2D y coordinate, normally calculated from meshgrid: x,y = np.meshgird((,)) :param D: diameter :return: a binary 2D array

Source code in PtyLab/utils/utils.py

def circ(x, y, D):
    """
    generate a binary array containing a circle on a 2D grid
    :param x: 2D x coordinate, normally calculated from meshgrid: x,y = np.meshgird((,))
    :param y: 2D y coordinate, normally calculated from meshgrid: x,y = np.meshgird((,))
    :param D: diameter
    :return: a binary 2D array
    """
    circle = (x**2 + y**2) < (D / 2) ** 2
    return circle

rect(arr, threshold=0.5)

generate a binary array containing a rectangle on a 2D grid :param x: 2D x coordinate, normally calculated from meshgrid: x,y = np.meshgird((,)) :param threshold: threshold value to binarilize the input array, default value 0.5 :return: a binary array

Source code in PtyLab/utils/utils.py

def rect(arr, threshold=0.5):
    """
    generate a binary array containing a rectangle on a 2D grid
    :param x: 2D x coordinate, normally calculated from meshgrid: x,y = np.meshgird((,))
    :param threshold: threshold value to binarilize the input array, default value 0.5
    :return: a binary array
    """
    arr = abs(arr)
    return arr < threshold

posit(x)

returns 0 when x negative

Source code in PtyLab/utils/utils.py

def posit(x):
    """
    returns 0 when x negative
    """
    r = (x + abs(x)) / 2
    # r[r<0]=0 #todo check which way is faster
    return r

fraccircshift(A, shiftsize)

fraccircshift expands numpy.roll to fractional shifts values, using linear interpolation. :param A: ndarray :param shiftsize: shift size in each dimension of A, len(shiftsize)==A.ndim.

Source code in PtyLab/utils/utils.py

def fraccircshift(A, shiftsize):
    """
    fraccircshift expands numpy.roll to fractional shifts values, using linear interpolation.
    :param A: ndarray
    :param shiftsize: shift size in each dimension of A, len(shiftsize)==A.ndim.
    """
    integer = np.floor(shiftsize).astype(int)  # integer portions of shiftsize
    fraction = shiftsize - integer
    dim = len(shiftsize)
    # the dimensions are treated one after another
    for n in np.arange(dim):
        intn = integer[n]
        fran = fraction[n]
        shift1 = intn
        shift2 = intn + 1
        # linear interpolation
        A = (1 - fran) * np.roll(A, shift1, axis=n) + fran * np.roll(A, shift2, axis=n)
    return A

cart2pol(x, y)

Transform Cartesian to polar coordinates :param x: :param y: :return:

Source code in PtyLab/utils/utils.py

def cart2pol(x, y):
    """
    Transform Cartesian to polar coordinates
    :param x:
    :param y:
    :return:
    """
    th = np.arctan2(y, x)
    r = np.hypot(x, y)
    return th, r

orthogonalizeModes(p, method=None)

Imposes orthogonality through singular value decomposition :return:

Source code in PtyLab/utils/utils.py

def orthogonalizeModes(p, method=None):
    """
    Imposes orthogonality through singular value decomposition
    :return:
    """
    # orthogonolize modes only for npsm and nosm which are lcoated and indices 1, 2
    xp = getArrayModule(p)

    if method == "snapShots":
        try:
            p, normalizedEigenvalues, V = _snapShotModes(p, xp)
        except Exception as e:
            # cuSOLVER can fail for reasons that have nothing to do with this
            # data -- a CUDA install that cannot load libcusolver, or a memory
            # pool that has left it no room to allocate its workspace. Fall back
            # to the host rather than losing the reconstruction, but say loudly
            # what went wrong: this costs a device round trip every call, so it
            # is not something to run for a whole reconstruction unnoticed.
            logger.warning(
                "Orthogonalizing modes on the CPU rather than the GPU: %s: %s",
                type(e).__name__,
                e,
            )
            p, normalizedEigenvalues, V = _snapShotModes(asNumpyArray(p), np)
        return xp.asarray(p), normalizedEigenvalues, V

    else:
        U, s, V = xp.linalg.svd(
            p.reshape(p.shape[0], p.shape[1] * p.shape[2]), full_matrices=False
        )
        p = xp.dot(xp.diag(s), V).reshape(p.shape[0], p.shape[1], p.shape[2])
        normalizedEigenvalues = s**2 / xp.sum(s**2)

        return xp.asarray(p), normalizedEigenvalues, U.T.conj()

zernikeAberrations(Xp, Yp, D, z_coeff)

Compute the first 19 Zernike aberrations based on Zernike polynomials Based on https://en.wikipedia.org/wiki/Zernike_polynomials#OSA/ANSI_standard_indices

Xp,Yp - meshgrid coordinates D - radius within which to generate the zernike aberrations z_coeff - 19 element long list containing coefficients.

minimal example:

import matplotlib.pyplot as plt
import numpy as np

# create the circular dimensions which will define the size
# of a unit circle used for zernike aberration calculations
Xp,Yp = np.mgrid[-128:128, -128:128]
D = 128

# Get defocus aberration (4th index)
z_coeff = np.zeros(19)
z_coeff[4] = 3
Z = zernikeAberrations(Xp,Yp,D,z_coeff)

# plot the polynoial
plt.figure(1)
plt.imshow(np.angle(Z))
plt.show()

Source code in PtyLab/utils/utils.py

def zernikeAberrations(Xp, Yp, D, z_coeff):
    """
    Compute the first 19 Zernike aberrations based on Zernike polynomials
    Based on https://en.wikipedia.org/wiki/Zernike_polynomials#OSA/ANSI_standard_indices

    Xp,Yp - meshgrid coordinates
    D - radius within which to generate the zernike aberrations
    z_coeff - 19 element long list containing coefficients.

    minimal example:

        import matplotlib.pyplot as plt
        import numpy as np

        # create the circular dimensions which will define the size
        # of a unit circle used for zernike aberration calculations
        Xp,Yp = np.mgrid[-128:128, -128:128]
        D = 128

        # Get defocus aberration (4th index)
        z_coeff = np.zeros(19)
        z_coeff[4] = 3
        Z = zernikeAberrations(Xp,Yp,D,z_coeff)

        # plot the polynoial
        plt.figure(1)
        plt.imshow(np.angle(Z))
        plt.show()
    """

    aperture = circ(Xp, Yp, D)
    angle = np.double(np.arctan2(Yp, Xp)) * aperture
    p = np.double(np.hypot(Xp, Yp)) * aperture
    p = p / np.max(p)

    Z = dict()
    Z[0] = z_coeff[0]  # pistom
    Z[1] = z_coeff[1] * 4 ** (1 / 2.0) * p * np.cos(angle)
    # tip
    Z[2] = z_coeff[2] * 4 ** (1 / 2.0) * p * np.sin(angle)
    # tilt
    Z[3] = z_coeff[3] * 3 ** (1 / 2.0) * (2 * p**2 - 1)
    # defocus
    Z[4] = z_coeff[4] * 6 ** (1 / 2.0) * (p**2) * np.sin(2 * angle)
    # astigmatism
    Z[5] = z_coeff[5] * 6 ** (1 / 2.0) * (p**2) * np.cos(2 * angle)
    # astigmatism
    Z[6] = z_coeff[6] * 8 ** (1 / 2.0) * (3 * p**3 - 2 * p) * np.sin(angle)
    # coma
    Z[7] = z_coeff[7] * 8 ** (1 / 2.0) * (3 * p**3 - 2 * p) * np.cos(angle)
    # coma
    Z[8] = z_coeff[8] * 8 ** (1 / 2.0) * (p**3) * np.sin(3 * angle)
    # trefoil
    Z[9] = z_coeff[9] * 8 ** (1 / 2.0) * (p**3) * np.cos(3 * angle)
    # trefoil
    Z[10] = z_coeff[10] * 5 ** (1 / 2.0) * (6 * p**4 - 6 * p**2 + 1)
    # spherical
    Z[11] = (
        z_coeff[11] * 10 ** (1 / 2.0) * (4 * p**4 - 3 * p**2) * np.cos(2.0 * angle)
    )
    # 2nd astigmatism
    Z[12] = (
        z_coeff[12] * 10 ** (1 / 2.0) * (4 * p**4 - 3 * p**2) * np.sin(2.0 * angle)
    )
    # 2nd astigmatism
    Z[13] = z_coeff[13] * 10 ** (1 / 2.0) * (p**4) * np.cos(4.0 * angle)
    Z[14] = z_coeff[14] * 10 ** (1 / 2.0) * (p**4) * np.sin(4.0 * angle)
    Z[15] = (
        z_coeff[15]
        * 12 ** (1 / 2.0)
        * (10 * p**5 - 12 * p**3 + 3 * p)
        * np.cos(angle)
    )
    Z[16] = (
        z_coeff[16]
        * 12 ** (1 / 2)
        * (10 * p**5 - 12 * p**3 + 3 * p)
        * np.sin(angle)
    )
    Z[17] = z_coeff[17] * 12 ** (1 / 2) * (5 * p**5 - 4 * p**3) * np.cos(3 * angle)
    Z[18] = z_coeff[18] * 12 ** (1 / 2) * (5 * p**5 - 4 * p**3) * np.sin(3 * angle)

    return aperture * np.exp(1j * np.sum(list(Z.values())))

p2bin(im, binningFactor)

perform binning at a factor of power of 2, return binned image and the indices for before and after binning. :Params im: input image for binning :Params binningFactor: must be power of 2 in the current implementation :return:

Source code in PtyLab/utils/utils.py

def p2bin(im, binningFactor):
    """
    perform binning at a factor of power of 2, return binned image and the indices for before and after binning.
    :Params im: input image for binning
    :Params binningFactor: must be power of 2 in the current implementation
    :return:
    """
    M, N = im.shape
    if np.mod(binningFactor, 2) != 0 and binningFactor != 1:
        raise ValueError("binning factor needs to be a power of 2")
    if np.mod(M, binningFactor) != 0 or np.mod(N, binningFactor) != 0:
        raise ValueError(
            "#rows and #columns of reference need to be divided by binningFactor!"
        )

    if binningFactor != 1:
        for k in range(1, int(np.log2(binningFactor))):
            im_binned = bin2(im)
        im_binned_ind = range(im_binned.size)
        im_ind = np.arange(M * N).reshape(M // binningFactor, N, binningFactor)
        im_ind = np.stack(im_ind, axis=1).reshape(M, N)
    else:
        im_binned = im
        im_binned_ind = range(im_binned.size)
        im_ind = im_binned_ind
    return im_binned, im_ind, im_binned_ind

bin2(X)

perform 2-by-2 binning. :Params X: input 2D image for binning return: Y: output 2D image after 2-by-2 binning

Source code in PtyLab/utils/utils.py

def bin2(X):
    """
    perform 2-by-2 binning.
    :Params X: input 2D image for binning
    return: Y: output 2D image after 2-by-2 binning
    """
    # simple 2-fold binning
    m, n = X.shape
    Y = np.sum(X.reshape(2, m // 2, 2, n // 2), axis=(0, 2))
    return Y

visualisation

hsv2rgb(hsv)

Convert a 3D hsv np.ndarray to rgb (5 times faster than colorsys). https://stackoverflow.com/questions/27041559/rgb-to-hsv-python-change-hue-continuously h,s should be a numpy arrays with values between 0.0 and 1.0 v should be a numpy array with values between 0.0 and 255.0 :param hsv: np.ndarray of shape (x,y,3) :return: hsv2rgb returns an array of uints between 0 and 255.

Source code in PtyLab/utils/visualisation.py

def hsv2rgb(hsv: np.ndarray) -> np.ndarray:
    """
    Convert a 3D hsv np.ndarray to rgb (5 times faster than colorsys).
    https://stackoverflow.com/questions/27041559/rgb-to-hsv-python-change-hue-continuously
    h,s should be a numpy arrays with values between 0.0 and 1.0
    v should be a numpy array with values between 0.0 and 255.0
    :param hsv: np.ndarray of shape (x,y,3)
    :return: hsv2rgb returns an array of uints between 0 and 255.
    """
    xp = getArrayModule(hsv)
    rgb = xp.empty_like(hsv)
    rgb[..., 3:] = hsv[..., 3:]
    h, s, v = hsv[..., 0], hsv[..., 1], hsv[..., 2]
    i = (h * 6.0).astype("uint8")
    f = (h * 6.0) - i
    p = v * (1.0 - s)
    q = v * (1.0 - s * f)
    t = v * (1.0 - s * (1.0 - f))
    i = i % 6
    conditions = [s == 0.0, i == 1, i == 2, i == 3, i == 4, i == 5, i == i]
    rgb[..., 0] = xp.select(conditions, [v, q, p, p, t, v, v])  # , default=v)
    rgb[..., 1] = xp.select(conditions, [v, v, v, q, p, p, t])  # , default=t)
    rgb[..., 2] = xp.select(conditions, [v, p, t, v, v, q, p])  # , default=p)
    return rgb.astype("uint8")

complex2rgb(u, amplitudeScalingFactor=1, force_numpy=True, center_phase=False)

Preparation function for a complex plot, converting a 2D complex array into an rgb array :param u: a 2D complex array :return: an rgb array for complex plot

Source code in PtyLab/utils/visualisation.py

def complex2rgb(u, amplitudeScalingFactor=1, force_numpy=True, center_phase=False):
    """
    Preparation function for a complex plot, converting a 2D complex array into an rgb array
    :param u: a 2D complex array
    :return: an rgb array for complex plot
    """
    # hue (normalize angle)
    # if u is on the GPU, remove it as we can toss it now.
    xp = getArrayModule(u)
    # u = asNumpyArray(u)
    if center_phase:
        N = u.shape[-1]
        phexp = xp.sum(u[..., N // 3 : 2 * N // 3, N // 3 : 2 * N // 3], axis=(-2, -1))
        u = u * phexp.conj() / (abs(phexp) + 1e-9)
    h = xp.angle(u)
    h = (h + np.pi) / (2 * np.pi)
    # saturation  (ones)
    s = xp.ones_like(h)
    # value (normalize brightness to 8-bit)
    v = xp.abs(u)
    if amplitudeScalingFactor == "2sigma":
        ASF = v.mean() + 2 * np.std(v)
        ASF = ASF / v.max()
    elif amplitudeScalingFactor is None:
        ASF = 1.0 / v.max()
        amplitudeScalingFactor = ASF
    else:
        ASF = amplitudeScalingFactor

    if ASF != 1 and amplitudeScalingFactor != "2sigma":
        v[v > amplitudeScalingFactor * np.max(v)] = amplitudeScalingFactor * np.max(v)
    v = v / (xp.max(v) + xp.finfo(float).eps) * (2**8 - 1)

    hsv = xp.dstack([h, s, v])
    rgb = hsv2rgb(hsv)
    if isGpuArray(rgb) and force_numpy:
        rgb = rgb.get()
    return rgb

complex2rgb_vectorized(probe, **kwargs)

Turn complex image into rgb for every line.

The individual images are all autoscaled, so you cannot compare them.

Source code in PtyLab/utils/visualisation.py

def complex2rgb_vectorized(probe, **kwargs):
    """Turn complex image into rgb for every line.

    The individual images are all autoscaled, so you cannot compare them.
    """
    xp = getArrayModule(probe)
    original_shape = probe.shape
    probe = probe.reshape(-1, *probe.shape[-2:])
    probe_rgb = xp.array([complex2rgb(p, force_numpy=False, **kwargs) for p in probe])
    probe_rgb = probe_rgb.reshape(original_shape + (3,))
    return probe_rgb

plotExtent(pixelSize, axisUnit, shape)

Extent for imshow, expressed in axisUnit.

Real-space axes run from zero, as they always have. Reciprocal axes are centred on zero frequency instead, which is where the pupil sits.

:param pixelSize: sample spacing of the array, in SI units :param str axisUnit: any key of unitRatio :param shape: shape of the array that is plotted :return: [left, right, bottom, top] for imshow

Source code in PtyLab/utils/visualisation.py

def plotExtent(pixelSize, axisUnit, shape):
    """
    Extent for imshow, expressed in axisUnit.

    Real-space axes run from zero, as they always have. Reciprocal axes are
    centred on zero frequency instead, which is where the pupil sits.

    :param pixelSize: sample spacing of the array, in SI units
    :param str axisUnit: any key of unitRatio
    :param shape: shape of the array that is plotted
    :return: [left, right, bottom, top] for imshow
    """
    step = pixelSize * unitRatio[axisUnit]
    width, height = step * shape[1], step * shape[0]
    if axisUnit.startswith("1/"):
        return [-width / 2, width / 2, height / 2, -height / 2]
    return [0, width, height, 0]

complexPlot(rgb, ax=None, pixelSize=1, axisUnit='pixel')

Plot a 2D complex plot (hue for phase, brightness for amplitude). Input array need to be prepared by using the complex2rgb function. :param rgb: a rgb array that is converted from a 2D complex np.ndarray by using complex2rgb :param ax: Optional axis to plot in :param pixelSize: pixelSize in x and y, to display the physical dimension of the plot :param str axisUnit: Options: default 'pixel', 'm', 'cm', 'mm', 'um', and the reciprocal '1/m', '1/mm', '1/um' for Fourier-space quantities :return: An hsv plot

Source code in PtyLab/utils/visualisation.py

def complexPlot(rgb, ax=None, pixelSize=1, axisUnit="pixel"):
    """
    Plot a 2D complex plot (hue for phase, brightness for amplitude). Input array need to be prepared by using
    the complex2rgb function.
    :param rgb: a rgb array that is converted from a 2D complex np.ndarray by using complex2rgb
    :param ax: Optional axis to plot in
    :param pixelSize: pixelSize in x and y, to display the physical dimension of the plot
    :param str axisUnit: Options: default 'pixel', 'm', 'cm', 'mm', 'um', and the
        reciprocal '1/m', '1/mm', '1/um' for Fourier-space quantities
    :return: An hsv plot
    """

    if not ax:
        fig, ax = plt.subplots()
    extent = plotExtent(pixelSize, axisUnit, rgb.shape)

    im = ax.imshow(rgb, extent=extent, interpolation=None)
    ax.set_ylabel(axisUnit)
    ax.set_xlabel(axisUnit)

    divider = make_axes_locatable(ax)
    cax = divider.append_axes("right", size="5%", pad=0.1)

    norm = mpl.colors.Normalize(vmin=-np.pi, vmax=np.pi)
    scalar_mappable = mpl.cm.ScalarMappable(norm=norm, cmap=mpl.cm.hsv)
    scalar_mappable.set_array([])
    cbar = plt.colorbar(scalar_mappable, ax=ax, cax=cax, ticks=[-np.pi, 0, np.pi])
    cbar.ax.set_yticklabels([r"$-\pi$", "0", r"$\pi$"])
    return im

modeTile(P, normalize=True)

Tile 3D data into a single 2D array :param P: A complex np.ndarray :param normalize: normalize each mode individually :param pixelSize: pixelSize in x and y, to display the physical dimension of the plot :return: A big array with flattened modes

Source code in PtyLab/utils/visualisation.py

def modeTile(P, normalize=True):
    """
    Tile 3D data into a single 2D array
    :param P: A complex np.ndarray
    :param normalize: normalize each mode individually
    :param pixelSize: pixelSize in x and y, to display the physical dimension of the plot
    :return: A big array with flattened modes
    """
    if P.ndim == 3 and P.shape[0] > 1:
        if normalize:
            maxs = np.max(abs(P), axis=(-1, -2)) + 1e-6
            P = (P.T / maxs).T
        S = P.shape[0]
        s = math.ceil(np.sqrt(S))
        if s > np.sqrt(S):
            P = np.pad(P, ((0, s**2 - S), (0, 0), (0, 0)), "constant")
        P = P[: s**2, ...]
        P = P.reshape((s, s) + P.shape[1:]).transpose(
            (1, 2, 0, 3) + tuple(range(4, P.ndim + 1))
        )
        P = P.reshape((s * P.shape[1], s * P.shape[3]) + P.shape[4:])
    elif P.ndim == 4 and P.shape[0] > 1:
        if normalize:
            maxs = np.max(abs(P), axis=(-1, -2)) + 1e-6
            P = (P.T / maxs.T).T
        P = np.swapaxes(P, 1, 2).reshape(
            P.shape[0] * P.shape[2], P.shape[1] * P.shape[3]
        )
    else:
        P = np.squeeze(P)
    return P

hsvplot(u, ax=None, pixelSize=1, axisUnit='pixel', amplitudeScalingFactor=1)

perform complex plot :param ax :param pixelSize, default 1 :param axisUnit, default 'pixel', options: 'm', 'cm', 'mm', 'um' return: a complex plot

Source code in PtyLab/utils/visualisation.py

def hsvplot(u, ax=None, pixelSize=1, axisUnit="pixel", amplitudeScalingFactor=1):
    """
    perform complex plot
    :param ax
    :param pixelSize, default 1
    :param axisUnit, default 'pixel', options: 'm', 'cm', 'mm', 'um'
    return: a complex plot
    """
    u = np.squeeze(asNumpyArray(u))
    rgb = complex2rgb(u, amplitudeScalingFactor=amplitudeScalingFactor)
    complexPlot(rgb, ax, pixelSize, axisUnit)

hsvmodeplot(P, ax=None, normalize=True, pixelSize=1, axisUnit='pixel', amplitudeScalingFactor=1)

Place multi complex images in a square grid and use hsvplot to display :param P: A complex np.ndarray :param normalize: normalize each mode individually :param pixelSize: pixelSize in x and y, to display the physical dimension of the plot :return: a tiled complex plot

Source code in PtyLab/utils/visualisation.py

def hsvmodeplot(
    P, ax=None, normalize=True, pixelSize=1, axisUnit="pixel", amplitudeScalingFactor=1
):
    """
    Place multi complex images in a square grid and use hsvplot to display
    :param P: A complex np.ndarray
    :param normalize: normalize each mode individually
    :param pixelSize: pixelSize in x and y, to display the physical dimension of the plot
    :return: a tiled complex plot
    """

    Q = modeTile(np.squeeze(asNumpyArray(P)), normalize=normalize)
    hsvplot(
        Q,
        ax=ax,
        pixelSize=pixelSize,
        axisUnit=axisUnit,
        amplitudeScalingFactor=amplitudeScalingFactor,
    )

setColorMap()

create the colormap for diffraction data (the same as matlab) return: customized matplotlib colormap

Source code in PtyLab/utils/visualisation.py

def setColorMap():
    """
    create the colormap for diffraction data (the same as matlab)
    return: customized matplotlib colormap
    """
    colors = [
        (1, 1, 1),
        (0, 0.0875, 1),
        (0, 0.4928, 1),
        (0, 1, 0),
        (1, 0.6614, 0),
        (1, 0.4384, 0),
        (0.8361, 0, 0),
        (0.6505, 0, 0),
        (0.4882, 0, 0),
    ]

    n = 255  # Discretizes the interpolation into n bins
    cm = LinearSegmentedColormap.from_list("cmap", colors, n)
    return cm

show3Dslider(A, colormap='diffraction')

show a 3D plot with a slider.

In a Jupyter notebook an inline ipywidgets slider is used. In a script the interactive pyqtgraph viewer is used.

:param A: a 3D array :param colormap: matplotlib colormap, default, customized colormap for plotting diffraction data return: a pyqtgraph plot

Source code in PtyLab/utils/visualisation.py

def show3Dslider(A, colormap="diffraction"):
    """
    show a 3D plot with a slider.

    In a Jupyter notebook an inline ipywidgets slider is used.
    In a script the interactive pyqtgraph viewer is used.

    :param A: a 3D array
    :param colormap: matplotlib colormap, default, customized colormap for plotting diffraction data
    return: a pyqtgraph plot
    """
    print(A.min(), A.max())

    # resolve colormap once — matplotlib LinearSegmentedColormap works for both paths
    if colormap == "diffraction":
        cmap = setColorMap()
    else:
        cmap = mpl.cm.get_cmap(colormap)

    if _is_notebook():
        import ipywidgets as widgets
        from IPython.display import display
        import matplotlib.pyplot as plt

        # Create output widget for displaying figure
        out = widgets.Output()

        def update_frame(change):
            # Create new figure for each frame update
            # Handle both old API (event dict) and new API (direct value)
            if isinstance(change, dict):
                frame = int(change['new'])
            else:
                frame = int(change)
            fig, ax = plt.subplots(1, 1, figsize=(6, 6))
            im = ax.imshow(A[frame], cmap=cmap, origin="lower")
            plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
            ax.axis("off")
            plt.tight_layout()

            with out:
                out.clear_output(wait=True)
                display(fig)
                plt.close(fig)

        # Create slider widget
        slider = widgets.IntSlider(
            min=0,
            max=A.shape[0] - 1,
            step=1,
            value=0,
            description="Frame"
        )

        # Link slider to update function
        slider.observe(update_frame, names="value")

        # Display slider and output
        display(widgets.VBox([slider, out]))

        # Initial display
        update_frame(0)
    else:
        app = pg.mkQApp()
        imv = pg.ImageView(view=pg.PlotItem())
        imv.setWindowTitle("Close to proceed")
        imv.setImage(A)

        # set the colormap
        positions = np.linspace(0, 1, cmap.N)
        colors = [(np.array(cmap(i)[:-1]) * 255).astype("int") for i in positions]
        imv.setColorMap(pg.ColorMap(pos=positions, color=colors))
        imv.show()
        app.exec_()