Utils

utils.bmadx.bmad_to_cheetah_coords(bmad_coords: Tensor, p0c: Tensor, mc2: Tensor) → Tensor

Transforms Bmad coordinates to Cheetah coordinates.

Parameters:

  • bmad_coords – 6-dimensional particle vectors in Bmad coordinates.

  • p0c – Reference momentum in eV/c.

utils.bmadx.bmad_to_cheetah_z_pz(z: Tensor, pz: Tensor, p0c: Tensor, mc2: float) → tuple[Tensor]

Transforms Bmad longitudinal coordinates to Cheetah coordinates and computes reference energy.

Parameters:

  • z – Bmad longitudinal coordinate (c*delta_t).

  • pz – Bmad longitudinal momentum (delta_E/p0c).

  • p0c – Reference momentum in eV/c.

  • mc2 – Particle mass in eV/c^2.

utils.bmadx.calculate_quadrupole_coefficients(k1: Tensor, length: Tensor, rel_p: Tensor, eps: float = 2.220446049250313e-16) → tuple[list[Tensor], list[Tensor]]

Returns 2x2 transfer matrix elements aij and the coefficients to calculate the change in z position.

NOTE: Accumulated error due to machine epsilon.

Parameters:

  • k1 – Quadrupole strength (k1 > 0 ==> defocus).

  • length – Quadrupole length.

  • rel_p – Relative momentum P/P0.

  • eps – Machine precision epsilon, default to double precision.

Returns:

Tuple of transfer matrix elements and coefficients. a11, a12, a21, a22: Transfer matrix elements. c1, c2, c3: Second order derivatives of z such that

z = c1 * x_0^2 + c2 * x_0 * px_0 + c3 * px_0^2.

utils.bmadx.cheetah_to_bmad_coords(cheetah_coords: Tensor, ref_energy: Tensor, mc2: Tensor) → Tensor

Transforms Cheetah coordinates to Bmad coordinates.

Parameters:

  • cheetah_coords – 7-dimensional particle vectors in Cheetah coordinates.

  • ref_energy – Reference energy in eV.

utils.bmadx.cheetah_to_bmad_z_pz(tau: Tensor, delta: Tensor, ref_energy: Tensor, mc2: float) → Tensor

Transforms Cheetah longitudinal coordinates to Bmad coordinates and computes p0c.

Parameters:

  • tau – Cheetah longitudinal coordinate (c*delta_t).

  • delta – Cheetah longitudinal momentum (delta_E/p0c).

  • ref_energy – Reference energy in eV.

  • mc2 – Particle mass in eV/c^2.

utils.bmadx.cosc(x)

cosc(x) = (cos(x)-1)/x^2 = -1/2 [sinc(x/2)]^2

utils.bmadx.low_energy_z_correction(pz: Tensor, p0c: Tensor, mc2: Tensor, ds: Tensor) → Tensor

Corrects the change in z-coordinate due to speed < speed_of_light.

Parameters:

  • pz – Particle longitudinal momentum.

  • p0c – Reference particle momentum in eV.

  • mc2 – Particle mass in eV.

  • ds – Drift length.

Returns:

dz=(ds-d_particle) + ds*(beta - beta_ref)/beta_ref

utils.bmadx.offset_particle_set(x_offset: Tensor, y_offset: Tensor, tilt: Tensor, x_lab: Tensor, px_lab: Tensor, y_lab: Tensor, py_lab: Tensor) → list[Tensor]

Transforms particle coordinates from lab to element frame.

Parameters:

  • x_offset – Element x-coordinate offset.

  • y_offset – Element y-coordinate offset.

  • tilt – Tilt angle (rad).

  • x_lab – x-coordinate in lab frame.

  • px_lab – x-momentum in lab frame.

  • y_lab – y-coordinate in lab frame.

  • py_lab – y-momentum in lab frame.

Returns:

x, px, y, py coordinates in element frame.

