AMReX 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 native AMReX particle loader: A faster, more direct way to load and plot particle data without the yt overhead. (TODO: Requires testing in 3D geometries.)

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

Setup: Imports and Data Downloads

Imports

import flekspy
from flekspy.util import download_testfile, unit_one
from flekspy.util.transformations import create_field_transform
from flekspy.util.gmm import generate_synthetic_gmm_data, compare_gmm_results
from flekspy.amrex import AMReXParticle

import numpy as np
import matplotlib.pyplot as plt

# Set display for XArrays
import xarray as xr

xr.set_options(
    display_expand_coords=False,
    display_expand_data=False,
    display_expand_data_vars=False,
)

<xarray.core.options.set_options at 0x7f7f05801e80>

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")

Using the Native AMReX Particle Loader

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 AMReXParticle. This is much faster than the yt loader with a simple interface.

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

AMReXParticle 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 by providing the x_range, y_range, and z_range. If any of those is neglected, there is no constraint in that dimension.

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 Plot

Phase space plotting is achieved via plot_phase:

fig, ax = pd.plot_phase("vy", "vz", bins=(50, 32))
plt.show()

Adding 1D distributions via marginals=True:

pd.plot_phase("velocity_x", "velocity_y", bins=(32, 32), marginals=True)
plt.show()

Customization:

# Define the magnetic field vector for the transformation
B = np.array([1.0, 1.0, 0.0])
# Define an electric field vector\n
E = np.array([0.0, 1.0, 0.0])

# Create the transformation function using the utility
# We pass the original component names from the data header
b_field_transform = create_field_transform(B, pd.header.real_component_names)
eb_field_transform = create_field_transform(
    B, pd.header.real_component_names, e_field=E
)

fig, ax = plt.subplots(2, 2, figsize=(8, 6), constrained_layout=True)

pd.plot_phase(
    "vx",
    "vy",
    ax=ax[0, 0],
    add_colorbar=False,
    use_kde=True,
    kde_grid_size=30,
    kde_bandwidth=0.0002,
    title="KDE with GMM Estimate",
)

pd.plot_phase(
    "x",
    "vz",
    ax=ax[0, 1],
    add_colorbar=False,
    x_range=(-0.004, 0.004),
    y_range=(-0.005, 0.005),
    plot_zero_lines=False,
    ylabel=r"$v_z [m/s]$",
    title="x-vz",
)

pd.plot_phase(
    "v_parallel",
    "v_perp",
    ax=ax[1, 0],
    transform=b_field_transform,
    add_colorbar=False,
    plot_zero_lines=False,
    x_range=(-0.004, 0.004),
    y_range=(-0.005, 0.005),
    xlabel=r"$v_{\parallel}$",
    ylabel=r"$v_{\perp}$",
    title="B-field Aligned Velocities",
)

pd.plot_phase(
    "v_B",
    "v_E",
    ax=ax[1, 1],
    transform=eb_field_transform,
    add_colorbar=False,
    plot_zero_lines=False,
    x_range=(-0.004, 0.004),
    y_range=(-0.005, 0.005),
    xlabel=r"$v_B$",
    ylabel=r"$v_{E, \perp}$",
    title="E-B Aligned Velocities",
)
# GMM fit
gmm = pd.fit_gmm(1, "velocity_x", "velocity_y")
pd.plot_gmm_fit(gmm, ax=ax[0, 0], scale=1.0)

gmm2 = pd.fit_gmm(1, "v_B", "v_E", transform=eb_field_transform)
pd.plot_gmm_fit(gmm2, ax=ax[1, 1], scale=1.0)

plt.show()

Pairplot:

pd.pairplot()
plt.show()

pd.pairplot(corner=True)
plt.show()

Orthogonal intersecting planes:

pd.plot_intersecting_planes("x", "vx", "vy")
plt.show()

3D scattering:

xmin, xmax, ymin, ymax, zmin, zmax = -0.001, 0.001, -0.001, 0.001, 0.0, 0.01
hist_range = [[xmin, xmax], [ymin, ymax], [zmin, zmax]]
pd.plot_phase_3d("vx", "vy", "vz", normalize=True, hist_range=hist_range)
plt.show()

Matplotlib does not have real 3D rendering, which makes it hard for visualizing 3D phase plots.

GMM Fitting

