From 4fa72eccc503da3ce76bfb8369207df59374c8d5 Mon Sep 17 00:00:00 2001
From: Noe Brucy <noe.brucy@cea.fr>
Date: Thu, 22 Jul 2021 15:38:16 +0200
Subject: [PATCH] [plotter] add sinks overlays

---
 plotter.py | 44 +++++++++++++++++++++++++++++++++-----------
 1 file changed, 33 insertions(+), 11 deletions(-)

diff --git a/plotter.py b/plotter.py
index 25fcc8b..ef8cbb1 100644
--- a/plotter.py
+++ b/plotter.py
@@ -21,6 +21,7 @@ from numpy.polynomial.polynomial import polyfit
 from scipy import optimize
 from scipy.ndimage.filters import gaussian_filter1d
 from scipy.stats import linregress
+import pandas as pd
 
 if os.environ.get("DISPLAY", "") == "":
     print("No display found. Using non-interactive Agg backend")
@@ -189,6 +190,7 @@ class Plotter(Aggregator, BaseProcessor):
         self.gen_simus()
 
         self.save = None
+        self.current_snap = None
 
     def gen_simus(self):
         self.simulations = {}
@@ -336,7 +338,7 @@ class Plotter(Aggregator, BaseProcessor):
                 real_ax = ax[i]
             except TypeError as e:
                 if ax is None:
-                    fig, real_ax = plt.subplots(1, 1)
+                    _, real_ax = plt.subplots(1, 1)
                 elif not_array_error(e):
                     real_ax = ax
                 else:
@@ -344,12 +346,14 @@ class Plotter(Aggregator, BaseProcessor):
 
             # Find plot save
             if from_cells or rule.kind == "cells":
-                if not os.path.exists(self.snaps[run][num].cells_filename):
+                self.current_snap = self.snaps[run][num]
+                if not os.path.exists(self.current_snap.cells_filename):
                     self.snaps[run][num].load_cells()
                     self.snaps[run][num].unload_cells()
                     save = tables.open_file(self.snaps[run][num].cells_filename)
             elif rule.kind == "snapshot":
-                save = tables.open_file(self.snaps[run][num].filename)
+                self.current_snap = self.snaps[run][num]
+                save = tables.open_file( self.current_snap.filename)
             else:
                 save = tables.open_file(self.study.filename, "r")
 
@@ -786,18 +790,36 @@ class Plotter(Aggregator, BaseProcessor):
         )
 
     def _overlay_particles(
-        self, ax_los, im_extent, unit_space=U.pc, center_space=False, **kwargs
+        self,
+        ax_los,
+        im_extent,
+        unit_space=U.pc,
+        center_space=False,
+        parts = True,
+        sinks = False,
+        
+        **kwargs
     ):
         """
         Add an overlay with particles data
         """
 
-        # Open particle HDF5 file
-        filename = self.save.get_node("/hdf5/write_particles").read()[0].decode()
-        hdf5_parts = tables.open_file(filename, "r")
-        part_pos = hdf5_parts.get_node("/data/pos").read()
-        hdf5_parts.close()
-        unit_length = self.save.root._v_attrs["unit_length"]
+        unit_length = self.current_snap.info["unit_length"]
+
+        if sinks:
+            self.current_snap.sinks()
+            sinks = pd.DataFrame(self.current_snap.get_value("/datasets/sinks"))
+            part_pos = sinks[["x", "y", "z"]].values
+            mass = sinks.M
+
+        if parts:
+            # Open particle HDF5 filetype_from_ext
+            filename = self.save.get_node("/hdf5/write_particles").read()[0].decode()
+            hdf5_parts = tables.open_file(filename, "r")
+            part_pos = hdf5_parts.get_node("/data/pos").read()
+            mass =  hdf5_parts.get_node("/data/mass").read()
+            mass *= self.current_snap.info["unit_mass"].express(U.Msun)
+            hdf5_parts.close()
 
         # index of the horizontal axis
         ih = self._ax_nb[self._axes_h[ax_los]]
@@ -831,7 +853,7 @@ class Plotter(Aggregator, BaseProcessor):
         part_v = part_v[mask]
 
         # Scatter plot
-        plt.scatter(part_h, part_v, **kwargs)
+        plt.scatter(part_h, part_v, s=mass/5e3, **kwargs)
 
     def _overlay_speed(
         self, ax_los, im_extent, unit=U.km_s, unit_coeff=1.0, key_v=None, **kwargs