Adds feature "plot current sheet"

changelog/143.feature.rst

@@ -0,0 +1 @@
+Added a method to plot a current sheet via :meth:`~sunkit_pyvista.plotter.SunpyPlotter.plot_current_sheet`.

examples/field_lines.py → examples/field_lines_current_sheets.py

@@ -1,11 +1,14 @@
 """
-======================================
-Plotting Field Lines from sunkit-magex
-======================================
+=========================================================
+Plotting field lines and current sheets from sunkit-magex
+=========================================================
 
 ``sunkit-pyvista`` can be used to plot field lines from ``sunkit-magex``.
 
 This example requires the `streamtracer <https://streamtracer.readthedocs.io/en/stable/>`__ package.
+
+This example uses specific sample data from ``sunkit-magex`` to demonstrate the plotting of field lines and current sheets.
+If you want to adapt this, you have to make sure that the input magnetic field data has the same date as the AIA image.
 """
 
 import astropy.units as u
@@ -24,33 +27,16 @@
 from sunkit_pyvista.sample import LOW_RES_AIA_193
 
 ###############################################################################
-# We will be using an AIA 193 image from the sunpy sample data as the base image.
-
-# sphinx_gallery_defer_figures
-
-# Start by creating a plotter
-plotter = SunpyPlotter()
-
-# Plot a map
-plotter.plot_map(LOW_RES_AIA_193, clip_interval=[1, 99] * u.percent, assume_spherical_screen=False)
-# Add an arrow to show the solar rotation axis
-plotter.plot_solar_axis()
-
-###############################################################################
-# We now need to do the magnetic field extrapolation using ``sunkit-magex``.
-
-# sphinx_gallery_defer_figures
+# We will start by doing the magnetic field extrapolation using ``sunkit-magex``.
 
 # We load a gong_map from sunkit-magex
 gong_fname = get_gong_map()
 gong_map = sunpy.map.Map(gong_fname)
-# Now we plot the Gong Map to fill in the farside.
-plotter.plot_map(gong_map, cmap="hmimag", clip_interval=[1, 99] * u.percent)
 
-# Create points spaced between lat={-90, 90} degrees
+# Create points spaced between lat={-89, 89} degrees
 lat = np.linspace(-np.pi / 2, np.pi / 2, 32, endpoint=False)
 # Create 32 points spaced between long={0, 360} degrees
-lon = np.linspace(0, 2 * np.pi, 32, endpoint=False)
+lon = np.linspace(1, 2 * np.pi, 32, endpoint=False)
 # Make a 2D grid from these 1D points
 lat, lon = np.meshgrid(lat, lon, indexing="ij")
 # Create lon, lat and radial coordinate values by using a sunkit-magex
@@ -65,6 +51,21 @@
 tracer = tracing.FortranTracer()
 field_lines = tracer.trace(seeds, output_)
 
+###############################################################################
+# We will be using an AIA 193 image from the sunpy sample data as the base image.
+
+# sphinx_gallery_defer_figures
+
+# Start by creating a plotter
+plotter = SunpyPlotter()
+
+# Plot the AIA  map
+plotter.plot_map(LOW_RES_AIA_193, clip_interval=[1, 99] * u.percent, assume_spherical_screen=False)
+# Add an arrow to show the solar rotation axis
+plotter.plot_solar_axis()
+# Now we plot the Gong Map to fill in the farside.
+plotter.plot_map(gong_map, cmap="hmimag", clip_interval=[1, 99] * u.percent)
+
 ###############################################################################
 # We can also specify a color function while plotting the field lines.
 # This function takes a single field line, and returns a color either
@@ -80,6 +81,8 @@ def my_fline_color_func(field_line):
 
 # Plotting the field lines
 plotter.plot_field_lines(field_lines, color_func=my_fline_color_func)