To demonstrate the capability of the GMM fitting and parameter extraction, we can create a synthetic dataset of particles with known properties, fit a GMM to it, and verify that the extracted parameters match the input parameters.

from scipy.stats import multivariate_normal
from matplotlib.colors import LogNorm
from sklearn.mixture import GaussianMixture

# Define original parameters for comparison
params_isotropic = [
    {"center": np.array([0.0, 0.0]), "temp": 0.5},
    {"center": np.array([1.0, 2.0]), "temp": 0.2},
]

# Define the parameters for two synthetic Maxwellian distributions
components_isotropic = [
    {
        "n_samples": 1500,
        "mean": params_isotropic[0]["center"],
        "cov": np.array(
            [[params_isotropic[0]["temp"], 0], [0, params_isotropic[0]["temp"]]]
        ),
    },
    {
        "n_samples": 1000,
        "mean": params_isotropic[1]["center"],
        "cov": np.array(
            [[params_isotropic[1]["temp"], 0], [0, params_isotropic[1]["temp"]]]
        ),
    },
]

# Generate the synthetic data using the utility function
synthetic_data = generate_synthetic_gmm_data(components_isotropic, random_state=0)

print(f"Shape of the synthetic dataset: {synthetic_data.shape}")

Shape of the synthetic dataset: (2500, 2)

Fit a GMM to the synthetic data with 2 components:

gmm_synthetic = GaussianMixture(n_components=2, random_state=0)
gmm_synthetic.fit(synthetic_data)

GaussianMixture(n_components=2, random_state=0)

In a Jupyter environment, please rerun this cell to show the HTML representation or trust the notebook.
On GitHub, the HTML representation is unable to render, please try loading this page with nbviewer.org.

Compare the results. Since the synthetic data was generated with \(m=1\) and \(k_B=1\), the temperature is numerically equal to the variance, so we can compare directly.

compare_gmm_results(params_isotropic, gmm_synthetic, isotropic=True)

--- Original Synthetic Data Parameters ---
  Population 1: Center = [0.0, 0.0], Temperature = 0.5
  Population 2: Center = [1.0, 2.0], Temperature = 0.2

--- GMM Extracted Parameters ---
  Component 1: Center = [-0.03793, -0.04825], v_th_sq = 0.47898
  Component 2: Center = [1.00069, 2.0008], v_th_sq = 0.21018

Finally, we visualize the result. The plot below shows a 2D histogram of the synthetic particle data, with the fitted GMM components overlaid as red dashed ellipses. The ellipses represent the 1, 2, and 3-sigma contours of the Gaussian distributions, clearly showing how the GMM has identified the two distinct populations in the data.

fig, ax = plt.subplots(figsize=(8, 6), constrained_layout=True)

# Plot the 2D histogram of the synthetic data
ax.hist2d(
    synthetic_data[:, 0],
    synthetic_data[:, 1],
    bins=50,
    norm=LogNorm(),
    cmap="turbo",
    density=True,
)
ax.set_title("Synthetic Data with GMM Fit Overlay")
ax.set_xlabel("vx")
ax.set_ylabel("vy")

# Create a meshgrid for contour plotting
x = np.linspace(synthetic_data[:, 0].min(), synthetic_data[:, 0].max(), 100)
y = np.linspace(synthetic_data[:, 1].min(), synthetic_data[:, 1].max(), 100)
X, Y = np.meshgrid(x, y)
pos = np.dstack((X, Y))

# Overlay the GMM contours
for i in range(gmm_synthetic.n_components):
    mean = gmm_synthetic.means_[i]
    cov = gmm_synthetic.covariances_[i]
    rv = multivariate_normal(mean, cov)
    levels = rv.pdf(mean) * np.exp(-0.5 * np.array([1.0, 2.0, 3.0]) ** 2)
    ax.contour(X, Y, rv.pdf(pos), levels=np.sort(levels), colors="red", linestyles="--")

ax.set_aspect("equal", "box")
plt.show()

Anisotropic Example

Now, let’s do the same for an anisotropic distribution, where the temperature is different in different directions. This is common in magnetized plasmas, where temperature is often described in terms of \(T_\parallel\) and \(T_\perp\) with respect to the magnetic field.

# Define original parameters for comparison
params_anisotropic = [
    {
        "center": np.array([0.5, 0.5]),
        "temp_parallel": 0.8,
        "temp_perp": 0.2,
    }
]

