From 0cc77d69afef39afc3603e0784131babb3be72ca Mon Sep 17 00:00:00 2001
From: Noe Brucy <noe.brucy@ens-rennes.fr>
Date: Mon, 15 Jun 2020 14:23:27 +0200
Subject: [PATCH] Add B streamline plot

---
 plotter.py       | 38 ++++++++++++++++++++++++++++++++++++
 postprocessor.py | 50 ++++++++++++++++++++++++++++++++++++------------
 2 files changed, 76 insertions(+), 12 deletions(-)

diff --git a/plotter.py b/plotter.py
index b27ab65..ea0893f 100644
--- a/plotter.py
+++ b/plotter.py
@@ -521,6 +521,32 @@ class Plotter(Aggregator, BaseProcessor):
             coordinates="figure",
         )
 
+    def _overlay_B(self, ax_los):
+        ax_h = self._axes_h[ax_los]
+        ax_v = self._axes_v[ax_los]
+        dmap_Bh_node = self.save.get_node("/maps/B_h_{}".format(ax_los))
+        dmap_Bh = dmap_Bh_node.read()
+        dmap_Bv = self.save.get_node("/maps/B_v_{}".format(ax_los)).read()
+
+        vel_red = self.pp_params.plot.vel_red
+        radius = self.save.root.maps._v_attrs.radius
+        center = self.save.root.maps._v_attrs.center
+        lbox = self.save.root._v_attrs.lbox
+
+        map_Bh_red = dmap_Bh[::vel_red, ::vel_red]  # take only a subset of velocities
+        map_Bv_red = dmap_Bv[::vel_red, ::vel_red]
+        nh = map_Bh_red.shape[0]
+        nv = map_Bv_red.shape[1]
+        vec_h = (
+            np.arange(nh) * 2.0 / nh * radius - radius + center[0] + radius / nh
+        ) * lbox
+        vec_v = (
+            np.arange(nv) * 2.0 / nv * radius - radius + center[1] + radius / nv
+        ) * lbox
+        hh, vv = np.meshgrid(vec_h, vec_v)
+
+        P.streamplot(hh, vv, map_Bh_red, map_Bv_red)
+
     def _plot_radial(self, name, ax_los, label=None, xlog=False, ylog=False):
 
         radial_bins = self.save.get_node("/radial/radial_bins_" + ax_los).read()
@@ -945,6 +971,18 @@ class Plotter(Aggregator, BaseProcessor):
                 "Density slice with speed overlay",
                 dependencies=["rho", "speed_h", "speed_v"],
             ),
+            "rho_B": PlotRule(
+                self,
+                partial(
+                    self._plot_map,
+                    "rho",
+                    label=r"$\rho$",
+                    unit=cst.Msun_pc3,
+                    overlays=[self._overlay_B],
+                ),
+                "Density slice with magnetic field overlay",
+                dependencies=["rho", "B_h", "B_v"],
+            ),
             "jeans_ratio": PlotRule(
                 self,
                 partial(
diff --git a/postprocessor.py b/postprocessor.py
index aa5727e..f36198d 100644
--- a/postprocessor.py
+++ b/postprocessor.py
@@ -313,21 +313,33 @@ class PostProcessor(HDF5Container):
         datamap_rho = slicing.SliceMap(self._amr, self._cam[ax_los], self._rho_op, z=z)
         return (datamap_rho.map).T
 
-    def _speed_h(self, ax_los, z=0.0):
-        vh_op = ScalarOperator(
-            lambda dset: dset["vel"][:, self._ax_nb[self._axes_h[ax_los]]],
-            self._ro.info["unit_velocity"],
+    def _vector_h(self, name, unit, ax_los, z=0.0):
+        h_op = ScalarOperator(
+            lambda dset: dset[name][:, self._ax_nb[self._axes_h[ax_los]]],
+            self._ro.info[unit],
         )
-        dmap_vh = slicing.SliceMap(self._amr, self._cam[ax_los], vh_op, z=z).map.T
-        return dmap_vh
+        dmap_h = slicing.SliceMap(self._amr, self._cam[ax_los], h_op, z=z).map.T
+        return dmap_h
+
+    def _vector_v(self, name, unit, ax_los, z=0.0):
+        v_op = ScalarOperator(
+            lambda dset: dset[name][:, self._ax_nb[self._axes_v[ax_los]]],
+            self._ro.info[unit],
+        )
+        dmap_v = slicing.SliceMap(self._amr, self._cam[ax_los], v_op, z=z).map.T
+        return dmap_v
+
+    def _speed_h(self, ax_los, z=0.0):
+        return self._vector_h("vel", "unit_velocity", ax_los, z)
 
     def _speed_v(self, ax_los, z=0.0):
-        vv_op = ScalarOperator(
-            lambda dset: dset["vel"][:, self._ax_nb[self._axes_v[ax_los]]],
-            self._ro.info["unit_velocity"],
-        )
-        dmap_vv = slicing.SliceMap(self._amr, self._cam[ax_los], vv_op, z=z).map.T
-        return dmap_vv
+        return self._vector_v("vel", "unit_velocity", ax_los, z)
+
+    def _B_h(self, ax_los, z=0.0):
+        return self._vector_h("Br", "unit_mag", ax_los, z)
+
+    def _B_v(self, ax_los, z=0.0):
+        return self._vector_v("Br", "unit_mag", ax_los, z)
 
     def _temperature(self, ax_los, z=0.0):
         P_op = ScalarOperator(lambda dset: dset["P"], self._ro.info["unit_pressure"])
@@ -654,6 +666,20 @@ class PostProcessor(HDF5Container):
                 "/maps",
                 unit=self.info["unit_velocity"],
             ),
+            "B_h": Rule(
+                self,
+                self._speed_h,
+                "Horizontal slice of the magnetic field wrt the line of sight",
+                "/maps",
+                unit=self.info["unit_mag"],
+            ),
+            "B_v": Rule(
+                self,
+                self._speed_v,
+                "Vertical  slice  of the magnetic field wrt the line of sight",
+                "/maps",
+                unit=self.info["unit_mag"],
+            ),
             "T": Rule(
                 self,
                 self._temperature,