Skip to content

Particles

pmd_beamphysics.ParticleGroup 

ParticleGroup(h5=None, data=None)

Particle Group class

Initialized on on openPMD beamphysics particle group:

  • h5: open h5 handle, or str that is a file
  • data: raw data

The required bunch data is stored in .data with keys

  • np.array: x, px, y, py, z, pz, t, status, weight
  • str: species

where:

  • x, y, z are positions in units of [m]
  • px, py, pz are momenta in units of [eV/c]
  • t is time in [s]
  • weight is the macro-charge weight in [C], used for all statistical calulations.
  • species is a proper species name: 'electron', etc.

Optional data:

  • np.array: id

where id is a list of unique integers that identify the particles.

Derived data can be computed as attributes:

  • .gamma, .beta, .beta_x, .beta_y, .beta_z: relativistic factors [1].
  • .r, .theta: cylidrical coordinates [m], [1]
  • .pr, .ptheta: momenta in the radial and angular coordinate directions [eV/c]
  • .Lz: angular momentum about the z axis [m*eV/c]
  • .energy : total energy [eV]
  • .kinetic_energy: total energy - mc^2 in [eV].
  • .higher_order_energy: total energy with quadratic fit in z or t subtracted [eV]
  • .p: total momentum in [eV/c]
  • .mass: rest mass in [eV]
  • .xp, .yp: Slopes \(x' = dx/dz = p_x/p_z\) and \(y' = dy/dz = p_y/p_z\) [1].

Normalized transvere coordinates can also be calculated as attributes:

  • .x_bar, .px_bar, .y_bar, .py_bar in [sqrt(m)]

The normalization is automatically calculated from the covariance matrix. See functions in .statistics for more advanced usage.

Their cooresponding amplitudes are:

.Jx, .Jy [m]

where Jx = (x_bar^2 + px_bar^2 )/2.

The momenta are normalized by the mass, so that <Jx> = norm_emit_x and similar for y.

Statistics of any of these are calculated with:

  • .min(X)
  • .max(X)
  • .ptp(X)
  • .avg(X)
  • .std(X)
  • .cov(X, Y, ...)
  • .histogramdd(X, Y, ..., bins=10, range=None)

with a string X as the name any of the properties above.

Useful beam physics quantities are given as attributes:

  • .norm_emit_x
  • .norm_emit_y
  • .norm_emit_4d
  • .higher_order_energy_spread
  • .average_current

Twiss parameters, including dispersion, for the \(x\) or \(y\) plane:

  • .twiss(plane='x', fraction=0.95, p0C=None)

For convenience, plane='xy' will calculate twiss for both planes.

Twiss matched particles, using a simple linear transformation:

  • .twiss_match(self, beta=None, alpha=None, plane='x', p0c=None, inplace=False)

The weight is required and must sum to > 0. The sum of the weights is in .charge. This can also be set: .charge = 1.234 # C, will rescale the .weight array

All attributes can be accessed with brackets: [key]

Additional keys are allowed for convenience: ['min_prop'] will return .min('prop') ['max_prop'] will return .max('prop') ['ptp_prop'] will return .ptp('prop') ['mean_prop'] will return .avg('prop') ['sigma_prop'] will return .std('prop') ['cov_prop1__prop2'] will return .cov('prop1', 'prop2')[0,1]

Units for all attributes can be accessed by:

  • .units(key)

Particles are often stored at the same time (i.e. from a t-based code), or with the same z position (i.e. from an s-based code.) Routines:

  • drift_to_z(z0)
  • drift_to_t(t0)

help to convert these. If no argument is given, particles will be drifted to the mean. Related properties are:

  • .in_t_coordinates returns True if all particles have the same \(t\) corrdinate
  • .in_z_coordinates returns True if all particles have the same \(z\) corrdinate

Convenient plotting is provided with:

  • .plot(...)
  • .slice_plot(...)

Use help(ParticleGroup.plot), etc. for usage.

Source code in pmd_beamphysics/particles.py 
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
def __init__(self, h5=None, data=None):
    if h5 and data:
        # TODO:
        # Allow merging or changing some array with extra data
        raise NotImplementedError("Cannot init on both h5 and data")

    if h5:
        # Allow filename
        if isinstance(h5, (str, pathlib.Path)):
            fname = os.path.expandvars(h5)
            assert os.path.exists(fname), f"File does not exist: {fname}"

            with File(fname, "r") as hh5:
                pp = particle_paths(hh5)
                assert len(pp) == 1, f"Number of particle paths in {h5}: {len(pp)}"
                data = load_bunch_data(hh5[pp[0]])

        elif isinstance(h5, File):
            pp = particle_paths(h5)
            assert len(pp) == 1, f"Number of particle paths in {h5}: {len(pp)}"
            data = load_bunch_data(h5[pp[0]])
        else:
            # Try dict
            data = load_bunch_data(h5)
    else:
        # Fill out data. Exclude species.
        data = full_data(data)
        species = list(set(data["species"]))

        # Allow for empty data (len=0). Otherwise, check species.
        if len(species) >= 1:
            assert len(species) == 1, f"mixed species are not allowed: {species}"
            data["species"] = species[0]

    self._settable_array_keys = [
        "x",
        "px",
        "y",
        "py",
        "z",
        "pz",
        "t",
        "status",
        "weight",
    ]
    # Optional data
    for k in ["id"]:
        if k in data:
            self._settable_array_keys.append(k)

    self._settable_scalar_keys = ["species"]
    self._settable_keys = self._settable_array_keys + self._settable_scalar_keys
    # Internal data. Only allow settable keys
    self._data = {k: data[k] for k in self._settable_keys}