+# Plotting the current sheet
+plotter.plot_current_sheet(output_)
 
 # Set the camera coordinate to view the plot correctly
 camera_coord = SkyCoord(

sunkit_pyvista/plotter.py

@@ -9,6 +9,7 @@
 from astropy.visualization import AsymmetricPercentileInterval
 from astropy.visualization.wcsaxes import Quadrangle
 from matplotlib import colors
+from sunkit_magex.pfss.coords import strum2cart
 from sunpy.coordinates import HeliocentricInertial, Helioprojective
 from sunpy.coordinates.utils import get_rectangle_coordinates
 from sunpy.map.maputils import all_corner_coords_from_map
@@ -509,6 +510,31 @@ def color_func(field_line):
 
         self._add_mesh_to_dict(block_name="field_lines", mesh=spline)
 
+    def plot_current_sheet(self, pfss_out, **kwargs):
+        """
+        Plot current sheet "(Br=0)".
+
+        Parameters
+        ----------
+        pfss_out : `sunkit_magex.pfss.output.Output`
+            Magnetic field calculated by `sunkit_magex`.
+        **kwargs :
+            Keyword arguments are passed to `pyvista.Plotter.add_mesh`.
+        """
+        sc_vect = pfss_out.grid.sc
+        pc_vect = np.insert(pfss_out.grid.pc, 0, pfss_out.grid.pc[-1])
+        rg_vect = pfss_out.grid.rg
+        coord = SkyCoord((pfss_out.input_map.meta["crval1"] * u.deg).to(u.rad), 0.0 * u.rad, frame=HeliocentricInertial)
+        coord = coord.transform_to(self.coordinate_frame)
+        bc_r = np.insert(pfss_out.bc[0], 0, pfss_out.bc[0][-1], axis=0)
+        [s_arr, p_arr, r_arr] = np.meshgrid(sc_vect, pc_vect, rg_vect)
+        x_arr, y_arr, z_arr = strum2cart(r_arr, s_arr, p_arr + coord.lon.to(u.rad).value)
+        pfss_out_pv = pv.StructuredGrid(x_arr, y_arr, z_arr)
+        pfss_out_pv["Br"] = bc_r.ravel("F")
+        isos_br = pfss_out_pv.contour(isosurfaces=1, rng=[0, 0])
+        self.plotter.add_mesh(isos_br, **kwargs)
+        self._add_mesh_to_dict(block_name="current_sheet", mesh=isos_br)
+
     def save(self, filepath, *, overwrite=False):
         """
         Save all the meshes.

sunkit_pyvista/tests/test_magex.py

@@ -1,3 +1,7 @@
+"""
+This file contains tests for any methods that use sunkit-magex.
+"""
+
 import astropy.constants as const
 import astropy.units as u
 import matplotlib.pyplot as plt
@@ -10,6 +14,8 @@
 
 from sunkit_pyvista import SunpyPlotter
 
+pv.OFF_SCREEN = True
+
 
 def test_field_lines_figure(aia171_test_map, plotter, verify_cache_image):
     pfss = pytest.importorskip("sunkit_magex.pfss")
@@ -26,7 +32,7 @@
     lat, lon = np.meshgrid(lat, lon, indexing="ij")
     lat, lon = lat.ravel() * u.rad, lon.ravel() * u.rad
     radius = 1.2
-    tracer = tracing.PythonTracer()
+    tracer = tracing.FortranTracer()
     input_ = pfss.Input(gong_map, nrho, rss)
     output_ = pfss.pfss(input_)
     seeds = SkyCoord(lon, lat, radius * const.R_sun, frame=gong_map.coordinate_frame)
@@ -42,7 +48,7 @@
     plotter.show(cpos=(0, 1, 0), before_close_callback=verify_cache_image)
 
 
-def test_field_lines_and_color_func(plotter):
+def test_field_lines_and_color_func(plotter, verify_cache_image):
     pfss = pytest.importorskip("sunkit_magex.pfss")
 
     from sunkit_magex.pfss import tracing
@@ -57,7 +63,7 @@
     lat, lon = np.meshgrid(lat, lon, indexing="ij")
     lat, lon = lat.ravel() * u.rad, lon.ravel() * u.rad
     radius = 1.2
-    tracer = tracing.PythonTracer()
+    tracer = tracing.FortranTracer()
     input_ = pfss.Input(gong_map, nrho, rss)
     output_ = pfss.pfss(input_)
     seeds = SkyCoord(lon, lat, radius * const.R_sun, frame=gong_map.coordinate_frame)
@@ -72,3 +78,19 @@
 
     plotter = SunpyPlotter()
     plotter.plot_field_lines(field_lines, color_func=color_func)
+    plotter.show(before_close_callback=verify_cache_image)
+
+
+def test_current_sheet_figure(plotter, verify_cache_image):
+    pfss = pytest.importorskip("sunkit_magex.pfss")
+
+    from sunkit_magex.pfss.sample_data import get_gong_map
+
+    gong_fname = get_gong_map()
+    gong_map = smap.Map(gong_fname)
+    nrho = 35
+    rss = 2.5
+    input_ = pfss.Input(gong_map, nrho, rss)
+    output_ = pfss.pfss(input_)
+    plotter.plot_current_sheet(output_)
+    plotter.show(before_close_callback=verify_cache_image)

sunkit_pyvista/tests/test_plotting.py

@@ -1,19 +1,19 @@
 """
-This file contains figure comparison tests.
+This file contains tests for the main plotting routines.
 """
 
 import astropy.constants as const
 import astropy.units as u
 import pytest
-import pyvista
+import pyvista as pv
 import sunpy.data.test as test
 import sunpy.map as smap
 from astropy.coordinates import SkyCoord
 from sunpy.coordinates import frames
 
 from sunkit_pyvista import SunpyPlotter
 
-pyvista.OFF_SCREEN = True
+pv.OFF_SCREEN = True
 
 
 @pytest.fixture()