Merge pull request #12 from nicolasaunai/plot

add plot to maps

mpmaps/mpmaps.py

@@ -5,6 +5,7 @@
 from platformdirs import user_data_dir
 from .globals import grids
 import os
+import matplotlib.pyplot as plt
 
 
 from spok.models import planetary as smp
@@ -15,6 +16,29 @@
 
 class MPMap:
     def __init__(self, **kwargs):
+        """
+        Keyword arguments:
+        ------------------
+
+        clock: float, (default:-90)
+                IMF clock angle in degrees, 0 is due north
+
+        cone: float, (default: 55)
+                IMF cone angle in degrees, 0 is radial IMF
+
+        tilt: float, (default: 0)
+                Dipole tilt axis
+
+        bimf: float, (default: 5)
+                IMF amplitude in nT
+
+        nws: float (default 5)
+                solar wind density, in cm^-3
+
+        mp_thick: float, (default 800)
+                Magnetopause thickness in km
+                only used for computing the current density
+        """
         self._grid_path = os.path.join(user_data_dir(), "mpmaps")
         self._clock = kwargs.get("clock", -90)
         self._cone = kwargs.get("cone", 55)
@@ -214,6 +238,57 @@ def reconnection_rate(self, rec_angle="max_rate"):
         R = k * u / v
         return R
 
+    def plot(self, **kwargs):
+        """
+        Keyword arguments
+        -----------------
+
+        value : string ("shear_angle":default, "reconnection_rate", "current_density")
+                the quantity that is plotted
+
+        xlim : tuple, lower bound of the plot area
+        ylim : tuple, upper bound of the plot area
+        filename : string, file name to save the figure on disk
+
+        other keywork arguments: see MPMap
+
+        example : mp.plot(value="shear_angle", tilt=14, xlim=(-18,18), ylim=(-18,18))
+        """
+
+        if "ax" in kwargs:
+            ax = kwargs["ax"]
+            fig = ax.get_figure()
+        else:
+            fig, ax = plt.subplots()
+
+        msh = smp.Magnetosheath()
+        phi = np.arange(0, np.pi * 2 + 0.1, 0.1)
+        theta = np.pi / 2
+        _, y, z = msh.magnetopause(theta, phi)
+        ax.plot(y, z, ls="--", color="k")
+
+        xmin, xmax = kwargs.get("xlim", (self.Y.min(), self.Y.max()))
+        ymin, ymax = kwargs.get("ylim", (self.Z.min(), self.Z.max()))
+
+        ax.set_xlim((xmin, xmax))
+        ax.set_ylim((ymin, ymax))
+        ax.set_xlabel(r"$Y/R_e$")
+        ax.set_ylabel(r"$Z/R_e$")
+
+        self.set_parameters(**kwargs)
+        val = kwargs.get("value", "shear_angle")
+        sa = getattr(self, val)()
+
+        ax.pcolormesh(self.Y, self.Z, sa, cmap="jet")
+        ax.axvline(0, ls="--", color="k")
+        ax.axhline(0, ls="--", color="k")
+        ax.set_aspect("equal")
+
+        if "filename" in kwargs:
+            fig.savefig(kwargs["filename"])
+
+        return fig, ax
+
     def _find_rec_angle_max_rate(self):
         def _dRR_dtheta_local(theta, params):
             n1 = params[:, 0] * 1e6 * cst.m_p