essos.dynamics¶

Attributes¶

`_VMEC_GUIDING_CENTER_MODELS`
`_GUIDING_CENTER_COLLISION_MODELS`
`_AXIS_REGION`

Classes¶

`Particles`
`LevelsetStoppingCriterion`	Stop tracing when a signed-distance level set is crossed.
`Tracing`

Functions¶

`gc_to_fullorbit`(field, initial_xyz, initial_vparallel, ...)	Computes full orbit positions for given guiding center positions,
`GuidingCenterCollisionsDiffusionMu`(→ jax.numpy.ndarray)
`GuidingCenterCollisionsDriftMuStratonovich`(...)
`GuidingCenterCollisionsDriftMuIto`(→ jax.numpy.ndarray)
`GuidingCenterCollisionsDiffusion`(→ jax.numpy.ndarray)
`GuidingCenterCollisionsDrift`(→ jax.numpy.ndarray)
`_gc_quantities`(field, points)	Guiding-center field quantities, fused when the field provides them.
`GuidingCenter`(→ jax.numpy.ndarray)
`LorentzCollisionsDiffusion`(→ jax.numpy.ndarray)
`LorentzCollisionsDrift`(→ jax.numpy.ndarray)
`Lorentz`(→ jax.numpy.ndarray)
`FieldLine`(→ jax.numpy.ndarray)
`FieldLineArclength`(→ jax.numpy.ndarray)	Trace the same field line with physical arclength as the parameter.
`FieldLineToroidal`(→ jax.numpy.ndarray)	Trace a flux-coordinate field with toroidal angle as the parameter.
`_fill_stopped_trajectories`(trajectories, criteria, args)	Hold each trajectory at its last point inside all level sets.
`_to_axis_regular`(y)	Map (s, theta, ...) to (sqrt(s) cos theta, sqrt(s) sin theta, ..., 0).
`_from_axis_regular`(y)	Map (u, w, ..., Theta) back to (s, theta, ...), with theta in [0, 2 pi).
`_axis_regular`(vector_field)	Express a VMEC guiding-center vector field in a chart that is regular on the axis.
`_vmec_boundary_event`(t, y, args, **kwargs)	LCFS event for VMEC guiding centers traced in (u, w).
`trace_field_lines`(field, initial_conditions, *[, ...])	Trace field lines by toroidal angle or physical arclength.
`connection_length`(field, initial_conditions, wall, *, ...)	Connection length and wall strike points of field lines.

Module Contents¶

essos.dynamics.gc_to_fullorbit(field, initial_xyz, initial_vparallel, total_speed, mass, charge, phase_angle_full_orbit=0)¶

Computes full orbit positions for given guiding center positions, parallel speeds, and total velocities using JAX for efficiency.

The full-orbit start satisfies x - b x v / Omega = X with the signed gyrofrequency Omega = charge |B| / mass, so the guiding center of the returned state is the requested point for either sign of the charge.

class essos.dynamics.Particles(initial_xyz=None, initial_vparallel_over_v=None, charge=ALPHA_PARTICLE_CHARGE, mass=ALPHA_PARTICLE_MASS, energy=FUSION_ALPHA_PARTICLE_ENERGY, min_vparallel_over_v=-1, max_vparallel_over_v=1, field=None, initial_vxvyvz=None, initial_xyz_fullorbit=None, phase_angle_full_orbit=0)¶

charge = 3.204353268e-19¶

mass = 6.69509884346e-27¶

energy = 5.639661751679999e-13¶

initial_xyz¶

nparticles¶

initial_xyz_fullorbit = None¶

initial_vxvyvz = None¶

phase_angle_full_orbit = 0¶

particle_index¶

random_keys¶

total_speed¶

initial_vparallel¶

initial_vperpendicular¶

to_full_orbit(field)¶

join(other, field=None)¶

classmethod InitializeParticlesAroundSurfaceAxis(surface, n_particles, distance_from_axis=0.0, charge=ALPHA_PARTICLE_CHARGE, mass=ALPHA_PARTICLE_MASS, energy=FUSION_ALPHA_PARTICLE_ENERGY, min_vparallel_over_v=-1, max_vparallel_over_v=1, field=None, random_seed=42, n_arc_samples=1000, boundary_surface=None, distance_mode='absolute', boundary_bisection_steps=32)¶

Initialize particles randomly distributed around/along a magnetic axis extracted from a surface.

Parameters:

surface – SurfaceRZFourier object to extract axis from
n_particles – Number of particles to initialize
distance_from_axis – Perpendicular distance (in Frenet frame) from the axis (0.0 for particles on axis, >0 for particles around axis). If distance_mode=’fraction_to_boundary’, this is interpreted as a fraction in [0, 1] of the local axis-to-boundary distance.
charge – Particle charge (default: alpha particle charge)
mass – Particle mass (default: alpha particle mass)
energy – Particle kinetic energy
min_vparallel_over_v – Minimum parallel velocity fraction
max_vparallel_over_v – Maximum parallel velocity fraction
field – Magnetic field object (for converting to full orbit if needed)
random_seed – Seed for random number generation
n_arc_samples – Number of samples for arc-length parametrization
boundary_surface – Optional surface used as geometric boundary when distance_mode=’fraction_to_boundary’.
distance_mode – ‘absolute’ or ‘fraction_to_boundary’.
boundary_bisection_steps – Number of bisection iterations used to find axis-to-boundary distance along each particle direction.

