flekspy.tp.test_particles

Attributes

logger

EARTH_RADIUS_KM

Classes

Indices

Defines constant indices for test particles.

FLEKSTP

A class that is used to read and plot test particles. Each particle ID consists of

Functions

interpolate_at_times(→ polars.DataFrame)

Interpolates multiple numeric columns of a DataFrame at specified time points.

plot_integrated_energy(df[, outname])

Plots 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[source]

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 | List[str], iDomain: int = 0, iSpecies: int = 0, unit: str = 'planetary', mass: float = proton_mass, charge: float = elementary_charge, iListStart: int = 0, iListEnd: int = -1, readAllFiles: bool = False, use_cache: bool = False)[source]

Bases: object

A class that is used to read and plot test particles. Each particle ID consists of a CPU index, a particle index on each CPU, and a location index. By default, 7 real numbers saved for each step: time + position + velocity. Additional field information are also stored if available.

This class is a lazy, iterable container. It avoids loading all data into memory at once, making it efficient for large datasets. You can access particle trajectories using standard container operations.

Parameters:

dirs (str) – the path to the test particle dataset.

Examples: >>> tp = FLEKSTP(“res/run1/PC/test_particles”, iSpecies=1) >>> len(tp) 10240 >>> trajectory = tp[0] >>> tp.plot_trajectory(tp.IDs[3]) >>> tp.save_trajectory_to_csv(tp.IDs[5]) >>> ids, pData = tp.read_particles_at_time(0.0, doSave=False) >>> f = tp.plot_location(pData)

use_cache = False
unit = 'planetary'
_trajectory_cache
mass
charge
iSpecies = 0
pfiles = []
indextotime = []
particle_locations: Dict[Tuple[int, int], List[Tuple[str, int]]]
IDs
filetime = []
__repr__()[source]
__len__()[source]
__iter__()[source]
__getitem__(key)[source]
getIDs()[source]
read_particle_list(filename: str) → Dict[Tuple[int, int], int][source]

Read and return a list of the particle IDs.

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

Get the first record stored in one file.

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

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.

Returns:

a numpy array of tuples contains the particle IDs. pData: a numpy real array with the particle weight, location and velocity.

Return type:

ids

Examples: >>> ids, pData = pt.read_particles_at_time(3700, doSave=True)

save_trajectory_to_csv(pID: Tuple[int, int], filename: str = None, shiftTime: bool = False, scaleTime: bool = False) → None[source]

Save the trajectory of a particle to a csv file.

Parameters:

  • pID – particle ID.

  • shiftTime (bool) – If set to True, set the initial time to be 0.

  • scaleTime (bool) – If set to True, scale the time into [0,1] range.

Example: >>> tp.save_trajectory_to_csv((3,15))

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

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

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

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

Parameters:

  • pID – particle ID

  • index – The index of the record to be returned. 0: first record. -1: last record (default).

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

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

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

Return the initial conditions of a test particle.

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

Return the final conditions of a test particle.

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

Return the test particles whose initial conditions satisfy the requirement set by the user defined function f_select. The first argument of f_select is the particle ID, and the second argument is the ID of a particle.

Examples: >>> from flekspy.tp import Indices >>> def f_select(tp, pid): >>> pData = tp.read_initial_condition(pid) >>> inTime = pData[Indices.TIME] < 3601 >>> inRegion = pData[Indices.X] > 20 >>> return inTime and inRegion >>> >>> pselected = tp.select_particles(f_select) >>> tp.plot_trajectory(list(pselected.keys())[1])

get_kinetic_energy(vx, vy, vz)[source]
get_kinetic_energy_change_rate(pt_lazy: polars.LazyFrame) → polars.Series[source]

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

get_pitch_angle(pID)[source]
static _get_pitch_angle_lazy(lf: polars.LazyFrame) → polars.Series[source]

Calculates the pitch angle from a LazyFrame.

static get_pitch_angle_from_v_b(vx, vy, vz, bx, by, bz)[source]
get_first_adiabatic_invariant(pt_lazy: polars.LazyFrame) → polars.Series[source]

