flekspy.tp.test_particles

class flekspy.tp.test_particles.FLEKSTP(dirs: str | List[str], iDomain: int = 0, iSpecies: int = 0, 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

_trajectory_cache

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], unit: str = 'planetary', 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])

static get_kinetic_energy(vx, vy, vz, mass=proton_mass)[source]

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(pID, mass=proton_mass)[source]: Calculates the 1st adiabatic invariant of a particle.

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(pID: Tuple[int, int]) → 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(pID: Tuple[int, int], unit='planetary', mass=proton_mass, charge=elementary_charge) → 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_gyroradius_to_curvature_ratio(pID: Tuple[int, int], unit='planetary', mass=proton_mass, charge=elementary_charge) → polars.Series[source]: Calculates the ratio of the particle’s gyroradius to the magnetic field’s radius of curvature.

get_gradient_drift(pID: Tuple[int, int], unit='planetary', mass=proton_mass, charge=elementary_charge) → 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]

plot_trajectory(pID: Tuple[int, int], *, mass=proton_mass, type='quick', xaxis='t', yaxis='x', ax=None, **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)

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.