Coils and fieldline tracing

In this notebook we define magnetic filament coils in GVEC and perform fieldline tracing using the coil field.

The notebook covers the following topics:

• defining coils in gvec

• visualizing coils

• evaluating coil fields

• fieldline tracing using coil fields

• Poincaré plots

Defining coils in GVEC

In GVEC all coils are discretized via line segments.

To define a coil in GVEC one needs,

• vertices of the line segments in Cartesian coordinates

• current in the coil

(This is similar to the coils format as used, e.g., by MAKEGRID).

Since this might be a common use-case, let us create a CoilSet, from a coils-formatted file as mentioned above. Here, we utilize the NCSX coils openly available from the STELLOPT examples.

In GVEC, several Coils can be combined to a CoilSet, so that the total field from all coils in the CoilSet can be easily evaluated. To define all the relevant coils, we extract the vertices and currents from the coils.NCSX file via pandas:

Extracting coils

# set the number of OpenMP threads before importing gvec
import os

os.environ["OMP_NUM_THREADS"] = "1"

import pandas as pd
import numpy as np
from gvec.coils import Coil, CoilSet

from urllib.request import urlretrieve

# get the coils from the STELLOPT examples
url = "https://princetonuniversity.github.io/STELLOPT/examples/coils.c09r00"
filename = "coils.NCSX"
urlretrieve(url, filename)

# read the coil file
coil_data = pd.read_csv(
    "coils.NCSX",
    sep=r"\s+",
    usecols=[0, 1, 2, 3],
    skiprows=3,
    names=["x", "y", "z", "I"],
)[:-1]

# find the separation of the coils
idx_coil_end = np.where(coil_data["I"] == 0.0)

# extract the currents
coil_currents = [coil_data["I"][i - 1] for i in idx_coil_end[0]]

# extract the vertices
coil_points = np.array(coil_data).astype(float)[:, :-1].T

# gather all coils to combine them into a CoilSet
coils = []
offset = 0
for i, end_idx in enumerate(idx_coil_end[0]):
    # create a single Coil
    coils.append(
        Coil(coil_points[:, offset : end_idx + 1], coil_currents[i])
    )
    offset = end_idx + 1

coil_set = CoilSet(coils)

Saving and loading coils

Once coils have been converted into GVEC objects, one can more easily save and load them:

coil_set.save("NCSX_coils.nc")
del coil_set
coil_set = CoilSet.load("NCSX_coils.nc")

Accessing individual coils in a CoilSet

When creating a CoilSet we can also assign names to the individual coils by passing a list of strings via the coil_names argument, if no names are passed the coils will just be numbered. Individual coils in the CoilSet can be accessed via this name. For example, to access the first coil in our CoilSet we just call:

coil_set["coil_0"]

Coil(points=401, current=6.52e+05 A)

Visualizing coils

Once we have a Coil or CoilSet, they can be easily visualized using plotly:

coil_set.plot();

To visualize one individual coil, just select it from the CoilSet.

coil_set["coil_0"].plot();

Evaluating the coil-field

To evaluate the magnetic field \(\mathbf{B}_{\text{coils}}\) at specific points, one has to pass the Cartesian coordinates of the evaluation points to the eval_B routine of the Coil or CoilSet object. The returned object is an xarray.Dataset, that contains the evaluation points and the corresponding \(\mathbf{B}_{\text{coils}}\) again in Cartesian coordinates, that is

\(\mathbf{B}_{\text{coils}} = (B_x, B_y, B_z)\).

Note that the evaluation points can also be the pos quantity from a GVEC evaluation data-set.

As previously mentioned, coils are discretized via filamentary line segments and therefore the magnetic field can be obtained using a compact expressions for the Biot–Savart fields of a filamentary segment. This expression is implemented using FORTRAN and is parallelized over the coil segments (so do not forget to set OMP_NUM_THREADS)

# define some positions where we want to know B_coils
n_eval_pos = 500
eval_pos = np.zeros((3, n_eval_pos))
eval_pos[0, :] = np.linspace(1.2, 1.62, n_eval_pos)

coil_set.eval_B(eval_pos)

<xarray.Dataset> Size: 24kB
Dimensions:  (xyz: 3, points: 500)
Coordinates:
  * xyz      (xyz) <U1 12B 'x' 'y' 'z'
Dimensions without coordinates: points
Data variables:
    B        (xyz, points) float64 12kB -8.825e-16 1.114e-15 ... 0.2179 0.2171
    pos      (xyz, points) float64 12kB 1.2 1.201 1.202 1.203 ... 0.0 0.0 0.0

xarray.Dataset

