imaging.py 11.1 KB
Newer Older
1
2
3
4
5
6
"""
Defines classes and functions to manage reading and writing files associated
with absorption imaging.
"""
from __future__ import annotations
import numpy as np
whooie's avatar
whooie committed
7
import numpy.linalg as la
8
import lmfit
9
10
import pathlib
import toml
11
import warnings
12
import PIL.Image
whooie's avatar
whooie committed
13
import lib.plotdefs as pd
14
import lib.dataio as io
15

whooie's avatar
whooie committed
16
17
18
H = 6.626070040e-34
C = 2.99792458e+8

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

whooie's avatar
whooie committed
21
class ImagingData:
22
    def __init__(self, outdir: pathlib.Path, arrays: dict[str, np.ndarray],
23
            config: dict[str, ...]=None, comments: str=None,
24
            results: dict[str, ...]=None):
25
        self.outdir = outdir
26
        self.arrays = arrays
27
        self.config = dict() if config is None else config
whooie's avatar
whooie committed
28
        self.comments = str() if comments is None else comments
29
        self.results = results
30

31
32
    def save(self, target: pathlib.Path=None, overwrite: bool=False,
            printflag: bool=True):
33
34
        T = self.outdir if target is None else target
        if printflag: print(f"[imaging] Saving data to '{T}':")
35
        if not T.is_dir():
whooie's avatar
whooie committed
36
            if printflag: print(f"[imaging] mkdir {T}")
37
            T.mkdir(parents=True, exist_ok=True)
whooie's avatar
whooie committed
38

39
        arrays_file = io.fresh_filename(T.joinpath("arrays.npz"), overwrite)
whooie's avatar
whooie committed
40
        if printflag: print("[imaging] save arrays")
41
42
        #np.savez_compressed(arrays_file, **self.arrays)
        np.savez(arrays_file, **self.arrays)
whooie's avatar
whooie committed
43

44
        config_file = io.fresh_filename(T.joinpath("config.toml"), overwrite)
whooie's avatar
whooie committed
45
        if printflag: print("[imaging] save camera config")
46
47
        with config_file.open('w') as outfile:
            toml.dump(self.config, outfile)
whooie's avatar
whooie committed
48

49
        comments_file = io.fresh_filename(T.joinpath("comments.txt"), overwrite)
whooie's avatar
whooie committed
50
        if printflag: print("[imaging] save comments")
51
        comments_file.write_text(self.comments)
whooie's avatar
whooie committed
52

53
        if self.results is not None:
54
            results_file = io.fresh_filename(
55
                T.joinpath("results.toml"), overwrite)
whooie's avatar
whooie committed
56
            if printflag: print("[imaging] save results")
57
58
            with results_file.open('w') as outfile:
                toml.dump(self.results, outfile)
whooie's avatar
whooie committed
59
        #if printflag: print("  Done.")
60
61
62
        return self

    @staticmethod
63
64
65
66
    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}':")
67
        if not target.is_dir():
68
            raise Exception(f"Target '{target}' does not exist")
whooie's avatar
whooie committed
69
70

        arrays_file = target.joinpath("arrays.npz")
whooie's avatar
whooie committed
71
        if printflag: print("[imaging] load arrays")
72
        if not arrays_file.is_file():
73
            raise Exception(f"Arrays file '{arrays_file}' does not exist")
74
        arrays = dict(np.load(arrays_file))
whooie's avatar
whooie committed
75
76

        config_file = target.joinpath("config.toml")
whooie's avatar
whooie committed
77
        if printflag: print("[imaging] load camera config")
78
        if not config_file.is_file():
79
            raise Exception(f"Config file '{config_file}' does not exist")
whooie's avatar
whooie committed
80
        config = toml.load(config_file)
whooie's avatar
whooie committed
81
82

        comments_file = target.joinpath("comments.txt")
whooie's avatar
whooie committed
83
        if printflag and comments_file.is_file(): print("[imaging] load comments")
84
        comments = comments_file.read_text() \
whooie's avatar
whooie committed
85
86
87
                if comments_file.is_file() else str()

        results_file = target.joinpath("results.toml")
whooie's avatar
whooie committed
88
        if printflag and results_file.is_file(): print("[imaging] load results")
89
        results = results_file.read_text() \
whooie's avatar
whooie committed
90
                if results_file.is_file() else None
whooie's avatar
whooie committed
91
        #if printflag: print("  Done.")
92
        return ImagingData(outdir, arrays, config, comments, results)
93

94
    def render_arrays(self, target: pathlib.Path=None, printflag: bool=True):
95
        T = self.outdir if target is None else target
96
        T = T.joinpath("images")
whooie's avatar
whooie committed
97
        if printflag: print(f"[imaging] Render images to {T}:")
98
        if not T.is_dir():
whooie's avatar
whooie committed
99
            if printflag: print(f"[imaging] mkdir {T}")
100
            T.mkdir(parents=True, exist_ok=True)
101
        for label, array in self.arrays.items():