Calculates the 1st adiabatic invariant of a particle. The output units depend on the input data’s units: - “planetary” (e.g., velocity in km/s, B-field in nT): result is in [1e9 J/T]. - “SI” (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[source]

Calculates the magnetic field magnitude.

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

Calculates the magnetic field curvature vector and adds it to the DataFrame. κ = (b ⋅ ∇)b Depending on the selected units, output curvature may be - “planetary”: [1/RE] - “SI”: [1/m]

get_ExB_drift(pt_lazy: polars.LazyFrame) → polars.DataFrame[source]

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[source]

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”: [km/s] - “SI”: [m/s]

get_adiabaticity_parameter(pt_lazy: polars.LazyFrame) → polars.Series[source]

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[source]

Calculates the gradient of the magnetic field magnitude.

get_gradient_drift(pt_lazy: polars.LazyFrame) → polars.DataFrame[source]

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”: [km/s] - “SI”: [m/s]

get_polarization_drift(pt_lazy: polars.LazyFrame) → polars.DataFrame[source]

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”: [km/s] - “SI”: [m/s]

get_betatron_acceleration(pt, mu)[source]

Calculates 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

Parameters:

  • pt – A Polars LazyFrame containing the particle trajectory. It must include columns for time, velocity (vx, vy, vz), magnetic field (bx, by, bz), and the magnetic field gradient tensor (e.g., ‘dbxdx’, ‘dbydx’, etc.).

  • mu – A Polars Series containing the magnetic moment (first adiabatic invariant) of the particle.

Returns:

A new Polars LazyFrame with added intermediate columns and the final ‘dW_betatron’ column representing the rate of energy change in [eV/s].

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

Computes the change of energy of a single particle based on the 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. The first term on the right hand side is the parallel acceleration, the second term is the Betatron acceleration, and the third term is one type of Fermi acceleration.

Parameters:

pID (Tuple[int, int]) – The particle ID (cpu, id).

Returns:

A Polars DataFrame with the time, the three energy change components, and the total, in [eV/s].

integrate_drift_accelerations(pid: tuple[int, int])[source]

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

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

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

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

Analyzes and plots the energy changes for each term in the guiding center approximation.

This method computes the parallel, Betatron, and Fermi accelerations using get_energy_change_guiding_center. It also calculates the total kinetic energy change and treats the difference between the kinetic energy change and the sum of the guiding center terms (dW_total) as the non-adiabatic term.

The method generates a plot with four subplots: 1. dW_parallel: Energy change due to parallel electric fields. 2. dW_betatron: Energy change due to the Betatron effect. 3. dW_fermi: Energy change due to Fermi acceleration. 4. dW_total and Non-adiabatic term: The sum of the above terms compared

with the non-adiabatic heating component.

Parameters:

  • pID (Tuple[int, int]) – The particle ID (cpu, id).

  • outname (str, optional) – If provided, the plot is saved to this filename instead of being shown. Defaults to None.

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

Analyzes a specific drift for a particle, plotting its velocity, the electric field, the energy change rate, and the integrated energy change.

Parameters:

  • pID (tuple[int, int]) – The particle ID (cpu, id).

  • drift_type (str) – The type of drift to analyze. Supported options are: ‘ExB’, ‘gradient’, ‘curvature’, ‘polarization’.

  • outname (str, optional) – If provided, the plot is saved to this filename instead of being shown. Defaults to None.

find_shock_crossing_time(pid, b_threshold_factor=2.5, verbose=False)[source]

Finds the shock crossing time for a single particle.

The shock is identified by finding the first rate of change in the magnetic field magnitude that exceeds a threshold, which signifies a rapid transition between the upstream and downstream regions.

Parameters:

  • pid – particle index.

  • b_threshold_factor (float) – A multiplier for the standard deviation of the B-field derivative. A larger value makes the detection less sensitive to minor fluctuations. Defaults to 2.5.

  • verbose (bool) – If True, prints diagnostic information. Defaults to False.

Returns:

The time of the shock crossing in seconds. Returns None if

no significant crossing is detected based on the criteria.

Return type:

float or None

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

Analyzes particles to find their state upstream and downstream of a shock.

This function iterates through a list of particle IDs. For each particle, it first identifies the shock crossing time. It then calculates specific upstream and downstream time points based on this crossing. Finally, it interpolates the particle’s full state (position, velocity, fields) at these two points and collects the results.

Parameters:

  • pids (list) – A list of particle IDs (e.g., [(0, 1), (0, 2), …]) to process.

  • delta_t_up (float) – The time in seconds before the shock crossing to define the upstream point. Defaults to 20.0.

  • delta_t_down (float) – The time in seconds after the shock crossing to define the downstream point. Defaults to 40.0.

  • b_threshold_factor (float) – The sensitivity factor for shock detection, passed to find_shock_crossing_time. Defaults to 2.5.

  • verbose (bool) – If True, prints progress and individual shock detection times. Defaults to False.

Returns:

A tuple containing two Polars DataFrames:

  • The first DataFrame contains the states of all valid particles at their respective upstream times.

  • The second DataFrame contains the states of all valid particles at their respective downstream times.

Each DataFrame includes the original particle ID (pid_rank, pid_idx), the shock crossing time (t_cross), and the interpolated physical quantities. Returns (None, None) if no particles with a valid shock crossing are found.

Return type:

tuple[pl.DataFrame, pl.DataFrame]

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

Finds the de Hoffmann-Teller frame velocity and the shock normal vector using the method from Sonnerup et al. [2006], which minimizes the residual electric field.

Parameters:

  • upstream_df (pl.DataFrame) – DataFrame with upstream particle states.

  • downstream_df (pl.DataFrame) – DataFrame with downstream particle states.

Returns:

A tuple containing:

  • V_HT (np.ndarray | None): The de Hoffmann-Teller velocity vector in [km/s] if successful, otherwise None.

  • shock_normal (np.ndarray | None): The estimated shock normal vector if successful, otherwise None.

Return type:

tuple[np.ndarray | None, np.ndarray | None]

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

Analyzes a particle’s trajectory in the de Hoffmann-Teller (HT) frame.

This method performs the following steps: 1. Finds the shock crossing and determines the upstream and downstream states. 2. Calculates the de Hoffmann-Teller velocity (V_HT) and the shock normal. 3. Transforms the particle’s velocity and the electric/magnetic fields

into the HT frame.

  1. In this frame, the energy gain is a direct measure of non-ideal acceleration. It calculates and plots this energy gain.

  2. Generates a summary plot of the analysis.

Parameters:

  • pID (Tuple[int, int]) – The particle ID (cpu, id) to analyze.

  • outname (str, optional) – If provided, the plot is saved to this filename. Defaults to None (displays plot).

  • verbose (bool, optional) – If True, prints diagnostic information. Defaults to False.

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)[source]