Returns:

Particles object with initial positions distributed around the axis

essos.dynamics.GuidingCenterCollisionsDiffusionMu(t, initial_condition, args) → jax.numpy.ndarray¶

essos.dynamics.GuidingCenterCollisionsDriftMuStratonovich(t, initial_condition, args) → jax.numpy.ndarray¶

essos.dynamics.GuidingCenterCollisionsDriftMuIto(t, initial_condition, args) → jax.numpy.ndarray¶

essos.dynamics.GuidingCenterCollisionsDiffusion(t, initial_condition, args) → jax.numpy.ndarray¶

essos.dynamics.GuidingCenterCollisionsDrift(t, initial_condition, args) → jax.numpy.ndarray¶

essos.dynamics._gc_quantities(field, points)¶

Guiding-center field quantities, fused when the field provides them.

Fields that do not override MagneticField.gc_quantities (for example Vmec, or a field that is not a pytree) use their individual methods; the choice is made at trace time.

essos.dynamics.GuidingCenter(t, initial_condition, args) → jax.numpy.ndarray¶

essos.dynamics.LorentzCollisionsDiffusion(t, initial_condition, args) → jax.numpy.ndarray¶

essos.dynamics.LorentzCollisionsDrift(t, initial_condition, args) → jax.numpy.ndarray¶

essos.dynamics.Lorentz(t, initial_condition, args) → jax.numpy.ndarray¶

essos.dynamics.FieldLine(t, initial_condition, field) → jax.numpy.ndarray¶

essos.dynamics.FieldLineArclength(t, initial_condition, field) → jax.numpy.ndarray¶: Trace the same field line with physical arclength as the parameter.

essos.dynamics.FieldLineToroidal(t, initial_condition, field) → jax.numpy.ndarray¶: Trace a flux-coordinate field with toroidal angle as the parameter.

essos.dynamics._fill_stopped_trajectories(trajectories, criteria, args)¶: Hold each trajectory at its last point inside all level sets.

essos.dynamics._VMEC_GUIDING_CENTER_MODELS¶

essos.dynamics._GUIDING_CENTER_COLLISION_MODELS¶

essos.dynamics._AXIS_REGION = 0.01¶

essos.dynamics._to_axis_regular(y)¶: Map (s, theta, …) to (sqrt(s) cos theta, sqrt(s) sin theta, …, 0).

essos.dynamics._from_axis_regular(y)¶: Map (u, w, …, Theta) back to (s, theta, …), with theta in [0, 2 pi).

essos.dynamics._axis_regular(vector_field)¶

Express a VMEC guiding-center vector field in a chart that is regular on the axis.

The state is (u, w, …, Theta) with (u, w) = sqrt(s) (cos a, sin a) and theta = a + Theta. (s, theta) is singular on the magnetic axis and (u, w) is not, so orbits cross it instead of stopping there. Of the poloidal rotation dtheta/dt, the fraction c = s / (s + _AXIS_REGION) goes to Theta and the rest rotates (u, w): du/dt = u/(2s) ds/dt - (1 - c) w dtheta/dt, dw/dt = w/(2s) ds/dt + (1 - c) u dtheta/dt, dTheta/dt = c dtheta/dt, all regular on the axis. Away from it, a fixed-step solver then advances theta as an angle, not by rotating (u, w) with a truncation error that spirals orbits outward. Drift vectors and diffusion matrices transform alike; the map involves position only and the position has no noise, so Ito and Stratonovich forms need no extra drift.

essos.dynamics._vmec_boundary_event(t, y, args, **kwargs)¶: LCFS event for VMEC guiding centers traced in (u, w).

class essos.dynamics.LevelsetStoppingCriterion(classifier, maximum_distance=0.0)¶

Stop tracing when a signed-distance level set is crossed.

classifier must be positive inside its reference surface. A positive maximum_distance permits tracing that far outside the surface before stopping.

classifier¶

maximum_distance¶

__call__(t, y, args, **kwargs)¶

class essos.dynamics.Tracing(trajectories_input=None, initial_conditions=None, times_to_trace=None, field=None, electric_field=None, model=None, maxtime: float = 1e-07, timestep: int = 1e-08, rtol=1e-07, atol=1e-07, particles=None, condition=None, species=None, tag_gc=1.0, boundary=None, rejected_steps=None, solver=None, stopping_criteria=None, progress=False, devices=None, max_steps=1000000)¶

stopping_criteria = None¶

devices¶

rejected_steps = 100¶

model = None¶

initial_conditions = None¶

times_to_trace = None¶