# Define parameters for a synthetic anisotropic distribution
components_anisotropic = [
    {
        "n_samples": 2500,
        "mean": params_anisotropic[0]["center"],
        "cov": np.array(
            [
                [params_anisotropic[0]["temp_parallel"], 0],
                [0, params_anisotropic[0]["temp_perp"]],
            ]
        ),
    }
]

# Generate the synthetic data
anisotropic_data = generate_synthetic_gmm_data(
    components_anisotropic, shuffle=False, random_state=42
)

Fit the GMM:

from sklearn.mixture import GaussianMixture

gmm_anisotropic = GaussianMixture(n_components=1, random_state=0)
gmm_anisotropic.fit(anisotropic_data)

GaussianMixture(random_state=0)

In a Jupyter environment, please rerun this cell to show the HTML representation or trust the notebook.
On GitHub, the HTML representation is unable to render, please try loading this page with nbviewer.org.

compare_gmm_results(params_anisotropic, gmm_anisotropic, isotropic=False)

--- Original Synthetic Data Parameters ---
  Population 1: Center = [0.5, 0.5], Temp Parallel = 0.8, Temp Perp = 0.2

--- GMM Extracted Parameters ---
  Component 1: Center = [0.48075, 0.49185], v_th_sq_parallel = 0.78808, v_th_sq_perp = 0.20246

Plot the 2D histogram of the synthetic anisotropic data:

fig, ax = plt.subplots(figsize=(8, 6), constrained_layout=True)

ax.hist2d(
    anisotropic_data[:, 0],
    anisotropic_data[:, 1],
    bins=50,
    norm=LogNorm(),
    cmap="turbo",
    density=True,
)
ax.set_title("Anisotropic Synthetic Data with GMM Fit Overlay")
ax.set_xlabel(r"$v_{\parallel}$")
ax.set_ylabel(r"$v_{\perp}$")

# Create a meshgrid for contour plotting
x = np.linspace(anisotropic_data[:, 0].min(), anisotropic_data[:, 0].max(), 100)
y = np.linspace(anisotropic_data[:, 1].min(), anisotropic_data[:, 1].max(), 100)
X, Y = np.meshgrid(x, y)
pos = np.dstack((X, Y))

# Overlay the GMM contours
for i in range(gmm_anisotropic.n_components):
    mean = gmm_anisotropic.means_[i]
    cov = gmm_anisotropic.covariances_[i]
    rv = multivariate_normal(mean, cov)
    levels = rv.pdf(mean) * np.exp(-0.5 * np.array([1.0, 2.0, 3.0]) ** 2)
    ax.contour(X, Y, rv.pdf(pos), levels=np.sort(levels), colors="red", linestyles="--")

ax.set_aspect("equal", "box")
plt.show()

Core Population Extraction

We can also automatically extract the core population from the suprathermal population. This assumes the core is a Maxwellian distribution centered at the peak phase space density. The method extract_core_population fits a GMM to find the core center and thermal velocity, then applies a cutoff (default 3 * v_th) to separate the populations.

# Extract core population using 2D velocity
core_mask, supra_mask = pd.extract_core_population(
    ["velocity_x", "velocity_y"], v_dim=2, v_th=0.0001, cutoff=3.0
)

print(f"Core particles: {np.sum(core_mask)}")
print(f"Suprathermal particles: {np.sum(supra_mask)}")

# Visualize the separation in 2D (vx vs vy)
fig, ax = plt.subplots(figsize=(8, 6), constrained_layout=True)

# Get raw data for plotting
vx = pd.rdata[:, pd.header.real_component_names.index("velocity_x")]
vy = pd.rdata[:, pd.header.real_component_names.index("velocity_y")]

ax.scatter(vx[core_mask], vy[core_mask], s=1, c="tab:blue", label="Core", alpha=0.5)
ax.scatter(vx[supra_mask], vy[supra_mask], s=1, c="tab:red", label="Suprathermal", alpha=0.5)

ax.set_xlabel(r"$v_x$")
ax.set_ylabel(r"$v_y$")
ax.set_title("Core vs Suprathermal Extraction")
ax.legend()
plt.show()

Core particles: 109702
Suprathermal particles: 146721

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     ] 2026-01-04 16:24:44,369 Parameters: current_time              = 2.066624095642792

yt : [INFO     ] 2026-01-04 16:24:44,370 Parameters: domain_dimensions         = [64 40  1]