Plots the trajectory and velocities of the particle pID.

Example: >>> tp.plot_trajectory((3,15))

plot_location(pData: numpy.ndarray)[source]

Plot the location of particles pData.

Examples: >>> 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[source]

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[source]

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][source]

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])[source]

Plots the verification results.

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

Verifies the guiding center model by comparing it against the full particle trajectory.

This method performs the following steps: 1. Calculates a “true” guiding center trajectory by applying a low-pass

filter (moving average) to the full particle trajectory, smoothing out the gyromotion.

  1. Calculates a “predicted” guiding center trajectory by numerically integrating the guiding center velocity, which is the sum of the parallel velocity and all perpendicular drift velocities (E x B, gradient, curvature, and polarization).

  2. Generates a plot comparing the full trajectory, the “true” GC trajectory, and the “predicted” GC trajectory for each coordinate (X, Y, Z).

A close overlap between the “true” and “predicted” trajectories validates the guiding center approximation and the drift calculations.

Parameters:

  • pID (Tuple[int, int]) – The ID of the particle to analyze.

  • smoothing_gyro_periods (float) – The size of the moving average window in units of gyro-periods. Defaults to 1.0.

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

Interpolates multiple numeric columns of a DataFrame at specified time points.

Parameters:

  • df – The input Polars DataFrame or LazyFrame.

  • times_to_interpolate – A list of time points (floats or ints) at which to interpolate.

Returns:

A new DataFrame containing the interpolated rows for each specified time.

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

Plots integrated energy quantities as a function of time.

Parameters:

  • df (pl.DataFrame) – A Polars DataFrame containing a time column and one or more integrated energy columns.

  • outname (str) – If not None, save the plot to file.