maxtime = 1e-07¶

timestep = 1e-08¶

rtol = 1e-07¶

atol = 1e-07¶

_trajectories = None¶

particles = None¶

species = None¶

tag_gc = 1.0¶

_axis_regular¶

_has_boundary_event = False¶

max_steps = 1000000¶

progress = False¶

progress_meter¶

solver = None¶

total_particles_unresolved¶

trace()¶

_vector_field(vector_field)¶

property trajectories¶

energy()¶

v_perp()¶

to_vtk(filename)¶

plot(ax=None, show=True, axis_equal=True, n_trajectories_plot=5, **kwargs)¶

loss_fraction_BioSavart(boundary)¶

Memory-efficient boundary loss fraction evaluation.

Uses flattened single vmap instead of nested double vmap to reduce memory usage by ~80% while maintaining accuracy.

Parameters:: boundary – SurfaceClassifier for boundary evaluation
Returns:: Cumulative loss fraction over time total_particles_lost: Total number of particles lost lost_times: Time of loss for each particle
Return type:: loss_fractions

loss_fraction(r_max=1.0)¶

Cumulative loss fraction of a flux-coordinate trace.

A particle is lost at the first saved time with s >= r_max, or with a non-finite state, which is what the LCFS event leaves after it stops a trace. The default r_max is that LCFS.

loss_fraction_BioSavart_collisions(boundary)¶

Memory-efficient boundary loss fraction for collision models.

Optimized version using flattened vmap.

loss_fraction_collisions(r_max=1.0)¶: As loss_fraction(), with the energy and position of each lost particle at its last finite saved state.

poincare_plot(shifts=[jnp.pi / 2], orientation='toroidal', length=1, ax=None, show=True, color=None, **kwargs)¶

Plot Poincare sections from Cartesian trajectories. :param shifts: Apply a linear shift to dependent data. Default is [pi/2]. :type shifts: list, optional :param orientation: ‘toroidal’ - find time values when toroidal angle = shift [0, 2pi].

‘z’ - find time values where z coordinate = shift. Default is ‘toroidal’.

Parameters:

length (float, optional) – A way to shorten data. 1 - plot full length, 0.1 - plot 1/10 of data length. Default is 1.
ax (matplotlib.axes._subplots.AxesSubplot, optional) – Matplotlib axis to plot on. Default is None.
show (bool, optional) – Whether to display the plot. Default is True.
color – "time", one Matplotlib color, or one color per trajectory.
**kwargs – Additional keyword arguments for plotting.

Toroidal crossings are found from the unwrapped Cartesian azimuth, so the branch cut at phi=0 does not create or discard intersections.

_tree_flatten()¶

classmethod _tree_unflatten(aux_data, children)¶

essos.dynamics.trace_field_lines(field, initial_conditions, *, toroidal_turns=None, length=None, samples=1000, tolerance=1e-07, stopping_criteria=None, progress=True, label='field lines', devices=None)¶

Trace field lines by toroidal angle or physical arclength.

Specify exactly one of toroidal_turns or length. Toroidal tracing is intended for fields represented in flux coordinates; Cartesian coil fields use arclength, so multiplying the magnetic field does not change the traced distance. samples includes both endpoints. The returned Tracing object provides trajectories, event flags, plotting, and Poincare sections.

Parameters:

field – ESSOS-compatible magnetic field.
initial_conditions – One seed per row, in the field’s coordinates.
toroidal_turns – Number of full toroidal turns to follow.
length – Physical arclength to follow for a Cartesian field.
samples – Number of saved points along each line.
tolerance – Relative and absolute adaptive-integration tolerance.
stopping_criteria – Optional event callable or sequence of callables.
progress – Show Diffrax’s terminal progress bar.
label – Text printed before compilation and after completion; set to None to suppress these two messages.
devices – Optional explicit sequence of JAX devices.

essos.dynamics.connection_length(field, initial_conditions, wall, *, max_length, tolerance=1e-08, max_steps=100000)¶

Connection length and wall strike points of field lines.

Each seed is followed along +B and -B by physical arclength until it crosses the wall or reaches max_length. The crossing is located by Diffrax event root finding, so the strike point is exact up to tolerance rather than limited by a sampling interval. The result is differentiable with respect to the seeds and field parameters.

Parameters:

field – ESSOS-compatible Cartesian magnetic field.
initial_conditions – Seeds of shape (n, 3). Seeds on or outside the wall return zero length and hit true.
wall – Object with evaluate_xyz(xyz) (e.g. SurfaceClassifier) or a callable wall(xyz), positive inside the wall and zero on it.
max_length – Cap on the length followed in each direction.
tolerance – Relative and absolute integration and root-finding tolerance.
max_steps – Maximum adaptive steps per direction; a line that exhausts them returns nan length and hit false.

Returns:

Dict with lengths (n, 2) (forward, backward), connection_length (n,) (their sum), strike_points (n, 2, 3) (end points; the wall hit when hit is true) and hit (n, 2) booleans.