imaging.py 12.4 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
whooie's avatar
whooie committed
14
from lib.plotdefs import S
15
import lib.dataio as io
16

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

whooie's avatar
whooie committed
20
21
22
ROI = [932, 729, 155, 125] # [x, y, w, h]
ROI_s = (S[ROI[1]:ROI[1] + ROI[3]], S[ROI[0]:ROI[0] + ROI[2]])

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

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

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

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

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

53
        comments_file = io.fresh_filename(T.joinpath("comments.txt"), overwrite)
whooie's avatar
whooie committed
54
        if printflag: print("[imaging] save comments")
55
        comments_file.write_text(self.comments)
whooie's avatar
whooie committed
56

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

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

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

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

        comments_file = target.joinpath("comments.txt")
whooie's avatar
whooie committed
87
        if printflag and comments_file.is_file(): print("[imaging] load comments")
88
        comments = comments_file.read_text() \
whooie's avatar
whooie committed
89
90
91
                if comments_file.is_file() else str()

        results_file = target.joinpath("results.toml")
whooie's avatar
whooie committed
92
        if printflag and results_file.is_file(): print("[imaging] load results")
93
94
        with results_file.open('r') as infile:
            results = toml.load(infile)
whooie's avatar
whooie committed
95
        #if printflag: print("  Done.")
96
        return ImagingData(outdir, arrays, config, comments, results)
97

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

111
    def compute_results(self, printflag: bool=True):
whooie's avatar
whooie committed
112
113
114
        raise NotImplementedError

class AbsorptionData(ImagingData):
115
116
117
118
119
120
121
    @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)

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

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
162
163
164
    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

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

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

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

def compute_mot_number(image, QE, gain, exposure_time, solid_angle,
225
        detuning, intensity_parameter, k=None):
whooie's avatar
whooie committed
226
227
228
229
230
231
232
    # H, W = image.shape
    # K_h = 3.0
    # slice_h = S[int(H / K_h) : int((K_h - 1) * H / K_h)]
    # K_w = 3.0
    # slice_w = S[int(W / K_w) : int((K_w - 1) * W / K_w)]
    # im = image[slice_h, slice_w]
    im = image[ROI_s]
233
234
235
    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()
236
    electron_rate = im.astype(np.float64).sum()
237
238
    photon_rate = (
        electron_rate
239
        / 10**(gain / 20) # originally 10, testing out at 20
240
241
242
243
        / QE
        / solid_angle
        / exposure_time
    )
whooie's avatar
whooie committed
244
245
246
247
248
249
    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)
250
    return N
whooie's avatar
whooie committed
251

whooie's avatar
whooie committed
252
def lls_fit_gaussian(A, dA, k=None): # deprecated
whooie's avatar
whooie committed
253
254
255
256
257
258
259
260
261
262
263
    """
    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)
264
265
266
267
268
269
270
271
272
273
274
275
    if k is not None:
        (pd.Plotter()
            .imshow(
                100 * np.exp(-X**2 / Sx - Y**2 / Sy),
                extent=[x.min(), x.max(), y.max(), y.min()],
                cmap="rainbow"
            )
            .savefig(f"gaussian_{k}.png")
            .close()
        )
        print(i0, j0, A.shape)
        print(np.sqrt(1 / Sx), np.sqrt(1 / Sy))
whooie's avatar
whooie committed
276
    return np.sqrt(1 / Sx), np.sqrt(1 / Sy)
277

278
def lmfit_gaussian(data, dA, k=None):
Yb Tweezer's avatar
Yb Tweezer committed
279
280
281
    """
    Assumes uniform, square dA.
    """
whooie's avatar
whooie committed
282
283
    D = data[ROI_s]
    i0, j0 = [k // 2 for k in D.shape]
284
    x0 = np.sqrt(dA) * j0
whooie's avatar
whooie committed
285
    x = np.sqrt(dA) * np.arange(D.shape[1])
286
    y0 = np.sqrt(dA) * i0
whooie's avatar
whooie committed
287
    y = np.sqrt(dA) * np.arange(D.shape[0])
288
    X, Y = np.meshgrid(x, y)
289
    
whooie's avatar
whooie committed
290
291
292
    downsample = 1
    sampler = S[::downsample]
    _D = D[sampler, sampler]
293
294
    _X = X[sampler, sampler]
    _Y = Y[sampler, sampler]
295

296
    def model(params: lmfit.Parameters):
297
298
299
300
301
        A = params["A"]
        x0 = params["x0"]
        sx = params["sx"]
        y0 = params["y0"]
        sy = params["sy"]
302
303
        B = params["B"]
        return A * np.exp(-((_X - x0) / sx)**2 - ((_Y - y0) / sy)**2) + B
304

305
306
    def residual(params: lmfit.Parameters):
        m = model(params)
whooie's avatar
whooie committed
307
        return ((_D - m)**2).flatten()
308
309
310
311

    params = lmfit.Parameters()
    params.add("A", value=data.max())
    params.add("x0", value=x0)
312
    params.add("sx", value=(x.max() - x.min()) / 12)
313
    params.add("y0", value=y0)
314
315
    params.add("sy", value=(y.max() - y.min()) / 12)
    params.add("B", value=data.mean())
316
    fit = lmfit.minimize(residual, params)
317
318
    if not fit.success:
        raise Exception("Error in Gaussian fit")
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
    if k is not None:
        (pd.Plotter()
            .imshow(
                model(fit.params),
                extent=[x.min(), x.max(), y.max(), y.min()],
                cmap="rainbow"
            )
            .savefig(f"gaussian_{k}.png")
            .close()
        )
        print(fit.params["A"])
        print(fit.params["x0"])
        print(fit.params["sx"])
        print(fit.params["y0"])
        print(fit.params["sy"])
334
335
    return float(fit.params["sx"]), float(fit.params["sy"])