desc.particles.trace_particles

desc.particles.trace_particles(field, initializer, model, ts, params=None, rtol=1e-05, atol=1e-05, max_steps=None, min_step_size=1e-08, bounds=None, solver=Tsit5(scan_kind='bounded'), adjoint=RecursiveCheckpointAdjoint(), chunk_size=None, options=None, throw=True, return_aux=False)Source

Trace charged particles in an equilibrium or external magnetic field.

For jit friendly version of this function or to have more control over the integration, see desc.particles._trace_particles.

Parameters:

  • field (MagneticField or Equilibrium) – Source of magnetic field to integrate

  • initializer (AbstractParticleInitializer) – Particle initializer

  • model (AbstractTrajectoryModel) – Trajectory model to integrate with.

  • ts (array-like) – Strictly increasing array of time values where output will be recorded.

  • params (dict, optional) – Parameters of the field object, needed for automatic differentiation. Defaults to field.params_dict.

  • rtol (float, optional) – Relative and absolute tolerances for PID stepsize controller.

  • atol (float, optional) – Relative and absolute tolerances for PID stepsize controller.

  • max_steps (int) – Maximum number of steps for whole integration. This will be passed to the diffrax.diffeqsolve function. Defaults to (ts[-1] - ts[0]) / min_step_size

  • min_step_size (float) – minimum step size (in t) that the integration can take. Defaults to 1e-8

  • bounds (array of shape(3, 2), optional) – Bounds for particle trajectory during tracing. Trajectories that leave this region will be stopped, and points outside the region will be shown with NaNs. When tracing an Equilibrium, the default bounds are set to [[0, 1], [-inf, inf], [-inf, inf]] for rho, theta, zeta coordinates (LCFS). When tracing a MagneticField, the default bounds are set to [[0, inf], [-inf, inf], [-inf, inf]] for R, phi, Z coordinates (no bounds).

  • solver (diffrax.AbstractSolver, optional) – diffrax Solver object to use in integration. Defaults to Tsit5(scan_kind=’bounded’), a RK45 explicit solver.

  • adjoint (diffrax.AbstractAdjoint, optional) – How to take derivatives of the trajectories. RecursiveCheckpointAdjoint supports reverse mode AD and tends to be the most efficient. For forward mode AD use diffrax.ForwardMode().

  • chunk_size (int, optional) – Chunk size for integration over particles. If None (default), the integration will be done over all particles at once without chunking.

  • options (dict, optional) –

    Additional keyword arguments to pass to the field computation,
    • iotaProfile

      Iota profile of the Equilibrium, if not already assigned.

    • source_grid: Grid

      Source grid to use for field computation during Biot-Savart.

  • throw (bool, optional) – Whether to throw an error if the integration fails. If False, will return NaN for the points where the integration failed. Defaults to True.

  • return_aux (bool, optional) – Whether to return auxiliary information from the integrator. If True, will return a tuple (x, v, aux) where aux consists ts, stats and result from diffrax.diffeqsolve. Defaults to False. ts may become useful if SaveAt(steps=True) is used (see _trace_particles). Note that x, v and ts will be padded with NaNs to max_steps size in that case.

Returns:

  • x (ndarray, shape(num_particles, num_timesteps, 3)) – Position of each particle at each requested time, in either r,phi,z or rho,theta,zeta depending model.frame.

  • v (ndarray) – Velocity of each particle at specified times. The exact number of columns will depend on model.vcoords.