yt : [INFO     ] 2026-01-04 16:24:44,370 Parameters: domain_left_edge          = [-0.016 -0.01   0.   ]

yt : [INFO     ] 2026-01-04 16:24:44,371 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

<xarray.Dataset> Size: 19MB
Dimensions:              (X: 64, Y: 40, particle_index: 256423)
Coordinates: (4)
Dimensions without coordinates: particle_index
Data variables: (13)
Attributes:
    cut_norm:  z
    cut_loc:   0.5 code_length
    time:      2.066624095642792 code_time
    nstep:     -1
    filename:  /home/runner/work/flekspy/flekspy/docs/data/3d_particle_region...

xarray.Dataset

• Dimensions:
  • X: 64
  • Y: 40
  • particle_index: 256423
• Coordinates: (4)
  • X
    (X)
    float64
    -0.016 -0.01549 ... 0.01549 0.016

    array([-0.016   , -0.015492, -0.014984, -0.014476, -0.013968, -0.01346 ,
           -0.012952, -0.012444, -0.011937, -0.011429, -0.010921, -0.010413,
           -0.009905, -0.009397, -0.008889, -0.008381, -0.007873, -0.007365,
           -0.006857, -0.006349, -0.005841, -0.005333, -0.004825, -0.004317,
           -0.00381 , -0.003302, -0.002794, -0.002286, -0.001778, -0.00127 ,
           -0.000762, -0.000254,  0.000254,  0.000762,  0.00127 ,  0.001778,
            0.002286,  0.002794,  0.003302,  0.00381 ,  0.004317,  0.004825,
            0.005333,  0.005841,  0.006349,  0.006857,  0.007365,  0.007873,
            0.008381,  0.008889,  0.009397,  0.009905,  0.010413,  0.010921,
            0.011429,  0.011937,  0.012444,  0.012952,  0.01346 ,  0.013968,
            0.014476,  0.014984,  0.015492,  0.016   ])

  • Y
    (Y)
    float64
    -0.01 -0.009487 ... 0.009487 0.01

    array([-0.01    , -0.009487, -0.008974, -0.008462, -0.007949, -0.007436,
           -0.006923, -0.00641 , -0.005897, -0.005385, -0.004872, -0.004359,
           -0.003846, -0.003333, -0.002821, -0.002308, -0.001795, -0.001282,
           -0.000769, -0.000256,  0.000256,  0.000769,  0.001282,  0.001795,
            0.002308,  0.002821,  0.003333,  0.003846,  0.004359,  0.004872,
            0.005385,  0.005897,  0.00641 ,  0.006923,  0.007436,  0.007949,
            0.008462,  0.008974,  0.009487,  0.01    ])

  • particle_position_x
    (particle_index)
    float64
    -0.01575 -0.01576 ... 0.01558

    unyt_array([-0.01575008, -0.01575525, -0.01587573, ...,  0.01554996,
            0.01552487,  0.01557616], shape=(256423,), units='code_length')

  • particle_position_y
    (particle_index)
    float64
    -0.00987 -0.009944 ... 0.009694

    unyt_array([-0.00986998, -0.00994404, -0.00975847, ...,  0.00981498,
            0.00953845,  0.00969414], shape=(256423,), units='code_length')

