Methods for particle data§

The methods in this page are accessible to a data object when:

<data_obj>.attrs['pic_data_type'] == 'part'#(1)!

<data_obj> may be either ozzy.DataArray or ozzy.Dataset

bin_into_grid §

bin_into_grid(axes_ds, t_var='t', w_var='q', r_var=None)

Bin particle data into a grid (density distribution).

Parameters:

Name	Type	Description	Default
`axes_ds` §	`Dataset`	Dataset containing grid axes information. Tip The axis information can be created for example with: `import ozzy as oz nx = 200 ny = 150 xlims = (0.0, 30.0) ylims = (-4.0, 4.0) axes_ds = oz.Dataset( coords={ "x1": oz.utils.axis_from_extent(nx, xlims), "x2": oz.utils.axis_from_extent(ny, ylims), }, pic_data_type = "grid")` Or it can be obtained from an existing grid data object with: `# fields may be an existing Dataset or DataArray axes_ds = fields.coords` Note about axis attributes By default, the `long_name` and `units` attributes of the resulting grid axes are taken from the original particle Dataset. But these attributes are overriden if they are passed along with the `axes_ds` Dataset.	required
`t_var` §	`str`	Name of the time dimension in the input datasets.	`'t'`
`w_var` §	`str`	Name of the variable representing particle weights or particle charge.	`'q'`
`r_var` §	`str \| None`	Name of the variable representing particle radial positions. If provided and if one of the `axes_ds` coordinates is `r_var`, the particle weights are divided by this variable.	`None`

Returns:

Name	Type	Description
`parts`	`Dataset`	Dataset containing the charge density distribution on the grid.

Raises:

Type	Description
`KeyError`	If no spatial dimensions are found in the input `axes_ds`.
`ValueError`	If the `axes_ds` argument does not contain grid data.

Notes

The binned density data is multiplied by a factor that ensures that the total volume integral of the density corresponds to the sum of all particle weights \(Q_w\). If \(w\) is each particle's weight variable and \(N_p\) is the total number of particles, then \(Q_w\) is defined as:

\[ Q_w = \sum_i^{N_p} w_i \]

Note that different simulation codes have different conventions in terms of what \(Q_w\) corresponds to.

Examples:

Usage

import ozzy as oz
import numpy as np

# Create a sample particle dataset
particles = oz.Dataset(
    {
        "x1": ("pid", np.random.uniform(0, 10, 10000)),
        "x2": ("pid", np.random.uniform(0, 5, 10000)),
        "q": ("pid", np.ones(10000)),
    },
    coords={"pid": np.arange(10000)},
    attrs={"pic_data_type": "part"}
)

# Create axes for binning
axes = oz.Dataset(
    coords={
        "x1": oz.utils.axis_from_extent(100, (0.0, 10.0)),
        "x2": oz.utils.axis_from_extent(50, (0.0, 5.0)),
    },
    pic_data_type = "grid",
)

# Bin particles into grid (Cartesian geometry)
grid_data = particles.ozzy.bin_into_grid(axes)

# Example 2: Using a different weight variable
particles["w"] = ("pid", np.random.uniform(0.5, 1.5, 10000))
grid_data_weighted = particles.ozzy.bin_into_grid(axes, w_var="w")

# Example 3: Axisymmetric geometry
grid_data_axisym = particles.ozzy.bin_into_grid(axes, r_var="x2")

# Example 4: Time-dependent data
time_dependent_particles = particles.expand_dims(dim={"t": [0, 1, 2]})
time_dependent_grid = time_dependent_particles.ozzy.bin_into_grid(axes)

get_emittance §

get_emittance(
    norm_emit=True,
    axisym=False,
    p_all_vars=["p1", "p2", "p3"],
    x_var="x2",
    p_var="p2",
    w_var="q",
)

Calculate the RMS beam emittance.

Computes the normalized or geometric RMS emittance based on particle positions and momenta. For axisymmetric beams, returns the Lapostolle emittance (see Notes below).

Warning

This method assumes that the particle dimension is "pid".

Parameters:

Name	Type	Description	Default
`norm_emit` §	`bool`	Whether to calculate normalized emittance (multiplied by \(\left< \beta \gamma \right>\)).	`True`
`axisym` §	`bool`	If `True`, calculate Lapostolle emittance for axisymmetric beams	`False`
`p_all_vars` §	`list[str]`	List of names of momentum components. Note The components should be sorted, e.g. `["px", "py", "pz"]`. If `axisym=True`, only the two first components will be adopted.	`['p1', 'p2', 'p3']`
`x_var` §	`str`	Variable name for position coordinate in Dataset that should be used for emittance calculation	`'x2'`
`p_var` §	`str`	Variable name for momentum coordinate in Dataset that should be used for emittance calculation. This argument is only relevant when `axisym=False`, otherwise it is set to `p_all_vars[1]`.	`'p2'`
`w_var` §	`str`	Variable name for particle weights in Dataset	`'q'`

