flekspy.tp.test_particles

Attributes

logger
EARTH_RADIUS_KM

Classes

Indices	Defines constant indices for test particles.
FLEKSTP	Read and analyze particle trajectories stored in FLEKS test-particle files.

Functions

interpolate_at_times(→ polars.DataFrame)	Interpolate multiple numeric columns at specified time points.
plot_integrated_energy(df[, outname])	Plot integrated energy quantities as a function of time.

Module Contents

flekspy.tp.test_particles.logger

flekspy.tp.test_particles.EARTH_RADIUS_KM = 6378

class flekspy.tp.test_particles.Indices

Bases: enum.IntEnum

Defines constant indices for test particles.

TIME = 0

X = 1

Y = 2

Z = 3

VX = 4

VY = 5

VZ = 6

BX = 7

BY = 8

BZ = 9

EX = 10

EY = 11

EZ = 12

DBXDX = 13

DBXDY = 14

DBXDZ = 15

DBYDX = 16

DBYDY = 17

DBYDZ = 18

DBZDX = 19

DBZDY = 20

DBZDZ = 21

class flekspy.tp.test_particles.FLEKSTP(dirs: str, iDomain: int = 0, iSpecies: int = 0, unit: str = 'planetary', mass: float = proton_mass, charge: float = elementary_charge, iListStart: int = 0, iListEnd: int = -1, use_cache: bool = False)

Bases: object

Read and analyze particle trajectories stored in FLEKS test-particle files.

Each particle ID consists of a CPU index and a local particle index. The class exposes a lazy iterable interface for loading trajectory data without materializing the full dataset in memory.

use_cache = False

unit = 'planetary'

_trajectory_cache

mass

charge

iSpecies = 0

pfiles = []

particle_locations: Dict[Tuple[int, int], List[Tuple[str, int]]]

IDs

filetime = []

__repr__()

__len__()

__iter__()

__getitem__(key)

getIDs()

read_particle_list(filename: str) → Dict[Tuple[int, int], int]

Read and return a list of the particle IDs.

_read_the_first_record(filename: str) → List[float] | None

Get the first record stored in one file.

read_particles_at_time(time: float, doSave: bool = False) → Tuple[numpy.ndarray, numpy.ndarray]

Get the information of all the particles at a given time. If doSave, save to a CSV file with the name “particles_t***.csv”.

Note that the time tags in filetime do not include the last saved time. The function returns a tuple (ids, pData), where ids stores particle IDs and pData stores the particle weight, location, and velocity data.

save_trajectory(pID: Tuple[int, int], filename: str = None, shiftTime: bool = False, scaleTime: bool = False, format: str = 'csv') → None

Save the trajectory of a particle to a file. The output name defaults to a generated trajectory file name unless a custom filename is provided. Set shiftTime to reset the initial time to 0 and scaleTime to normalize it into the range [0, 1]. Supported formats are "csv" and "parquet".

save_trajectories(pIDs: List[Tuple[int, int]] | List[int], filename: str = 'trajectories.h5') → None

Save the trajectories of multiple particles to a single HDF5 file.

This helper can accept either a list of particle IDs or a list of integer indices. The output is written to an HDF5 archive with one dataset per particle.

_get_particle_raw_data(pID: Tuple[int, int]) → numpy.ndarray

Reads all raw trajectory data for a particle across multiple files.

_read_particle_record(pID: Tuple[int, int], index: int = -1) → list | None

Return a specific record of a test particle given its ID.

Use index=0 for the first record and index=-1 for the final record.

read_particle_trajectory(pID: Tuple[int, int]) → polars.LazyFrame

Return the trajectory of a test particle as a polars LazyFrame.

get_column_names() → List[str]

read_initial_condition(pID: Tuple[int, int]) → list | None

Return the initial conditions of a test particle.

read_final_condition(pID: Tuple[int, int]) → list | None

Return the final conditions of a test particle.

select_particles(f_select: Callable = None) → List[Tuple[int, int]]

Return the test particles whose initial conditions satisfy the requirement set by the user defined function f_select. The function receives the particle container and a particle ID, and should return a boolean value.

get_kinetic_energy(vx, vy, vz)

get_kinetic_energy_change_rate(pt_lazy: polars.LazyFrame) → polars.Series

Calculates the rate of change of kinetic energy in [eV/s].