Attributes

pmd_beamphysics.ParticleGroup.Jx property 
Jx

Normalized amplitude J in the x-px plane

pmd_beamphysics.ParticleGroup.Jy property 
Jy

Normalized amplitude J in the y-py plane

pmd_beamphysics.ParticleGroup.Lz property 
Lz

Angular momentum around the z axis in m*eV/c Lz = x * py - y * px

pmd_beamphysics.ParticleGroup.average_current property 
average_current

Simple average current = charge / dt in [A], with dt = (max_t - min_t) If particles are in \(t\) coordinates, will trydt = (max_z - min_z)*c_light*beta_z

pmd_beamphysics.ParticleGroup.beta property 
beta

Relativistic beta

pmd_beamphysics.ParticleGroup.beta_x property writable 
beta_x

Relativistic beta, x component

pmd_beamphysics.ParticleGroup.beta_y property writable 
beta_y

Relativistic beta, y component

pmd_beamphysics.ParticleGroup.beta_z property writable 
beta_z

Relativistic beta, z component

pmd_beamphysics.ParticleGroup.charge property writable 
charge

Total charge in C

pmd_beamphysics.ParticleGroup.data property 
data

Internal data dict

pmd_beamphysics.ParticleGroup.energy property 
energy

Total energy in eV

pmd_beamphysics.ParticleGroup.gamma property writable 
gamma

Relativistic gamma

pmd_beamphysics.ParticleGroup.higher_order_energy property 
higher_order_energy

Fits a quadratic (order=2) to the Energy vs. time, and returns the energy with this subtracted.

pmd_beamphysics.ParticleGroup.higher_order_energy_spread property 
higher_order_energy_spread

Legacy syntax to compute the standard deviation of higher_order_energy.

pmd_beamphysics.ParticleGroup.id property writable 
id

id integer

pmd_beamphysics.ParticleGroup.in_t_coordinates property 
in_t_coordinates

Returns True if all particles have the same t coordinate

pmd_beamphysics.ParticleGroup.in_z_coordinates property 
in_z_coordinates

Returns True if all particles have the same z coordinate

pmd_beamphysics.ParticleGroup.kinetic_energy property 
kinetic_energy

Kinetic energy in eV

pmd_beamphysics.ParticleGroup.mass property 
mass

Rest mass in eV

pmd_beamphysics.ParticleGroup.n_alive property 
n_alive

Number of alive particles, defined by status == 1

pmd_beamphysics.ParticleGroup.n_dead property 
n_dead

Number of alive particles, defined by status != 1

pmd_beamphysics.ParticleGroup.n_particle property 
n_particle

Total number of particles. Same as len

pmd_beamphysics.ParticleGroup.norm_emit_4d property 
norm_emit_4d

Normalized emittance in the xy planes (4D)

pmd_beamphysics.ParticleGroup.norm_emit_x property 
norm_emit_x

Normalized emittance in the x plane

pmd_beamphysics.ParticleGroup.norm_emit_y property 
norm_emit_y

Normalized emittance in the x plane

pmd_beamphysics.ParticleGroup.p property 
p

Total momemtum in eV/c

pmd_beamphysics.ParticleGroup.pr property 
pr

Momentum in the radial direction in eV/c r_hat = cos(theta) xhat + sin(theta) yhat pr = p dot r_hat

pmd_beamphysics.ParticleGroup.ptheta property 
ptheta

Momentum in the polar theta direction. theta_hat = -sin(theta) xhat + cos(theta) yhat ptheta = p dot theta_hat Note that Lz = r*ptheta

pmd_beamphysics.ParticleGroup.px property writable 
px

px coordinate in [eV/c]

pmd_beamphysics.ParticleGroup.px_bar property 
px_bar

Normalized px in units of sqrt(m)

pmd_beamphysics.ParticleGroup.py property writable 
py

py coordinate in [eV/c]

pmd_beamphysics.ParticleGroup.py_bar property 
py_bar

Normalized py in units of sqrt(m)

pmd_beamphysics.ParticleGroup.pz property writable 
pz

pz coordinate in [eV/c]

pmd_beamphysics.ParticleGroup.r property 
r

Radius in the xy plane: r = sqrt(x^2 + y^2) in m

pmd_beamphysics.ParticleGroup.species property writable 
species

species string

pmd_beamphysics.ParticleGroup.species_charge property 
species_charge

Species charge in C

pmd_beamphysics.ParticleGroup.status property writable 
status

status coordinate in [1]

pmd_beamphysics.ParticleGroup.t property writable 
t

t coordinate in [s]

pmd_beamphysics.ParticleGroup.theta property 
theta

Angle in xy plane: theta = arctan2(y, x) in radians

pmd_beamphysics.ParticleGroup.weight property writable 
weight

weight coordinate in [C]

pmd_beamphysics.ParticleGroup.x property writable 
x

