imaging.py

"""
Defines classes and functions to manage reading and writing files associated
with absorption imaging.
"""
from __future__ import annotations
import numpy as np
import numpy.linalg as la
import lmfit
import pathlib
import toml
import warnings
import PIL.Image
import lib.plotdefs as pd
import lib.dataio as io

H = 6.626070040e-34
C = 2.99792458e+8

DEF_DX = 3.45 # pixel size for Flir Grasshopper [um]

class ImagingData:
    def __init__(self, outdir: pathlib.Path, arrays: dict[str, np.ndarray],
            config: dict[str, ...]=None, comments: str=None,
            results: dict[str, ...]=None):
        self.outdir = outdir
        self.arrays = arrays
        self.config = dict() if config is None else config
        self.comments = str() if comments is None else comments
        self.results = results

    def save(self, target: pathlib.Path=None, overwrite: bool=False,
            printflag: bool=True):
        T = self.outdir if target is None else target
        if printflag: print(f"[imaging] Saving data to '{T}':")
        if not T.is_dir():
            if printflag: print(f"[imaging] mkdir {T}")
            T.mkdir(parents=True, exist_ok=True)

        arrays_file = io.fresh_filename(T.joinpath("arrays.npz"), overwrite)
        if printflag: print("[imaging] save arrays")
        #np.savez_compressed(arrays_file, **self.arrays)
        np.savez(arrays_file, **self.arrays)

        config_file = io.fresh_filename(T.joinpath("config.toml"), overwrite)
        if printflag: print("[imaging] save camera config")
        with config_file.open('w') as outfile:
            toml.dump(self.config, outfile)

        comments_file = io.fresh_filename(T.joinpath("comments.txt"), overwrite)
        if printflag: print("[imaging] save comments")
        comments_file.write_text(self.comments)

        if self.results is not None:
            results_file = io.fresh_filename(
                T.joinpath("results.toml"), overwrite)
            if printflag: print("[imaging] save results")
            with results_file.open('w') as outfile:
                toml.dump(self.results, outfile)
        #if printflag: print("  Done.")
        return self

    @staticmethod
    def load(target: pathlib.Path, outdir: pathlib.Path=None,
            printflag: bool=True):
        outdir = target if outdir is None else target
        if printflag: print(f"[imaging] Loading data from target '{target}':")
        if not target.is_dir():
            raise Exception(f"Target '{target}' does not exist")

        arrays_file = target.joinpath("arrays.npz")
        if printflag: print("[imaging] load arrays")
        if not arrays_file.is_file():
            raise Exception(f"Arrays file '{arrays_file}' does not exist")
        arrays = dict(np.load(arrays_file))

        config_file = target.joinpath("config.toml")
        if printflag: print("[imaging] load camera config")
        if not config_file.is_file():
            raise Exception(f"Config file '{config_file}' does not exist")
        config = toml.load(config_file)

        comments_file = target.joinpath("comments.txt")
        if printflag and comments_file.is_file(): print("[imaging] load comments")
        comments = comments_file.read_text() \
                if comments_file.is_file() else str()

        results_file = target.joinpath("results.toml")
        if printflag and results_file.is_file(): print("[imaging] load results")
        results = results_file.read_text() \
                if results_file.is_file() else None
        #if printflag: print("  Done.")
        return ImagingData(outdir, arrays, config, comments, results)

    def render_arrays(self, target: pathlib.Path=None, printflag: bool=True):
        T = self.outdir if target is None else target
        T = T.joinpath("images")
        if printflag: print(f"[imaging] Render images to {T}:")
        if not T.is_dir():
            if printflag: print(f"[imaging] mkdir {T}")
            T.mkdir(parents=True, exist_ok=True)
        for label, array in self.arrays.items():
            if printflag: print(f"[imaging] render '{label}'")
            PIL.Image.fromarray(array).save(T.joinpath(label + ".png"))
        #if printflag: print("  Done.")
        return self

    def compute_results(self, printflag: bool=True):
        raise NotImplementedError

class AbsorptionData(ImagingData):
    @staticmethod
    def load(target: pathlib.Path, outdir: pathlib.Path=None,
            printflag: bool=True):
        data = ImagingData.load(target, outdir, printflag)
        return AbsorptionData(
            data.outdir, data.arrays, data.config, data.comments, data.results)

    def compute_results(self, dA=(3.45)**2, printflag: bool=True):
        _dA = self.config.get("bin_size", 1)**2 \
                * (DEF_DX**2 if dA is None else dA)
        if printflag: print("[imaging] Compute results")
        OD = compute_od(
            self.arrays["shadow"], self.arrays["bright"], self.arrays["dark"])
        self.arrays["od"] = OD
        sx, sy = lmfit_gaussian(OD, _dA)
        self.results = dict() if self.results is None else self.results
        self.results.update({"sx": sx, "sy": sy})
        return self

    def render_arrays(self, target: pathlib.Path=None, dx=None,
            printflag: bool=True):
        T = self.outdir if target is None else target
        T = T.joinpath("images")
        _dx = self.config.get("bin_size", 1) * (DEF_DX if dx is None else dx)
        if printflag: print(f"[imaging] Render images to {T}:")
        if not T.is_dir():
            if printflag: print(f"[imaging] mkdir {T}")
            T.mkdir(parents=True, exist_ok=True)
        for label, array in self.arrays.items():
            if printflag: print(f"[imaging] render '{label}'")
            if label == "od":
                warnings.filterwarnings("ignore")
                H, W = array.shape
                (pd.Plotter()
                    .imshow(array,
                        cmap=pd.colormaps["vibrant"], vmin=-0.1, vmax=0.15,
                        extent=[0, W * _dx, H * _dx, 0]
                    )
                    .colorbar()
                    #.set_clim(0, array.max())
                    .set_xlabel("[$\\mu$m]")
                    .set_ylabel("[$\\mu$m]")
                    .savefig(T.joinpath(label + ".png"))
                    .close()
                )
                warnings.resetwarnings()
            else:
                PIL.Image.fromarray(array).save(T.joinpath(label + ".png"))
        #if printflag: print("  Done.")
        return self

