utilities_lv65

utilities_lv65.py

import numpy as np


### Requested pixel2q and q2pixel ####
def pixel2q(pixel, E_ev, detz_m, pixelsize_m=75e-6):
    """
    returns q in reciprocal nm
    E_ev: photon Energy in eV
    detz_m: detector distance in m
    pixelsize_m: detector pixelsize in m
    """
    qstep_nm = pixelsize_m / detz_m * (2 * np.pi) / (1240 / E_ev)
    return pixel * qstep_nm


def q2pixel(q_rnm, E_ev, detz_m, pixelsize_m=75e-6):
    """
    return pixel on the detector/in the correlation
    q_rnm: q in reciprocal nm
    E_ev: photon Energy in eV
    detz_m: detector distance in m
    pixelsize_m: detector pixelsize in m
    """
    qstep_nm = pixelsize_m / detz_m * (2 * np.pi) / (1240 / E_ev)
    return q_rnm / qstep_nm


### some other utilities ###
def radial_profile(data, center=None, calcStd=False, os=1):
    """
    calculates a ND radial profile of data around center. will ignore nans
    calStd: calculate standard deviation, return tuple of (profile, std)
    os: oversample by a factor. With default 1 the stepsize will be 1 pixel, with 2 it will be .5 pixels etc.
    """
    if center is None:
        center = np.array(data.shape) // 2
    if len(center) != data.ndim:
        raise TypeError("center should be of length data.ndim")
    center = np.array(center)[tuple([slice(len(center))] + data.ndim * [None])]
    ind = np.indices(data.shape)
    r = (np.rint(os * np.sqrt(((ind - center) ** 2).sum(axis=0)))).astype(int)
    databin = np.bincount(r.ravel(), (np.nan_to_num(data)).ravel())
    nr = np.bincount(r.ravel(), ((~np.isnan(data)).astype(float)).ravel())
    with np.errstate(divide="ignore", invalid="ignore"):
        radialprofile = databin / nr
    if not calcStd:
        return radialprofile

    data2bin = np.bincount(r.ravel(), (np.nan_to_num(data ** 2)).ravel())
    with np.errstate(divide="ignore", invalid="ignore"):
        radial2profile = data2bin / nr
    std = np.sqrt(radial2profile - radialprofile ** 2)
    return radialprofile, std


def centered_part(arr, newshape):
    """
    Return the center newshape portion of the array.
    """
    newshape = np.asarray(newshape)
    currshape = np.array(arr.shape)
    newshape = np.minimum(newshape, currshape)
    startind = (np.ceil((currshape - newshape) / 2)).astype(int)
    endind = startind + newshape
    myslice = [slice(startind[k], endind[k]) for k in range(len(endind))]
    return arr[tuple(myslice)]