x coordinate in [m]

pmd_beamphysics.ParticleGroup.x_bar property 
x_bar

Normalized x in units of sqrt(m)

pmd_beamphysics.ParticleGroup.xp property 
xp

x slope px/pz (dimensionless)

pmd_beamphysics.ParticleGroup.y property writable 
y

y coordinate in [m]

pmd_beamphysics.ParticleGroup.y_bar property 
y_bar

Normalized y in units of sqrt(m)

pmd_beamphysics.ParticleGroup.yp property 
yp

y slope py/pz (dimensionless)

pmd_beamphysics.ParticleGroup.z property writable 
z

z coordinate in [m]

Functions

pmd_beamphysics.ParticleGroup.assign_id 
assign_id()

Assigns unique ids, integers from 1 to n_particle

Source code in pmd_beamphysics/particles.py 
369
370
371
372
373
374
375
376
def assign_id(self):
    """
    Assigns unique ids, integers from 1 to n_particle

    """
    if "id" not in self._settable_array_keys:
        self._settable_array_keys.append("id")
    self.id = np.arange(1, self["n_particle"] + 1)
pmd_beamphysics.ParticleGroup.avg 
avg(key)

Statistical average

Source code in pmd_beamphysics/particles.py 
616
617
618
619
620
621
def avg(self, key):
    """Statistical average"""
    dat = self[key]  # equivalent to self.key for accessing properties above
    if np.isscalar(dat):
        return dat
    return np.average(dat, weights=self.weight)
pmd_beamphysics.ParticleGroup.bunching 
bunching(wavelength)

Calculate the normalized bunching parameter, which is the magnitude of the complex sum of weighted exponentials at a given point.

The formula for bunching is given by:

\[ B(z, \lambda) = \frac{\left|\sum w_i e^{i k z_i}\right|}{\sum w_i} \]

where: - \( z \) is the position array, - \( \lambda \) is the wavelength, - \( k = \frac{2\pi}{\lambda} \) is the wave number, - \( w_i \) are the weights.

Parameters:

Name Type Description Default
wavelength float

Wavelength of the wave.

 required

Returns:

Type Description
complex

The normalized bunching parameter.

Raises:

Type Description
ValueError

If wavelength is not a positive number.
Source code in pmd_beamphysics/particles.py 
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
def bunching(self, wavelength):
    r"""
    Calculate the normalized bunching parameter, which is the magnitude of the
    complex sum of weighted exponentials at a given point.

    The formula for bunching is given by:

    $$
    B(z, \lambda) = \frac{\left|\sum w_i e^{i k z_i}\right|}{\sum w_i}
    $$

    where:
    - \( z \) is the position array,
    - \( \lambda \) is the wavelength,
    - \( k = \frac{2\pi}{\lambda} \) is the wave number,
    - \( w_i \) are the weights.

    Parameters
    ----------
    wavelength : float
        Wavelength of the wave.


    Returns
    -------
    complex
        The normalized bunching parameter.

    Raises
    ------
    ValueError
        If `wavelength` is not a positive number.
    """

    if self.in_z_coordinates:
        # Approximate z
        z = self.t * self.avg("beta_z") * c_light
    else:
        z = self.z

    return statistics.bunching(z, wavelength, weight=self.weight)
pmd_beamphysics.ParticleGroup.copy 
copy()

Returns a deep copy

Source code in pmd_beamphysics/particles.py 
1251
1252
1253
def copy(self):
    """Returns a deep copy"""
    return deepcopy(self)
pmd_beamphysics.ParticleGroup.cov 
cov(*keys)

Covariance matrix from any properties

Example: P = ParticleGroup(h5) P.cov('x', 'px', 'y', 'py')

Source code in pmd_beamphysics/particles.py 
631
632
633
634
635
636
637
638
639
640
641
def cov(self, *keys):
    """
    Covariance matrix from any properties

    Example:
    P = ParticleGroup(h5)
    P.cov('x', 'px', 'y', 'py')

    """
    dats = np.array([self[key] for key in keys])
    return np.cov(dats, aweights=self.weight)
pmd_beamphysics.ParticleGroup.delta 
delta(key)

Attribute (array) relative to its mean

Source code in pmd_beamphysics/particles.py 
598
599
600
def delta(self, key):
    """Attribute (array) relative to its mean"""
    return self[key] - self.avg(key)
pmd_beamphysics.ParticleGroup.drift 
drift(delta_t)

Drifts particles by time delta_t

Source code in pmd_beamphysics/particles.py 
842
843
844
845
846
847
848
849
def drift(self, delta_t):
    """
    Drifts particles by time delta_t
    """
    self.x = self.x + self.beta_x * c_light * delta_t
    self.y = self.y + self.beta_y * c_light * delta_t
    self.z = self.z + self.beta_z * c_light * delta_t
    self.t = self.t + delta_t
pmd_beamphysics.ParticleGroup.drift_to_t 
drift_to_t(t=None)

Drifts all particles to the same t

If no z is given, particles will be drifted to the average t

Source code in pmd_beamphysics/particles.py 
859
860
861
862
863
864
865
866
867
868
869
870
def drift_to_t(self, t=None):
    """
    Drifts all particles to the same t

    If no z is given, particles will be drifted to the average t
    """
    if t is None:
        t = self.avg("t")
    dt = t - self.t
    self.drift(dt)
    # Fix t to be exactly this value
    self.t = np.full(self.n_particle, t)
