FLEKS Python Visualization Toolkit: AMReX Data

flekspy is a Python package for processing FLEKS data. This notebook focuses on handling data in the AMReX format, which is used for both field and particle data. We will cover two main ways to load and analyze this data:

  1. Using the yt-based loader: This is the primary method, leveraging the power of the yt library for slicing, plotting, and analyzing both field and particle data.

  2. Using the experimental native AMReX particle loader: A faster, more direct way to load and plot particle data without the yt overhead.

1. Setup: Imports and Data Downloads

Imports

import flekspy
from flekspy.util import download_testfile, unit_one
from flekspy.amrex import AMReXParticleData

import numpy as np
import matplotlib.pyplot as plt
import yt

Downloading Demo Data

We’ll download two different AMReX datasets to demonstrate various features.

# Dataset 1: General purpose field and particle data
url_1 = "https://raw.githubusercontent.com/henry2004y/batsrus_data/master/3d_particle.tar.gz"
download_testfile(url_1, "data")

# Dataset 2: Smaller particle dataset for more specific examples
url_2 = "https://raw.githubusercontent.com/henry2004y/batsrus_data/master/fleks_particle_small.tar.gz"
download_testfile(url_2, "data")

2. Analyzing Field and PIC Particle Data with the yt Loader

The primary way to load AMReX data is with flekspy.load, which uses yt on the backend. This returns a yt dataset object, giving you access to the full power of the yt analysis framework.

ds = flekspy.load("data/3d*amrex", use_yt_loader=True)

yt : [INFO     ] 2025-10-09 21:08:52,671 Parameters: current_time              = 2.066624095642792

yt : [INFO     ] 2025-10-09 21:08:52,671 Parameters: domain_dimensions         = [64 40  1]

yt : [INFO     ] 2025-10-09 21:08:52,672 Parameters: domain_left_edge          = [-0.016 -0.01   0.   ]

yt : [INFO     ] 2025-10-09 21:08:52,673 Parameters: domain_right_edge         = [0.016 0.01  1.   ]

Slicing and Plotting Field Data

We can easily take a 2D slice of the 3D 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,)

Flekspy provides convenient wrappers for creating plots from these slices.

f, axes = dc.plot("By", pcolor=True)
dc.add_stream(axes[0], "Bx", "By", color="w")
dc.add_contour(axes[0], "Bx")

[2025-10-09 21:08:53,124] [flekspy.util.data_container] [INFO] Plots at z = 0.5 code_length
_images/cc1e873f723b6b90e9d9c1994bb3565502c2d706ed05d6691f4fdeec34258ba1.png

For more control, you can also extract the data and plot it directly with Matplotlib.

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()
_images/93d0dc0d945eae85923e2a4208d5d7405f570220f4e552fd4a0c8395a833ca60.png

Visualizing Phase Space Distributions

We can also analyze the particle data. yt allows us to create derived fields, which are useful for weighting particles. Here, we create a unit_one field that simply assigns a value of 1 to each particle.

ds.add_field(
    ds.pvar("unit_one"),
    function=unit_one,
    sampling_type="particle",
    units="dimensionless",
)

Now we can create a phase space plot. We’ll select a spatial region and plot the distribution of y- and z-velocities (p_uy, p_uz), weighted by the particle weight (p_w).

x_field = "p_uy"
y_field = "p_uz"
z_field = "p_w"
xleft = [-0.016, -0.01, 0.0]
xright = [0.016, 0.01, 1.0]

zmin, zmax = 1e-5, 2e-3

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_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")
pp.set_log(pp.fields[0], True)
pp.show()

Selecting and Plotting Particles in Geometric Regions

yt makes it easy to select particles within various geometric shapes.

Sphere

Plot the spatial location and velocity space of particles within a sphere.

sp = ds.sphere(center=[0, 0, 0], radius=1)

# Plot particle locations
pp_loc = ds.plot_particles("p_x", "p_y", "p_w", region=sp, unit_type="si", x_bins=64, y_bins=64)
pp_loc.show()

# Plot phase space
pp_phase = ds.plot_phase("p_uy", "p_uz", "p_w", region=sp, unit_type="si", x_bins=64, y_bins=64)
pp_phase.show()

yt : [INFO     ] 2025-10-09 21:08:54,629 xlim = -0.016000 0.016000

yt : [INFO     ] 2025-10-09 21:08:54,629 ylim = -0.010000 0.010000

yt : [INFO     ] 2025-10-09 21:08:54,631 xlim = -0.016000 0.016000

yt : [INFO     ] 2025-10-09 21:08:54,632 ylim = -0.010000 0.010000

yt : [INFO     ] 2025-10-09 21:08:54,635 Splatting (('particles', 'p_w')) onto a 800 by 800 mesh using method 'ngp'


Disk

Plot the spatial location and velocity space of particles within a disk.

disk = ds.disk(center=[0, 0, 0], normal=[0, 0, 1], radius=0.5, height=1.0)

# Plot particle locations
pp_loc = ds.plot_particles("p_x", "p_y", "p_w", region=disk, unit_type="si", x_bins=64, y_bins=64)
pp_loc.show()

