Merge pull request #84 from dstansby/plutotracing

Enable pluto field line tracing

CHANGELOG.rst

@@ -4,6 +4,11 @@ Changelog
 Version 0.4.0
 -------------
 
+New features
+~~~~~~~~~~~~
+- It is now possible trace magnetic field lines through PLUTO data.
+  The code for doing this is exactly the same as for MAS data.
+
 Bug fixes
 ~~~~~~~~~
 - Removed a couple of un-needed dependencies from the package.

psipy/conftest.py

@@ -61,4 +61,4 @@ def pluto_model():
     return pluto.PLUTOOutput(get_pluto_directory())
 
 
-fixture_union("all_models", [mas_model, pluto_model])
+fixture_union("model", [mas_model, pluto_model])

psipy/model/base.py

@@ -1,5 +1,8 @@
 import abc
 from pathlib import Path
+from typing import Optional
+
+import xarray as xr
 
 from .variable import Variable
 
@@ -105,6 +108,37 @@ def get_unit(self, var):
         factor : float
         """
 
+    @abc.abstractmethod
+    def get_runit(self):
+        """
+        Return the units for the radial coordiate.
+
+        Returns
+        -------
+        unit : `astropy.units.Unit`
+        """
+
+    @abc.abstractmethod
+    def cell_corner_b(self, t_idx: Optional[int] = None) -> xr.DataArray:
+        """
+        Get the magnetic field vector at the cell corners.
+
+        Parameters
+        ----------
+        t_idx : int, optional
+            If more than one timestep is present in the loaded model, a
+            timestep index at which to get the vectors must be provided.
+
+        Returns
+        -------
+        xarray.DataArray
+
+        Notes
+        -----
+        The phi limits go from 0 to 2pi inclusive, with the vectors at phi=0
+        equal to the vectors at phi=2pi.
+        """
+
     # Properties start here
     @property
     def loaded_variables(self):

psipy/model/mas.py

@@ -71,24 +71,6 @@ def __str__(self):
         return f"MAS output in directory {self.path}\n" + super().__str__()
 
     def cell_corner_b(self, t_idx: Optional[int] = None) -> xr.DataArray:
-        """
-        Get the magnetic field vector at the cell corners.
-
-        Parameters
-        ----------
-        t_idx : int, optional
-            If more than one timestep is present in the loaded model, a
-            timestep index at which to get the vectors must be provided.
-
-        Returns
-        -------
-        xarray.DataArray
-
-        Notes
-        -----
-        The phi limits go from 0 to 2pi inclusive, with the vectors at phi=0
-        equal to the vectors at phi=2pi.
-        """
         if not set(["br", "bt", "bp"]) <= set(self.variables):
             raise RuntimeError("MAS output must have the br, bt, bp variables loaded")
 

psipy/model/pluto.py

@@ -1,4 +1,6 @@
 import astropy.units as u
+import numpy as np
+import xarray as xr
 
 from psipy.io import get_pluto_variables, read_pluto_files
 from .base import ModelOutput
@@ -34,3 +36,29 @@ def get_variables(self):
 
     def load_file(self, var):
         return read_pluto_files(self.path, var)
+
+    def cell_corner_b(self, t_idx: int = None) -> xr.DataArray:
+        if not set(["Bx1", "Bx2", "Bx3"]) <= set(self.variables):
+            raise RuntimeError(
+                "PLUTO output must have the BX1, Bx2, Bx3 variables loaded"
+            )
+
+        r_coords = self["Bx1"].r_coords
+        t_coords = self["Bx1"].theta_coords
+        p_coords = self["Bx1"].phi_coords
+
+        br = self["Bx1"].data.isel(time=t_idx or 0)
+        bt = self["Bx2"].data.isel(time=t_idx or 0)
+        bp = self["Bx3"].data.isel(time=t_idx or 0)
+
+        # Add an extra layer of cells around phi=2pi for the tracer
+        br = np.concatenate((br, br[0:1, :, :]), axis=0)
+        bt = np.concatenate((bt, bt[0:1, :, :]), axis=0)
+        bp = np.concatenate((bp, bp[0:1, :, :]), axis=0)
+        new_pcoords = np.append(p_coords, p_coords[0:1] + 2 * np.pi)
+
+        return xr.DataArray(
+            np.stack([bp, bt, br], axis=-1),
+            dims=["phi", "theta", "r", "component"],
+            coords=[new_pcoords, t_coords, r_coords, ["bp", "bt", "br"]],
+        )

psipy/tracing/tests/test_tracing.py

@@ -1,14 +1,13 @@
 import astropy.units as u
 import numpy as np
 
-from psipy.data import sample_data
-from psipy.model import MASOutput
+from psipy.model import MASOutput, PLUTOOutput
 from psipy.tracing import FieldLines, FortranTracer
 
 
-def test_tracer(mas_model):
+def test_tracer(model):
     # Simple smoke test of field line tracing
-    bs = mas_model.cell_corner_b()
+    bs = model.cell_corner_b()
     # Fake data to be unit vectors pointing in radial direction
     bs.loc[..., "bp"] = 0
     bs.loc[..., "bt"] = 0
@@ -17,32 +16,37 @@ def test_tracer(mas_model):
     def cell_corner_b(self):
         return bs
 
-    mas_model.cell_corner_b = cell_corner_b
+    model.cell_corner_b = cell_corner_b
 
     tracer = FortranTracer()
 
     r = 40 * u.R_sun
     lat = 0 * u.deg
     lon = 0 * u.deg
 
-    flines = tracer.trace(mas_model, lon=lon, lat=lat, r=r)
+    flines = tracer.trace(model, lon=lon, lat=lat, r=r)
     assert len(flines) == 1
 
     # Check that with auto step size, number of steps is close to number of
     # radial coordinates
-    assert len(bs.coords["r"]) == 140
-    assert flines[0].xyz.shape == (139, 3)
+    if isinstance(model, MASOutput):
+        assert len(bs.coords["r"]) == 140
+        assert flines[0].xyz.shape == (139, 3)
+    elif isinstance(model, PLUTOOutput):
+        assert len(bs.coords["r"]) == 141
+        assert flines[0].xyz.shape == (140, 3)
 
     tracer = FortranTracer(step_size=0.5)
-    flines = tracer.trace(mas_model, lon=lon, lat=lat, r=r)
-    # Check that changing step size has an effect
-    assert flines[0].xyz.shape == (278, 3)
+    flines = tracer.trace(model, lon=lon, lat=lat, r=r)
+
+    if isinstance(model, MASOutput):
+        assert flines[0].xyz.shape == (278, 3)
+    elif isinstance(model, PLUTOOutput):
+        assert flines[0].xyz.shape == (280, 3)
 
 
-def test_fline_io(mas_model, tmpdir):
+def test_fline_io(model, tmpdir):
     # Test saving and loading field lines
-    mas_path = sample_data.mas_sample_data()
-    model = MASOutput(mas_path)
     tracer = FortranTracer()
 
     r = 40 * u.R_sun

psipy/visualization/pyvista.py

@@ -7,7 +7,7 @@
 from vtkmodules.vtkCommonCore import vtkCommand
 from vtkmodules.vtkRenderingCore import vtkCellPicker
 
-from psipy.model import MASOutput
+from psipy.model import ModelOutput
 
 if TYPE_CHECKING:
     from psipy.tracing import FortranTracer
@@ -27,7 +27,7 @@ class MASPlotter:
     plotter : pyvista.Plotter
     """
 
-    def __init__(self, mas_output: MASOutput):
+    def __init__(self, mas_output: ModelOutput):
         self.pvplotter = pv.Plotter()
         self.mas_output = mas_output
         self.tracer: Optional[FortranTracer] = None