Allows for a color function to be specified to plot_field_lines()

changelog/70.feature.rst

@@ -0,0 +1 @@
+Allows user to specify color via a color function to :meth:`~sunkit_pyvista.plotter.SunpyPlotter.plot_field_lines`.

examples/field_lines.py

@@ -5,8 +5,10 @@
 
 sunkit-pyvista can be used to plot field lines from `pfsspy`.
 """
+import matplotlib.pyplot as plt
 import numpy as np
 import pfsspy
+from matplotlib import colors
 from pfsspy import tracing
 from pfsspy.sample_data import get_gong_map
 
@@ -55,7 +57,18 @@
                  frame=gong_map.coordinate_frame)
 field_lines = tracer.trace(seeds, output_)
 
-# We plot the field lines
-plotter.plot_field_lines(field_lines)
+
+# We can also specify a color function while plotting the field lines.
+# This function takes a single field line, and returns a color either
+# in the form of a string, (r,g,b) or (r,g,b,a) tuple.
+# In this case we use a Matplotlib norm and colormap to return a tuple of RGBA values.
+def my_fline_color_func(field_line):
+    norm = colors.LogNorm(vmin=1, vmax=1000)
+    cmap = plt.get_cmap('viridis')
+    return cmap(norm(np.abs(field_line.expansion_factor)))
+
+
+# Plotting the field lines
+plotter.plot_field_lines(field_lines, color_func=my_fline_color_func)
 
 plotter.show()

sunkit_pyvista/plotter.py

@@ -373,24 +373,46 @@ def plot_quadrangle(self, bottom_left, top_right=None, width: u.deg = None,
         self.plotter.add_mesh(quad_block, color=color, **kwargs)
         self._add_mesh_to_dict(block_name='quadrangles', mesh=quad_block)
 
-    def plot_field_lines(self, field_lines, **kwargs):
+    def plot_field_lines(self, field_lines, color_func=None, **kwargs):
         """
         Plot magnetic field lines from `pfsspy`.
 
         Parameters
         ----------
         field_lines : `pfsspy.fieldline.FieldLines`
             Field lines to be plotted.
+        color_func : function
+            Function to get the color for each field line.
+            If not given, defaults to showing closed field lines in black,
+            and open field lines in blue (positive polarity) or red (negative polarity).
+            The function must have the signature::
+
+                def color_func(field_line: pfsspy.fieldline.FieldLine) -> color:
+
+             Where ``color`` is any color that `pyvista` recognises
+             (e.g. a RGBA tuple, a RGB tuple, a color string)
         **kwargs :
             Keyword arguments are handed to `pyvista.Plotter.add_mesh`.
         """
+        if not color_func:
+            def color_func(field_line):
+                color = {0: 'black', -1: 'tab:blue', 1: 'tab:red'}.get(field_line.polarity)
+                return colors.to_rgb(color)
+
         field_line_meshes = pv.MultiBlock([])
         for field_line in field_lines:
             grid = self._coords_to_xyz(field_line.coords.ravel())
             field_line_mesh = pv.StructuredGrid(grid[:, 0], grid[:, 1], grid[:, 2])
-            color = {0: 'black', -1: 'tab:blue', 1: 'tab:red'}.get(field_line.polarity)
+            color = color_func(field_line)
+            opacity = 1
+            if isinstance(color, tuple):
+                color = list(color)
+                if len(color) == 4:
+                    opacity = color[3]
+                    color = color[:3]
+
             field_line_mesh.add_field_array([color], 'color')
-            self.plotter.add_mesh(field_line_mesh, color=color, **kwargs)
+            self.plotter.add_mesh(field_line_mesh, color=color, opacity=opacity, **kwargs)
             field_line_meshes.append(field_line_mesh)
 
         self._add_mesh_to_dict(block_name='field_lines', mesh=field_line_meshes)

sunkit_pyvista/tests/test_plotting.py

@@ -1,10 +1,12 @@
 """
 This file contains figure comparison tests.
 """
+import matplotlib.pyplot as plt
 import numpy as np
 import pfsspy
 import pytest
 import pyvista
+from matplotlib import colors
 from pfsspy import tracing
 from pfsspy.sample_data import get_gong_map
 
@@ -78,6 +80,11 @@ def test_field_lines_figure(aia171_test_map, plotter, verify_cache_image):
                      frame=gong_map.coordinate_frame)
     field_lines = tracer.trace(seeds, output_)
 
+    def color_function(field_line):
+        norm = colors.LogNorm(vmin=1, vmax=1000)
+        cmap = plt.get_cmap('magma')
+        return cmap(norm(np.abs(field_line.expansion_factor)))
+
     plotter.plot_map(aia171_test_map)
-    plotter.plot_field_lines(field_lines)
+    plotter.plot_field_lines(field_lines, color_function)
     plotter.show(cpos=(0, 1, 0), before_close_callback=verify_cache_image)

sunkit_pyvista/tests/test_pyvista.py

@@ -1,9 +1,12 @@
+
 import pathlib
 
+import matplotlib.pyplot as plt
 import numpy as np
 import pfsspy
 import pytest
 import pyvista as pv
+from matplotlib import colors
 from pfsspy import tracing
 from pfsspy.sample_data import get_gong_map
 
@@ -131,7 +134,7 @@ def test_multi_block(plotter):
     assert plotter.all_meshes['solar_axis'][0].n_points == 101
 
 
-def test_field_lines(aia171_test_map, plotter):
+def test_field_lines_and_color_func(aia171_test_map, plotter):
     gong_fname = get_gong_map()
     gong_map = smap.Map(gong_fname)
     nrho = 35
@@ -152,6 +155,14 @@ def test_field_lines(aia171_test_map, plotter):
     assert isinstance(plotter.all_meshes['field_lines'][0],
                       pv.core.composite.MultiBlock)
 
+    def color_func(field_line):
+        norm = colors.LogNorm(vmin=1, vmax=1000)
+        cmap = plt.get_cmap('viridis')
+        return cmap(norm(np.abs(field_line.expansion_factor)))
+
+    plotter = SunpyPlotter()
+    plotter.plot_field_lines(field_lines, color_func=color_func)
+
 
 def test_save_and_load(aia171_test_map, plotter, tmp_path):
     plotter.plot_map(aia171_test_map)