• Data variables: (13)
  • Bx
    (X, Y)
    float64
    2.245e+05 2.235e+05 ... 2.247e+05

    unyt_array([[224465.92167695, 223511.3594672 , 224188.31487282, ...,
            224410.22217894, 223618.42356061, 224447.7408606 ],
           [223933.67751213, 223981.94052489, 224434.84378174, ...,
            223799.59921274, 222412.93468601, 225301.42823948],
           [223376.74980711, 224719.842702  , 223908.0040976 , ...,
            224688.9042025 , 222852.5160104 , 225387.37712019],
           ...,
           [225145.71592373, 223375.5474363 , 225656.78270949, ...,
            223993.25928821, 225547.75297624, 222709.63650977],
           [223827.66818206, 225067.89685698, 224506.52337642, ...,
            223641.83159904, 224879.89021393, 224195.32849444],
           [223481.16404539, 225052.95193819, 223564.68854986, ...,
            224701.52539833, 223904.26482222, 224725.18441804]],
          shape=(64, 40), units='code_magnetic')

  • By
    (X, Y)
    float64
    811.5 811.0 ... 1.644e+03 2.063e+03

    unyt_array([[  811.54203824,   811.03928808,   371.52324842, ...,
              275.57541332,  1156.27306893,  1301.06444771],
           [ -716.32440201,  -654.15854473,  -931.75247168, ...,
            -2026.82789956, -1091.41371436,  -884.40359742],
           [-2102.89933806, -2346.03966744, -2279.50823346, ...,
              591.04185444, -1673.25864492, -2405.79896696],
           ...,
           [ 3692.21822119,  3513.8050763 ,  3372.41743103, ...,
             3939.60190766,  4739.40759212,  3894.30234193],
           [ 3620.30412781,  3424.41559369,  3023.71315135, ...,
             3339.10345523,  3558.58822225,  3585.86425009],
           [ 1845.05218209,  2402.38977168,  3759.87195937, ...,
             1947.90358107,  1644.35040647,  2062.84384674]], shape=(64, 40), units='code_magnetic')

  • Bz
    (X, Y)
    float64
    -1.081e+04 ... -1.507e+04

    unyt_array([[-10806.32527664,  -9667.99753249, -12670.86423788, ...,
            -11484.68943187, -11293.62180832, -10483.65835111],
           [-10957.06174223, -11633.89932035, -12829.24450963, ...,
             -9899.22827234, -13412.56189913, -10215.90081238],
           [-11461.57769511, -10959.79971088,  -9612.14866188, ...,
            -11234.18569548, -10833.80754336, -10215.37719531],
           ...,
           [-10597.83994477, -12805.00259909,  -8589.69935516, ...,
            -11324.35021517,  -8780.01479233, -10830.02923898],
           [ -9623.08060758, -10959.11325371,  -9609.34303382, ...,
            -10546.23795463,  -8017.2310553 , -15121.33771394],
           [-12739.01502721, -10497.79542355, -11607.1965854 , ...,
            -13558.38232816,  -9440.99839003, -15065.22348579]],
          shape=(64, 40), units='code_magnetic')

  • Ex
    (X, Y)
    float64
    -0.3597 -1.162 ... 0.02884 -0.6913

    unyt_array([[-0.3596532 , -1.16208451, -0.46147923, ..., -0.03412608,
            -1.90650998, -1.73923923],
           [ 0.8080756 ,  0.26497417,  0.26816024, ...,  0.79046229,
            -0.4931928 ,  0.00705748],
           [ 0.60278065,  0.64179747, -1.48451264, ..., -0.52272708,
            -2.15070739, -0.87191115],
           ...,
           [ 0.62597684, -0.23157669,  0.3038241 , ...,  0.90546155,
             0.63842718,  0.39326574],
           [-0.22493516, -0.85246413,  0.69883731, ...,  2.13671645,
             2.68147884,  0.96930665],
           [-1.27249402, -1.01073755, -0.36535417, ...,  1.16289407,
             0.02883647, -0.69133293]], shape=(64, 40), units='code_magnetic*code_velocity')

  • Ey
    (X, Y)
    float64
    -2.757 -3.439 ... -3.208 -2.571

    unyt_array([[-2.75736511, -3.43869233, -3.38487312, ..., -3.98125864,
            -2.61109881, -2.12552311],
           [-2.47630574, -2.89356259, -3.59112206, ..., -3.67031166,
            -3.47062178, -2.16188577],
           [-2.67296463, -3.32475296, -3.55444452, ..., -3.02380257,
            -3.29298641, -2.39430971],
           ...,
           [-2.96382969, -1.81339858, -0.87390519, ..., -2.92939119,
            -3.88255793, -3.07691515],
           [-3.32560662, -3.42380214, -2.24088479, ..., -3.36355245,
            -3.76370397, -2.39832328],
           [-4.25307213, -4.29802304, -3.46127942, ..., -4.45019606,
            -3.20830486, -2.57093152]], shape=(64, 40), units='code_magnetic*code_velocity')

  • Ez
    (X, Y)
    float64
    10.05 8.677 8.577 ... 8.214 8.859

    unyt_array([[10.04808593,  8.67715126,  8.57707251, ...,  8.27933475,
             9.26091944, 10.33592784],
           [ 9.73864734,  8.63404384,  8.62647961, ...,  8.31378477,
             8.93794768, 10.14241813],
           [ 9.68415013,  9.56862392,  9.4095804 , ...,  8.9888938 ,
             8.70280589,  9.15658924],
           ...,
           [ 7.56837574,  7.58648581,  7.80563171, ...,  7.19534765,
             7.84981302,  7.9899229 ],
           [ 7.5151505 ,  8.56152046,  8.78726949, ...,  7.92919407,
             8.02132381,  7.85520397],
           [ 9.2983856 ,  9.23132522,  9.27564268, ...,  8.06094185,
             8.21364844,  8.85940931]], shape=(64, 40), units='code_magnetic*code_velocity')

  • particle_cpu
    (particle_index)
    float64
    0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0

    unyt_array([0., 0., 0., ..., 0., 0., 0.], shape=(256423,), units='(dimensionless)')

  • particle_id
    (particle_index)
    float64
    2.16e+05 3.997e+04 ... 2.54e+05

    unyt_array([216050.,  39972., 216020., ..., 255785., 254230., 254010.],
          shape=(256423,), units='(dimensionless)')

  • particle_supid
    (particle_index)
    float64
    1.0 1.0 1.0 1.0 ... 1.0 1.0 1.0 1.0

    unyt_array([1., 1., 1., ..., 1., 1., 1.], shape=(256423,), units='(dimensionless)')

  • particle_velocity_x
    (particle_index)
    float64
    6.393e-05 0.0001933 ... 0.000691

    unyt_array([6.39300997e-05, 1.93264974e-04, 1.05126409e-04, ...,
           3.06691515e-04, 1.55426098e-04, 6.91048536e-04], shape=(256423,), units='code_velocity')

  • particle_velocity_y
    (particle_index)
    float64
    -0.0001576 -0.0004673 ... 0.0006013

    unyt_array([-1.57582877e-04, -4.67309560e-04,  5.68051228e-05, ...,
            1.83998720e-04, -3.44191338e-04,  6.01267985e-04], shape=(256423,), units='code_velocity')

  • particle_velocity_z
    (particle_index)
    float64
    0.4787 -0.2406 ... 0.1597 0.4071

    unyt_array([ 0.4786543 , -0.24063388,  0.83048468, ...,  0.0132917 ,
            0.15966753,  0.40709236], shape=(256423,), units='code_velocity')

  • particle_weight
    (particle_index)
    float64
    9.342e+21 9.342e+21 ... 9.342e+21

    unyt_array([9.34162108e+21, 9.34162108e+21, 9.34162108e+21, ...,
           9.34162108e+21, 9.34162108e+21, 9.34162108e+21], shape=(256423,), units='code_mass')