102
            if printflag: print(f"[imaging] render '{label}'")
103
            PIL.Image.fromarray(array).save(T.joinpath(label + ".png"))
whooie's avatar
whooie committed
104
        #if printflag: print("  Done.")
105
106
        return self

107
    def compute_results(self, printflag: bool=True):
whooie's avatar
whooie committed
108
109
110
        raise NotImplementedError

class AbsorptionData(ImagingData):
111
112
113
114
115
116
117
    @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)

118
    def compute_results(self, dA=(3.45)**2, printflag: bool=True):
119
120
        _dA = self.config.get("bin_size", 1)**2 \
                * (DEF_DX**2 if dA is None else dA)
whooie's avatar
whooie committed
121
        if printflag: print("[imaging] Compute results")
122
123
        OD = compute_od(
            self.arrays["shadow"], self.arrays["bright"], self.arrays["dark"])
whooie's avatar
whooie committed
124
        self.arrays["od"] = OD
Yb Tweezer's avatar
Yb Tweezer committed
125
        sx, sy = lmfit_gaussian(OD, _dA)
whooie's avatar
whooie committed
126
127
128
129
        self.results = dict() if self.results is None else self.results
        self.results.update({"sx": sx, "sy": sy})
        return self

130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
    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

162
def compute_od(shadow: np.ndarray, bright: np.ndarray, dark: np.ndarray):
whooie's avatar
typo    
whooie committed
163
    with np.errstate(divide="ignore", invalid="ignore"):
whooie's avatar
typo    
whooie committed
164
        T = (shadow - dark) / (bright - dark)
Yb Tweezer's avatar
Yb Tweezer committed
165
        H, W = T.shape
whooie's avatar
whooie committed
166
        OD = -np.log(T)
167
    OD[~np.isfinite(OD)] = 0.0
whooie's avatar
whooie committed
168
    return OD
169

whooie's avatar
whooie committed
170
class FluorescenceData(ImagingData):
171
172
173
174
175
176
177
    @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)

178
179
    def compute_results(self, subtract_bkgd: bool=True, dA: float=3.45**2,
            printflag: bool=True):
whooie's avatar
whooie committed
180
181
        _dA = self.config.get("bin_size", 1)**2 \
                * (DEF_DX**2 if dA is None else dA)
whooie's avatar
whooie committed
182
        if printflag: print("[imaging] Compute results")
183
184
185
186
187
188
189
190
191
192
193
        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,
194
            k=None
195
        ) if subtract_bkgd and "background" in self.arrays.keys() else 0.0
196
        for k, (label, array) in enumerate(self.arrays.items()):
whooie's avatar
whooie committed
197
198
            dtype_info = np.iinfo(array.dtype)
            if array.max() >= (dtype_info.max - dtype_info.bits):
199
200
                print(
                    f"[imaging] WARNING: image '{label}' may contain clipping")
201
            N = compute_mot_number(
202
                array.astype(np.float64),
203
                **compute_mot_number_params,
204
                k=None
205
            )
Yb Tweezer's avatar
Yb Tweezer committed
206
            sx, sy = lls_fit_gaussian(array, _dA)
207
            self.results = dict() if self.results is None else self.results
208
209
            self.results.update({
                label: {
210
211
                    "N": float(N) - float(N_bkgd)
                        if label != "background" else float(N),
212
213
                    "sx": float(sx),
                    "sy": float(sy)
214
                }
215
            })
whooie's avatar
whooie committed
216
217
218
        return self

def compute_mot_number(image, QE, gain, exposure_time, solid_angle,
219
        detuning, intensity_parameter, k=None):
whooie's avatar
whooie committed
220
    H, W = image.shape
Yb Tweezer's avatar
Yb Tweezer committed
221
    K_h = 3.0
222
223
224
225
    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]
226
227
228
    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()
229
    electron_rate = im.astype(np.float64).sum()
230
231
232
233
234
235
236
    photon_rate = (
        electron_rate
        / 10**(gain / 10)
        / QE
        / solid_angle
        / exposure_time
    )
whooie's avatar
whooie committed
237
238
239
240
241
242
    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)
243
    return N
whooie's avatar
whooie committed
244

Yb Tweezer's avatar
Yb Tweezer committed
245
def lls_fit_gaussian(A, dA):
whooie's avatar
whooie committed
246
247
248
249
250
251
252
253
254
255
256
257
    """
    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)
258

Yb Tweezer's avatar
Yb Tweezer committed
259
260
261
262
def lmfit_gaussian(data, dA):
    """
    Assumes uniform, square dA.
    """
263
264
265
266
267
268
269
270
    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)

271
    def model(params: lmfit.Parameters):
272
273
274
275
276
277
278
        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)

279
280
    def residual(params: lmfit.Parameters):
        m = model(params)
281
282
283
284
285
286
287
288
        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)
289
    fit = lmfit.minimize(residual, params)
290
291
292
293
    if not fit.success:
        raise Exception("Error in Gaussian fit")
    return float(fit.params["sx"]), float(fit.params["sy"])