• Dimensions:
  • xyz: 3
  • points: 500
• Coordinates: (1)
  • xyz
    (xyz)
    <U1
    'x' 'y' 'z'

    array(['x', 'y', 'z'], dtype='<U1')

• Data variables: (2)
  • B
    (xyz, points)
    float64
    -8.825e-16 1.114e-15 ... 0.2171

    long_name :

    magnetic coil field

    symbol :

    \mathbf{B}_C

    array([[-8.82540568e-16,  1.11412615e-15, -8.87311058e-16, ...,
            -1.70870262e-16, -1.64365049e-16, -2.50667542e-16],
           [-2.44689296e-01, -2.28506459e-01, -2.12098244e-01, ...,
             1.52464239e+00,  1.52340113e+00,  1.52216281e+00],
           [ 7.85722346e-01,  7.91616169e-01,  7.97482698e-01, ...,
             2.18612986e-01,  2.17859606e-01,  2.17108561e-01]],
          shape=(3, 500))

  • pos
    (xyz, points)
    float64
    1.2 1.201 1.202 ... 0.0 0.0 0.0

    long_name :

    position vector

    symbol :

    \mathbf{x}

    array([[1.2       , 1.20084168, 1.20168337, ..., 1.61831663, 1.61915832,
            1.62      ],
           [0.        , 0.        , 0.        , ..., 0.        , 0.        ,
            0.        ],
           [0.        , 0.        , 0.        , ..., 0.        , 0.        ,
            0.        ]], shape=(3, 500))

• Indexes: (1)
  • xyz
    PandasIndex

    PandasIndex(Index(['x', 'y', 'z'], dtype='object', name='xyz'))

• Attributes: (0)

Alternatively one can directly evaluate \(|\mathbf{B}_{\text{coils}}|\), via eval_mod_B

import matplotlib.pyplot as plt

ds_B = coil_set.eval_mod_B(eval_pos)

# plot the coil-field magnitude along a line
fig, ax = plt.subplots()
ax.plot(ds_B.pos.sel(xyz="x"), ds_B.mod_B)
ax.set(xlabel="x", ylabel=r"$B_{\text{coils}}$");

The same syntax holds for a single Coil

coil_set["coil_0"].eval_B(eval_pos)

<xarray.Dataset> Size: 24kB
Dimensions:  (xyz: 3, points: 500)
Coordinates:
  * xyz      (xyz) <U1 12B 'x' 'y' 'z'
Dimensions without coordinates: points
Data variables:
    B        (xyz, points) float64 12kB 1.093 1.097 1.101 ... 0.03239 0.03215
    pos      (xyz, points) float64 12kB 1.2 1.201 1.202 1.203 ... 0.0 0.0 0.0

xarray.Dataset

• Dimensions:
  • xyz: 3
  • points: 500
• Coordinates: (1)
  • xyz
    (xyz)
    <U1
    'x' 'y' 'z'

    array(['x', 'y', 'z'], dtype='<U1')

• Data variables: (2)
  • B
    (xyz, points)
    float64
    1.093 1.097 ... 0.03239 0.03215

    long_name :

    magnetic coil field

    symbol :

    \mathbf{B}_C

    array([[ 1.09341322,  1.09736568,  1.10123127, ..., -0.02585256,
            -0.02619852, -0.02654328],
           [-0.46695895, -0.45959147, -0.45210725, ...,  0.4336591 ,
             0.43334811,  0.43303859],
           [ 0.17105316,  0.174376  ,  0.17768632, ...,  0.03262589,
             0.03238751,  0.0321499 ]], shape=(3, 500))

  • pos
    (xyz, points)
    float64
    1.2 1.201 1.202 ... 0.0 0.0 0.0

    long_name :

    position vector

    symbol :

    \mathbf{x}

    array([[1.2       , 1.20084168, 1.20168337, ..., 1.61831663, 1.61915832,
            1.62      ],
           [0.        , 0.        , 0.        , ..., 0.        , 0.        ,
            0.        ],
           [0.        , 0.        , 0.        , ..., 0.        , 0.        ,
            0.        ]], shape=(3, 500))

• Indexes: (1)
  • xyz
    PandasIndex

    PandasIndex(Index(['x', 'y', 'z'], dtype='object', name='xyz'))

• Attributes: (0)

Fieldline tracing in Cartesian coordinates

• fieldline tracing in \(x,y,z\) avoids issues for geometries beyond cylindrical coordinates (e.g. in G-Frame)

• fieldline tracing core routine trace_fieldlines

• trace_fieldlines wraps scipy.integrate.solve_ivp