pmd_beamphysics.ParticleGroup.fractional_split 
fractional_split(fractions, key)

Split particles based on a given array key and a list of specified fractions or a single fraction.

Parameters:

Name Type Description Default
fractions float or list of float

A fraction or a list of fractions used for splitting the particles. All values must be between 0 and 1 (exclusive).

 required
key str

The attribute of particles to be used for sorting and splitting (e.g., 'z' for longitudinal position).

 required

Returns:

Name Type Description
particle_groups list of ParticleGroup

A list of ParticleGroup objects, each representing a subset of particles based on the specified fractions.
Description

This function splits the given group of particles into multiple subsets based on the provided attribute (e.g., position). The splits are determined such that each specified fraction of the total particle weights is separated. The function first sorts the particles by the specified key, computes the cumulative sum of weights, and determines the split values. It then returns a list of ParticleGroup objects representing the split subsets.

Source code in pmd_beamphysics/particles.py 
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
def fractional_split(self, fractions: Union[float, int, list], key: str):
    """
    Split particles based on a given array key and a list of specified fractions or a single fraction.

    Parameters
    ----------
    fractions : float or list of float
        A fraction or a list of fractions used for splitting the particles. All values must be between 0 and 1 (exclusive).

    key : str
        The attribute of particles to be used for sorting and splitting (e.g., 'z' for longitudinal position).

    Returns
    -------
    particle_groups : list of ParticleGroup
        A list of ParticleGroup objects, each representing a subset of particles based on the specified fractions.

    Description
    -----------
    This function splits the given group of particles into multiple subsets based on the provided attribute (e.g., position).
    The splits are determined such that each specified fraction of the total particle weights is separated.
    The function first sorts the particles by the specified key, computes the cumulative sum of weights,
    and determines the split values. It then returns a list of ParticleGroup objects representing the split subsets.

    """

    # Ensure fractions is a list
    if isinstance(fractions, (float, int)):
        fractions = [fractions]

    # Validate fraction values
    if any(f <= 0 or f >= 1 for f in fractions):
        raise ValueError("All fraction values must be between 0 and 1 (exclusive)")

    # Sort particles by the specified key
    ixs = np.argsort(self[key])
    sorted_particles = self[ixs]

    # Sorted weights
    ws = sorted_particles.weight
    total_weight = np.sum(ws)
    cw = np.cumsum(ws) / total_weight

    # Use vectorized searchsorted to determine split indices
    fractions = np.array(fractions)
    split_indices = np.searchsorted(cw, fractions, side="right")

    # Create ParticleGroup subsets for each split
    particle_groups = []
    previous_index = 0
    for isplit in split_indices:
        particle_groups.append(sorted_particles[previous_index:isplit])
        previous_index = isplit

    # Add the remaining particles to the last group
    if previous_index < len(sorted_particles):
        particle_groups.append(sorted_particles[previous_index:])

    return particle_groups
pmd_beamphysics.ParticleGroup.from_bmad classmethod 
from_bmad(bmad_dict)

Convert Bmad particle data as a dict to ParticleGroup data.

See: ParticleGroup.to_bmad or particlegroup_to_bmad

Parameters:

Name Type Description Default
bmad_data

Dict with keys: 'x' 'px' 'y' 'py' 'z' 'pz', 'charge' 'species', 'tref' 'state'

 required

Returns:

Type Description
ParticleGroup
Source code in pmd_beamphysics/particles.py 
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
@classmethod
@functools.wraps(bmad.bmad_to_particlegroup_data)
def from_bmad(cls, bmad_dict):
    """
    Convert Bmad particle data as a dict
    to ParticleGroup data.

    See: ParticleGroup.to_bmad or particlegroup_to_bmad

    Parameters
    ----------
    bmad_data: dict
        Dict with keys:
        'x'
        'px'
        'y'
        'py'
        'z'
        'pz',
        'charge'
        'species',
        'tref'
        'state'

    Returns
    -------
    ParticleGroup
    """
    data = bmad.bmad_to_particlegroup_data(bmad_dict)
    return cls(data=data)
pmd_beamphysics.ParticleGroup.higher_order_energy_calc 
higher_order_energy_calc(order=2)

Fits a polynmial with order order to the Energy vs. time, , and returns the energy with this subtracted.

Source code in pmd_beamphysics/particles.py 
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
def higher_order_energy_calc(self, order=2):
    """
    Fits a polynmial with order `order` to the Energy vs. time, , and returns the energy with this subtracted.
    """
    # order=2
    if self.std("z") < 1e-12:
        # must be at a screen. Use t
        t = self.t
    else:
        # All particles at the same time. Use z to calc t
        t = self.z / c_light
    energy = self.energy

    best_fit_coeffs = np.polynomial.polynomial.polyfit(t, energy, order)
    best_fit = np.polynomial.polynomial.polyval(t, best_fit_coeffs)
    return energy - best_fit
pmd_beamphysics.ParticleGroup.histogramdd 
histogramdd(*keys, bins=10, range=None)

Wrapper for numpy.histogramdd, but accepts property names as keys.

