Some improvements

disk_postprocess.py

@@ -5,11 +5,10 @@ import os
 import pymses
 import numpy as np
 import matplotlib as mpl
-
+if os.environ.get('DISPLAY','') == '':
-if os.environ.get("DISPLAY", "") == "":
+    print('No display found. Using non-interactive Agg backend')
-    print("No display found. Using non-interactive Agg backend")
+    mpl.use('Agg')
-    mpl.use("Agg")
+from matplotlib.patches import Polygon
-
 import pylab as P
 import glob as glob
 
@@ -27,6 +26,9 @@ from pymses.analysis import Camera, raytracing, slicing, splatting
 from pymses.filters import CellsToPoints
 from pymses.analysis import ScalarOperator, FractionOperator, MaxLevelOperator
 from scipy.stats import gaussian_kde
+import matplotlib.patches as mpatches
+from matplotlib.collections import PatchCollection
+
 
 
 import module_extract as me
@@ -122,19 +124,9 @@ def compute_image_data(path, num, radius=0.5,
         return
 
     # Prepare saving data
-    if "T" in images and not "rho" in images:
+    maps_disk = {'time' : time, 'im_extent' : im_extent,
-        images.append("rho")
+                 'center' : center, 'radius' : radius, 'lbox' : lbox,
-    if "Q" in images and not "coldens" in images:
+                 'images': images, 'axes_los': axes_los}
-        images.append("coldens")
-    maps_disk = {
-        "time": time,
-        "im_extent": im_extent,
-        "center": center,
-        "radius": radius,
-        "lbox": lbox,
-        "images": images,
-        "axes_los": axes_los,
-    }
 
     # Prepare raytracing
     rho_op = ScalarOperator(lambda dset: dset["rho"], ro.info["unit_density"])
@@ -344,6 +336,11 @@ def plot_maps(path, num,
             ):
                 continue
 
+            if (image == 'jeans'):
+                maps_disk['jeans_' + ax_los] = np.sqrt(np.pi * maps_disk['T_' + ax_los] / maps_disk['rho_' + ax_los])
+            if (image == 'jeans_ratio'):
+                maps_disk['jeans_ratio_' + ax_los] = maps_disk['jeans_' + ax_los] * 2**(maps_disk['levels_' + ax_los])
+
             map_disk = maps_disk[image + '_' + ax_los]
 
             if image == 'Q' :
@@ -353,6 +350,13 @@ def plot_maps(path, num,
                               cmap='RdBu',
                               norm=mpl.colors.LogNorm(),
                               vmin=0.01, vmax=100.)
+            elif image == 'jeans_ratio' :
+                im = P.imshow(map_disk,
+                              extent=im_extent,
+                              origin='lower',
+                              cmap='RdBu',
+                              norm=mpl.colors.LogNorm(),
+                              vmin=0.1, vmax=1000.)
             else:
                 im = P.imshow(
                     map_disk,
@@ -396,7 +400,7 @@ def plot_maps(path, num,
                     cont.levels = cont.levels + 1
 
                     P.clabel(cont,
-                             levels_ar[levels_ar < 11][1:],
+                             cont.levels[cont.levels < 11],
                              inline=1, fontsize=8., fmt='%1d')
             elif image == 'rho':
                 cbar.set_label(r'$\rho$ (code)')
@@ -421,6 +425,8 @@ def plot_maps(path, num,
                 cbar.set_label(r'$T (code)$')
             elif image == 'Q':
                 cbar.set_label(r'$Q$')
+            elif image == 'jeans':
+                cbar.set_label(r'Jeans\'s lenght')
             else:
                 cbar.set_label(image)
 
@@ -805,20 +811,67 @@ def plot_disk_prop(
             P.close()
 
 
-def disk_pdf(
+def sigma_pdf(fluct_maps, mask_flat, put_title=False, title='',  nb_bin_hist=50, tag=''):
-    path,
+    # Sigma-PDF
-    num,
+    dcoldens = np.log10(fluct_maps['coldens']).flatten()
-    maps_disk,
+
-    pos_star=[1.0, 1.0],
+    nb_cells = np.sum(mask_flat)
-    force=False,
+    P.grid()
-    interactive=False,
+    P.yscale('log')
-    nb_bin_hist=50,
+    P.ylim([0.5 / nb_cells, 1.])
-    tag="",
+    P.xlabel(r'$\log(\Sigma / \bar{\Sigma})$')
-    rad_min=0.075,
+    P.ylabel(r'$\mathcal{P}_\Sigma$')
-    rad_max=0.3,
+    if put_title:
-    put_title=True,
+        P.title(title)
-    do_speed=True,
+    values, edges, _ = P.hist(dcoldens[mask_flat],
-):
+                              nb_bin_hist, range=(-1, 3),
+                              weights = np.ones(nb_cells) / nb_cells)
+    centers = 0.5 * (edges[1:] + edges[:-1])
+
+    # Variance
+    var = np.var(dcoldens[mask_flat])
+    # Compute the slope of the right part of the histogramm
+    mask_fit = (centers > 0.) & (centers < 1.25) & (values > 0)
+    if (np.sum(mask_fit > 0)):
+        (a, b, rho, _, stderr) = linregress(centers[mask_fit], np.log10(values[mask_fit]))
+        P.plot(centers, 10**(a*centers + b), '--', linewidth=2)
+        print("pdf a=%e, b=%e, rho=%e, var=%e"% (a, b, rho, var))
+        try :
+            beta = int(tag.split('_')[1][4:])
+        except ValueError :
+            beta = 0
+        fit = {'beta': beta,'slope': a, 'origin': b, 'correlation': rho, 'stderr':stderr, 'var':var}
+
+    return fit
+
+def plot_dcsdrho(fluct_maps, mask_flat, put_title=False, title='',  nb_bin_hist=50, tag=''):
+    drho  = fluct_maps['rho'].flatten()
+    dcs   = fluct_maps['cs'].flatten()
+
+    nb_points = np.sum(mask_flat)
+    size_hist = int(np.sqrt(nb_points))
+
+    P.hist2d(np.log10(drho[mask_flat]), np.log10(dcs[mask_flat]), (size_hist, size_hist),  norm=mpl.colors.LogNorm())
+    P.xlabel(r'$\log(\rho / \bar{\rho})$')
+    P.ylabel(r'$\log(c_s / \bar{c_s})$')
+    P.ylim([-0.6, 0.6])
+    P.xlim([-1, 1.])
+
+    (a, b, rho, _, stderr) = linregress(np.log10(drho[mask_flat]), np.log10(dcs[mask_flat]))
+    P.plot(np.log10(drho[mask_flat]), a*np.log10(drho[mask_flat]) + b, '--', linewidth=2)
+    print("cs/rho a=%e, b=%e, rho=%e"% (a, b, rho))
+
+    try :
+        beta = int(tag.split('_')[1][4:])
+    except ValueError :
+        beta = 0
+    fit = {'beta': beta,'slope': a, 'origin': b, 'correlation': rho, 'stderr':stderr}
+
+def disk_pdf(path, num, maps_disk,
+             pos_star=[1., 1.], force=False, interactive=False,
+             nb_bin_hist=50, tag='',
+             rad_min=0.075, rad_max=0.3, put_title=True,
+             do_speed=True):
 
     # Load property file
     name_prop = path + "/prop_disk_" + str(num).zfill(5) + ".save"
@@ -853,9 +906,11 @@ def disk_pdf(
     print("maps file loaded")
 
     # Properties
-    time = prop_disk["time"]
+    time        =  prop_disk['time']
-    im_extent = maps_disk["im_extent"]
+    im_extent   =  maps_disk['im_extent']
-    title = tag.split("_")[1] + " t=" + str(time)[0:5] + " (code)"
+    title       =  None
+    if (put_title):
+        title   =  tag.split('_')[1] + ' t='+ str(time)[0:5] +' (code)'
 
     # Load coldens
     coldens_map = maps_disk["coldens_z"]
@@ -883,16 +938,14 @@ def disk_pdf(
 
     # Mask selecting the zone of interest
     mask_map = (rr > rad_min) & (rr < rad_max)
-    mask_flat = mask_map.flatten()
 
     # Additionnal maps
     rho_map = maps_disk["rho_z"]
     cs_map = np.sqrt(maps_disk["T_z"])
 
-    vx_map = maps_disk["vx_z"]
+    vx_map = maps_disk['vx_z']
-    vy_map = maps_disk["vy_z"]
+    vy_map = maps_disk['vy_z']
     v_map  = np.sqrt(vx_map**2 + vy_map**2)
-    v_kepl = np.sqrt(1.0 / rr)
     xx_star = xx - pos_star[0]
     yy_star = yy - pos_star[1]
     vrad_map = (vx_map * xx_star + vy_map * yy_star) / rr
@@ -931,45 +984,29 @@ def disk_pdf(
         avg_maps[cur_map] = np.reshape(avg_flat, rr.shape)
 
         #  Select zone of interest
-        avg_maps[cur_map][np.logical_not(mask_map)] = np.nan
+        # avg_maps[cur_map][np.logical_not(mask_map)] = np.nan
 
         # Compute fluctuation
         fluct_maps[cur_map] = map_arr / avg_maps[cur_map]
 
+    if not interactive:
+        plot_fluctuations_map(path, num, maps, fluct_maps, avg_maps, im_extent, mask_map, put_title, title,
+                             nb_bin_hist, tag, prop_disk)
+
+    # Save on disk
+    f = open(name_prop,'w')
+    pickle.dump(prop_disk, f)
+    f.close()
+
+    return (xx, yy, fluct_maps)
+
+def plot_fluctuations_map(path, num, maps, fluct_maps, avg_maps, im_extent, mask_map, put_title, title,
+                          nb_bin_hist, tag, prop_disk, interactive=False):
+
+    mask_flat = mask_map.flatten()
+
     # Sigma-PDF
-    dcoldens = np.log10(fluct_maps['coldens']).flatten()
+    fit = sigma_pdf(fluct_maps, mask_flat, put_title, title, nb_bin_hist, tag)
-
-    nb_cells = np.sum(mask_flat)
-    P.grid()
-    P.yscale('log')
-    P.ylim([0.5 / nb_cells, 1.])
-    P.xlabel(r'$\log(\Sigma / \bar{\Sigma})$')
-    P.ylabel(r'$\mathcal{P}_\Sigma$')
-    if put_title:
-        P.title(title)
-    values, edges, _ = P.hist(
-        dcoldens[mask_flat],
-        nb_bin_hist,
-        range=(-1, 3),
-        weights=np.ones(nb_cells) / nb_cells,
-    )
-    centers = 0.5 * (edges[1:] + edges[:-1])
-
-    # Variance
-    var = np.var(dcoldens[mask_flat])
-    # Compute the slope of the right part of the histogramm
-    mask_fit = (centers > 0.) & (centers < 1.25) & (values > 0)
-    if (np.sum(mask_fit > 0)):
-        (a, b, rho, _, stderr) = linregress(centers[mask_fit], np.log10(values[mask_fit]))
-        P.plot(centers, 10**(a*centers + b), '--', linewidth=2)
-        print("a=%e, b=%e, rho=%e, var=%e"% (a, b, rho, var))
-        try :
-            beta = int(tag.split('_')[1][4:])
-        except ValueError :
-            beta = 0
-        fit = {'beta': beta,'slope': a, 'origin': b, 'correlation': rho, 'stderr':stderr, 'var':var}
-        prop_disk['fit'] = fit
-        
     if interactive:
         P.figure()
     else:
@@ -999,13 +1036,10 @@ def disk_pdf(
         P.savefig(path + '/drho_hist_' + tag + '_' + str(num).zfill(5) + out_ext)
         P.close()
 
-
-
     # Derived quantities
+    dcoldens  = np.log10(fluct_maps['coldens']).flatten()
     dcs       = fluct_maps['cs'].flatten()
     dv        = fluct_maps['v'].flatten()
-    dvaz_kepl = abs(maps['vaz'] - v_kepl) / v_kepl
-    fluct_maps['vaz_kepl'] = dvaz_kepl
     dmach = abs(maps['v'] - avg_maps['v']) / maps['cs']
     dmach[dmach > 10.] = 10.
     fluct_maps['mach'] = dmach
@@ -1017,31 +1051,10 @@ def disk_pdf(
     prop_disk['dmach_mean'] = dmach_mean
     print("dmach_mean = {}".format(dmach_mean))
 
-    
     # Fluctuations plots
 
     # dcs = f(drho)
-    P.hist2d(
+    prop_disk['fit_cs'] = plot_dcsdrho(fluct_maps, mask_flat, put_title, title, nb_bin_hist, tag)
-        np.log10(drho[mask_flat]),
-        np.log10(dcs[mask_flat]),
-        (1000, 1000),
-        norm=mpl.colors.LogNorm(),
-    )
-    P.xlabel(r"$\log(\rho / \bar{\rho})$")
-    P.ylabel(r"$\log(c_s / \bar{c_s})$")
-    P.ylim([-0.6, 0.6])
-    P.xlim([-1, 1.])
-
-    (a, b, rho, _, stderr) = linregress(np.log10(drho[mask_flat]), np.log10(dcs[mask_flat]))
-    P.plot(np.log10(drho[mask_flat]), a*np.log10(drho[mask_flat]) + b, '--', linewidth=2)
-    print("cs/rho a=%e, b=%e, rho=%e"% (a, b, rho))
-
-    try :
-        beta = int(tag.split('_')[1][4:])
-    except ValueError :
-        beta = 0
-    fit = {'beta': beta,'slope': a, 'origin': b, 'correlation': rho, 'stderr':stderr}
-    prop_disk['fit_cs'] = fit
 
     P.colorbar()
     if interactive:
@@ -1096,15 +1109,6 @@ def disk_pdf(
                           vmin=-2.,
                           vmax=2.)
             label = r'$log(\rho/\bar{\rho})$'
-        elif cur_map == 'vaz_kepl':
-            im = P.imshow(fluct_map,
-                          extent=im_extent,
-                          origin='lower',
-                          cmap='RdBu_r',
-                          norm=mpl.colors.LogNorm(),
-                          vmax=1.,
-                          vmin=0.01)
-            label = r'$|v_\varphi - v_{kepl}|/v_{kepl}$'
         elif cur_map == 'vaz':
             im = P.imshow(fluct_map,
                           extent=im_extent,
@@ -1135,12 +1139,6 @@ def disk_pdf(
             P.savefig(name_im)
             P.close()
 
-    # Save on disk
-    f = open(name_prop,'w')
-    pickle.dump(prop_disk, f)
-    f.close()
-
-
 
 
 
@@ -1541,3 +1539,142 @@ def clump_study(path, num, path_out, tag, force=False, rho_thres=1e3, lvl_thres=
     # P.imshow(np.log10(m['coldens_z']), extent=m['im_extent'], cmap='plasma')
 
     # P.savefig(path_out + '/clumps_size_' + tag + '_' +  str(num).zfill(5) + out_ext)
+
+
+
+class interactive_plots:
+
+    def onbuttonrelease(self, event):
+        """Deal with click events"""
+        button = ['left','middle','right']
+        toolbar = P.get_current_fig_manager().toolbar
+        if toolbar.mode=='zoom rect' and event.inaxes == self.ax_col:
+            print("zooming ")
+            xlim = self.ax_col.get_xlim()
+            ylim = self.ax_col.get_ylim()
+            self.reset_mask()
+        elif self.add_mask and event.inaxes == self.ax_col:
+            self.plot_side()
+            P.draw()
+
+    def onbuttonpress(self, event):
+        """Deal with click events"""
+        button = ['left','middle','right']
+        toolbar = P.get_current_fig_manager().toolbar
+        if toolbar.mode!='':
+            print("You clicked on something, but toolbar is in mode {:s}.".format(toolbar.mode))
+        print(self.add_mask)
+        if self.add_mask and toolbar.mode=='' and event.inaxes == self.ax_col:
+            ix, iy = event.xdata, event.ydata
+            print("Add patch {}, {}".format(ix, iy))
+            xlim = self.ax_col.get_xlim()
+            ylim = self.ax_col.get_ylim()
+            radius = 0.05 * min(abs(xlim[1] - xlim[0]), abs(ylim[1] - ylim[0]))
+            circle = mpatches.Circle([ix, iy], radius,  color='black', alpha=0.1,  ec="none")
+            self.circles.append(circle)
+            self.ax_col.add_artist(circle)
+            self.ax_col.draw_artist(circle)
+            self.patch_mask = self.patch_mask | ((self.xx - ix)**2 + (self.yy -iy)**2 < radius**2)
+            #self.plot_side()
+
+    def onkeypress(self, event):
+        """whenever a key is pressed"""
+        if not event.inaxes:
+            return
+        if event.key == 't':
+            self.add_mask = not self.add_mask
+            print("Add mode is {}".format(self.add_mask))
+        elif event.key == 'r':
+            self.reset_mask()
+
+    def plot_side(self):
+        if (self.add_mask):
+            mask = (self.patch_mask & self.mask).flatten()
+        else:
+            mask = self.mask.flatten()
+        self.ax_gamma.clear()
+        P.sca(self.ax_gamma)
+        plot_dcsdrho(self.fluct_maps, mask, tag=self.tag)
+
+        self.ax_pdf.clear()
+        P.sca(self.ax_pdf)
+        sigma_pdf(self.fluct_maps, mask, tag=self.tag, nb_bin_hist=self.args.pdf_nb_bin)
+
+    def reset_mask(self):
+        xlim = self.ax_col.get_xlim()
+        ylim = self.ax_col.get_ylim()
+        self.mask = (self.xx >= xlim[0]) & (self.xx <= xlim[1]) & (self.yy >= ylim[0]) & (self.yy <= ylim[1])
+        self.patch_mask = np.full(self.mask.shape, False)
+        for circle in self.circles:
+            circle.remove()
+        self.circles = []
+        self.plot_side()
+        P.draw()
+
+
+
+    def __init__(self, args, path, num,
+                 maps_disk=None,
+                 tag='',
+                 set_lim=True) :
+        """
+        Interactive plotting
+
+        Parameters
+        ----------
+        num       output number
+        path      path of the pipeline output
+
+        """
+        self.args = args
+        self.add_mask = False
+        self.circles = []
+        self.tag = tag
+
+        path_out = path
+
+        # Load maps file
+        print("load maps file")
+        name_maps = path + '/maps_disk' + '_' + tag + '_' + format(num,'05') + '.save'
+
+        if maps_disk is None:
+            if (len(glob.glob(name_maps)) == 0):
+                raise IOError('no pickle file for disk maps {}. Run make_image_disk() first'.format(name_maps))
+        f = open(name_maps,'r')
+        maps_disk = pickle.load(f)
+        f.close()
+        print("maps file loaded")
+
+        im_extent   = maps_disk['im_extent']
+
+        fig = P.figure();
+        self.ax_col = P.subplot(1, 2, 1)
+        coldens = maps_disk['coldens_z']
+        im = self.ax_col.imshow(coldens,
+                                extent=im_extent,
+                                origin='lower',
+                                norm=mpl.colors.LogNorm())
+        if set_lim:
+            im.set_clim(0.01, 100)
+        self.ax_col.set_xlabel(r'$x$')
+        self.ax_col.set_ylabel(r'$y$')
+
+        self.xx, self.yy, self.fluct_maps = disk_pdf(path, num, maps_disk, tag=self.tag, force=True, put_title=False, interactive=True)
+        coord_flat = zip(self.xx.flatten(), self.yy.flatten())
+
+        self.ax_gamma = P.subplot(2, 2, 2)
+        self.ax_pdf = P.subplot(2,2,4)
+
+        xlim = self.ax_col.get_xlim()
+        ylim = self.ax_col.get_ylim()
+        self.mask = (self.xx >= xlim[0]) & (self.xx <= xlim[1]) & (self.yy >= ylim[0]) & (self.yy <= ylim[1])
+        self.patch_mask = np.full(self.mask.shape, False)
+
+        self.plot_side()
+
+        fig.canvas.mpl_connect('button_release_event', self.onbuttonrelease)
+        fig.canvas.mpl_connect('button_press_event', self.onbuttonpress)
+        fig.canvas.mpl_connect('key_press_event', self.onkeypress)
+
+        P.tight_layout()
+        P.show()

pipeline_disk.py

@@ -109,26 +109,23 @@ parser.add_argument("--print_outputs",
                     help="print names of outputs",
                     action='store_true')
 
-parser.add_argument(
+parser.add_argument("--fft",
-    "--fft", help="use quick and dirty fft rendering", action="store_true"
+                    help="use quick and dirty fft rendering",
-)
+                    action='store_true')
-parser.add_argument(
+parser.add_argument("--images", nargs='*',
-    "--images",
+                    default=['coldens', 'rho', 'speed', 'Q', 'T'],
-    nargs="*",
+                    choices=['coldens', 'rho', 'speed', 'Q', 'T', 'levels', 'cpu', 'jeans', 'jeans_ratio'])
-    default=["coldens", "rho", "speed", "Q", "T"],
+parser.add_argument("--axes", nargs='*',
-    choices=["coldens", "rho", "speed", "Q", "T", "levels", "cpu"],
+                    default=['x', 'y', 'z'],
-)
+                    choices=['x', 'y', 'z'])
-parser.add_argument(
+parser.add_argument("--zoom",
-    "--axes", nargs="*", default=["x", "y", "z"], choices=["x", "y", "z"]
+                    help="zoom",
-)
+                    type=float,
-parser.add_argument("--zoom", help="zoom", type=float, default=2.0)
+                    default=2.)
-parser.add_argument(
+parser.add_argument("-ms", "--mapsize",
-    "-ms",
-    "--mapsize",
                     help="size of the maps created in he map mode (in pixel)",
                     type=int,
-    default=1024,
+                    default=1024)
-)
 
 
 parser.add_argument("--nb_bin",
@@ -221,8 +218,23 @@ me.P = dp.P
 # List of id that were successfully computed
 run_succeded = {}
 
+
+# Care for dependencies
+images = args.images
+if 'jeans_ratio' in images and not 'jeans' in images :
+     images = ['jeans'] + images
+if 'jeans_ratio' in images and not 'levels' in images :
+     images.append('levels')
+if 'jeans' in images and not 'T' in images :
+     images.append('T')
+if ('T' in images or 'jeans' in images) and not 'rho' in images :
+     images.append('rho')
+if 'Q' in images and not 'coldens' in images :
+     images.append('coldens')
+
+# Go through all runs
 for run in runs:
-    path_suffix = project + "/" + run
+     path_suffix = project + '/' + run
      path_in = storage_in + path_suffix
      path_out = storage_out + path_suffix
 
@@ -234,8 +246,8 @@ for run in runs:
      if not os.path.exists(path_out):
           os.makedirs(path_out)
      try:
-            copy(path_in + "/disk.nml", path_out)
+          copy(path_in + '/disk.nml', path_out)
-            copy(path_in + "/output_00001/compilation.txt", path_out)
+          copy(path_in + '/output_00001/compilation.txt', path_out)
      except:
           pass
 
@@ -271,7 +283,7 @@ for run in runs:
                                                        map_size=args.mapsize,
                                                        force=args.force_redo,
                                                        axes_los=args.axes,
-                                                   images=args.images,
+                                                       images=images,
                                                        pos_star=np.array([args.x, args.y, args.z]),
                                                        fft=args.fft)
                          print("[{}, {}] maps computed".format(run, num))
@@ -280,7 +292,7 @@ for run in runs:
                          maps_disk = dp.plot_maps(path_out, num,
                                                   maps_disk=maps_disk,
                                                   axes_los=args.axes,
-                                             images=args.images,
+                                                  images=images,
                                                   tag=tag_run,
                                                   force=args.force_redo,
                                                   interactive=args.interactive,
@@ -322,12 +334,15 @@ for run in runs:
                          print("[{}, {}] computing clumps".format(run, num))
                          dp.clump_study(path_in, num, path_out, tag_run)
                          print("[{}, {}] clumps computed".format(run, num))
+                    if args.interactive:
+                         print("[{}, {}] Interactive session".format(run, num))
+                         maps_disk = dp.interactive_plots(args, path_out, num, tag=tag_run)
                     # If we are here, success !
                     success = True
                     run_succeded[run].append(num)
                except (ValueError, IOError, KeyError) as e:
                     print(e)
-                if args.watch and failures < args.allowed_failures:
+                    if(args.watch and failures < args.allowed_failures):
                          failures = failures + 1
                          print("Unable to proceed for run {} output {}. Trying again in {} s ({} tries remaining)".format(run, num, args.waiting_time, args.allowed_failures - failures))
                          time.sleep(args.waiting_time)
@@ -372,7 +387,7 @@ if args.compare:
                                pdf=args.pdf or args.cpdf)
           except (KeyError, IOError) as e:
                print(e)
-                if args.skip:
+               if (args.skip):
                     pass
                else:
                     raise