Data-type-specific methods§

These methods can be accessed via the standard accessor way, i.e.:

DataArrayDataset

xarray.DataArray.ozzy.<method>

xarray.Dataset.ozzy.<method>

Note for developers

The following methods are defined within specific mixin classes: GridMixin (grid_mixin.py) and PartMixin (part_mixin.py).

Grid data§

GridMixin §

Mixin class for operations on grid-like data objects.

The methods in this class are accessible to a data object when <data_obj>.attrs['pic_data_type'] == 'grid'.

coords_from_extent §

coords_from_extent(mapping)

Add coordinates to DataArray | Dataset based on axis extents.

For each axis name and extent tuple in the mapping, get the axis values and assign them to a new coordinate in the data object.

Parameters:

Name	Type	Description	Default
`mapping` §	`dict[str, tuple[float, float]]`	Dictionary mapping axis names to (min, max) extents	required

Returns:

Name	Type	Description
`obj`	`Same type as self._obj`	Object with added coordinate values

Examples:

Example 1

import ozzy as oz
da = oz.DataArray(np.zeros((4,3)), dims=['x', 'y'], pic_data_type='grid')
mapping = {'x': (0, 1), 'y': (-1, 2)}
da = da.ozzy.coords_from_extent(mapping)

get_bin_edges §

get_bin_edges(time_dim='t')

Get bin edges along each spatial axis.

Calculates bin edges from coordinate values. This is useful for binning operations (see example below).

Parameters:

Name	Type	Description	Default
`time_dim` §	`str`	Name of time coordinate	`'t'`

Returns:

Type	Description
`list[ndarray]`	List of bin edges for each spatial axis

Examples:

Using numpy.histogramdd

import numpy as np
bin_edges = axes_ds.ozzy.get_bin_edges('t')
dist, edges = np.histogramdd(part_coords, bins=bin_edges, weights=ds_i[wvar])

get_space_dims §

get_space_dims(time_dim='t')

Get names of spatial dimensions.

Returns coordinate names that are not the time dimension.

Parameters:

Name	Type	Description	Default
`time_dim` §	`str`	Name of time coordinate	`'t'`

Returns:

Type	Description
`list[str]`	Spatial coordinate names

Examples:

Example 1

import ozzy as oz
ds = oz.Dataset(...)
spatial_dims = ds.ozzy.get_space_dims('t')
print(spatial_dims)

Particle data§

PartMixin §

Mixin class for operations on particle-like data objects.

The methods in this class are accessible to a data object when <data_obj>.attrs['pic_data_type'] == 'part'.

bin_into_grid §

bin_into_grid(
    axes_ds, time_dim="t", weight_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
`time_dim` §	`str`	Name of the time dimension in the input datasets.	`'t'`
`weight_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, 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, weight_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,
    all_pvars=["p1", "p2", "p3"],
    xvar="x2",
    pvar="p2",
    wvar="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`
`all_pvars` §	`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']`
`xvar` §	`str`	Variable name for position coordinate in Dataset that should be used for emittance calculation	`'x2'`
`pvar` §	`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 `all_pvars[1]`.	`'p2'`
`wvar` §	`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, xvar="r", all_pvars=["px", "pr"])
# Returns Dataset with normalized emittance in k_p^(-1) rad

get_energy_spectrum §

get_energy_spectrum(
    axis_ds=None, nbins=None, enevar="ene", wvar="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 `enevar` as a coordinate. If `None`, `nbins` must be provided. Note If the label and unit attributes exist in `axis_ds[enevar]` (`'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`
`enevar` §	`str`	Name of the energy variable in the dataset, default is `"ene"`.	`'ene'`
`wvar` §	`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, enevar="p1", wvar="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,
    axisym=False,
    r_var="x2",
    time_dim="t",
    weight_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`
`axisym` §	`bool`	Whether geometry is 2D cylindrical (axisymmetric), in which case the particle weights are divided by the radial coordinate (`r_var`).	`False`
`r_var` §	`str`	Name of the radial coordinate. This argument is ignored if `axisym = False`.	`'x2'`
`time_dim` §	`str`	Name of the time dimension in the input datasets.	`'t'`
`weight_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,
    all_pvars=["p1", "p2", "p3"],
    min_count=None,
    slice_var="x1_box",
    xvar="x2",
    pvar="p2",
    wvar="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`
`all_pvars` §	`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'`
`xvar` §	`str`	Variable name for the transverse position coordinate that should be used for emittance calculation.	`'x2'`
`pvar` §	`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 `all_pvars[1]`.	`'p2'`
`wvar` §	`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", xvar="r", all_pvars=["px","pr"])
# Returns Dataset with normalized emittance in k_p^(-1) rad

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

Data-type-specific methods§

Grid data§

GridMixin §

coords_from_extent §

mapping §

get_bin_edges §

time_dim §

get_space_dims §

time_dim §

Particle data§

PartMixin §

bin_into_grid §

axes_ds §

time_dim §

weight_var §

r_var §

get_emittance §

norm_emit §

axisym §

all_pvars §

xvar §

pvar §

wvar §

get_energy_spectrum §

axis_ds §

nbins §

enevar §

wvar §

get_phase_space §

vars §

extents §

nbins §

axisym §

r_var §

time_dim §

weight_var §

get_slice_emittance §

axis_ds §

nbins §

norm_emit §

axisym §

all_pvars §

min_count §

slice_var §

xvar §

pvar §

wvar §

mean_std §

vars §

axes_ds §

expand_time §

axisym §

sample_particles §

n §

`mapping` §

`time_dim` §

`time_dim` §

`axes_ds` §

`time_dim` §

`weight_var` §

`r_var` §

`norm_emit` §

`axisym` §

`all_pvars` §

`xvar` §

`pvar` §

`wvar` §

`axis_ds` §

`nbins` §

`enevar` §

`wvar` §

`vars` §

`extents` §

`nbins` §

`axisym` §

`r_var` §

`time_dim` §

`weight_var` §

`axis_ds` §

`nbins` §

`norm_emit` §

`axisym` §

`all_pvars` §

`min_count` §

`slice_var` §

`xvar` §

`pvar` §

`wvar` §

`vars` §

`axes_ds` §

`expand_time` §

`axisym` §

`n` §