utils.bmadx.offset_particle_unset(x_offset: Tensor, y_offset: Tensor, tilt: Tensor, x_ele: Tensor, px_ele: Tensor, y_ele: Tensor, py_ele: Tensor) → list[Tensor]

Transforms particle coordinates from element to lab frame.

Parameters:

  • x_offset – Element x-coordinate offset.

  • y_offset – Element y-coordinate offset.

  • tilt – Tilt angle (rad).

  • x_ele – x-coordinate in element frame.

  • px_ele – x-momentum in element frame.

  • y_ele – y-coordinate in element frame.

  • py_ele – y-momentum in element frame.

Returns:

x, px, y, py coordinates in lab frame.

utils.bmadx.particle_rf_time(z, pz, p0c, mc2)

Returns rf time of Particle p.

utils.bmadx.sinc(x)

sinc(x) = sin(x)/x.

utils.bmadx.sqrt_one(x)

Routine to calculate Sqrt[1+x] - 1 to machine precision.

utils.bmadx.track_a_drift(length: Tensor, x_in: Tensor, px_in: Tensor, y_in: Tensor, py_in: Tensor, z_in: Tensor, pz_in: Tensor, p0c: Tensor, mc2: Tensor) → tuple[Tensor]

Exact drift tracking used in different elements.

utils.device.is_mps_available_and_functional()

Check if MPS is available and functional (for GitHub Actions).

utils.elementwise_linspace.elementwise_linspace(start: Tensor, end: Tensor, steps: int) → Tensor

Generate a tensor of linearly spaced values between two tensors element-wise.

Parameters:

  • start – Any-dimensional tensor of the starting value for the set of points.

  • end – Any-dimensional tensor of the ending value for the set of points.

  • steps – Size of the last dimension of the constructed tensor.

Returns:

A tensor of shape start.shape + (steps,) containing steps linearly spaced values between each pair of elements in start and end.

utils.kde.kde_histogram_1d(x: Tensor, bins: Tensor, bandwidth: Tensor, weights: Tensor | None = None, epsilon: Tensor | float = 1e-10) → Tensor

Estimate the histogram using KDE of the input tensor.

The computation uses kernel density estimation which requires a bandwidth (smoothing) parameter.

Parameters:

  • x – Input tensor to compute the histogram with shape \((B, D)\).

  • bins – The number of bins to use the histogram \((N_{bins})\).

  • bandwidth – Gaussian smoothing factor with shape shape (1,).

  • weights – Weights of the input tensor of shape \((B, N)\).

  • epsilon – A scalar, for numerical stability. Default: 1e-10.

Returns:

Computed 1d histogram of shape \((B, N_{bins})\).

Examples:
>>> x = torch.rand(1, 10)
>>> bins = torch.torch.linspace(0, 255, 128)
>>> hist = kde_histogram_1d(x, bins, bandwidth=torch.tensor(0.9))
>>> hist.shape
torch.Size([1, 128])
utils.kde.kde_histogram_2d(x1: Tensor, x2: Tensor, bins1: Tensor, bins2: Tensor, bandwidth: Tensor, weights: Tensor | None = None, epsilon: float | Tensor = 1e-10) → Tensor

Estimate the 2D histogram of the input tensor.

The computation uses kernel density estimation which requires a bandwidth (smoothing) parameter.

This is a modified version of the kornia.enhance.histogram implementation.

Parameters:

  • x1 – Input tensor to compute the histogram with shape \((B, D1)\).

  • x2 – Input tensor to compute the histogram with shape \((B, D2)\).

  • bins – Bin coordinates.

  • bandwidth – Gaussian smoothing factor with shape shape (1,).

  • weights – Weights of the input tensor of shape \((B, N)\).

  • epsilon – A scalar, for numerical stability. Default: 1e-10.

Returns:

Computed histogram of shape \((B, N_{bins}, N_{bins})\).

