Crate rustitude_core

Rustitude

Demystifying Amplitude Analysis with Rust and Python

The rustitude-core crate aims to implement common amplitude analysis techniques in Rust with bindings to Python. This crate does not include the Python bindings, see the GitHub repo for more information on the Python API.

The three core principles of rustitude-core are:

1. Parallelization over events is automatically handeled by a Manager.
2. Amplitudes are written to do as much work as possible ahead of time, and evaluations use caching as much as possible automatically.
3. Developers just need to implement the Node trait to write a new amplitude, everything else is handled by the crate.

Table of Contents

• Dataset Structure
• Creating a New Amplitude
• Combining Amplitudes into Models
• Managing Parameters
• Evaluating Likelihoods

Dataset Structure

A Dataset is essentially just a wrapper for a Vec of Events. The current Event structure is as follows:

pub struct Event {
    pub index: usize,                    // Position of event within dataset
    pub weight: f64,                     // Event weight
    pub beam_p4: FourMomentum,           // Beam four-momentum
    pub recoil_p4: FourMomentum,         // Recoil four-momentum
    pub daughter_p4s: Vec<FourMomentum>, // Four-momenta of final state particles sans recoil
    pub eps: Vector3<f64>,               // Beam polarization vector
}

In the Rust API, we can create Datasets from ROOT files as well as Parquet files. ROOT file reading is done through oxyroot - This still has some issues, and large files or files with user metadata might fail to load. The alternative Parquet format can be obtained from a ROOT file by using a conversion script like the one provided here. By default, we expect all of the Event fields to be mirrored as the following branches:

