From 2c1c690ba8425384c3a4c31d993ce16c6f289a22 Mon Sep 17 00:00:00 2001
From: Noe Brucy <noe.brucy@cea.fr>
Date: Tue, 19 Jan 2021 10:08:42 +0100
Subject: [PATCH] [plotter] add particle scatter overlay

---
 plotter.py       | 27 +++++++++++++++++++++++++++
 postprocessor.py |  5 +++++
 2 files changed, 32 insertions(+)

diff --git a/plotter.py b/plotter.py
index 262d572..a9a5e59 100644
--- a/plotter.py
+++ b/plotter.py
@@ -689,6 +689,32 @@ class Plotter(Aggregator, BaseProcessor):
             **kwargs,
         )
 
+    def _overlay_particles(self, ax_los, im_extent, **kwargs):
+        """
+        Add an overlay with particles data
+        """
+
+        # Open particle HDF5 file
+        filename = self.save.get_node("/hdf5/particles").read()
+        hdf5_parts = tables.open(filename, "r")
+        part_pos = hdf5_parts.get_node("/data/pos").read()
+
+        # index of the horizontal axis
+        ih = self._ax_nb[self._axes_h[ax_los]]
+        # index of the vertical axis
+        iv = self._ax_nb[self._axes_v[ax_los]]
+
+        # horizontal coordinates
+        part_h = part_pos[:, ih]
+        part_v = part_pos[:, iv]
+
+        # Renormalize
+        part_h = im_extent[0] + (im_extent[1] - im_extent[0]) * part_h
+        part_v = im_extent[2] + (im_extent[3] - im_extent[2]) * part_v
+
+        # Scatter plot
+        plt.scatter(part_h, part_v)
+
     def _overlay_speed(
         self, ax_los, im_extent, unit=U.km_s, unit_coeff=1.0, key_v=None, **kwargs
     ):
@@ -1739,6 +1765,7 @@ class Plotter(Aggregator, BaseProcessor):
             "speed": self._overlay_speed,
             "levels": self._overlay_levels,
             "contour": self._overlay_contour,
+            "particles": self._overlay_particles,
         }
 
         super(Plotter, self).def_rules()
diff --git a/postprocessor.py b/postprocessor.py
index 99e2ccd..7ad1df3 100644
--- a/postprocessor.py
+++ b/postprocessor.py
@@ -213,6 +213,10 @@ class PostProcessor(HDF5Container):
             self.path_out + "/cells_" + tag_name + format(num, "05") + ".h5"
         )
 
+        self.parts_filename = (
+            self.path_out + "/parts_" + tag_name + format(num, "05") + ".h5"
+        )
+
         self.pspec_filename = (
             self.path_out + "/pspec_" + tag_name + format(num, "05") + ".h5"
         )
@@ -1160,6 +1164,7 @@ class PostProcessor(HDF5Container):
         """Ensure particles are written in the hdf5 file"""
         self.load_parts()
         self.unload_parts()
+        return self.parts_filename
 
     def _filaments(self):