Examples:
>>> x1 = torch.rand(2, 32)
>>> x2 = torch.rand(2, 32)
>>> bins = torch.torch.linspace(0, 255, 128)
>>> hist = kde_histogram_2d(x1, x2, bins, bandwidth=torch.tensor(0.9))
>>> hist.shape
torch.Size([2, 128, 128])
utils.physics.compute_relativistic_factors(energy: Tensor, particle_mass_eV: Tensor) → tuple[Tensor, Tensor, Tensor]

Computes the relativistic factors gamma, inverse gamma squared and beta for particles.

Parameters:

  • energy – Energy in eV.

  • particle_mass_eV – Mass of the particle in eV.

Returns:

gamma, igamma2, beta.

utils.statistics.match_distribution_moments(samples: Tensor, target_mu: Tensor, target_cov: Tensor, weights: Tensor | None = None) → Tensor

Match the first and second moments of a sample distribution to a target distribution.

Parameters:

  • samples – Input samples of shape (…, num_samples, num_features).

  • target_mu – Mean of the target distribution of shape (…, num_features).

  • target_cov – Covariance of the target distribution of shape (…, num_features, num_features).

  • weights – Weights for the samples of shape (…, num_samples).

Returns:

Transformed samples.

utils.statistics.unbiased_weighted_covariance(inputs1: Tensor, inputs2: Tensor, weights: Tensor, dim: int = None) → Tensor

Compute the unbiased weighted covariance of two tensors.

Parameters:

  • input1 – Input tensor 1 of shape (…, sample_size).

  • input2 – Input tensor 2 of shape (…, sample_size).

  • weights – Weights tensor of shape (…, sample_size).

  • dim – Dimension along which to compute the covariance.

Returns:

Unbiased weighted covariance of shape (…, 2, 2).

utils.statistics.unbiased_weighted_covariance_matrix(inputs: Tensor, weights: Tensor) → Tensor

Compute the unbiased weighted covariance matrix of a tensor.

Parameters:

  • inputs – Input tensor of shape (…, sample_size, num_features).

  • weights – Weights tensor of shape (…, sample_size).

Returns:

Unbiased weighted covariance matrix.

utils.statistics.unbiased_weighted_std(inputs: Tensor, weights: Tensor, dim: int = None) → Tensor

Compute the unbiased weighted standard deviation of a tensor.

Parameters:

  • inputs – Input tensor.

  • weights – Weights tensor.

  • dim – Dimension along which to compute the standard deviation.

Returns:

Unbiased weighted standard deviation.

utils.statistics.unbiased_weighted_variance(inputs: Tensor, weights: Tensor, dim: int = None) → Tensor

Compute the unbiased weighted variance of a tensor.

Parameters:

  • inputs – Input tensor.

  • weights – Weights tensor.

  • dim – Dimension along which to compute the variance.

Returns:

Unbiased weighted variance.

class utils.unique_name_generator.UniqueNameGenerator(prefix: str)

Generates a unique name given a prefix.

utils.vector.squash_index_for_unavailable_dims(index: tuple, shape: tuple) → tuple

Helper function to create an index that works on only partially broadcasted tensors.

You might need this when you vectorized one parameter of your model as (n,) and another as (m, 1). Fully broadcasted, this would result in vector shapes of (m, n). Some functions in Cheetah expect an index into the fully broadcasted vector shape. However, sometimes this index is passed onto a function that only handles a result from the model that was only affected by one of these vectorizations, e.g., the second one. This function lets you squash the vector index meant for a vectorised result of shape (m, n) into a vector index that works on a result with vector shape (m, 1) or (n,).

Example: The input vector shapes are (a) (3,) and (b) (2, 1), resulting in a broadcasted shape of (2, 3). A valid vector index for the fully broadcasted model would be (1, 2). If the result is only affected by vectorisation (b) with shape (2, 1), this function would squash the second dimension of the index to 0, resulting in (1, 0). Similarly, if the result is only affected by vectorisation (a) with shape (3,), this function would remove the first dimension of the index, resulting in (2,).

Parameters:

  • index – The index to squash.

  • shape – The shape of the result that was only affected by part of the vectorisations.

Returns:

The squashed index that is valid for indexing into a tensor of shape.