Returns:

Type	Description
`Dataset`	Dataset containing the calculated emittance and particle counts for each data point. The emittance variable is named `"emit_norm"` if `norm_emit=True`, otherwise `"emit"`. Also includes a `"counts"` variable with particle counts.

Notes

The geometric emittance along a given transverse dimension \(i\) is calculated according to:

\(\epsilon_i = \sqrt{\left<x_i^2\right> \left<{x'_i}^2\right> - \left(x_i x'_i\right)^2}\)

where \(x_i\) is the particle position, and \(x'_i \approx p_i / p_\parallel\) is the trace for relativistic particles with longitudinal momentum \(p_\parallel\) and transverse momentum \(p_i \ll p_\parallel\). The angle brackets denote a weighted average over particles.

The normalized emittance (norm_emit=True, default) is calculated as:

\(\epsilon_{N,i} = \left< \beta \gamma \right> \ \epsilon_i\)

where \(\beta \gamma = \left| \vec{p} \right| / (m_\mathrm{sp} c)\).

For a 2D cylindrical, axisymmetric geometry (axisym=True) this function returns the Lapostolle emittance¹^,², i.e.:

\(\epsilon = 4 \ \epsilon_i\)

Examples:

Calculate normalized emittance in 2D cyl. geometry

import ozzy as oz
import numpy as np

# Create a sample particle dataset
particles = oz.Dataset(
    {
        "x": ("pid", np.random.uniform(0, 10, 10000)),
        "r": ("pid", np.random.uniform(0, 5, 10000)),
        "px": ("pid", np.random.uniform(99, 101, 10000)),
        "pr": ("pid", np.random.uniform(-2e-4, 2e-4, 10000)),
        "q": ("pid", np.ones(10000)),
    },
    coords={"pid": np.arange(10000)},
    attrs={"pic_data_type": "part"}
)

emittance = particles.ozzy.get_emittance(axisym=True, x_var="r", p_all_vars=["px", "pr"])
# Returns Dataset with normalized emittance in k_p^(-1) rad

get_energy_spectrum §

get_energy_spectrum(
    axis_ds=None, nbins=None, ene_var="ene", w_var="q"
)

Calculate the energy spectrum of particles.

This method computes a histogram of particle energy, binning the energy values and summing the associated charge or weighting variable in each bin.

Parameters:

Name	Type	Description	Default
`axis_ds` §	`Dataset or None`	Dataset containing the energy axis to use for binning. Must have `ene_var` as a coordinate. If `None`, `nbins` must be provided. Note If the label and unit attributes exist in `axis_ds[ene_var]` (`'long_name'` and `'units'`, respectively), these attributes are adopted for the output dataset.	`None`
`nbins` §	`int or None`	Number of bins to use for the energy axis. Only used if `axis_ds` is `None`. If `None`, `axis_ds` must be provided.	`None`
`ene_var` §	`str`	Name of the energy variable in the dataset, default is `"ene"`.	`'ene'`
`w_var` §	`str`	Name of the weighting variable (typically charge) in the dataset, default is `"q"`.	`'q'`

Returns:

Type	Description
`Dataset`	A new dataset containing the energy spectrum with the following variables: - The weighting variable (e.g., `"q"`) containing the histogram of charge per energy bin - `"counts"` containing the number of particles in each energy bin

Notes

The absolute value of the weighting variable is used for the calculation.

Examples:

Basic usage with number of bins

import numpy as np
import ozzy as oz

# Create a sample particle dataset
rng = np.random.default_rng()
ds = oz.Dataset(
    {
        "ene": ("pid", rng.normal(100, 5, size=10000) ),
        "q": ("pid", rng.random(10000)),
    },
    coords={"pid": np.arange(10000)},
    attrs={"pic_data_type": "part"}
)

# Get energy spectrum
spectrum = ds.ozzy.get_energy_spectrum(nbins=100)
# Plot the result
spectrum.q.plot()

Using a custom energy axis

import numpy as np
import ozzy as oz

# Create a sample particle dataset
rng = np.random.default_rng()
ds = oz.Dataset(
    {
        "p1": ("pid", rng.lognormal(3.0, 1.0, size=10000) ),
        "weight": ("pid", rng.random(10000)),
    },
    coords={"pid": np.arange(10000)},
    attrs={"pic_data_type": "part"}
)

# Create a custom logarithmic energy axis
energy_axis = np.logspace(-1, 3, 50)  # 50 points from 0.1 to 1000
axis_ds = oz.Dataset(coords={"p1": energy_axis}, pic_data_type="grid")
axis_ds["p1"].attrs["long_name"] = r"$p_1$"
axis_ds["p1"].attrs["units"] = r"$m_\mathrm{sp} c$"

# Get energy spectrum using this axis
spectrum = ds.ozzy.get_energy_spectrum(axis_ds=axis_ds, ene_var="p1", w_var="weight")
# Plot the result
spectrum["weight"].plot(marker=".")
# Spectrum now contains the summed weights in each logarithmic energy bin

get_phase_space §

get_phase_space(
    vars,
    extents=None,
    nbins=200,
    r_var=None,
    t_var="t",
    w_var="q",
)

Generate a phase space grid from particle data.

Creates a gridded dataset by depositing particle quantities onto a 2D phase space.

Parameters:

Name	Type	Description	Default
`vars` §	`list[str]`	Variables to deposit onto phase space.	required
`extents` §	`dict[str, tuple[float, float]]`	Minimum and maximum extent for each variable. If not specified, extents are calculated from the data.	`None`
`nbins` §	`int \| dict[str, int]`	Number of bins for each variable. If `int`, the same number of bins is used for all variables.	`200`
`r_var` §	`str \| None`	Name of the variable representing particle radial positions. If provided and if one of `vars` is `r_var`, the particle weights are divided by this variable.	`None`
`t_var` §	`str`	Name of the time dimension in the input datasets.	`'t'`
`w_var` §	`str`	Name of the variable representing particle weights or particle charge.	`'q'`

Returns:

Type	Description
`Dataset`	Dataset with phase space data.

Examples:

Transverse phase space

import ozzy as oz
import numpy as np

# Create a sample particle dataset
ds = oz.Dataset(
    {
        "x1": ("pid", np.random.rand(10000)),
        "x2": ("pid", np.random.rand(10000)),
        "p1": ("pid", np.random.rand(10000)),
        "p2": ("pid", np.random.rand(10000)),
        "q": ("pid", np.ones(10000)),
    },
    coords={"pid": np.arange(10000)},
    pic_data_type="part",
    data_origin="ozzy",
)

ds_ps = ds.ozzy.get_phase_space(['p2', 'x2'], nbins=100)

get_slice_emittance §

get_slice_emittance(
    axis_ds=None,
    nbins=None,
    norm_emit=True,
    axisym=False,
    p_all_vars=["p1", "p2", "p3"],
    min_count=None,
    slice_var="x1_box",
    x_var="x2",
    p_var="p2",
    w_var="q",
)

Calculate the RMS slice emittance of particle data.

This method computes the slice emittance by binning particles along a specified variable and calculating the RMS emittance (normalized or geometric) for each slice. For axisymmetric beams, returns the Lapostolle emittance (see Notes below).

Warning

This method assumes that the particle dimension is "pid".

Parameters:

Name	Type	Description	Default
`axis_ds` §	`Dataset or None`	Dataset containing coordinate information for binning. If provided, bin edges are extracted from this dataset. Either `axis_ds` or `nbins` must be specified. Note If the label and unit attributes exist in `axis_ds[slice_var]` (`'long_name'` and `'units'`, respectively), these attributes are adopted for the output dataset.	`None`
`nbins` §	`int or None`	Number of bins to use for slicing. Either `axis_ds` or `nbins` must be specified.	`None`
`norm_emit` §	`bool`	Whether to calculate normalized emittance (multiplied by \(\left< \beta \gamma \right>\)).	`True`
`axisym` §	`bool`	Whether to apply Lapostolle factor of 4.	`False`
`p_all_vars` §	`list[str]`	List of names of momentum components. Note The components should be sorted, e.g. `["px", "py", "pz"]`. If `axisym=True`, only the two first components will be adopted.	`['p1', 'p2', 'p3']`
`min_count` §	`int or None`	Minimum number of particles required in each bin for valid calculation.	`None`
`slice_var` §	`str`	Variable name to use for slicing/binning the particles.	`'x1_box'`
`x_var` §	`str`	Variable name for the transverse position coordinate that should be used for emittance calculation.	`'x2'`
`p_var` §	`str`	Variable name for the transverse momentum coordinate that should be used for emittance calculation. This argument is only relevant when `axisym=False`, otherwise it is set to `p_all_vars[1]`.	`'p2'`
`w_var` §	`str`	Variable name for the particle weights/charges.	`'q'`

Returns:

Type	Description
`Dataset`	Dataset containing the calculated slice emittance and particle counts per bin. The emittance variable is named `"slice_emit_norm"` if `norm_emit=True`, otherwise `"slice_emit"`. Also includes a `"counts"` variable with particle counts per bin.

Notes

Particles are binned along the specified slice_var variable, and the emittance is computed for each binned ensemble.

The geometric emittance along a given transverse dimension \(i\) is calculated according to:

\(\epsilon_i = \sqrt{\left<x_i^2\right> \left<{x'_i}^2\right> - \left(x_i x'_i\right)^2}\)

where \(x_i\) is the particle position, and \(x'_i \approx p_i / p_\parallel\) is the trace for relativistic particles with longitudinal momentum \(p_\parallel\) and transverse momentum \(p_i \ll p_\parallel\). The angle brackets denote a weighted average over particles.

The normalized emittance (norm_emit=True, default) is calculated as:

\(\epsilon_{N,i} = \left< \beta \gamma \right> \ \epsilon_i\)

where \(\beta \gamma = \left| \vec{p} \right| / (m_\mathrm{sp} c)\).

For a 2D cylindrical, axisymmetric geometry (axisym=True) this function returns the Lapostolle emittance¹^,², i.e.:

\(\epsilon = 4 \ \epsilon_i\)

Examples:

Calculate normalized slice emittance in 2D cyl. geometry

import ozzy as oz
import numpy as np

# Create a sample particle dataset
particles = oz.Dataset(
    {
        "x": ("pid", np.random.uniform(0, 10, 10000)),
        "r": ("pid", np.random.uniform(0, 5, 10000)),
        "px": ("pid", np.random.uniform(99, 101, 10000)),
        "pr": ("pid", np.random.uniform(-2e-4, 2e-4, 10000)),
        "q": ("pid", np.ones(10000)),
    },
    coords={"pid": np.arange(10000)},
    attrs={"pic_data_type": "part"}
)

# Longitudinal axis along which to bin
axis = oz.utils.axis_from_extent(500, (0,10))
axis_ds = oz.Dataset({"x": axis}, pic_data_type = "grid")

emittance = particles.ozzy.get_slice_emittance(axis_ds=axis_ds, axisym=True, slice_var="x", x_var="r", p_all_vars=["px","pr"])
# Returns Dataset with normalized emittance in k_p^(-1) rad

get_weighted_median §

get_weighted_median(var, w_var='q', t_var='t')

Calculate the weighted median of a variable in the particle dataset.

This method computes the median of var weighted by the values in w_var, for each value in the time variable t_var (if the t_var dimension exists).

Parameters:

Name	Type	Description	Default
`var` §	`str`	Name of the variable for which to calculate the weighted median.	required
`w_var` §	`str`	Name of the weighting variable, by default `"q"`. The absolute value of this variable is used for weighting.	`'q'`
`t_var` §	`str`	Name of the time dimension to iterate over, by default `"t"`. If this dimension exists in the dataset, a weighted median is calculated for each time step.	`'t'`

Returns:

Type	Description
`DataArray`	DataArray containing the weighted median value(s). If `t_var` is present in the dataset, the result will have the same `t_var` dimension.

Notes

The weighted median is calculated by: 1. Sorting the data according to the variable of interest 2. Computing the cumulative sum of weights 3. Finding the point where the cumulative sum of weights reaches half of the total weight

For an odd number of observations, the midpoint value is used directly. For an even number, the average of the two middle values is used.

Examples:

Basic usage with particle energy

import numpy as np
import ozzy as oz

# Create a sample particle dataset
rng = np.random.default_rng(seed=42)
ds = oz.Dataset(
    {
        "energy": ("pid", rng.normal(100, 20, size=1000)),
        "q": ("pid", rng.random(1000)),
    },
    coords={"pid": np.arange(1000)},
    pic_data_type = "part",
)

# Calculate the weighted median of energy
median_energy = ds.ozzy.get_weighted_median(var="energy")
print(f"Weighted median energy: {median_energy.values:.2f}")
# Weighted median energy: ~100.00 (exact value will vary)

Time-dependent weighted median

import numpy as np
import ozzy as oz

# Create a sample particle dataset with time dimension
rng = np.random.default_rng(seed=42)
times = np.linspace(0, 10, 5)
energies = np.zeros((5, 100))

# Create time-dependent energies
for i, t in enumerate(times):
    energies[i] = rng.normal(100 + t*10, 20, size=100)

ds = oz.Dataset(
    {
        "energy": (["t", "pid"], energies),
        "q": (["t", "pid"], rng.random((5, 100))),
    },
    coords={
        "t": times,
        "pid": np.arange(100)
    },
    pic_data_type = "part",
)

# Calculate the time-dependent weighted median of energy
median_energy = ds.ozzy.get_weighted_median(var="energy")

# Plot the result
median_energy.plot()
# The plot will show the weighted median energy increasing over time

mean_std §

mean_std(vars, axes_ds, expand_time=True, axisym=False)

Calculate mean and standard deviation of variables.

Bins the particle data onto the grid specified by axes_ds and calculates the mean and standard deviation for each bin.

Parameters:

Name	Type	Description	Default
`vars` §	`str \| list[str]`	The variable(s) for which to calculate statistics.	required
`axes_ds` §	`Dataset \| DataArray \| Coordinates`	Data object containing the axes to use for the calculation (as xarray coordinates). Tip The axes object can be taken from an existing Dataset or DataArray via `axes_ds = <data_obj>.coords`.	required
`expand_time` §	`bool`	If `True`, statistics are calculated separately for each timestep.	`True`
`axisym` §	`bool`	If `True`, azimuthal symmetry is assumed.	`False`

Returns:

Type	Description
`Dataset`	Dataset containing the calculated mean and standard deviation of the particle variables.

Examples:

Get mean and std of 'x2' and 'p2'

import ozzy as oz
import numpy as np

# Create a sample particle dataset
ds = oz.Dataset(
    {
        "x1": ("pid", np.random.rand(10000)),
        "x2": ("pid", np.random.rand(10000)),
        "p1": ("pid", np.random.rand(10000)),
        "p2": ("pid", np.random.rand(10000)),
        "q": ("pid", np.ones(10000)),
    },
    coords={"pid": np.arange(10000)},
    pic_data_type="part",
    data_origin="ozzy",
)

# Create axes for binning
axes_ds = oz.Dataset(
    coords={
        "x1": np.linspace(0, 1, 21),
    },
    pic_data_type="grid",
    data_origin="ozzy",
)

# Calculate mean and standard deviation
ds_mean_std = ds.ozzy.mean_std(["x2", "p2"], axes_ds)

sample_particles §

sample_particles(n)

Downsample a particle Dataset by randomly choosing particles.

Parameters:

Name	Type	Description	Default
`n` §	`int`	Number of particles to sample.	required

Returns:

Type	Description
`Dataset`	Dataset with sampled particles.

Examples:

Sample 1000 particles

import ozzy as oz
import numpy as np


# Create a sample particle dataset
ds = oz.Dataset(
    {
        "x1": ("pid", np.random.rand(10000)),
        "x2": ("pid", np.random.rand(10000)),
        "p1": ("pid", np.random.rand(10000)),
        "p2": ("pid", np.random.rand(10000)),
        "q": ("pid", np.ones(10000)),
    },
    coords={"pid": np.arange(10000)},
    pic_data_type="part",
    data_origin="ozzy",
)

# Sample 1000 particles
ds_small = ds.ozzy.sample_particles(1000)
print(len(ds_small.pid))
# 1000

# Try to sample more particles than available
ds_all = ds.ozzy.sample_particles(20000)
# WARNING: number of particles to be sampled is larger than total particles. Proceeding without any sampling.
print(len(ds_all.pid))
# 10000

Methods for particle data§

bin_into_grid §

axes_ds §

t_var §

w_var §

r_var §

get_emittance §

norm_emit §

axisym §

p_all_vars §

x_var §

p_var §

w_var §

get_energy_spectrum §

axis_ds §

nbins §

ene_var §

w_var §

get_phase_space §

vars §

extents §

nbins §

r_var §

t_var §

w_var §

get_slice_emittance §

axis_ds §

nbins §

norm_emit §

axisym §

p_all_vars §

min_count §

slice_var §

x_var §

p_var §

w_var §

get_weighted_median §

var §

w_var §

t_var §

mean_std §

vars §

axes_ds §

expand_time §

axisym §

sample_particles §

n §

`axes_ds` §

`t_var` §

`w_var` §

`r_var` §

`norm_emit` §

`axisym` §

`p_all_vars` §

`x_var` §

`p_var` §

`w_var` §

`axis_ds` §

`nbins` §

`ene_var` §

`w_var` §

`vars` §

`extents` §

`nbins` §

`r_var` §

`t_var` §

`w_var` §

`axis_ds` §

`nbins` §

`norm_emit` §

`axisym` §

`p_all_vars` §

`min_count` §

`slice_var` §

`x_var` §

`p_var` §

`w_var` §

`var` §

`w_var` §

`t_var` §

`vars` §

`axes_ds` §

`expand_time` §

`axisym` §

`n` §