def bin(ndarray, new_shape, operation="sum"):
    """
    # https://stackoverflow.com/a/29042041
    bin an ndarray
    ndarray: nd-array
    new_shape: shape to bin to. shape of ndarray has to be integer multiple of new_shape along each dimension
    operation: string. sum, mean, max, or min. operation to use
    """
    ops = ["sum", "mean", "max", "min"]
    operation = operation.lower()
    if operation not in ops:
        raise ValueError("Operation not supported.")
    if ndarray.ndim != len(new_shape):
        raise ValueError("Shape mismatch: {} -> {}".format(ndarray.shape, new_shape))
    compression_pairs = [(d, c // d) for d, c in zip(new_shape, ndarray.shape)]
    flattened = [i for p in compression_pairs for i in p]
    ndarray = ndarray.reshape(flattened)
    for i in range(len(new_shape)):
        op = getattr(ndarray, operation)
        ndarray = op(-1 * (i + 1))
    return ndarray


def binFactor(ndarray, n, operation="mean"):
    """
    rebin an ndarray
    parameters:
    ndarray: nd-array
    n: scalar or list, factor to bin with. if scalar, same factor along all dimensions
    operation: string. sum, mean, max, or min. operation to use
    if n doesnt divide arr along one dimensions, last elements of n will silently be dropped
    """
    if not (np.size(n) == 1 or np.size(n) == np.ndim(ndarray)):
        raise ValueError("n should be scalar or of same length as ndarray has dimensions")
    newshape = np.maximum(np.array(ndarray.shape) // n, 1)
    return bin(
        ndarray[tuple((slice(None, s) for s in n * newshape))],
        newshape,
        operation,
    )


### the q/pixel plotting functions ###
import matplotlib.pyplot as plt
from matplotlib.patches import Circle


def plot_2d(image, mask_center=True, E_ev=None, detz_m=None, pixelsize_m=75e-6, ax=None, colorbar=True, rings=(0.125, 0.25, 0.5, 0.75, 1), **kwargs):
    """
    plots an image.
    mask_center (default true): mask the center cross
    if E_ev, detz_m and pixelsize_m are all not None, q circles will be added
    ax: axis to plot to, if None, gca()
    colorbar: add colorbar
    rings: relative radius of the q circles to add
    other kwargs will be passed to matshow, e.g. vmin/vmax/cmap
    returns the axis
    """
        
    if mask_center:
        i = image.copy()
        centered_part(i, (3, 1))[:] = np.nan
        centered_part(i, (1, 3))[:] = np.nan
    else:
        i = image
    if ax is None:
        ax = plt.gca()
    plot = ax.matshow(i, **kwargs)
    if not any(x is None for x in (E_ev, pixelsize_m, detz_m)):
        maxq = pixel2q(max(i.shape) / 2, E_ev, detz_m, pixelsize_m)
        labeldir = 1
        for r in rings:
            q = np.round(r * maxq, int(np.ceil(1 - np.log10(maxq))))
            p = q2pixel(q, E_ev, detz_m, pixelsize_m)
            circ = Circle(
                np.array(i.shape) // 2,
                p,
                fill=False,
                linestyle=":",
                alpha=1,
            )
            ax.add_patch(circ)
            ax.text(i.shape[0] / 2 + p / np.sqrt(2), i.shape[0] / 2 + labeldir * p / np.sqrt(2), f'{q}{"$nm^{-1}$"}', fontsize="x-small", va="center")
            labeldir *= -1
    Nx, Ny = i.shape
    ax.set_xticks([*range(0, Nx, Nx // 4), Nx - 1])
    ax.set_yticks([*range(0, Ny, Ny // 4), Ny - 1])
    ax.xaxis.set_major_formatter(lambda x, pos: f"{int(x-Nx/2)}")
    ax.yaxis.set_major_formatter(lambda y, pos: f"{int(y-Ny/2)}")
    if colorbar:
        plt.colorbar(plot, ax=ax)
    return ax


def plot_radialprofile(image, mask_center=True, E_ev=None, detz_m=None, pixelsize_m=75e-6, ax=None, **kwargs):
    """
    plots the radial profile of image.
    mask_center (default true): mask the center cross
    if E_ev, detz_m and pixelsize_m are all not None, a q-axis will be added
    ax: axis to plot to, if None, gca()
    returns the axis
    """
    if mask_center:
        i = image.copy()
        centered_part(i, (3, 1))[:] = np.nan
        centered_part(i, (1, 3))[:] = np.nan
    else:
        i = image
    if ax is None:
        ax = plt.gca()
    r = radial_profile(i)
    pixel_x = np.arange(0, len(r))
    ax.plot(pixel_x, r, **kwargs)
    if not any(x is None for x in (E_ev, pixelsize_m, detz_m)):
        q_x = pixel2q(pixel_x, E_ev, detz_m, pixelsize_m)
        ax2 = ax.twiny()
        ax2.plot(q_x, np.repeat(np.nan, len(q_x)))
        ax2.set_xlabel("$nm^{-1}$")
    ax.set_xlabel("px")
    return ax


def plot_lineprofiles(image, mask_center=True, E_ev=None, detz_m=None, pixelsize_m=75e-6, ax=None, **kwargs):
    """
    plots a horizontal and vertical line profile of image
    mask_center (default true): mask the center cross
    if E_ev, detz_m and pixelsize_m are all not None, a q-axis will be added
    ax: axis to plot to, if None, gca()
    returns the axis
    """
    if mask_center:
        i = image.copy()
        centered_part(i, (3, 1))[:] = np.nan
        centered_part(i, (1, 3))[:] = np.nan
    else:
        i = image
    if ax is None:
        ax = plt.gca()
    pixel_x, pixel_y = (np.arange(-s / 2, s / 2) for s in i.shape)
    lv = centered_part(i, (1, i.shape[1])).ravel()
    lh = centered_part(i, (i.shape[0], 1)).ravel()
    ax.plot(pixel_x, lv, label="vertical", **kwargs)
    ax.plot(pixel_y, lh, label="horizontal", **kwargs)
    if not any(x is None for x in (E_ev, pixelsize_m, detz_m)):
        q_x = pixel2q(pixel_x, E_ev, detz_m, pixelsize_m)
        ax2 = ax.twiny()
        ax2.plot(q_x, np.repeat(np.nan, len(q_x)))
        ax2.set_xlabel("$nm^{-1}$")
    ax.set_xlabel("px")

    return ax