Branch Name	Data Type	Notes
Weight	Float32
E_Beam	Float32
Px_Beam	Float32
Py_Beam	Float32
Pz_Beam	Float32
E_FinalState	[Float32]	[recoil, daughter #1, daughter #2, …]
Px_FinalState	[Float32]	[recoil, daughter #1, daughter #2, …]
Py_FinalState	[Float32]	[recoil, daughter #1, daughter #2, …]
Pz_FinalState	[Float32]	[recoil, daughter #1, daughter #2, …]
EPS	[Float32]	[$P_\gamma \cos(\Phi)$, $P_\gamma \sin(\Phi)$, $0.0$] for linear polarization with magnitude $P_\gamma$ and angle $\Phi$

A Parquet file with these columns can be loaded with the following:

use rustitude_core::prelude::*;
fn main() -> Result<(), RustitudeError> {
    let dataset = Dataset::from_parquet("path/to/file.parquet")?;
    println!("{}", dataset.events()[0]); // print first event
}

Because the beam is often directed along the $z$-axis, there is an alternative way to store the EPS vector without a new branch (for linear polarization. The $x$ and $y$ components of EPS can be stored as Px_Beam and Py_Beam respectively, and the format can be loaded using Dataset::from_parquet_eps_in_beam.

Creating a New Amplitude

To make a new amplitude, we will first create a new struct and then implement Node. Let’s start with a trivial example, an amplitude which returns a complex scalar. This particular amplitude is already implemented as a convenience struct called ComplexScalar.

use rustitude_core::prelude::*
struct ComplexScalar;
impl Node for ComplexScalar {
    fn calculate(&self, parameters: &[f64], _event: &Event) -> Result<Complex64, RustitudeError> {
        Ok(Complex64::new(parameters[0], parameters[1]))
    }

    fn parameters(&self) -> Vec<String> {
        vec!["real".to_string(), "imag".to_string()]
    }
}

For a second example, we can look at the precalculation feature. Here’s an Dalitz-like amplitude for the $\omega$ particle:

use rayon::prelude::*;
use rustitude_core::prelude::*;

#[derive(Default)]
pub struct OmegaDalitz {
    dalitz_z: Vec<f64>,
    dalitz_sin3theta: Vec<f64>,
    lambda: Vec<f64>,
}

impl Node for OmegaDalitz {
    fn precalculate(&mut self, dataset: &Dataset) -> Result<(), RustitudeError> {
        (self.dalitz_z, (self.dalitz_sin3theta, self.lambda)) = dataset
            .events
            .read()
            .par_iter()
            .map(|event| {
                let pi0 = event.daughter_p4s[0];
                let pip = event.daughter_p4s[1];
                let pim = event.daughter_p4s[2];
                let omega = pi0 + pip + pim;

                let dalitz_s = (pip + pim).m2();
                let dalitz_t = (pip + pi0).m2();
                let dalitz_u = (pim + pi0).m2();

                let m3pi = (2.0 * pip.m()) + pi0.m();
                let dalitz_d = 2.0 * omega.m() * (omega.m() - m3pi);
                let dalitz_sc = (1.0 / 3.0) * (omega.m2() + pip.m2() + pim.m2() + pi0.m2());
                let dalitz_x = f64::sqrt(3.0) * (dalitz_t - dalitz_u) / dalitz_d;
                let dalitz_y = 3.0 * (dalitz_sc - dalitz_s) / dalitz_d;

                let dalitz_z = dalitz_x * dalitz_x + dalitz_y * dalitz_y;
                let dalitz_sin3theta = f64::sin(3.0 * f64::asin(dalitz_y / f64::sqrt(dalitz_z)));

                let pip_omega = pip.boost_along(&omega);
                let pim_omega = pim.boost_along(&omega);
                let pi_cross = pip_omega.momentum().cross(&pim_omega.momentum());

                let lambda = (4.0 / 3.0) * f64::abs(pi_cross.dot(&pi_cross))
                    / ((1.0 / 9.0) * (omega.m2() - (2.0 * pip.m() + pi0.m()).powi(2)).powi(2));

                (dalitz_z, (dalitz_sin3theta, lambda))
            })
            .unzip();
        Ok(())
    }

    fn calculate(&self, parameters: &[f64], event: &Event) -> Result<Complex64, RustitudeError> {
        let dalitz_z = self.dalitz_z[event.index];
        let dalitz_sin3theta = self.dalitz_sin3theta[event.index];
        let lambda = self.lambda[event.index];
        let alpha = parameters[0];
        let beta = parameters[1];
        let gamma = parameters[2];
        let delta = parameters[3];
        Ok(f64::sqrt(f64::abs(
            lambda
                * (1.0
                    + 2.0 * alpha * dalitz_z
                    + 2.0 * beta * dalitz_z.powf(3.0 / 2.0) * dalitz_sin3theta
                    + 2.0 * gamma * dalitz_z.powi(2)
                    + 2.0 * delta * dalitz_z.powf(5.0 / 2.0) * dalitz_sin3theta),
        ))
        .into())
    }

    fn parameters(&self) -> Vec<String> {
        vec![
            "alpha".to_string(),
            "beta".to_string(),
            "gamma".to_string(),
            "delta".to_string(),
        ]
    }
}

Combining Amplitudes into Models

We can use several operations to modify and combine amplitudes. Since amplitudes yield complex values, the following convenience methods are provided: real, and imag give the real and imaginary part of the amplitude, respectively. Additionally, amplitudes can be added and multiplied together using operator overloading. All sums are interpreted as coherent sums, and products with these coherent sums are distributed. Incoherent sums are handled at the Model level.

To incoherently sum two Amplitudes, say amp1 and amp2, we would first assume that we actually want the absolute square of the given term (or write our amplitude as the square root of what we really want), and then include them both in our model:

use rustitude_core::prelude::*;
// Define amp1/amp2: Amplitude here...
let model = Model::new(vec![amp1.as_cohsum(), amp2.as_cohsum()])

To reiterate, this would yield something like $\left|\text{amp}_1\right|^2 + \left|\text{amp}_2\right|^2$.

The Scalar, ComplexScalar, and PolarComplexScalar amplitudes all have convenience methods, scalar, cscalar, and pcscalar respectively. We then wrap the final expression in a Model which can manage all of the Parameters.

use rustitude_core::prelude::*;

#[derive(Default)]
pub struct OmegaDalitz { ... }
impl Node for OmegaDalitz { ... }

let complex_term = cscalar("my complex scalar");
let omega_dalitz = amplitude!("omega dalitz", OmegaDalitz::default());
let term = (complex_term * omega_dalitz).norm_sqr();
term.print_tree();
// [ norm sqr ]
//   ┗━[ * ]
//       ┣━ !my complex scalar(real, imag)
//       ┗━ !omega dalitz(alpha, beta, gamma, delta)
let model = Model::new(term);

Managing Parameters

Now that we have a model, we might want to constrain or fix parameters. Parameters are identified solely by their name and the name of the amplitude they are associated with. This means that two amplitudes with the same name will share parameters which also have the same name. If we want to intentionally set one parameter in a particular amplitude equal to another, we can use the Model::constrain. This will reduce the number of free parameters in the fit, and will yield a RustitudeError if either of the parameters is not found. Parameters can also be fixed and freed using Model::fix and Model::free respectively, and these methods are mirrored in Manager and ExtendedLogLikelihood for convenience.

Evaluating Likelihoods

If we wanted to obtain the negative log-likelihood for this particular amplitude, we need to link our Model to a Dataset. This is done using a Manager. Finally, two Managers may be combined into an ExtendedLogLikelihood. Both of these manager-like structs have an evaluate method that takes some parameters as a &[f64] (along with a usize for the number of threads to use for the ExtendedLogLikelihood).

use rustitude_core::prelude::*;

#[derive(Default)]
pub struct OmegaDalitz { ... }
impl Node for OmegaDalitz { ... }

let complex_term = cscalar("my complex scalar");
let omega_dalitz = amplitude!("omega dalitz", OmegaDalitz::default());
let term = (complex_term * omega_dalitz).norm_sqr();
let model = Model::new(term);
let dataset = Dataset::from_parquet("path/to/file.parquet").unwrap();
let dataset_mc = Dataset::from_parquet("path/to/monte_carlo_file.parquet").unwrap();
let nll = ExtendedLogLikelihood::new(
        Manager::new(&model, &dataset),
        Manager::new(&model, &dataset_mc)
    );
println!("NLL on 4 threads: {}", nll.evaluate(&nll.get_initial(), 4));

Modules

• amplitude
  The amplitude module contains structs and methods for defining and manipulating Amplitudes and Models
• dataset
  This module contains all the resources needed to load and examine datasets.
• errors
• four_momentum
  Four-Momentum
• manager
  This module contains methods to link Models with Datasets via a Manager::evaluate method. This module also holds a ExtendedLogLikelihood struct which holds two Managers and, as the name suggests, calculates an extended log-likelihood using a very basic method over data and (accepted) Monte-Carlo.
• prelude
  Recommended namespace for use and development.

Functions

• create_pool
  Creates a new thread pool.