Skip to content

Data-type-specific methods§

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

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 easily obtained from a grid Dataset (for example field data) with 

axes_ds = fields_ds.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": np.linspace(0, 10, 101),
        "x2": np.linspace(0, 5, 51),
    },
    attrs={"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(
    xvar="x2",
    pvar="p2",
    wvar="q",
    p_longit="p1",
    axisym=False,
)

Calculate normalized RMS beam emittance.

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

Parameters:

Name Type Description Default
xvar §
 str

Variable name for position coordinate in Dataset

 'x2'
pvar §
 str

Variable name for momentum coordinate in Dataset

 'p2'
wvar §
 str

Variable name for particle weights in Dataset

 'q'
p_longit §
 str

Variable name for longitudinal momentum in Dataset

 'p1'
axisym §
 bool

If True, calculate Lapostolle emittance for axisymmetric beams

 False

Returns:

Type Description
DataArray

Normalized emittance with attributes containing units and label information
Notes

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

\(\varepsilon_{N,i} = \gamma \sqrt{\left<x_i^2\right> \left<{x'_i}^2\right> - \left(x_i x'_i\right)^2}\)

where \(\gamma\) is the average Lorentz factor, \(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\).

For a 2D cylindrical, axisymmetric geometry this function returns the Lapostolle emittance1,2, i.e.:

\(\varepsilon_N = 4 \ \varepsilon_{N,i}\)

  1. J. D. Lawson, P. M. Lapostolle, and R. L. Gluckstern, Particle Accelerators 5, 61–65 (1973), ↩

  2. P. M. Lapostolle, IEEE Transactions on Nuclear Science 18, 1101–1104 (1971) ↩

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(xvar="r", pvar="pr", p_longit="px", axisym=True)
# Returns DataArray with normalized emittance in k_p^(-1) rad

get_phase_space §

get_phase_space(
    vars, extents=None, nbins=200, axisym=False, r_var="x2"
)

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'

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)

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