def compute_od(shadow: np.ndarray, bright: np.ndarray, dark: np.ndarray):
    with np.errstate(divide="ignore", invalid="ignore"):
        T = (shadow - dark) / (bright - dark)
        H, W = T.shape
        OD = -np.log(T)
    OD[~np.isfinite(OD)] = 0.0
    return OD

class FluorescenceData(ImagingData):
    @staticmethod
    def load(target: pathlib.Path, outdir: pathlib.Path=None,
            printflag: bool=True):
        data = ImagingData.load(target, outdir, printflag)
        return FluorescenceData(
            data.outdir, data.arrays, data.config, data.comments, data.results)

    def compute_results(self, subtract_bkgd: bool=True, dA: float=3.45**2,
            printflag: bool=True):
        _dA = self.config.get("bin_size", 1)**2 \
                * (DEF_DX**2 if dA is None else dA)
        if printflag: print("[imaging] Compute results")
        compute_mot_number_params = {
            "QE": 0.40,
            "gain": self.config["gain"],
            "exposure_time": self.config["exposure_time"] * 1e-6,
            "solid_angle": (0.0127**2 * np.pi) / (4 * np.pi * 0.125**2),
            "detuning": 50e6 * 2 * np.pi,
            "intensity_parameter": 1.0,
        }
        N_bkgd = compute_mot_number(
            self.arrays["background"].astype(np.float64),
            **compute_mot_number_params,
            k=None
        ) if subtract_bkgd and "background" in self.arrays.keys() else 0.0
        for k, (label, array) in enumerate(self.arrays.items()):
            dtype_info = np.iinfo(array.dtype)
            if array.max() >= (dtype_info.max - dtype_info.bits):
                print(
                    f"[imaging] WARNING: image '{label}' may contain clipping")
            N = compute_mot_number(
                array.astype(np.float64),
                **compute_mot_number_params,
                k=None
            )
            sx, sy = lls_fit_gaussian(array, _dA)
            self.results = dict() if self.results is None else self.results
            self.results.update({
                label: {
                    "N": float(N) - float(N_bkgd)
                        if label != "background" else float(N),
                    "sx": float(sx),
                    "sy": float(sy)
                }
            })
        return self

def compute_mot_number(image, QE, gain, exposure_time, solid_angle,
        detuning, intensity_parameter, k=None):
    H, W = image.shape
    K_h = 3.0
    slice_h = pd.S[int(H / K_h) : int((K_h - 1) * H / K_h)]
    K_w = 3.0
    slice_w = pd.S[int(W / K_w) : int((K_w - 1) * W / K_w)]
    im = image[slice_h, slice_w]
    if k is not None:
        pd.Plotter().imshow(image).savefig(f"image_{k}.png").close()
        pd.Plotter().imshow(im).savefig(f"img_{k}.png").close()
    electron_rate = im.astype(np.float64).sum()
    photon_rate = (
        electron_rate
        / 10**(gain / 10)
        / QE
        / solid_angle
        / exposure_time
    )
    Y = 29.1e6 * 2 * np.pi # transition linewidth; s^-1
    l = 399e-9 # transition wavelength; m
    I_sat = (Y * H * C * np.pi) / (3 * l**3)
    scatter_rate = lambda D, I: \
            (Y / 2) * I / (1 + 4 * (D / Y)**2 + I)
    N = photon_rate / scatter_rate(detuning, intensity_parameter)
    return N

def lls_fit_gaussian(A, dA):
    """
    Assumes uniform, square dA.
    """
    imax, jmax = np.where(A == A.max())
    i0, j0 = imax.mean(), jmax.mean()
    x = np.sqrt(dA) * (np.arange(A.shape[1]) - j0)
    y = np.sqrt(dA) * (np.arange(A.shape[0]) - i0)
    X, Y = np.meshgrid(x, y)
    Z = (-np.log(A / A.max())).flatten()
    M = np.array([X.flatten()**2, Y.flatten()**2]).T
    Sx, Sy = la.solve(M.T @ M, M.T @ Z)
    return np.sqrt(1 / Sx), np.sqrt(1 / Sy)

def lmfit_gaussian(data, dA):
    """
    Assumes uniform, square dA.
    """
    imax, jmax = np.where(data == data.max())
    i0, j0 = imax.mean(), jmax.mean()
    x0 = np.sqrt(dA) * j0
    x = np.sqrt(dA) * np.arange(data.shape[1])
    y0 = np.sqrt(dA) * i0
    y = np.sqrt(dA) * np.arange(data.shape[0])
    X, Y = np.meshgrid(x, y)

    def model(params: lmfit.Parameters):
        A = params["A"]
        x0 = params["x0"]
        sx = params["sx"]
        y0 = params["y0"]
        sy = params["sy"]
        return A * np.exp(-((X - x0) / sx)**2 - ((Y - y0) / sy)**2)

    def residual(params: lmfit.Parameters):
        m = model(params)
        return ((data - m)**2).flatten()

    params = lmfit.Parameters()
    params.add("A", value=data.max())
    params.add("x0", value=x0)
    params.add("sx", value=1)
    params.add("y0", value=y0)
    params.add("sy", value=1)
    fit = lmfit.minimize(residual, params)
    if not fit.success:
        raise Exception("Error in Gaussian fit")
    return float(fit.params["sx"]), float(fit.params["sy"])