FLEKS Python Visualization Toolkit
flekspy is a Python package for processing FLEKS data.
FLEKS data format
Field: *.out format or AMREX built-in format, whose directory name is assumed to end with “_amrex”
PIC particle: AMREX built-in format
Test particle: binary data format
Importing the package
import flekspy
Downloading demo data
If you don’t have FLEKS data to start with, you can download demo field data with the following:
from flekspy.util import download_testfile
url = "https://raw.githubusercontent.com/henry2004y/batsrus_data/master/batsrus_data.tar.gz"
download_testfile(url, "data")
Example test particle data can be downloaded as follows:
url = "https://raw.githubusercontent.com/henry2004y/batsrus_data/master/test_particles.tar.gz"
download_testfile(url, "data")
Example PIC particle data can be downloaded as follows:
url = "https://raw.githubusercontent.com/henry2004y/batsrus_data/master/3d_particle.tar.gz"
download_testfile(url, "data")
Loading data
flekspy.load is the interface to read files of all formats. It returns a different object for different formats.
file = "data/1d__raw_2_t25.60000_n00000258.out"
ds = flekspy.load(file)
ds
filename : data/1d__raw_2_t25.60000_n00000258.out
variables : ['x' 'Rho' 'Mx' 'My' 'Mz' 'Bx' 'By' 'Bz' 'Hyp' 'E' 'p' 'b1x' 'b1y' 'b1z'
'absdivB' 'r' 'g' 'cuty' 'cutz']
unit : normalized
nInstance : 1
npict : 1
time : 25.6
nIter : 258
ndim : 1
gencoord : False
grid : [256]
The variables can be accessed via dictionary keys:
p = ds.data["p"]
Extracting data
Variables at a series of locations can be extracted via extract_data, or at a single location via get_data.
import numpy as np
file = "data/z=0_fluid_region0_0_t00001640_n00010142.out"
ds = flekspy.load(file)
sat = np.array([[-28000.0, 0.0], [9000.0, 0.0]])
d = ds.extract_data(sat)
loc = np.array([0.0, 0.0])
d = ds.get_data(loc)
Visualizing fields
We provide convenient pp and pcolormesh method for IDL 1D and 2D outputs:
file = "data/1d__raw_2_t25.60000_n00000258.out"
ds = flekspy.load(file)
ds.plot("p", "Bx")
array([<Axes: xlabel='x', ylabel='p'>, <Axes: xlabel='x', ylabel='Bx'>],
dtype=object)
file = "data/z=0_fluid_region0_0_t00001640_n00010142.out"
ds = flekspy.load(file)
ds.pcolormesh("x", "y", "Bx", "By", "Bz", scale=False)
array([<Axes: title={'center': 'x'}, xlabel='X', ylabel='Y'>,
<Axes: title={'center': 'y'}, xlabel='X', ylabel='Y'>,
<Axes: title={'center': 'Bx'}, xlabel='X', ylabel='Y'>,
<Axes: title={'center': 'By'}, xlabel='X', ylabel='Y'>,
<Axes: title={'center': 'Bz'}, xlabel='X', ylabel='Y'>],
dtype=object)
Native AMReX format can be read using YT:
ds = flekspy.load("data/3d*amrex")
yt : [INFO ] 2025-03-18 17:27:21,056 Parameters: current_time = 2.066624095642792
yt : [INFO ] 2025-03-18 17:27:21,057 Parameters: domain_dimensions = [64 40 1]
yt : [INFO ] 2025-03-18 17:27:21,058 Parameters: domain_left_edge = [-0.016 -0.01 0. ]
yt : [INFO ] 2025-03-18 17:27:21,058 Parameters: domain_right_edge = [0.016 0.01 1. ]
We can take a slice at a given location, which works both for IDL and AMReX data:
dc = ds.get_slice("z", 0.5)
dc
variables : ['particle_cpu', 'particle_id', 'particle_position_x', 'particle_position_y', 'particle_supid', 'particle_velocity_x', 'particle_velocity_y', 'particle_velocity_z', 'particle_weight', 'Bx', 'By', 'Bz', 'Ex', 'Ey', 'Ez']
data range: [[unyt_quantity(-0.016, 'code_length'), unyt_quantity(0.016, 'code_length')], [unyt_quantity(-0.01, 'code_length'), unyt_quantity(0.01, 'code_length')], [0, 0]]
dimension : (256423,)
and plot the selected variables on the slice with colored contours, streamlines, and contour lines:
f, axes = dc.plot("By", pcolor=True)
dc.add_stream(axes[0], "Bx", "By", color="w")
dc.add_contour(axes[0], "Bx")
Plots at z = 0.5 code_length
Users can also obtain the data of a 2D plane, and visualize it with matplotlib. A complete example is given below:
import matplotlib.pyplot as plt
import numpy as np
ds = flekspy.load("data/3*amrex")
dc = ds.get_slice("z", 0.5)
f, axes = plt.subplots(
1, 2, figsize=(12, 4), constrained_layout=True, sharex=True, sharey=True
)
fields = ["By", "Bz"]
for ivar in range(2):
v = dc.evaluate_expression(fields[ivar])
vmin = v.min().v
vmax = v.max().v
ax = axes[ivar]
ax.set_title(fields[ivar], fontsize=16)
ax.set_ylabel("Y", fontsize=16)
ax.set_xlabel("X", fontsize=16)
c = ax.pcolormesh(dc.x.value, dc.y.value, np.array(v.T), cmap="turbo")
cb = f.colorbar(c, ax=ax, pad=0.01)
ax.set_xlim(np.min(dc.x.value), np.max(dc.x.value))
ax.set_xlim(np.min(dc.y.value), np.max(dc.y.value))
dc.add_stream(ax, "Bx", "By", density=0.5, color="w")
plt.show()
yt : [INFO ] 2025-03-18 17:27:22,158 Parameters: current_time = 2.066624095642792
yt : [INFO ] 2025-03-18 17:27:22,158 Parameters: domain_dimensions = [64 40 1]
yt : [INFO ] 2025-03-18 17:27:22,159 Parameters: domain_left_edge = [-0.016 -0.01 0. ]
yt : [INFO ] 2025-03-18 17:27:22,160 Parameters: domain_right_edge = [0.016 0.01 1. ]
Visualizing phase space distributions
from flekspy.util import unit_one
data_file = "data/3d_*amrex"
ds = flekspy.load(data_file)
# Add a user defined field. See yt document for more information about derived field.
ds.add_field(
ds.pvar("unit_one"),
function=unit_one,
sampling_type="particle",
units="dimensionless",
)
x_field = "p_uy"
y_field = "p_uz"
# z_field = "unit_one"
z_field = "p_w"
xleft = [-0.016, -0.01, 0.0]
xright = [0.016, 0.01, 1.0]
zmin, zmax = 1e-5, 2e-3
## Select and plot the particles inside a box defined by xleft and xright
region = ds.box(xleft, xright)
pp = ds.plot_phase(
x_field,
y_field,
z_field,
region=region,
unit_type="si",
x_bins=100,
y_bins=32,
domain_size=(xleft[0], xright[0], xleft[1], xright[1]),
)
pp.set_cmap(pp.fields[0], "turbo")
pp.set_font(
{
"size": 34,
"family": "DejaVu Sans",
}
)
pp.set_zlim(pp.fields[0], zmin=zmin, zmax=zmax)
pp.set_xlabel(r"$V_y$")
pp.set_ylabel(r"$V_z$")
pp.set_colorbar_label(pp.fields[0], "weight")
str_title = (
rf"x = [{xleft[0]:.1e}, {xright[0]:.1e}], y = [{xleft[1]:.1e}, {xright[1]:.1e}]"
)
pp.set_title(pp.fields[0], str_title)
pp.set_log(pp.fields[0], True)
pp.show()
# pp.save("test")
# pp.plots[("particles", z_field)].axes.xaxis.set_major_locator(
# ticker.MaxNLocator(4))
# pp.plots[("particles", z_field)].axes.yaxis.set_major_locator(
# ticker.MaxNLocator(4))
yt : [INFO ] 2025-03-18 17:27:23,021 Parameters: current_time = 2.066624095642792
yt : [INFO ] 2025-03-18 17:27:23,022 Parameters: domain_dimensions = [64 40 1]
yt : [INFO ] 2025-03-18 17:27:23,022 Parameters: domain_left_edge = [-0.016 -0.01 0. ]
yt : [INFO ] 2025-03-18 17:27:23,023 Parameters: domain_right_edge = [0.016 0.01 1. ]
Plot the location of particles that are inside a sphere
center = [0, 0, 0]
radius = 1
# Object sphere is defined in yt/data_objects/selection_objects/spheroids.py
sp = ds.sphere(center, radius)
pp = ds.plot_particles(
"p_x", "p_y", "p_w", region=sp, unit_type="si", x_bins=64, y_bins=64
)
pp.show()
# pp.save("test")
yt : [INFO ] 2025-03-18 17:27:23,857 xlim = -0.016000 0.016000
yt : [INFO ] 2025-03-18 17:27:23,858 ylim = -0.010000 0.010000
yt : [INFO ] 2025-03-18 17:27:23,860 xlim = -0.016000 0.016000
yt : [INFO ] 2025-03-18 17:27:23,861 ylim = -0.010000 0.010000
yt : [INFO ] 2025-03-18 17:27:23,864 Splatting (('particles', 'p_w')) onto a 800 by 800 mesh using method 'ngp'
Plot the velocity space of particles that are inside a sphere
pp = ds.plot_phase(
"p_uy", "p_uz", "p_w", region=sp, unit_type="si", x_bins=64, y_bins=64
)
pp.set_cmap(pp.fields[0], "turbo")
pp.set_font(
{
"size": 34,
"family": "DejaVu Sans",
}
)
pp.show()
Plot the location of particles that are inside a disk
center = [0, 0, 0]
normal = [0, 0, 1] # normal direction of the disk
radius = 0.5
height = 1.0
# Object sphere is defined in yt/data_objects/selection_objects/disk.py
disk = ds.disk(center, normal, radius, height)
pp = ds.plot_particles(
"p_x", "p_y", "p_w", region=disk, unit_type="si", x_bins=64, y_bins=64
)
pp.show()
yt : [INFO ] 2025-03-18 17:27:25,219 xlim = -0.016000 0.016000
yt : [INFO ] 2025-03-18 17:27:25,220 ylim = -0.010000 0.010000
yt : [INFO ] 2025-03-18 17:27:25,222 xlim = -0.016000 0.016000
yt : [INFO ] 2025-03-18 17:27:25,223 ylim = -0.010000 0.010000
yt : [INFO ] 2025-03-18 17:27:25,225 Splatting (('particles', 'p_w')) onto a 800 by 800 mesh using method 'ngp'
Plot the velocity space of particles that are inside a disk
pp = ds.plot_phase(
"p_uy", "p_uz", "p_w", region=disk, unit_type="si", x_bins=64, y_bins=64
)
pp.show()
Get the 2D array from the phase plot and reconstruct from scratch:
for var_name in pp.profile.field_data:
val = pp.profile.field_data[var_name]
x = pp.profile.x
y = pp.profile.y
plt.pcolormesh(x, y, val)
plt.show()
Plotting Test Particle Trajectories
Loading particle data
from flekspy import FLEKSTP
tp = FLEKSTP("data/test_particles", iSpecies=1)
Obtaining particle IDs
Test particle IDs consists of a CPU index and a particle index attached to the CPU.
pIDs = tp.getIDs()
Reading particle trajectory
traj = tp.read_particle_trajectory(pIDs[10])
print("time:\n")
print(traj[:, 0])
print("x:\n")
print(traj[:, 1])
time:
[0. 0.2865431 0.5730862 0.86495394 1.15735388 1.45110583
1.74436951 2.03866649]
x:
[-0.03138601 -0.03137759 -0.03136045 -0.0313433 -0.03132602 -0.03130875
-0.03129143 -0.03127424]
Getting initial location
x = tp.read_initial_location(pIDs[10])
print("time, X, Y, Z, Ux, Uy, Uz")
print(
f"{x[0]:.2e}, {x[1]:.2e}, {x[2]:.2e}, {x[3]:.2e}, {x[4]:.2e}, {x[5]:.2e}, {x[6]:.2e}"
)
time, X, Y, Z, Ux, Uy, Uz
0.00e+00, -3.14e-02, -1.86e-02, 0.00e+00, 5.87e-05, 3.70e-05, 3.98e-06
Plotting trajectory
tp.plot_trajectory(pIDs[0])
array([[<Axes: xlabel='x', ylabel='y'>, <Axes: xlabel='x', ylabel='z'>,
<Axes: xlabel='y', ylabel='z'>],
[<Axes: xlabel='time', ylabel='x'>,
<Axes: xlabel='time', ylabel='y'>,
<Axes: xlabel='time', ylabel='z'>],
[<Axes: xlabel='time', ylabel='Vx'>,
<Axes: xlabel='time', ylabel='Vy'>,
<Axes: xlabel='time', ylabel='Vz'>]], dtype=object)
Reading and visualizing all particle’s location at a given snapshot
ids, pData = tp.read_particles_at_time(0.0, doSave=False)
tp.plot_location(pData)
{'A': <Axes: label='A', xlabel='x', ylabel='y'>,
'B': <Axes: label='B', xlabel='x', ylabel='z'>,
'C': <Axes: label='C', xlabel='y', ylabel='z'>,
'D': <Axes3D: label='D', xlabel='x', ylabel='y', zlabel='z'>}
Selecting particles starting in a region
def f_select(tp, pid):
pData = tp.read_initial_location(pid)
inRegion = pData[FLEKSTP.ix_] > 0 and pData[FLEKSTP.iy_] > 0
return inRegion
pSelected = tp.select_particles(f_select)
Saving trajectories to CSV
tp.save_trajectory_to_csv(pIDs[0])