• Attributes: (5)

  cut_norm :

  z

  cut_loc :

  0.5 code_length

  time :

  2.066624095642792 code_time

  nstep :

  -1

  filename :

  /home/runner/work/flekspy/flekspy/docs/data/3d_particle_region0_1_t00000002_n00000007_amrex

flekspy provides convenient wrappers for creating plots from these slices. You can use query syntax to plot with value range limits:

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

plt.show()

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.fleks.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.values, dc.Y.values, np.array(v.T), cmap="turbo")
    cb = f.colorbar(c, ax=ax, pad=0.01)

    ax.set_xlim(np.min(dc.X.values), np.max(dc.X.values))
    ax.set_xlim(np.min(dc.Y.values), np.max(dc.Y.values))
    dc.fleks.add_stream(ax, "Bx", "By", density=0.5, color="w")

plt.show()

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     ] 2026-01-04 16:24:46,237 xlim = -0.016000 0.016000

yt : [INFO     ] 2026-01-04 16:24:46,237 ylim = -0.010000 0.010000

yt : [INFO     ] 2026-01-04 16:24:46,239 xlim = -0.016000 0.016000

yt : [INFO     ] 2026-01-04 16:24:46,239 ylim = -0.010000 0.010000

yt : [INFO     ] 2026-01-04 16:24:46,242 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     ] 2026-01-04 16:24:47,221 xlim = -0.016000 0.016000

yt : [INFO     ] 2026-01-04 16:24:47,222 ylim = -0.010000 0.010000

yt : [INFO     ] 2026-01-04 16:24:47,224 xlim = -0.016000 0.016000

yt : [INFO     ] 2026-01-04 16:24:47,224 ylim = -0.010000 0.010000

yt : [INFO     ] 2026-01-04 16:24:47,225 Splatting (('particles', 'p_w')) onto a 800 by 800 mesh using method 'ngp'

Advanced Example: Transforming Velocity Coordinates (WIP)

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

# 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()