• necessary inputs for trace_fieldlines are the starting positions of the fieldlines, the coil-set and the time for how long to trace the fieldlines

• provide surface normals and surface points to specify planes on which intersections with the fieldlines are tracked \(\rightarrow\) Poincaré plots

The fieldline equations that are solved are

\[ \begin{equation*} \frac{\mathbf{B}_{\text{coils}}(x,y,z)}{B_{\text{coils}}} = \dot{\mathbf{x}}. \end{equation*} \]

Where \(\mathbf{x}\) denotes the position vector.

In the example below we utilize the NCSX coilset to trace fieldlines.

Defining starting positions

As an initial setup we define the starting positions for our fieldlines. For this NCSX case cylindrical coordinates are more convenient to use, hence we specify the starts first in \(R,Z,\phi\) and then translate those into x,y,z.

n_fieldlines = 2  # number of initialized fieldlines

# define starting positions for the fieldlines in R/Z/phi
R_start = np.linspace(1.5, 1.56, n_fieldlines)
Z_start = np.zeros_like(R_start)
phis_start = np.zeros_like(R_start)

# translate R/Z/Phi starts into xyz
starts = np.zeros([3, n_fieldlines])
starts[0, :] = R_start * np.cos(phis_start)
starts[1, :] = R_start * np.sin(phis_start)
starts[2, :] = Z_start

Terminating tracing early

Sometimes it might be beneficial to terminate the fieldline tracing early, for example, when the fieldline leaves the relevant region. For such cases one can provide event routines to solve_ivp with the terminal attribute. In our case, we want to terminate the tracing if the fieldline crosses some \(R, Z\) limits. Below we show how such an event can be defined, for details see the solve_ivp documentation.

# limit the region where the fieldline is traced
trace_limits_R = [0.43600, 2.43600]
trace_limits_Z = [-1, 1]


# define a terminating event for the fieldline tracing
def terminal_event(t, xyz):
    R = np.sqrt(xyz[0] ** 2 + xyz[1] ** 2)
    if (trace_limits_R[0] <= R <= trace_limits_R[1]) and (
        trace_limits_Z[0] <= xyz[2] <= trace_limits_Z[1]
    ):
        return -1
    else:
        print(f"Terminate Fieldline: R={R:.2f}, Z={xyz[2]:.2f}")
        return 0


terminal_event.terminal = True

Automatic event creation via planes

To be able to plot Poincaré sections, we have to provide events to solve_ivp that occur when a fieldline crosses the relevant plane. In trace_fieldlines, such events can be automatically generated by providing surface normals (and optionally surface points). Here, we want to visualize \(\phi=const.\) cross-sections, thus we specify the corresponding surface normals.

n_planes = 4  # number of intersection planes for Poincaré plots

# define surface normals for of the intersection planes
phis = np.array([0, np.pi / 6, np.pi / 3, np.pi / 2])
surf_normals = [
    np.array([-np.sin(phi), np.cos(phi), 0]) for phi in phis
]

Start the fieldline tracing

• parallelized over fieldlines \(\rightarrow\) set n_jobs \(>1\)

• finer control via solve_ivp keyword arguments, atol, rtol, events etc.

• returns xarray.Datatree containing the relevant data grouped for each fieldline

from gvec.coils import trace_fieldlines

tree = trace_fieldlines(
    starts=starts,  # initial positions for the fieldlines
    coils=coil_set,  # Object for evaluating the B-field
    time=200,  # time to trace fieldlines
    surf_normals=surf_normals,  # intersection planes
    # surf_points = None # note that without surf_points the planes are expected to go through the origin
    # solve_ivp specific keywords
    events=[terminal_event],
    atol=1e-5,  # tune these tol. for more accurate results
    rtol=1e-4,
)

INFO    Tracing fieldlines using n_jobs=1
  0%|          | 0/2 [00:00<?, ?it/s]
 50%|█████     | 1/2 [00:18<00:18, 18.66s/it]
 50%|█████     | 1/2 [00:29<00:29, 30.00s/it]

---------------------------------------------------------------------------
KeyboardInterrupt                         Traceback (most recent call last)
Cell In[12], line 3
      1 from gvec.coils import trace_fieldlines
----> 3 tree = trace_fieldlines(
      4     starts=starts,  # initial positions for the fieldlines
      5     coils=coil_set,  # Object for evaluating the B-field
      6     time=200,  # time to trace fieldlines
      7     surf_normals=surf_normals,  # intersection planes
      8     # surf_points = None # note that without surf_points the planes are expected to go through the origin
      9     # solve_ivp specific keywords
     10     events=[terminal_event],
     11     atol=1e-5,  # tune these tol. for more accurate results
     12     rtol=1e-4,
     13 )

