lls_fit_gaussian is now deprecated, use lmfit_gaussian instead for more...

lls_fit_gaussian is now deprecated, use lmfit_gaussian instead for more accurate sx/sy; neaten up function args a bit

lib/imaging.py

@@ -86,8 +86,8 @@ class ImagingData:
 
         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
+        with results_file.open('r') as infile:
+            results = toml.load(infile)
         #if printflag: print("  Done.")
         return ImagingData(outdir, arrays, config, comments, results)
 
@@ -115,9 +115,8 @@ class AbsorptionData(ImagingData):
         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)
+    def compute_results(self, dA=DEF_DX**2, printflag: bool=True):
+        _dA = self.config.get("bin_size", 1)**2 * dA
         if printflag: print("[imaging] Compute results")
         OD = compute_od(
             self.arrays["shadow"], self.arrays["bright"], self.arrays["dark"])
@@ -175,10 +174,9 @@ class FluorescenceData(ImagingData):
         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,
+    def compute_results(self, subtract_bkgd: bool=True, dA: float=DEF_DX**2,
             printflag: bool=True):
-        _dA = self.config.get("bin_size", 1)**2 \
-                * (DEF_DX**2 if dA is None else dA)
+        _dA = self.config.get("bin_size", 1)**2 * dA
         if printflag: print("[imaging] Compute results")
         compute_mot_number_params = {
             "QE": 0.40,
@@ -203,7 +201,8 @@ class FluorescenceData(ImagingData):
                 **compute_mot_number_params,
                 k=None
             )
-            sx, sy = lls_fit_gaussian(array, _dA)
+            # sx, sy = lmfit_gaussian(array, _dA, label) # debug
+            sx, sy = lmfit_gaussian(array, _dA)
             self.results = dict() if self.results is None else self.results
             self.results.update({
                 label: {
@@ -242,7 +241,7 @@ def compute_mot_number(image, QE, gain, exposure_time, solid_angle,
     N = photon_rate / scatter_rate(detuning, intensity_parameter)
     return N
 
-def lls_fit_gaussian(A, dA):
+def lls_fit_gaussian(A, dA, k=None):
     """
     Assumes uniform, square dA.
     """
@@ -254,19 +253,36 @@ def lls_fit_gaussian(A, dA):
     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)
+    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))
     return np.sqrt(1 / Sx), np.sqrt(1 / Sy)
 
-def lmfit_gaussian(data, dA):
+def lmfit_gaussian(data, dA, k=None):
     """
     Assumes uniform, square dA.
     """
-    imax, jmax = np.where(data == data.max())
-    i0, j0 = imax.mean(), jmax.mean()
+    i0, j0 = [k // 2 for k in data.shape]
     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)
+    
+    downsample = 10
+    sampler = pd.S[::downsample]
+    _data = data[sampler, sampler]
+    _X = X[sampler, sampler]
+    _Y = Y[sampler, sampler]
 
     def model(params: lmfit.Parameters):
         A = params["A"]
@@ -274,20 +290,37 @@ def lmfit_gaussian(data, dA):
         sx = params["sx"]
         y0 = params["y0"]
         sy = params["sy"]
-        return A * np.exp(-((X - x0) / sx)**2 - ((Y - y0) / sy)**2)
+        B = params["B"]
+        return A * np.exp(-((_X - x0) / sx)**2 - ((_Y - y0) / sy)**2) + B
 
     def residual(params: lmfit.Parameters):
         m = model(params)
-        return ((data - m)**2).flatten()
+        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("sx", value=(x.max() - x.min()) / 12)
     params.add("y0", value=y0)
-    params.add("sy", value=1)
+    params.add("sy", value=(y.max() - y.min()) / 12)
+    params.add("B", value=data.mean())
     fit = lmfit.minimize(residual, params)
     if not fit.success:
         raise Exception("Error in Gaussian fit")
+    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"])
     return float(fit.params["sx"]), float(fit.params["sy"])