# Plot phase space
pp_phase = ds.plot_phase("p_uy", "p_uz", "p_w", region=disk, unit_type="si", x_bins=64, y_bins=64)
pp_phase.show()

yt : [INFO     ] 2025-10-09 21:08:55,745 xlim = -0.016000 0.016000

yt : [INFO     ] 2025-10-09 21:08:55,745 ylim = -0.010000 0.010000

yt : [INFO     ] 2025-10-09 21:08:55,748 xlim = -0.016000 0.016000

yt : [INFO     ] 2025-10-09 21:08:55,748 ylim = -0.010000 0.010000

yt : [INFO     ] 2025-10-09 21:08:55,749 Splatting (('particles', 'p_w')) onto a 800 by 800 mesh using method 'ngp'


Advanced Example: Slicing with Range Limits

Now we’ll use the second dataset to show a more advanced slicing feature. Inheriting from IDL procedures, you can pass strings to plot to limit the plotting range (note: this syntax is experimental).

ds2 = flekspy.load("data/fleks_particle_small/3d*amrex")
dc2 = ds2.get_slice("z", 0.001)

f, axes = dc2.plot("Bx>(2.2e5)<(3e5) Ex", figsize=(12, 6))
dc2.add_stream(axes[0], "Bx", "By", color="w")
dc2.add_contour(axes[1], "Bx", color="k")

yt : [INFO     ] 2025-10-09 21:08:56,904 Parameters: current_time              = 40.0

yt : [INFO     ] 2025-10-09 21:08:56,904 Parameters: domain_dimensions         = [16  8  6]

yt : [INFO     ] 2025-10-09 21:08:56,905 Parameters: domain_left_edge          = [-0.0064 -0.0032 -0.0024]

yt : [INFO     ] 2025-10-09 21:08:56,905 Parameters: domain_right_edge         = [0.0064 0.0032 0.0024]

[2025-10-09 21:08:57,129] [flekspy.util.data_container] [INFO] Plots at z = 0.001 code_length

/home/runner/work/flekspy/flekspy/src/flekspy/util/data_container.py:505: UserWarning: The following kwargs were not used by contour: 'color'
  ax.contour(self.x, self.y, v.T, *args, **kwargs)
_images/609017c7da0ba9eb0f406c5fd37d1e08f29a4d6a3f34e2ea134463af3279dd1a.png

Advanced Example: Transforming Velocity Coordinates (WIP)

This section demonstrates a work-in-progress for transforming particle velocities into a new coordinate system (e.g., field-aligned coordinates).

# l = [1, 0, 0]
# m = [0, 1, 0]
# n = [0, 0, 1]

# def _vel_l(field, data):
#     return (l[0] * data[("particles", "p_ux")] + l[1] * data[("particles", "p_uy")] + l[2] * data[("particles", "p_uz")])

# def _vel_m(field, data):
#     return (m[0] * data[("particles", "p_ux")] + m[1] * data[("particles", "p_uy")] + m[2] * data[("particles", "p_uz")])

# ds2.add_field(ds2.pvar("vel_l"), units="code_velocity", function=_vel_l, sampling_type="particle")
# ds2.add_field(ds2.pvar("vel_m"), units="code_velocity", function=_vel_m, sampling_type="particle")

# sp2 = ds2.sphere(center=[0, 0, 0], radius=0.001)

# x_field = ds2.pvar("vel_l")
# y_field = ds2.pvar("vel_m")
# z_field = ds2.pvar("p_w")

# pp = yt.create_profile(data_source=sp2, bin_fields=[x_field, y_field], fields=z_field, n_bins=[64, 64]).plot()
# pp.set_unit(x_field, "km/s")
# pp.set_unit(y_field, "km/s")
# pp.show()

3. Using the Native AMReX Particle Loader (Experimental)

For scenarios where you only need to load particle data and the overhead of yt is not desired, flekspy provides a direct, experimental loader called AMReXParticleData.

data_file = "data/3d_particle_region0_1_t00000002_n00000007_amrex"
pd = AMReXParticleData(data_file)
pd

AMReXParticleData from data/3d_particle_region0_1_t00000002_n00000007_amrex
Time: 2.066624095642792
Dimensions: 2
Domain Dimensions: [64, 40, 1]
Domain Edges: [-0.016, -0.01] to [0.016, 0.01]
Integer component names: ['particle_id', 'particle_cpu', 'supid']
Real component names: ['x', 'y', 'velocity_x', 'velocity_y', 'velocity_z', 'weight']
Particle data: Not loaded (access .idata or .rdata to load)

Particles within a given region can be extracted as follows:

x_range = (-0.008, 0.008)
y_range = (-0.005, 0.005)

rdata = pd.select_particles_in_region(x_range=x_range, y_range=y_range)

Phase space plotting is achieved via plot_phase:

fig, ax = pd.plot_phase("velocity_y", "velocity_z", bins=(100,32))
plt.show()
_images/49f8da1d8d2d56e06dd2e74362e21b3630e70a151bcbab601bb248f2de149293.png