Source code for torx.analysis.fast_tracer_m

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

from torx.decorators import AutoDocClass

# 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", "dop853")



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



    def __init__(self, equi):
        """Initialize fast trace class."""
        self.equi = equi
        # Build the trace function now, so that errors show here
        self.equi.make_trace_eq()


    @property
    def funcptr(self) -> int:
        """Return the address of the trace function of the equilibrium."""
        # Asked on each use, so changes of the equilibrium are traced
        return 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. If
            the integration fails, for example because the field line
            leaves the domain, the array is returned filled with NaN
            rather than the solver's partially written buffer.
        """
        if method not in _INTEGRATORS:
            raise ValueError(
                f"Unknown method {method!r}; "
                f"choose from {sorted(_INTEGRATORS)}."
            )
        import numbalsoda
        integrator = getattr(numbalsoda, method)
        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)
        if not success:
            usol[:] = np.nan

        return usol