Improve GUI

gui.py

@@ -170,6 +170,47 @@ class DraggableLine:
     def draw(self):
         self.parent.draw(update_map=False)
 
+    def clear(self):
+        for p in self.list_points:
+            p.point.remove()
+        self.list_points[1].line.remove()
+
+
+class FitterTool:
+    """
+    Flexible fitter tool
+    """
+
+    def __init__(self, ax, bounds=None):
+        self.ax = ax
+        self.bounds = bounds
+        self.bounds_selector = SpanSelector(
+            self.ax, self.onselect, "horizontal", useblit=True, span_stays=True
+        )
+        self.fitline = DraggableLine(self, self.ax_gamma, [0, 0], [1, 0.3], 0.05)
+
+        if update_map and self.gamma_3d_button.get_status()[1]:
+            lr = linregress(rho[rho > -4], cs[rho > -4])
+            (a, b, r, _, _) = lr
+            self.line_gamma.clear()
+            del self.line_gamma
+            rhomin, rhomax = np.min(rho), np.max(rho)
+            self.line_gamma = DraggableLine(
+                self,
+                self.ax_gamma,
+                [rhomin, a * rhomin + b],
+                [rhomax, a * rhomax + b],
+                0.05,
+            )
+            print("Gamma linregress : {}".format(lr))
+        else:
+            lps = self.line_gamma.list_points
+            xa, ya, xb, yb = lps[0].x, lps[0].y, lps[1].x, lps[1].y
+            a = (yb - ya) / (xb - xa)
+
+    def onselect(self, vmin, vmax):
+        self.bounds = (vmin, vmax)
+
 
 class InteractiveGUI:
     """
@@ -180,17 +221,13 @@ class InteractiveGUI:
         # amp is the current value of the slider
         self.rmin = self.srmin.val
         # update curve
-        self.draw()
+        self.clicked_reset(None)
-        # redraw canvas while idle
-        plt.draw()
 
     def update_rmax(self, val):
         # amp is the current value of the slider
         self.rmax = self.srmax.val
         # update curve
-        self.draw()
+        self.clicked_reset(None)
-        # redraw canvas while idle
-        plt.draw()
 
     def draw(self, first=False, update_map=True):
 
@@ -207,6 +244,7 @@ class InteractiveGUI:
                     self.fmap,
                     origin="lower",
                     cmap="RdBu_r",
+                    extent=self.im_extent,
                     norm=mpl.colors.LogNorm(),
                     vmin=1e-2,
                     vmax=1e2,
@@ -218,13 +256,9 @@ class InteractiveGUI:
 
             ## Gamma
             plt.sca(self.ax_gamma)
-            # (a, b, r, _, _) = linregress(self.frho_map[self.mask], self.fcs_map[self.mask])
 
-            if self.full_gamma_button.get_status()[0]:
+            if self.gamma_3d_button.get_status()[0]:
-                rho = np.log10(self.pp.cells["rho"])
+                rho, cs = self.rho3d[self.mask3d], self.P3d[self.mask3d]
-                cs = np.log10(np.sqrt(self.pp.cells["P"]))
-                mask_fin = np.isfinite(rho) & np.isfinite(cs)
-                rho, cs = rho[mask_fin], cs[mask_fin]
                 plt.xlabel(r"$\log(\rho)$")
                 plt.ylabel(r"$\log(P)$")
                 self.get_gamma = lambda a: a
@@ -253,18 +287,32 @@ class InteractiveGUI:
                     norm=mpl.colors.LogNorm(),
                     cmap=plt.get_cmap("plasma"),
                 )
-            plt.legend()
 
-        lps = self.line_gamma.list_points
+        if update_map and self.gamma_3d_button.get_status()[1]:
-        xa, ya, xb, yb = lps[0].x, lps[0].y, lps[1].x, lps[1].y
+            lr = linregress(rho[rho > 2], cs[rho > 2])
-        a = (yb - ya) / (xb - xa)
+            (a, b, r, _, _) = lr
+            self.line_gamma.clear()
+            del self.line_gamma
+            rhomin, rhomax = np.min(rho), np.max(rho)
+            self.line_gamma = DraggableLine(
+                self,
+                self.ax_gamma,
+                [rhomin, a * rhomin + b],
+                [rhomax, a * rhomax + b],
+                0.05,
+            )
+            print("Gamma linregress : {}".format(lr))
+        else:
+            lps = self.line_gamma.list_points
+            xa, ya, xb, yb = lps[0].x, lps[0].y, lps[1].x, lps[1].y
+            a = (yb - ya) / (xb - xa)
 
         self.ax_gamma.set_title("$\Gamma$ = {:.3g}".format(self.get_gamma(a)))
 
         ## PDF
         if first or update_map:
             plt.sca(self.ax_pdf)
-            nb_cells = np.sum(self.mask_flat.flatten())
+            nb_cells = np.sum(self.mask_flat)
             if first:
                 self.std_nb_cells = nb_cells
                 values, self.edges = np.histogram(
@@ -312,13 +360,20 @@ class InteractiveGUI:
         plt.sca(self.ax_profile)
         lps = self.line_profile.list_points
         xa, ya, xb, yb = lps[0].x, lps[0].y, lps[1].x, lps[1].y
-        xp, yp = line(int(xa), int(ya), int(xb), int(yb))
+        coor_pix = list(
-        rho_prof = self.frho_map[yp, xp]
+            np.asarray(
-        # Position on the line
+                (np.array([xa, ya, xb, yb]) * self.pp.pp_params.pymses.zoom + 0.5)
-
+                * self.shape[0],
-        # x = np.linspace(0, 1, len(rho_prof))
+                dtype=int,
+            )
+        )
+        xpp, ypp = line(*coor_pix)
+        rho_prof = self.rho_map[ypp, xpp]
+        xp = ((xpp / float(self.shape[0])) - 0.5) / self.pp.pp_params.pymses.zoom
+        yp = ((ypp / float(self.shape[1])) - 0.5) / self.pp.pp_params.pymses.zoom
+        print(xp, yp, xpp, ypp)
         x = np.sqrt(np.abs((xp - xa) * (xb - xa) + (yp - ya) * (yb - ya)))
-        x = (x - float(x[0])) / self.shape[0]  # /(x[-1] - x[0])
+        x = x - float(x[0])
         mask_fit_prof = (x >= self.fit_prof_vmin) & (x <= self.fit_prof_vmax)
         try:
             (a, b, r, _, _) = linregress(
@@ -328,14 +383,15 @@ class InteractiveGUI:
             print("Warning in linregress : {}".format(e))
             a, b, r = np.nan, np.nan, np.nan
 
-        xv, yv = float(xa) / self.shape[0], float(ya) / self.shape[1]
+        xv, yv = np.array([float(xa), float(ya)])
-        mask_vert = (xv >= self.pp.cells["pos"][:, 0]) & (
+        print(xv, yv)
-            xv < self.pp.cells["pos"][:, 0] + self.pp.cells["dx"]
+        mask_vert = (xv >= self.xy3d[:, 0]) & (
+            xv < self.xy3d[:, 0] + self.pp.cells["dx"]
         )
         mask_vert = (
             mask_vert
-            & (yv >= self.pp.cells["pos"][:, 1])
+            & (yv >= self.xy3d[:, 1])
-            & (yv < self.pp.cells["pos"][:, 1] + self.pp.cells["dx"])
+            & (yv < self.xy3d[:, 1] + self.pp.cells["dx"])
         )
         rho_vert = np.log10(self.pp.cells["rho"][mask_vert])
         z_vert = self.pp.cells["pos"][mask_vert][:, 2] - 0.5
@@ -364,39 +420,51 @@ class InteractiveGUI:
 
         plt.legend()
         plt.draw()
-        self.unselect_lasso()
+        self.unselect()
 
     def onselect(self, verts, select_data):
         # update curve
 
         path = Path(verts)
         self.mask_flat = self.mask.flatten()
+        if self.gamma_3d_button.get_status()[2]:
+            self.mask_flat = self.mask_flat | path.contains_points(select_data)
+        else:
+            self.mask_flat = self.mask_flat & path.contains_points(select_data)
 
-        self.mask_flat = self.mask_flat & path.contains_points(select_data)
         self.mask = self.mask_flat.reshape(self.shape)
+
+        if self.gamma_3d_button.get_status()[0]:
+            if self.gamma_3d_button.get_status()[2]:
+                self.mask3d = self.mask3d | path.contains_points(self.xy3d)
+            else:
+                self.mask3d = self.mask3d & path.contains_points(self.xy3d)
         self.draw()
 
     def clicked_select(self, val, selector, button):
         if not self.lasso:
             self.lasso = selector(
-                self.ax_fluct, partial(self.onselect, select_data=self.points)
+                self.ax_fluct, partial(self.onselect, select_data=self.xy)
-            )
-            self.lasso_gamma = selector(
-                self.ax_gamma, partial(self.onselect, select_data=self.rhocs)
             )
+            if not self.gamma_3d_button.get_status()[0]:
+                self.lasso_gamma = selector(
+                    self.ax_gamma, partial(self.onselect, select_data=self.rhocs)
+                )
             button.color = "0.55"
         else:
-            self.unselect_lasso()
+            self.unselect()
 
-    def unselect_lasso(self):
+    def unselect(self):
         self.lasso = False
         self.lasso_gamma = False
         self.lasso_button.color = "0.85"
         self.poly_button.color = "0.85"
 
     def clicked_reset(self, val):
-        self.mask = (self.rr > self.rmin) & (self.rr < self.rmax)
+        self.mask = (self.rr >= self.rmin) & (self.rr <= self.rmax)
         self.mask_flat = self.mask.flatten()
+        self.mask3d = np.isfinite(self.rho3d) & np.isfinite(self.P3d)
+        self.mask3d = self.mask3d & (self.rr3d >= self.rmin) & (self.rr3d <= self.rmax)
         self.lasso = False
         self.lasso_button.color = "0.85"
         self.draw()
@@ -431,21 +499,33 @@ class InteractiveGUI:
         self.pp.pdf_rho("z")
         self.pp.pdf_T("z")
         self.pp.load_cells()
+        self.im_extent = np.array(self.pp.get_attribute("/maps", "im_extent")) - 0.5
+
         self.datamap = self.pp.get_value("/maps/" + datamap_key)
+        self.rho_map = np.log10(self.pp.get_value("/maps/rho_z"))
         self.frho_map = np.log10(self.pp.get_value("/maps/fluct_rho_z"))
         self.T_map = self.pp.get_value("/maps/T_z")
         self.fcs_map = np.log10(np.sqrt(self.pp.get_value("/maps/fluct_T_z")))
-        self.rr = self.pp.get_value("/maps/rr_z")
 
         self.shape = self.datamap.shape
-        x = np.arange(self.shape[0])
+        x = np.linspace(self.im_extent[0], self.im_extent[1], self.shape[0])
-        y = np.arange(self.shape[1])
+        y = np.linspace(self.im_extent[2], self.im_extent[3], self.shape[1])
         self.xx, self.yy = np.meshgrid(x, y)
-        self.points = np.column_stack((self.xx.flatten(), self.yy.flatten()))
+        self.xy = np.column_stack((self.xx.flatten(), self.yy.flatten()))
+        self.rr = np.sqrt(self.xx ** 2 + self.yy ** 2)
         self.rhocs = np.column_stack((self.frho_map.flatten(), self.fcs_map.flatten()))
 
-        self.rmin = self.pp.pp_params.disk.rmin_pdf
+        self.rmin = self.pp.pp_params.disk.rmin_pdf / 2.0
-        self.rmax = self.pp.pp_params.disk.rmax_pdf
+        self.rmax = self.pp.pp_params.disk.rmax_pdf / 2.0
+        self.mask = (self.rr >= self.rmin) & (self.rr <= self.rmax)
+        self.mask_flat = self.mask.flatten()
+
+        self.rho3d = np.log10(self.pp.cells["rho"])
+        self.P3d = np.log10(self.pp.cells["P"])
+        self.xy3d = self.pp.cells["pos"][:, :2] - 0.5
+        self.rr3d = np.sqrt(np.sum((self.xy3d) ** 2, axis=1))
+        self.mask3d = np.isfinite(self.rho3d) & np.isfinite(self.P3d)
+        self.mask3d = self.mask3d & (self.rr3d >= self.rmin) & (self.rr3d <= self.rmax)
 
         self.ax_fluct = plt.axes([0.05, 0.6, 0.4, 0.35])
         self.ax_gamma = plt.axes([0.05, 0.15, 0.4, 0.35])
@@ -460,21 +540,26 @@ class InteractiveGUI:
             useblit=True,
             span_stays=True,
         )
-        self.line_profile = DraggableLine(
+        self.line_profile = DraggableLine(self, self.ax_fluct, [0, 0], [0.1, 0], 0.005)
-            self, self.ax_fluct, [100, 2000], [2000, 2000], 50
+        self.line_gamma = DraggableLine(self, self.ax_gamma, [0, 0], [1, 0.3], 0.05)
-        )
-        self.line_gamma = DraggableLine(self, self.ax_gamma, [0, 0], [1, 0.3], 0.03)
-        self.mask = (self.rr > self.rmin) & (self.rr < self.rmax)
-        self.mask_flat = self.mask.flatten()
 
-        ax_rmax = plt.axes([0.3, 0.03, 0.15, 0.02])
+        ax_rmax = plt.axes([0.05, 0.07, 0.15, 0.02])
-        self.srmax = Slider(ax_rmax, r"$r_{max}$", 0, 0.5, valinit=self.rmax)
+        self.srmax = Slider(
+            ax_rmax, r"$r_{max}$", 0, 0.7 / pp_params.pymses.zoom, valinit=self.rmax
+        )
         ax_rmin = plt.axes([0.05, 0.03, 0.15, 0.02])
-        self.srmin = Slider(ax_rmin, r"$r_{min}$", 0, 0.5, valinit=self.rmin)
+        self.srmin = Slider(
+            ax_rmin, r"$r_{min}$", 0, 0.7 / pp_params.pymses.zoom, valinit=self.rmin
+        )
         ax_lasso = plt.axes([0.6, 0.07, 0.19, 0.02])
         ax_poly = plt.axes([0.8, 0.07, 0.19, 0.02])
-        ax_full_gamma = plt.axes([0.1, 0.07, 0.19, 0.02])
+        ax_gamma_3d = plt.axes([0.3, 0.01, 0.19, 0.09])
-        self.full_gamma_button = CheckButtons(ax_full_gamma, ["Full $\Gamma$"], [True])
+        self.gamma_3d_button = CheckButtons(
+            ax_gamma_3d, ["3D $\Gamma$", "Fit  $\Gamma$", "Union"], [True, False, False]
+        )
+        self.gamma_3d_button.on_clicked(
+            lambda val: self.draw(first=False, update_map=True)
+        )
         self.lasso = False
         self.lasso_gamma = False
         self.lasso_button = Button(ax_lasso, "Lasso selector")
@@ -497,6 +582,6 @@ class InteractiveGUI:
         self.srmin.on_changed(self.update_rmin)
         self.srmax.on_changed(self.update_rmax)
 
-        self.draw(True)
+        self.draw(first=True, update_map=True)
 
         plt.show()

postprocessor.py

@@ -546,6 +546,38 @@ class PostProcessor(HDF5Container):
         pdf.attrs.var = np.var
         return True
 
+    # def kappa(self, mask):
+    #     mask_flat = mask.flatten()
+    #     fmap = self.get_value("/maps/fluct_coldens_z")
+    #     nb_cells = np.sum(self.mask_flat)
+    #     values, edges = np.histogram(np.log10(fmap.flatten()[mask_flat]),
+    #                                  self.pp_params.pdf.nb_bin,
+    #                                  weights = np.ones(nb_cells) / nb_cells)
+
+    #     centers = 0.5 * (edges[1:] + edges[:-1])
+    #     mask_fit = ((centers > self.pp.pp_params.pdf.xmin_fit) &
+    #                 (centers < self.pp.pp_params.pdf.xmax_fit) &
+    #                 (values > 0))
+    #     (a, b, r, p, err) = linregress(centers[mask_fit], np.log10(values[mask_fit]))
+    #     return a, err
+
+    # def gamma(self, mask):
+    #     rho3d = np.log10(self.cells["rho"])
+    #     P3d = np.log10(np.sqrt(self.cells["P"]))
+    #     xy3d = self.pp.cells["pos"][:, :2] - 0.5
+    #     self.rr3d = np.sqrt(np.sum((self.xy3d)**2, axis=1))
+    #     nb_cells = np.sum(self.mask_flat)
+    #     values, edges = np.histogram(np.log10(fmap.flatten()[mask_flat]),
+    #                                  self.pp_params.pdf.nb_bin,
+    #                                  weights = np.ones(nb_cells) / nb_cells)
+
+    #     centers = 0.5 * (edges[1:] + edges[:-1])
+    #     mask_fit = ((centers > self.pp.pp_params.pdf.xmin_fit) &
+    #                 (centers < self.pp.pp_params.pdf.xmax_fit) &
+    #              (values > 0))
+    #    (a, b, r, p, err) = linregress(centers[mask_fit], np.log10(values[mask_fit]))
+    #    return a, err
+
     def _sinks(self):
         csv_name = (
             self.path