Source code for torx.analysis.fast_tracer_m

"""Implementation of a class for fast fieldline tracing."""
import numpy as np

from torx.decorators import autodoc_class

from numba import cfunc
from numbalsoda import lsoda_sig, lsoda, dop853

# Available numbalsoda integrators. Both share the lsoda_sig cfunc signature
# and the same call/return interface, so make_trace_eq() works with either.
# lsoda is the default (adaptive stiff/non-stiff); dop853 is an explicit
# Runge-Kutta 8(5,3) that can be more accurate for non-stiff field lines.
_INTEGRATORS = {"lsoda": lsoda, "dop853": dop853}



@autodoc_class
class FastTracer():
    """Fast tracing class using numbalsoda package."""



    def __init__(self, equi):
        """Initialize fast trace class."""
        self.equi = equi
        self.funcptr = self.equi.make_trace_eq()




    def trace(self, r_start, z_start,
              phi_initial=0.0, phi_to_max=2*np.pi, rtol=1e-4, atol=1e-4,
              npoints=100, method="lsoda"):
        """
        Evaluate fast trace with numbalsoda.

        Parameters
        ----------
        r_start, z_start : float
            Starting normalized R and Z coordinates of the field line.
        phi_initial : float, optional
            Toroidal angle at the start of the trace (default 0.0).
        phi_to_max : float, optional
            Toroidal angle at the end of the trace (default 2*pi).
        rtol, atol : float, optional
            Relative and absolute integration tolerances (default 1e-4).
        npoints : int, optional
            Number of evenly spaced output points (default 100).
        method : {"lsoda", "dop853"}, optional
            numbalsoda integrator to use. "lsoda" (default) is adaptive
            stiff/non-stiff; "dop853" is an explicit Runge-Kutta 8(5,3)
            that can be more accurate for non-stiff field lines.

        Returns
        -------
        numpy.ndarray
            Array of shape (npoints, 3) with the (R, Z, phi) solution.
        """
        try:
            integrator = _INTEGRATORS[method]
        except KeyError:
            raise ValueError(
                f"Unknown method {method!r}; "
                f"choose from {sorted(_INTEGRATORS)}."
            ) from None
        u0 = np.array([r_start, z_start, phi_initial])
        t_eval = np.linspace(phi_initial, phi_to_max, npoints, endpoint=True)

        usol, success = integrator(self.funcptr, u0, t_eval,
                                   rtol=rtol, atol=atol)

        return usol