Code Design for Fast Particle Tracing

[f0uriest]

This is a meta issue for keeping track of things related to fast particle tracing.

There are currently 2 PRs (#1410, #734) that implement parts of this, but we should put some thought into a unifying design.

As I see it there are a few major parts:

Domain

• In vacuum fields provided by external coils/magnets etc, particles traced in $(R, \phi, Z)$ space
• In an equilibrium, possibly at finite beta, particles traced in $(\rho, \theta, \zeta)$ space

Model

• Full orbit
• Guiding center
• Guiding center + slowing down
• Guiding center + collisions (this turns the ODE into an SDE so probably not worth it?)
• Full orbit + slowing down (need to figure out equations for this?)

Initialization

• Manually specifying starting positions in configuration space
• Randomly sampled on curves/surfaces/volumes
• "Realistic" birth of alpha particles based on local fusion reaction rate

Gradients

• In general particle orbits can be chaotic, meaning AD may blow up
• Could simply clip gradients that get too large and rely on ensemble of non-chaotic trajectories for optimization
• Each particle is low dimensional (4-5) so could maybe also try least squares shadowing or something similar
• Use SPSA and SGD

Sketch of an API:

class AbstractTrajectoryModel(ABC):
    
    @abstractproperty
    def coordinates(self):
       """Phase space coordinates used by model, eg rho,theta,zeta,vpar,v..."""

    @abstractmethod
    def compute(self, x:jax.Array, field_or_eq) -> jax.Array:
        """Compute RHS of ODE with phase space position x."""


class AbstractParticleInitializer(ABC):
    """ABC for initial distribution of particles for tracing."""

    def __init__(self, m=4, q=2):
        self.m = m * proton_mass
        self.q = q * elementary_charge

    @abstractmethod
    def init_particles(self, N:int, model:AbstractTrajectoryModel) -> jax.Array:
        """Initialize a distribution of N particles"""
        pass




def trace_particles_vacuum(field, 
                           model:AbstractTrajectoryModel, 
                           initializer:AbstractParticleInitializer,
                           ...):

def trace_particles_equilibrium(eq, 
                                model:AbstractTrajectoryModel, 
                                initializer:AbstractParticleInitializer,
                                ...):

[joaopedrobiu6]

Hi! Nice idea, and I hope to be able to help you with what I can.

Some questions: Do you think that some of the code already done is reusable? How should this class be implemented? As an objective?

Finally, from my experience with tracers, to deal with numerical problems arising from small-scale dynamics, implementing a Boris pusher could be useful instead of only a regular solver for the Lorentz force equations!

[f0uriest]

Yeah I think most of the code we have will be re-used just moved around a bit. I started on branch rc/particles, see eg https://github.com/PlasmaControl/DESC/blob/rc/particles/desc/tracing.py

The idea is mainly to move some of the stuff thats in compute/objectives into a standalone module so that you can easily trace particles outside of optimization as well (eg for analysis, plotting etc), then the objective can just call those modules/functions.

[dpanici]

@f0uriest what is in rc/particles? what would the order of operations be for getting this done, would it be rc/particles, then retrofitting #734 and #1410 to use the API defined in rc/particles?

[f0uriest]

rc/particles I think includes everything from #734 and #1410 but with a unified API. rc/particles takes care of most of the things in this issue apart from the gradients, but it's all pretty rough and unfinished at the moment

[YigitElma]

I feel like instead of creating custom term classes for diffrax easier and less dependency future proof version would be to make the trace_particles() handle all the logic with an argument like mode. As I see in the rc/particles branch, there are only 2 ODEs, vacuum guiding center and guiding center with slowdown. There are slight differences for Equilibrium and MagneticField classes as well as the frame used. However, both frames can be computed from the same equation and taking dot products, so I think that is easily handleable.

I suggest just keeping particle initialization classes (since there can be more of them), have a single trace_particles() class like field_line_integrate() and pass in initialized particles (if we want to have this in optimization).

[YigitElma]

I will create a PR with my changes to rc/particles branch to have a code to talk on.