From 5b81d02558ca75b491aacb675b3131650d6e0f46 Mon Sep 17 00:00:00 2001 From: David Stansby Date: Sun, 23 Oct 2022 16:41:28 +0100 Subject: [PATCH 1/7] Enable pluto field line tracing --- examples/tracing/tracing_pluto.py | 48 +++++++++++++++++++++++++++++++ psipy/model/base.py | 34 ++++++++++++++++++++++ psipy/model/mas.py | 18 ------------ psipy/model/pluto.py | 31 ++++++++++++++++++++ psipy/visualization/pyvista.py | 4 +-- 5 files changed, 115 insertions(+), 20 deletions(-) create mode 100644 examples/tracing/tracing_pluto.py diff --git a/examples/tracing/tracing_pluto.py b/examples/tracing/tracing_pluto.py new file mode 100644 index 0000000..e240601 --- /dev/null +++ b/examples/tracing/tracing_pluto.py @@ -0,0 +1,48 @@ +""" +Field lines with pyvista +======================== +Visualising traced field lines with pyvista. + +pyvista is a Python package that provides a wrapper to the popular VTK library +for 3D visualisation. Unlike Matplotlib this is "true" 3D rendering, and is +much more performant in comparison. +""" +############################################################################### +# First, load the required modules. +import astropy.units as u +import numpy as np + +from psipy.data import sample_data +from psipy.model import PLUTOOutput +from psipy.tracing import FortranTracer +from psipy.visualization.pyvista import MASPlotter + +############################################################################### +# Load a set of PLUTO output files. +pluto_path = sample_data.pluto_sample_data() +model = PLUTOOutput(pluto_path) + +############################################################################### +# To trace field lines, start by creating a tracer. Then we create a set of +# seed points at which the field lines are drawn from. +tracer = FortranTracer() + +nseeds = 20 +r = np.ones(nseeds**2) * 0.5 * u.au +lat = np.linspace(-45, 45, nseeds**2, endpoint=False) * u.deg +lon = np.random.rand(nseeds**2) * 360 * u.deg + +# Do the tracing! +flines = tracer.trace(model, r=r, lat=lat, lon=lon) +print(flines) + +############################################################################### +# Visualise the results + +plotter = MASPlotter(model) +for fline in flines: + plotter.add_fline(fline) + +# Add a sphere at the inner boundary of the model +plotter.add_sphere(np.min(model["Bx1"].r_coords)) +plotter.show() diff --git a/psipy/model/base.py b/psipy/model/base.py index f2d3ede..23f8b04 100644 --- a/psipy/model/base.py +++ b/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): diff --git a/psipy/model/mas.py b/psipy/model/mas.py index 7786319..8c336bc 100644 --- a/psipy/model/mas.py +++ b/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") diff --git a/psipy/model/pluto.py b/psipy/model/pluto.py index af1dc4a..d3e307a 100644 --- a/psipy/model/pluto.py +++ b/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,32 @@ 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 + + print(self["Bx1"]) + print(self["Bx2"]) + + 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]) + + 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"]], + ) diff --git a/psipy/visualization/pyvista.py b/psipy/visualization/pyvista.py index d8fc0f4..f4f7376 100644 --- a/psipy/visualization/pyvista.py +++ b/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 From a717fd389a64535410b3d725d3920467f4d94a1a Mon Sep 17 00:00:00 2001 From: David Stansby Date: Thu, 27 Oct 2022 14:30:17 +0100 Subject: [PATCH 2/7] Apply tracing tests to all models --- psipy/conftest.py | 2 +- psipy/tracing/tests/test_tracing.py | 24 +++++++++++++----------- 2 files changed, 14 insertions(+), 12 deletions(-) diff --git a/psipy/conftest.py b/psipy/conftest.py index 328c3e0..abb636c 100644 --- a/psipy/conftest.py +++ b/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]) diff --git a/psipy/tracing/tests/test_tracing.py b/psipy/tracing/tests/test_tracing.py index 58eeda8..ea8117b 100644 --- a/psipy/tracing/tests/test_tracing.py +++ b/psipy/tracing/tests/test_tracing.py @@ -2,13 +2,13 @@ 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,7 +17,7 @@ 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() @@ -25,24 +25,26 @@ def cell_corner_b(self): 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 == (139, 3) tracer = FortranTracer(step_size=0.5) - flines = tracer.trace(mas_model, lon=lon, lat=lat, r=r) + flines = tracer.trace(model, lon=lon, lat=lat, r=r) # Check that changing step size has an effect assert flines[0].xyz.shape == (278, 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 From 3424d10d023adbf4f52673d4afb687642961b92c Mon Sep 17 00:00:00 2001 From: David Stansby Date: Thu, 27 Oct 2022 14:48:48 +0100 Subject: [PATCH 3/7] Add check on tracing phi coords --- psipy/tracing/tracing.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/psipy/tracing/tracing.py b/psipy/tracing/tracing.py index ae5aa72..cf61d8b 100644 --- a/psipy/tracing/tracing.py +++ b/psipy/tracing/tracing.py @@ -121,7 +121,6 @@ def _trace_from_grid(self, grid, seeds: np.ndarray, runit: u.Unit) -> FieldLines # Normalize step size to radial cell size rcoords = grid.zcoords - step_size = self.step_size * np.min(np.diff(rcoords)) - self.tracer = StreamTracer(max_steps, step_size) + self.tracer = StreamTracer(max_steps, self.step_size) self.tracer.trace(seeds, grid) return FieldLines(self.tracer.xs, runit) From 287ff10d2fd03a9d70cb7e4f59691ec1162027c5 Mon Sep 17 00:00:00 2001 From: David Stansby Date: Thu, 27 Oct 2022 14:49:22 +0100 Subject: [PATCH 4/7] Fix PLUTO phi coords --- psipy/model/pluto.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/psipy/model/pluto.py b/psipy/model/pluto.py index d3e307a..38d048f 100644 --- a/psipy/model/pluto.py +++ b/psipy/model/pluto.py @@ -47,9 +47,6 @@ def cell_corner_b(self, t_idx: int = None) -> xr.DataArray: t_coords = self["Bx1"].theta_coords p_coords = self["Bx1"].phi_coords - print(self["Bx1"]) - print(self["Bx2"]) - 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) @@ -58,7 +55,7 @@ def cell_corner_b(self, t_idx: int = None) -> xr.DataArray: 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]) + new_pcoords = np.append(p_coords, p_coords[0:1] + 2 * np.pi) return xr.DataArray( np.stack([bp, bt, br], axis=-1), From b5150444120ebfafd83bd6c3aecad9ba4635d2bb Mon Sep 17 00:00:00 2001 From: David Stansby Date: Thu, 27 Oct 2022 15:07:45 +0100 Subject: [PATCH 5/7] Working PLUTO tracing! --- psipy/tracing/tests/test_tracing.py | 10 ++++++---- psipy/tracing/tracing.py | 3 ++- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/psipy/tracing/tests/test_tracing.py b/psipy/tracing/tests/test_tracing.py index ea8117b..9f8bd50 100644 --- a/psipy/tracing/tests/test_tracing.py +++ b/psipy/tracing/tests/test_tracing.py @@ -1,7 +1,6 @@ import astropy.units as u import numpy as np -from psipy.data import sample_data from psipy.model import MASOutput, PLUTOOutput from psipy.tracing import FieldLines, FortranTracer @@ -35,12 +34,15 @@ def cell_corner_b(self): assert flines[0].xyz.shape == (139, 3) elif isinstance(model, PLUTOOutput): assert len(bs.coords["r"]) == 141 - assert flines[0].xyz.shape == (139, 3) + assert flines[0].xyz.shape == (140, 3) tracer = FortranTracer(step_size=0.5) flines = tracer.trace(model, lon=lon, lat=lat, r=r) - # Check that changing step size has an effect - assert flines[0].xyz.shape == (278, 3) + + 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(model, tmpdir): diff --git a/psipy/tracing/tracing.py b/psipy/tracing/tracing.py index cf61d8b..ae5aa72 100644 --- a/psipy/tracing/tracing.py +++ b/psipy/tracing/tracing.py @@ -121,6 +121,7 @@ def _trace_from_grid(self, grid, seeds: np.ndarray, runit: u.Unit) -> FieldLines # Normalize step size to radial cell size rcoords = grid.zcoords - self.tracer = StreamTracer(max_steps, self.step_size) + step_size = self.step_size * np.min(np.diff(rcoords)) + self.tracer = StreamTracer(max_steps, step_size) self.tracer.trace(seeds, grid) return FieldLines(self.tracer.xs, runit) From ed5412c77ae870263c59aab764c8f2a99f99db14 Mon Sep 17 00:00:00 2001 From: David Stansby Date: Thu, 27 Oct 2022 15:29:02 +0100 Subject: [PATCH 6/7] Remove pluto example --- examples/tracing/tracing_pluto.py | 48 ------------------------------- 1 file changed, 48 deletions(-) delete mode 100644 examples/tracing/tracing_pluto.py diff --git a/examples/tracing/tracing_pluto.py b/examples/tracing/tracing_pluto.py deleted file mode 100644 index e240601..0000000 --- a/examples/tracing/tracing_pluto.py +++ /dev/null @@ -1,48 +0,0 @@ -""" -Field lines with pyvista -======================== -Visualising traced field lines with pyvista. - -pyvista is a Python package that provides a wrapper to the popular VTK library -for 3D visualisation. Unlike Matplotlib this is "true" 3D rendering, and is -much more performant in comparison. -""" -############################################################################### -# First, load the required modules. -import astropy.units as u -import numpy as np - -from psipy.data import sample_data -from psipy.model import PLUTOOutput -from psipy.tracing import FortranTracer -from psipy.visualization.pyvista import MASPlotter - -############################################################################### -# Load a set of PLUTO output files. -pluto_path = sample_data.pluto_sample_data() -model = PLUTOOutput(pluto_path) - -############################################################################### -# To trace field lines, start by creating a tracer. Then we create a set of -# seed points at which the field lines are drawn from. -tracer = FortranTracer() - -nseeds = 20 -r = np.ones(nseeds**2) * 0.5 * u.au -lat = np.linspace(-45, 45, nseeds**2, endpoint=False) * u.deg -lon = np.random.rand(nseeds**2) * 360 * u.deg - -# Do the tracing! -flines = tracer.trace(model, r=r, lat=lat, lon=lon) -print(flines) - -############################################################################### -# Visualise the results - -plotter = MASPlotter(model) -for fline in flines: - plotter.add_fline(fline) - -# Add a sphere at the inner boundary of the model -plotter.add_sphere(np.min(model["Bx1"].r_coords)) -plotter.show() From 3289943ed84cd912df4d8a8c46a3168655f9e5b9 Mon Sep 17 00:00:00 2001 From: David Stansby Date: Thu, 27 Oct 2022 15:29:51 +0100 Subject: [PATCH 7/7] Add changelog entry --- CHANGELOG.rst | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 4fef05a..14e17e6 100644 --- a/CHANGELOG.rst +++ b/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.