get_pitch_angle(pID)

static _get_pitch_angle_lazy(lf: polars.LazyFrame) → polars.Series

Calculates the pitch angle from a LazyFrame.

static get_pitch_angle_from_v_b(vx, vy, vz, bx, by, bz)

get_first_adiabatic_invariant(pt_lazy: polars.LazyFrame) → polars.Series

Calculates the 1st adiabatic invariant of a particle.

The output units depend on the input data’s units:

• planetary units (e.g., velocity in km/s, B-field in nT): result is in [1e9 J/T].

• SI units (e.g., velocity in m/s, B-field in T): result is in [J/T].

static _calculate_bmag(df: polars.DataFrame | polars.LazyFrame) → polars.DataFrame | polars.LazyFrame

Calculates the magnetic field magnitude.

static _calculate_curvature(df: polars.DataFrame | polars.LazyFrame) → polars.DataFrame | polars.LazyFrame

Calculates the magnetic field curvature vector and adds it to the DataFrame. κ = (b ⋅ ∇)b.

Depending on the selected units, output curvature may be:

• planetary units: [1/RE]

• SI units: [1/m]

get_ExB_drift(pt_lazy: polars.LazyFrame) → polars.DataFrame

Calculates the convection drift velocity for a particle. v_exb = E x B / (B^2).

Assuming Earth’s planetary units, output drift velocity in [km/s].

get_curvature_drift(pt_lazy: polars.LazyFrame) → polars.DataFrame

Calculates the curvature drift velocity for a particle. v_c = (m * v_parallel^2 / (q*B^2)) * (B x κ).

Depending on the selected units, output drift velocity may be:

• planetary units: [km/s]

• SI units: [m/s]

get_adiabaticity_parameter(pt_lazy: polars.LazyFrame) → polars.Series

Calculates the adiabaticity parameter, defined as the ratio of the magnetic field’s radius of curvature to the particle’s gyroradius. When this parameter is >> 1, the motion is adiabatic.

static _calculate_gradient_b_magnitude(df: polars.DataFrame | polars.LazyFrame) → polars.DataFrame | polars.LazyFrame

Calculates the gradient of the magnetic field magnitude.

get_gradient_drift(pt_lazy: polars.LazyFrame) → polars.DataFrame

Calculates the gradient drift velocity for a particle. v_g = (μ / (q * B^2)) * (B x ∇|B|).

Depending on the selected units, output drift velocity may be:

• planetary units: [km/s]

• SI units: [m/s]

get_polarization_drift(pt_lazy: polars.LazyFrame) → polars.DataFrame

Calculates the polarization drift velocity for a particle. v_p = (m / (q * B^2)) * (dE_perp / dt).

Depending on the selected units, output drift velocity may be:

• planetary units: [km/s]

• SI units: [m/s]

get_betatron_acceleration(pt, mu)

Calculate the Betatron acceleration term from particle trajectory data.

The calculation follows the formula dW/dt = μ * (∂B/∂t), where the partial derivative is found using ∂B/∂t = dB/dt - v ⋅ ∇B. The input trajectory must include time, velocity, magnetic field, and magnetic gradient tensor columns.

get_energy_change_guiding_center(pID: Tuple[int, int]) → polars.DataFrame

Compute the change of energy of a single particle based on guiding-center theory.

The formula is given by dW/dt = q*E_parallel*v_parallel + mu*(∂B/∂t + u_E.∇B) + m*v_parallel^2*(u_E.κ), where W is the particle energy, B is the magnetic field magnitude, u_E is the E cross B drift, and κ is the magnetic field curvature.

integrate_drift_accelerations(pid: tuple[int, int])

Compute plasma drift velocities and the associated rate of energy change in [eV/s].

analyze_drifts(pid: tuple[int, int], outname=None, switchYZ=False)

Compute plasma drift velocities and the associated rate of energy change in [eV/s].

analyze_drifts_energy_change(pID: Tuple[int, int], outname=None)

Analyze and plot the energy changes for each term in the guiding-center approximation.

This method computes the parallel, Betatron, and Fermi acceleration terms using get_energy_change_guiding_center and compares them with the kinetic energy change rate. The residual difference between the total kinetic energy change and the guiding-center sum is interpreted as the non-adiabatic term.

analyze_drift(pID: tuple[int, int], drift_type: str, outname=None)