Computes the multidimensional histogram of keys. Internally uses weights.

Example: P.histogramdd('x', 'y', bins=50) Returns: np.array with shape 50x50, edge list

Source code in pmd_beamphysics/particles.py 
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
def histogramdd(self, *keys, bins=10, range=None):
    """
    Wrapper for numpy.histogramdd, but accepts property names as keys.

    Computes the multidimensional histogram of keys. Internally uses weights.

    Example:
        P.histogramdd('x', 'y', bins=50)
    Returns:
        np.array with shape 50x50, edge list

    """
    H, edges = np.histogramdd(
        np.array([self[k] for k in list(keys)]).T,
        weights=self.weight,
        bins=bins,
        range=range,
    )

    return H, edges
pmd_beamphysics.ParticleGroup.linear_point_transform 
linear_point_transform(mat3)

Applies a linear transformation to the particle's spatial coordinates and their conjugate momenta.

The spatial coordinates (x, y, z) are transformed using: [x', y', z'] = mat3 @ [x, y, z]

The conjugate momenta (px, py, pz) are transformed using: [px', py', pz'] = inv(mat3.T) @ [px, py, pz]

Parameters:

Name Type Description Default
mat3 array-like of shape (3, 3)

A 3x3 transformation matrix. Must be convertible to a NumPy array. Coordinates must be ordered as (x, y, z).

 required

Raises:

Type Description
ValueError

If mat3 is not a 3x3 matrix.
LinAlgError

If the transformation matrix is singular (i.e., not invertible).
Notes

This method modifies the object in-place.

Source code in pmd_beamphysics/particles.py 
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
def linear_point_transform(
    self, mat3: Union[np.ndarray, Sequence[Sequence[float]]]
) -> None:
    """
    Applies a linear transformation to the particle's spatial coordinates and
    their conjugate momenta.

    The spatial coordinates (x, y, z) are transformed using:
        [x', y', z'] = mat3 @ [x, y, z]

    The conjugate momenta (px, py, pz) are transformed using:
        [px', py', pz'] = inv(mat3.T) @ [px, py, pz]

    Parameters
    ----------
    mat3 : array-like of shape (3, 3)
        A 3x3 transformation matrix. Must be convertible to a NumPy array.
        Coordinates must be ordered as (x, y, z).

    Raises
    ------
    ValueError
        If `mat3` is not a 3x3 matrix.
    LinAlgError
        If the transformation matrix is singular (i.e., not invertible).

    Notes
    -----
    This method modifies the object in-place.
    """
    mat3 = np.asarray(mat3, dtype=float)
    if mat3.shape != (3, 3):
        raise ValueError("mat3 must be a 3x3 matrix.")

    # Grab original coordinates
    x, y, z = self.x, self.y, self.z
    px, py, pz = self.px, self.py, self.pz

    # Transform spatial coordinates: [x', y', z'] = mat3 @ [x, y, z]
    self.x = mat3[0, 0] * x + mat3[0, 1] * y + mat3[0, 2] * z
    self.y = mat3[1, 0] * x + mat3[1, 1] * y + mat3[1, 2] * z
    self.z = mat3[2, 0] * x + mat3[2, 1] * y + mat3[2, 2] * z

    # Transform momenta using inverse-transpose: [px', py', pz'] = inv(mat3.T) @ [px, py, pz]
    # Calculate inverse-transpose matrix elements
    inv_mat3_T = np.linalg.inv(mat3.T)
    self.px = inv_mat3_T[0, 0] * px + inv_mat3_T[0, 1] * py + inv_mat3_T[0, 2] * pz
    self.py = inv_mat3_T[1, 0] * px + inv_mat3_T[1, 1] * py + inv_mat3_T[1, 2] * pz
    self.pz = inv_mat3_T[2, 0] * px + inv_mat3_T[2, 1] * py + inv_mat3_T[2, 2] * pz
pmd_beamphysics.ParticleGroup.max 
max(key)

Maximum of any key

Source code in pmd_beamphysics/particles.py 
608
609
610
def max(self, key):
    """Maximum of any key"""
    return np.max(self[key])
pmd_beamphysics.ParticleGroup.min 
min(key)

Minimum of any key

Source code in pmd_beamphysics/particles.py 
604
605
606
def min(self, key):
    """Minimum of any key"""
    return np.min(self[key])  # was: getattr(self, key)
pmd_beamphysics.ParticleGroup.plot 
plot(key1='x', key2=None, bins=None, *, xlim=None, ylim=None, return_figure=False, tex=True, nice=True, ellipse=False, **kwargs)

1d or 2d density plot.

If one key is given, this will plot the density of that key. Example: .plot('x')

If two keys arg given, this will plot a 2d marginal plot. Example: .plot('x', 'px')

Parameters:

Name Type Description Default
particle_group

The object to plot

 required
key1

Key to bin on the x-axis

 'x'
key2

Key to bin on the y-axis.

 None
bins

Number of bins. If None, this will use a heuristic: bins = sqrt(n_particle/4)

 None
xlim

Manual setting of the x-axis limits. Note that these are in raw, unscaled units.

 None
ylim

Manual setting of the y-axis limits. Note that these are in raw, unscaled units.

 None
tex

Use TEX for labels

 True
nice

Scale to nice units

 True
ellipse

If True, plot an ellipse representing the 2x2 sigma matrix

 False
return_figure

If true, return a matplotlib.figure.Figure object

 False
**kwargs

Any additional kwargs to send to the the plot in: plt.subplots(**kwargs)

 {}

Returns:

Type Description
None or fig: matplotlib.figure.Figure

This only returns a figure object if return_figure=T, otherwise returns None
Source code in pmd_beamphysics/particles.py 
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
def plot(
    self,
    key1="x",
    key2=None,
    bins=None,
    *,
    xlim=None,
    ylim=None,
    return_figure=False,
    tex=True,
    nice=True,
    ellipse=False,
    **kwargs,
):
    """
    1d or 2d density plot.

    If one key is given, this will plot the density of that key.
    Example:
        .plot('x')

    If two keys arg given, this will plot a 2d marginal plot.
    Example:
        .plot('x', 'px')


    Parameters
    ----------
    particle_group: ParticleGroup
        The object to plot

    key1: str, default = 't'
        Key to bin on the x-axis

    key2: str, default = None
        Key to bin on the y-axis.

    bins: int, default = None
       Number of bins. If None, this will use a heuristic: bins = sqrt(n_particle/4)

    xlim: tuple, default = None
        Manual setting of the x-axis limits. Note that these are in raw, unscaled units.

    ylim: tuple, default = None
        Manual setting of the y-axis limits. Note that these are in raw, unscaled units.

    tex: bool, default = True
        Use TEX for labels

    nice: bool, default = True
        Scale to nice units

    ellipse: bool, default = True
        If True, plot an ellipse representing the
        2x2 sigma matrix

    return_figure: bool, default = False
        If true, return a matplotlib.figure.Figure object

    **kwargs
        Any additional kwargs to send to the the plot in: plt.subplots(**kwargs)


    Returns
    -------
    None or fig: matplotlib.figure.Figure
        This only returns a figure object if return_figure=T, otherwise returns None

    """

    if not key2:
        fig = density_plot(
            self, key=key1, bins=bins, xlim=xlim, tex=tex, nice=nice, **kwargs
        )
    else:
        fig = marginal_plot(
            self,
            key1=key1,
            key2=key2,
            bins=bins,
            xlim=xlim,
            ylim=ylim,
            tex=tex,
            nice=nice,
            ellipse=ellipse,
            **kwargs,
        )

    if return_figure:
        return fig
pmd_beamphysics.ParticleGroup.ptp 
ptp(key)

Peak-to-Peak = max - min of any key

Source code in pmd_beamphysics/particles.py 
612
613
614
def ptp(self, key):
    """Peak-to-Peak = max - min of any key"""
    return np.ptp(self[key])
pmd_beamphysics.ParticleGroup.rotate 
rotate(*, x_rot=0.0, y_rot=0.0, z_rot=0.0, order='zxy', xc=0.0, yc=0.0, zc=0.0)

Rotate the beam inplace by the specified angles first around the z axis, then the x axis, finally around the y axis.

Parameters:

Name Type Description Default
x_rot float

Rotation around the x axis, radians

 0.0
y_rot float

Rotation around the y axis, radians

 0.0
z_rot float

Rotation around the z axis, radians

 0.0
order str

A 3-character string specifying the rotation order (e.g., 'zxy', 'zyx'). Each character must be one of 'x', 'y', or 'z'.

 'zxy'
xc float

Center of rotation, x coordinate. Default to zero

 0.0
yc float

Center of rotation, x coordinate. Default to zero

 0.0
zc float

Center of rotation, x coordinate. Default to zero

 0.0
Notes

This method modifies the object in-place.

Source code in pmd_beamphysics/particles.py 
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
def rotate(
    self,
    *,
    x_rot: float = 0.0,
    y_rot: float = 0.0,
    z_rot: float = 0.0,
    order: str = "zxy",
    xc: float = 0.0,
    yc: float = 0.0,
    zc: float = 0.0,
) -> None:
    """
    Rotate the beam inplace by the specified angles first around the z axis, then the x axis, finally around the y axis.

    Parameters
    ----------
    x_rot : float, optional
        Rotation around the x axis, radians
    y_rot : float, optional
        Rotation around the y axis, radians
    z_rot : float, optional
        Rotation around the z axis, radians
    order : str
        A 3-character string specifying the rotation order (e.g., 'zxy', 'zyx').
        Each character must be one of 'x', 'y', or 'z'.
    xc : float
        Center of rotation, x coordinate. Default to zero
    yc : float
        Center of rotation, x coordinate. Default to zero
    zc : float
        Center of rotation, x coordinate. Default to zero

    Notes
    -----
    This method modifies the object in-place.
    """
    # Special case transformations
    if (x_rot == 0.0) and (y_rot == 0.0):
        return self.rotate_z(z_rot, xc=xc, yc=yc)
    if (x_rot == 0.0) and (z_rot == 0.0):
        return self.rotate_y(y_rot, xc=xc, zc=zc)
    if (z_rot == 0.0) and (y_rot == 0.0):
        return self.rotate_x(x_rot, yc=yc, zc=zc)

    # Shift beam
    self.x = self.x - xc
    self.y = self.y - yc
    self.z = self.z - zc

    # Perform rotation
    self.linear_point_transform(
        get_rotation_matrix(x_rot=x_rot, y_rot=y_rot, z_rot=z_rot, order=order)
    )

    # Shift back
    self.x = self.x + xc
    self.y = self.y + yc
    self.z = self.z + zc
pmd_beamphysics.ParticleGroup.rotate_x 
rotate_x(theta, yc=0.0, zc=0.0)

Rotate the object about the x-axis by angle theta (radians). Affects y–z coordinates and py–pz momenta.

Parameters:

Name Type Description Default
theta float

Angle of rotation, radians

 required
yc float

Center of rotation, y coordinate. Default to zero

 0.0
zc float

Center of rotation, z coordinate. Default to zero

 0.0
Source code in pmd_beamphysics/particles.py 
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
def rotate_x(
    self,
    theta: float,
    yc: float = 0.0,
    zc: float = 0.0,
) -> None:
    """
    Rotate the object about the x-axis by angle `theta` (radians).
    Affects y–z coordinates and py–pz momenta.

    Parameters
    ----------
    theta : float
        Angle of rotation, radians
    yc : float
        Center of rotation, y coordinate. Default to zero
    zc : float
        Center of rotation, z coordinate. Default to zero
    """
    # Shift beam
    self.y = self.y - yc
    self.z = self.z - zc

    c = np.cos(theta)
    s = np.sin(theta)

    y, z = self.y, self.z
    py, pz = self.py, self.pz

    self.y = c * y - s * z
    self.z = s * y + c * z

    self.py = c * py - s * pz
    self.pz = s * py + c * pz

    # Shift back
    self.y = self.y + yc
    self.z = self.z + zc
pmd_beamphysics.ParticleGroup.rotate_y 
rotate_y(theta, xc=0.0, zc=0.0)

Rotate the object about the y-axis by angle theta (radians). Affects z–x coordinates and pz–px momenta.

Parameters:

Name Type Description Default
theta float

Angle of rotation, radians

 required
xc float

Center of rotation, x coordinate. Default to zero

 0.0
zc float

Center of rotation, z coordinate. Default to zero

 0.0
Source code in pmd_beamphysics/particles.py 
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
def rotate_y(
    self,
    theta: float,
    xc: float = 0.0,
    zc: float = 0.0,
) -> None:
    """
    Rotate the object about the y-axis by angle `theta` (radians).
    Affects z–x coordinates and pz–px momenta.

    Parameters
    ----------
    theta : float
        Angle of rotation, radians
    xc : float
        Center of rotation, x coordinate. Default to zero
    zc : float
        Center of rotation, z coordinate. Default to zero
    """
    # Shift beam
    self.x = self.x - xc
    self.z = self.z - zc

    c = np.cos(theta)
    s = np.sin(theta)

    z, x = self.z, self.x
    pz, px = self.pz, self.px

    self.z = c * z - s * x
    self.x = s * z + c * x

    self.pz = c * pz - s * px
    self.px = s * pz + c * px

    # Shift beam
    self.x = self.x + xc
    self.z = self.z + zc
pmd_beamphysics.ParticleGroup.rotate_z 
rotate_z(theta, xc=0.0, yc=0.0)

Rotate the object about the z-axis by angle theta (radians). Affects x–y coordinates and px–py momenta.

Parameters:

Name Type Description Default
theta float

Angle of rotation, radians

 required
xc float

Center of rotation, x coordinate. Default to zero

 0.0
yc float

Center of rotation, y coordinate. Default to zero

 0.0
Source code in pmd_beamphysics/particles.py 
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
def rotate_z(
    self,
    theta: float,
    xc: float = 0.0,
    yc: float = 0.0,
) -> None:
    """
    Rotate the object about the z-axis by angle `theta` (radians).
    Affects x–y coordinates and px–py momenta.

    Parameters
    ----------
    theta : float
        Angle of rotation, radians
    xc : float
        Center of rotation, x coordinate. Default to zero
    yc : float
        Center of rotation, y coordinate. Default to zero
    """
    # Shift beam
    self.x = self.x - xc
    self.y = self.y - yc

    c = np.cos(theta)
    s = np.sin(theta)

    x, y = self.x, self.y
    px, py = self.px, self.py

    self.x = c * x - s * y
    self.y = s * x + c * y

    self.px = c * px - s * py
    self.py = s * px + c * py

    # Shift beam
    self.x = self.x + xc
    self.y = self.y + yc
pmd_beamphysics.ParticleGroup.slice_statistics 
slice_statistics(*keys, n_slice=100, slice_key=None)

Slice statistics

Source code in pmd_beamphysics/particles.py 
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
def slice_statistics(self, *keys, n_slice=100, slice_key=None):
    """
    Slice statistics

    """

    if slice_key is None:
        if self.in_t_coordinates:
            slice_key = "z"

        else:
            slice_key = "t"

    if slice_key in ("t", "delta_t"):
        density_name = "current"
    else:
        density_name = "density"

    keys = set(keys)
    keys.add("mean_" + slice_key)
    keys.add("ptp_" + slice_key)
    keys.add("charge")
    slice_dat = slice_statistics(
        self, n_slice=n_slice, slice_key=slice_key, keys=keys
    )

    slice_dat[density_name] = slice_dat["charge"] / slice_dat["ptp_" + slice_key]

    return slice_dat
pmd_beamphysics.ParticleGroup.std 
std(key)

Standard deviation (actually sample)

Source code in pmd_beamphysics/particles.py 
623
624
625
626
627
628
629
def std(self, key):
    """Standard deviation (actually sample)"""
    dat = self[key]
    if np.isscalar(dat):
        return 0
    avg_dat = self.avg(key)
    return np.sqrt(np.average((dat - avg_dat) ** 2, weights=self.weight))
pmd_beamphysics.ParticleGroup.twiss 
twiss(plane='x', fraction=1, p0c=None)

Returns Twiss and Dispersion dict.

plane can be:

'x', 'y', or 'xy'

Optionally a fraction of the particles, based on amplitiude, can be specified.

Source code in pmd_beamphysics/particles.py 
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
def twiss(self, plane="x", fraction=1, p0c=None):
    """
    Returns Twiss and Dispersion dict.

    plane can be:

    `'x'`, `'y'`, or `'xy'`

    Optionally a fraction of the particles, based on amplitiude, can be specified.
    """
    d = {}
    for p in plane:
        d.update(
            particle_twiss_dispersion(self, plane=p, fraction=fraction, p0c=p0c)
        )
    return d
pmd_beamphysics.ParticleGroup.twiss_match 
twiss_match(beta=None, alpha=None, plane='x', p0c=None, inplace=False)

Returns a ParticleGroup with requested Twiss parameters.

See: statistics.matched_particles

Source code in pmd_beamphysics/particles.py 
697
698
699
700
701
702
703
704
705
706
def twiss_match(self, beta=None, alpha=None, plane="x", p0c=None, inplace=False):
    """
    Returns a ParticleGroup with requested Twiss parameters.

    See: statistics.matched_particles
    """

    return matched_particles(
        self, beta=beta, alpha=alpha, plane=plane, inplace=inplace
    )
pmd_beamphysics.ParticleGroup.units 
units(key)

Returns the units of any key

Source code in pmd_beamphysics/particles.py 
393
394
395
def units(self, key):
    """Returns the units of any key"""
    return pg_units(key)
pmd_beamphysics.ParticleGroup.write 
write(h5, name=None)

Write particle data to an HDF5 file or group in openPMD format.

Parameters:

Name Type Description Default
h5 str or Path or File or Group

Target for writing: - str or pathlib.Path: path to a new HDF5 file to be created (opened in mode "w"). Environment variables in path strings are expanded (via os.path.expandvars) before file creation. - h5py.File: an already opened file handle; a new "particles" group is created. - h5py.Group: an existing group; assumed to be an appropriate location to write data. This is for advanced users who know what they are doing. Any other object is treated as a group-like handle compatible with write_pmd_bunch.

 required
name str

Name for the subgroup/bunch written inside the "particles" group (or provided group). If None, write_pmd_bunch will write directly to the "particles" group.

 None

Returns:

Type Description
None

Examples:

Write to a new file:

>>> P.write("beam.h5")

Write to an already opened file:

>>> import h5py
>>> with h5py.File("beam.h5", "w") as f:
...     P.write(f)

Write into an existing group and name the bunch:

>>> import h5py
>>> with h5py.File("beam.h5", "w") as f:
...     grp = f["particles"]
...     particles.write(grp, name="bunch1")
Source code in pmd_beamphysics/particles.py 
 985
 986
 987
 988
 989
 990
 991
 992
 993
 994
 995
 996
 997
 998
 999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
def write(self, h5, name=None) -> None:
    """
    Write particle data to an HDF5 file or group in openPMD format.

    Parameters
    ----------
    h5 : str or pathlib.Path or h5py.File or h5py.Group
        Target for writing:
        - str or pathlib.Path: path to a new HDF5 file to be created (opened in mode "w").
          Environment variables in path strings are expanded (via
          os.path.expandvars) before file creation.
        - h5py.File: an already opened file handle; a new "particles" group is created.
        - h5py.Group: an existing group; assumed to be an appropriate location to write data.
            This is for advanced users who know what they are doing.
            Any other object is treated as a group-like handle compatible with `write_pmd_bunch`.
    name : str, optional
        Name for the subgroup/bunch written inside the "particles" group (or provided group).
        If None, `write_pmd_bunch` will write directly to the "particles" group.

    Returns
    -------
    None

    Examples
    --------
    Write to a new file:
    >>> P.write("beam.h5")

    Write to an already opened file:
    >>> import h5py
    >>> with h5py.File("beam.h5", "w") as f:
    ...     P.write(f)

    Write into an existing group and name the bunch:
    >>> import h5py
    >>> with h5py.File("beam.h5", "w") as f:
    ...     grp = f["particles"]
    ...     particles.write(grp, name="bunch1")
    """
    if isinstance(h5, (str, pathlib.Path)):
        fname = os.path.expandvars(h5)
        g = File(fname, "w")
        pmd_init(g, basePath="/", particlesPath="particles")
        g = g.create_group("particles")
    elif isinstance(h5, File):
        pmd_init(h5, basePath="/", particlesPath="particles")
        g = h5.create_group("particles")
    else:
        g = h5

    write_pmd_bunch(g, self, name=name)