File ~/checkouts/readthedocs.org/user_builds/gvec/envs/latest/lib/python3.10/site-packages/gvec/coils.py:821, in trace_fieldlines(starts, coils, time, surf_normals, surf_points, n_jobs, return_solves, verbosity, **kwargs)
    819 else:  # serial over fieldlines but still OMP parallelization for magnetic field evaluation
    820     for i in tqdm(range(starts.shape[1])):
--> 821         solve = solve_ivp(
    822             fun=_push_cart,
    823             y0=starts[:, i],
    824             t_span=[0, time],
    825             **kwargs,
    826         )
    827         solves.append(solve)
    828 if return_solves:

File ~/checkouts/readthedocs.org/user_builds/gvec/envs/latest/lib/python3.10/site-packages/scipy/integrate/_ivp/ivp.py:655, in solve_ivp(fun, t_span, y0, method, t_eval, dense_output, events, vectorized, args, **options)
    653 status = None
    654 while status is None:
--> 655     message = solver.step()
    657     if solver.status == 'finished':
    658         status = 0

File ~/checkouts/readthedocs.org/user_builds/gvec/envs/latest/lib/python3.10/site-packages/scipy/integrate/_ivp/base.py:197, in OdeSolver.step(self)
    195 else:
    196     t = self.t
--> 197     success, message = self._step_impl()
    199     if not success:
    200         self.status = 'failed'

File ~/checkouts/readthedocs.org/user_builds/gvec/envs/latest/lib/python3.10/site-packages/scipy/integrate/_ivp/rk.py:144, in RungeKutta._step_impl(self)
    141 h = t_new - t
    142 h_abs = np.abs(h)
--> 144 y_new, f_new = rk_step(self.fun, t, y, self.f, h, self.A,
    145                        self.B, self.C, self.K)
    146 scale = atol + np.maximum(np.abs(y), np.abs(y_new)) * rtol
    147 error_norm = self._estimate_error_norm(self.K, h, scale)

File ~/checkouts/readthedocs.org/user_builds/gvec/envs/latest/lib/python3.10/site-packages/scipy/integrate/_ivp/rk.py:64, in rk_step(fun, t, y, f, h, A, B, C, K)
     62 for s, (a, c) in enumerate(zip(A[1:], C[1:]), start=1):
     63     dy = np.dot(K[:s].T, a[:s]) * h
---> 64     K[s] = fun(t + c * h, y + dy)
     66 y_new = y + h * np.dot(K[:-1].T, B)
     67 f_new = fun(t + h, y_new)

File ~/checkouts/readthedocs.org/user_builds/gvec/envs/latest/lib/python3.10/site-packages/scipy/integrate/_ivp/base.py:154, in OdeSolver.__init__.<locals>.fun(t, y)
    152 def fun(t, y):
    153     self.nfev += 1
--> 154     return self.fun_single(t, y)

File ~/checkouts/readthedocs.org/user_builds/gvec/envs/latest/lib/python3.10/site-packages/scipy/integrate/_ivp/base.py:23, in check_arguments.<locals>.fun_wrapped(t, y)
     22 def fun_wrapped(t, y):
---> 23     return np.asarray(fun(t, y), dtype=dtype)

File ~/checkouts/readthedocs.org/user_builds/gvec/envs/latest/lib/python3.10/site-packages/gvec/coils.py:789, in trace_fieldlines.<locals>._push_cart(t, R)
    788 def _push_cart(t, R):
--> 789     B = coils._eval_B_direct(R)
    790     B = B / np.sqrt(np.sum(B**2, axis=0))
    791     return B[:, 0]

File ~/checkouts/readthedocs.org/user_builds/gvec/envs/latest/lib/python3.10/site-packages/gvec/coils.py:565, in CoilSet._eval_B_direct(self, pos)
    559 for coil_name in self.coils:
    560     coil = self.coils[coil_name]
    561     _BS.biotsavart(
    562         n_positions=n_positions,
    563         xyz=np.asfortranarray(pos),
    564         n_points=coil.n_points,
--> 565         coil_points=np.asfortranarray(coil.coil_points),
    566         prefactor=coil.prefactor,
    567         b=B_aux,
    568     )
    569     B += B_aux
    570     B_aux *= 0.0

File ~/checkouts/readthedocs.org/user_builds/gvec/envs/latest/lib/python3.10/site-packages/xarray/core/common.py:303, in AttrAccessMixin.__getattr__(self, name)
    299 def __getattr__(self, name: str) -> Any:
    300     if name not in {"__dict__", "__setstate__"}:
    301         # this avoids an infinite loop when pickle looks for the
    302         # __setstate__ attribute before the xarray object is initialized
