Add option to fix z0

ColumbiaWillLab · Jul 22, 2019 · d3e5973 · d3e5973
1 parent 431b390
commit d3e5973
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 18 deletions.
diff --git a/config.py b/config.py
@@ -78,7 +78,26 @@ def atom_mass(self):
     @property
     def fix_theta(self):
         """If True, 2D Gaussian fit restricts theta to 0 (hor/ver)."""
-        return self.getboolean("fit", "fix_theta")
+        try:
+            return self.getboolean("fit", "fix_theta")
+        except configparser.NoOptionError:
+            return False
+
+    @fix_theta.setter
+    def fix_theta(self, val):
+        self["fit"]["fix_theta"] = str(val)
+
+    @property
+    def fix_z0(self):
+        """If True, 2D Gaussian fit restricts z0 to 0."""
+        try:
+            return self.getboolean("fit", "fix_z0")
+        except configparser.NoOptionError:
+            return False
+
+    @fix_z0.setter
+    def fix_z0(self, val):
+        self["fit"]["fix_z0"] = str(val)
 
     @property
     def roi(self):

diff --git a/gui/controls.py b/gui/controls.py
@@ -145,28 +145,40 @@ def __init__(self, master):
         super().__init__(master, text="Fit")
         self.master = master
 
-        self.fixtheta = tk.BooleanVar()
-        self.fixtheta.set(config.fix_theta)
-        fixtheta_btn = ttk.Checkbutton(
-            self,
-            text="Fix Theta",
-            variable=self.fixtheta,
-            command=self._toggle_fixtheta,
+        self.fix_z0 = tk.BooleanVar()
+        self.fix_z0.set(config.fix_z0)
+        fix_z0_btn = ttk.Checkbutton(
+            self, text="Fix z0", variable=self.fix_z0, command=self._toggle_fix_z0
         )
-        fixtheta_btn.grid(row=0, column=0, sticky="w")
-        self.fitvar = tk.BooleanVar()
-        self.fitvar.set(config.fit)
-        fitbtn = ttk.Checkbutton(
-            self, text="Enable Fitting", variable=self.fitvar, command=self._toggle_fit
+        fix_z0_btn.grid(row=0, column=0, sticky="w")
+
+        self.fix_theta = tk.BooleanVar()
+        self.fix_theta.set(config.fix_theta)
+        fix_theta_btn = ttk.Checkbutton(
+            self, text="Fix θ", variable=self.fix_theta, command=self._toggle_fix_theta
         )
-        fitbtn.grid(row=1, column=0, sticky="w")
+        fix_theta_btn.grid(row=0, column=1, sticky="w")
 
-    def _toggle_fixtheta(self):
-        config["fit"]["fix_theta"] = str(self.fixtheta.get())
+        self.enable_fit = tk.BooleanVar()
+        self.enable_fit.set(config.fit)
+        enable_fit_btn = ttk.Checkbutton(
+            self,
+            text="Enable Fitting",
+            variable=self.enable_fit,
+            command=self._toggle_fit,
+        )
+        enable_fit_btn.grid(row=1, column=0, sticky="w")
+
+    def _toggle_fix_theta(self):
+        config.fix_theta = self.fix_theta.get()
         config.save()
 
     def _toggle_fit(self):
-        config.fit = self.fitvar.get()
+        config.fit = self.enable_fit.get()
+
+    def _toggle_fix_z0(self):
+        config.fix_z0 = self.fix_z0.get()
+        config.save()
 
 
 class FitParams(ttk.Frame):

diff --git a/models/shots.py b/models/shots.py
@@ -154,7 +154,7 @@ def twoD_gaussian(self):
         model.set_param_hint(
             "theta", min=-np.pi / 4, max=np.pi / 4, vary=not config.fix_theta
         )
-        model.set_param_hint("z0", min=-1, max=1)
+        model.set_param_hint("z0", min=-1, max=1, vary=not config.fix_z0)
 
         result = model.fit(
             np.ravel(self.absorption[::5]),