Analyze a specific drift for a particle.

The method plots the time series of the drift velocity, the electric field components, the associated energy-change rate, and the integrated energy change. Supported drift types are 'ExB', 'gradient', 'curvature', and 'polarization'.

plot_work_energy_verification(pID: Tuple[int, int], outname=None)

Verify the work-energy theorem for a particle.

The method compares the kinetic-energy change rate with the electric-field work rate and then compares the integrated quantities over time.

find_shock_crossing_time(pid, b_threshold_factor=2.5, verbose=False)

Find the shock crossing time for a single particle.

The method identifies the first significant change in the magnetic-field magnitude derivative and returns the corresponding time in seconds. If no candidate crossing is found, None is returned.

get_shock_up_down_states(pids, delta_t_up=20.0, delta_t_down=40.0, b_threshold_factor=2.5, verbose=False)

Analyze particles to find their upstream and downstream states around a shock.

Each particle is checked for a valid shock-crossing time. The method then interpolates the full particle state at points before and after the shock and returns the resulting upstream and downstream DataFrames.

_get_HT_frame(upstream_df: polars.DataFrame, downstream_df: polars.DataFrame) → tuple[numpy.ndarray | None, numpy.ndarray | None]

Find the de Hoffmann-Teller frame velocity and shock normal vector.

This implementation follows the magnetic-coplanarity approach described by Sonnerup et al. and minimizes the residual electric field in the resulting transformed frame.

analyze_in_HT_frame(pID: Tuple[int, int], outname: str = None, verbose: bool = False)

Analyze a particle trajectory in the de Hoffmann-Teller frame.

The method identifies the shock crossing, computes the de Hoffmann-Teller velocity and shock normal, transforms the particle data into that frame, and plots the associated non-ideal energy gain.

plot_trajectory(pID: Tuple[int, int], *, fscaling=1, smoothing_window=None, t_start=None, t_end=None, dt=None, outname=None, shock_time=None, type='quick', xaxis='t', yaxis='x', switchYZ=False, splitYZ=False, ax=None, verbose=True, **kwargs)

Plots the trajectory and velocities of the particle pID.

Example: >>> from flekspy import FLEKSTP >>> tp = FLEKSTP(“res/run1/PC/test_particles”, iSpecies=1) >>> tp.plot_trajectory((3,15))

plot_location(pData: numpy.ndarray)

Plot the location of particles pData.

Examples: >>> from flekspy import FLEKSTP >>> tp = FLEKSTP(“res/run1/PC/test_particles”, iSpecies=1) >>> ids, pData = tp.read_particles_at_time(3700, doSave=True) >>> f = tp.plot_location(pData)

_calculate_true_gc_trajectory(pt: polars.DataFrame, smoothing_gyro_periods: float) → polars.DataFrame

Calculates the ‘true’ guiding center trajectory by smoothing.

static _integrate_velocity(v_series: polars.Series, initial_pos_series: polars.Series, dt_series: polars.Series) → polars.Series

Integrates a velocity series using the trapezoidal rule.

_calculate_predicted_gc_trajectory(pt: polars.DataFrame, pID: Tuple[int, int]) → Tuple[polars.Series, polars.Series, polars.Series]

Calculates the predicted guiding center trajectory from theory.

_plot_gc_verification(pt: polars.DataFrame, pos_gc_true: polars.DataFrame, pos_gc_pred: Tuple[polars.Series, polars.Series, polars.Series], pID: Tuple[int, int])

Plots the verification results.

verify_guiding_center_model(pID: Tuple[int, int], smoothing_gyro_periods=1.0)

Verify the guiding-center model against the full particle trajectory.

The method estimates a smoothed guiding-center path, predicts the path from the drift velocity terms, and plots the comparison for each spatial coordinate.

flekspy.tp.test_particles.interpolate_at_times(df: polars.DataFrame | polars.LazyFrame, times_to_interpolate: list[float]) → polars.DataFrame

Interpolate multiple numeric columns at specified time points.

The input DataFrame or LazyFrame is first converted to a regular time-ordered frame, then rows matching the requested interpolation times are returned.

flekspy.tp.test_particles.plot_integrated_energy(df: polars.DataFrame, outname=None, **kwargs)

Plot integrated energy quantities as a function of time.

The DataFrame should contain a time column and one or more integrated energy columns that will be overlaid on a single figure.