--> 303         for source in self._attr_sources:
    304             with suppress(KeyError):
    305                 return source[name]

File ~/checkouts/readthedocs.org/user_builds/gvec/envs/latest/lib/python3.10/site-packages/xarray/core/dataarray.py:913, in DataArray._attr_sources(self)
    910 @property
    911 def _attr_sources(self) -> Iterable[Mapping[Hashable, Any]]:
    912     """Places to look-up items for attribute-style access"""
--> 913     yield from self._item_sources
    914     yield self.attrs

File ~/checkouts/readthedocs.org/user_builds/gvec/envs/latest/lib/python3.10/site-packages/xarray/core/dataarray.py:922, in DataArray._item_sources(self)
    919 yield FilteredMapping(keys=self._coords, mapping=self.coords)
    921 # virtual coordinates
--> 922 yield FilteredMapping(keys=self.dims, mapping=self.coords)

File ~/checkouts/readthedocs.org/user_builds/gvec/envs/latest/lib/python3.10/site-packages/xarray/core/dataarray.py:843, in DataArray.dims(self)
    838     """Convert this variable to a pandas.Index. Only possible for 1D
    839     arrays.
    840     """
    841     return self.variable.to_index()
--> 843 @property
    844 def dims(self) -> tuple[Hashable, ...]:
    845     """Tuple of dimension names associated with this array.
    846 
    847     Note that the type of this property is inconsistent with
   (...)
    854     Dataset.dims
    855     """
    856     return self.variable.dims

KeyboardInterrupt: 

Poincaré plots

To visualize the Poincaré cross-sections, we have to extract the intersection positions of each fieldline with the specified planes. Intersection positions can be extracted for each fieldline via the event entry in the data set, numbered according to the order of the specified surface normals. Below we showcase a possible way to visualize the Poincaré sections using the returned data-tree.

fig, axs = plt.subplots(2, 2, figsize=(8, 8), tight_layout=True)
for fieldline in tree:
    # select the relevant fieldline
    ds = tree[fieldline]
    try:  # continue if the fieldline never intersected with any plane
        for i, ax in enumerate(axs.flatten()):
            # select intersection positions for the i-th plane
            events = ds[
                f"event_{i}"
            ]  # corresponds to surface_normal[i]

            # xyz --> RZphi
            R = np.sqrt(
                events.sel(xyz="x") ** 2 + events.sel(xyz="y") ** 2
            )
            ax.scatter(
                R[
                    ::2
                ],  # select every other point to just get one Poincaré section
                events.sel(xyz="z")[::2],
                marker=".",
                s=2,
            )
    except KeyError:
        continue

# add some plot labels and titles
for ax, phi in zip(axs.flatten(), phis):
    ax.set(
        xlabel="R / m",
        ylabel="Z / m",
        title=f"$\\phi$={phi / np.pi:.2f} $\\pi$",
    )
    ax.axis("equal")

Visualizing the fieldlines in 3D

With the data in the returned data-tree, we can also do some 3D visualizations to inspect the fieldlines and Poincaré sections with respect to the coils. Below is one example of how this could be achieved, again using the plotly package.

from plotly import graph_objects as go

# plot the coils
ax = coil_set.plot(
    show=False,
    line={"color": "black"},
    showlegend=False,
)
ax["layout"]["scene"]["aspectmode"] = "data"

for fieldline in tree:
    # plot the fieldlines
    ds = tree[fieldline]
    ax.add_trace(
        go.Scatter3d(
            x=ds.pos.sel(xyz="x"),
            y=ds.pos.sel(xyz="y"),
            z=ds.pos.sel(xyz="z"),
            mode="lines",
            opacity=0.2,
            line={"color": "green"},
            showlegend=False,
        )
    )
    # starting positions
    ax.add_trace(
        go.Scatter3d(
            x=starts[0, :],
            y=starts[1, :],
            z=starts[2, :],
            mode="markers",
            marker=dict(size=3, color="red"),
            showlegend=False,
        )
    )
    # plot the Poincaré sections in 3D
    try:
        for i in range(len(phis)):
            ax.add_trace(
                go.Scatter3d(
                    x=np.array(ds[f"event_{i}"].sel(xyz="x")),
                    y=np.array(ds[f"event_{i}"].sel(xyz="y")),
                    z=np.array(ds[f"event_{i}"].sel(xyz="z")),
                    mode="markers",
                    marker=dict(size=2, color="blue"),
                    showlegend=False,
                )
            )
    except KeyError:
        continue
ax.show();