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: f32,                     // 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<f32>,               // 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 with the ReadMethod::EPSInBeam method.

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::*;

#[derive(Clone)]
pub struct ComplexScalar;
impl<F: Field> Node<F> for ComplexScalar {
    fn calculate(&self, parameters: &[F], _event: &Event<F>) -> Result<Complex<F>, RustitudeError> {
        Ok(Complex::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, Clone)]
pub struct OmegaDalitz<F: Field> {
    dalitz_z: Vec<F>,
    dalitz_sin3theta: Vec<F>,
    lambda: Vec<F>,
}

impl<F: Field> Node<F> for OmegaDalitz<F> {
    fn precalculate(&mut self, dataset: &Dataset<F>) -> Result<(), RustitudeError> {
        (self.dalitz_z, (self.dalitz_sin3theta, self.lambda)) = dataset
            .events
            .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 = (F::TWO * pip.m()) + pi0.m();
                let dalitz_d = F::TWO * omega.m() * (omega.m() - m3pi);
                let dalitz_sc = (F::ONE / F::THREE) * (omega.m2() + pip.m2() + pim.m2() + pi0.m2());
                let dalitz_x = F::fsqrt(F::THREE) * (dalitz_t - dalitz_u) / dalitz_d;
                let dalitz_y = F::THREE * (dalitz_sc - dalitz_s) / dalitz_d;

                let dalitz_z = dalitz_x * dalitz_x + dalitz_y * dalitz_y;
                let dalitz_sin3theta = F::fsin(F::THREE * F::fasin(dalitz_y / F::fsqrt(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 = (F::FOUR / F::THREE) * F::fabs(pi_cross.dot(&pi_cross))
                    / ((F::ONE / F::NINE)
                        * (omega.m2() - (F::TWO * pip.m() + pi0.m()).fpowi(2)).fpowi(2));

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

    fn calculate(&self, parameters: &[F], event: &Event<F>) -> Result<Complex<F>, 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(F::fsqrt(F::fabs(
            lambda
                * (F::ONE
                    + F::TWO * alpha * dalitz_z
                    + F::TWO * beta * dalitz_z.fpowf(F::THREE / F::TWO) * dalitz_sin3theta
                    + F::TWO * gamma * dalitz_z.fpowi(2)
                    + F::TWO * delta * dalitz_z.fpowf(F::FIVE / F::TWO) * dalitz_sin3theta),
        ))
        .into())
    }

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

Note several of the generic features which allow this amplitude to be used with different numeric data types. Because it isn’t specifically written for 64-bit floats (f64s), we can conduct analyses that use the same code with 32-bit floats (f32s), which saves on memory and time while sacrificing a bit of precision. In fact, we can go a step further and conduct the majority of an analysis in 32-bit mode, switching over to 64-bit mode when we actually get near a solution and want the increased accuracy!

The Field trait contains a few mathematical constants like Field::PI() and Field::SQRT_2() as well as traits which implement most standard mathematical functions. See the Float trait for more details.

//! # 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. Models implicitly take the absolute square of each provided term in their constructor and add those results incoherently.

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!(amp1, amp2)

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::new("omega dalitz", OmegaDalitz::default());
let term = complex_term * omega_dalitz;
term.print_tree();
// [ norm sqr ]
//   ┗━[ * ]
//       ┣━ !my complex scalar(real, imag)
//       ┗━ !omega dalitz(alpha, beta, gamma, delta)
let model = model!(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 &[f32] (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 { ... }

fn main() -> Result<(), RustitudeError> {
    let complex_term = cscalar("my complex scalar");
    let omega_dalitz = Amplitude::new("omega dalitz", OmegaDalitz::default());
    let term = complex_term * omega_dalitz;
    let model = model!(term);
    let dataset = Dataset::from_parquet("path/to/file.parquet")?;
    let dataset_mc = Dataset::from_parquet("path/to/monte_carlo_file.parquet")?;
    let nll = ExtendedLogLikelihood::new(
        Manager::new(&model, &dataset),
        Manager::new(&model, &dataset_mc)
    );
    println!("NLL: {}", nll.evaluate(&nll.get_initial())?);
    Ok(())
}

Fitting Amplitudes to Data

Of course, the goal of all of this is to be able to construct a Model, load up a Dataset, create an ExtendedLogLikelihood, and fit the model to data. Here’s an example to show how that might be accomplished:

use ganesh::algorithms::NelderMead;
use ganesh::prelude::*;
use rustitude::gluex::harmonics::Zlm;
use rustitude::gluex::{
    resonances::BreitWigner,
    utils::{Frame, Reflectivity, Wave},
};
use rustitude::prelude::*;
fn main() -> Result<(), RustitudeError> {
    let a2_1320 = BreitWigner::new(&[0], &[1], 2).named("a2_1320");
    let a2_1700 = BreitWigner::new(&[0], &[1], 2).named("a2_1700");
    let pw_s_wave = piecewise_m("pw_s_wave", 40, (1.04, 1.72));
    let zlm_s0p = Zlm::new(Wave::S0, Reflectivity::Positive, Frame::Helicity).named("zlm_s0p");
    let zlm_s0n = Zlm::new(Wave::S0, Reflectivity::Negative, Frame::Helicity).named("zlm_s0n");
    let zlm_dn2p = Zlm::new(Wave::Dn2, Reflectivity::Positive, Frame::Helicity).named("zlm_dn2p");
    let zlm_dn1p = Zlm::new(Wave::Dn1, Reflectivity::Positive, Frame::Helicity).named("zlm_dn1p");
    let zlm_d0p = Zlm::new(Wave::D0, Reflectivity::Positive, Frame::Helicity).named("zlm_d0p");
    let zlm_d1p = Zlm::new(Wave::D1, Reflectivity::Positive, Frame::Helicity).named("zlm_d1p");
    let zlm_d2p = Zlm::new(Wave::D2, Reflectivity::Positive, Frame::Helicity).named("zlm_d2p");
    let zlm_dn2n = Zlm::new(Wave::Dn2, Reflectivity::Negative, Frame::Helicity).named("zlm_dn2n");
    let zlm_dn1n = Zlm::new(Wave::Dn1, Reflectivity::Negative, Frame::Helicity).named("zlm_dn1n");
    let zlm_d0n = Zlm::new(Wave::D0, Reflectivity::Negative, Frame::Helicity).named("zlm_d0n");
    let zlm_d1n = Zlm::new(Wave::D1, Reflectivity::Negative, Frame::Helicity).named("zlm_d1n");
    let zlm_d2n = Zlm::new(Wave::D2, Reflectivity::Negative, Frame::Helicity).named("zlm_d2n");
    let pos_d_wave = zlm_dn2p + zlm_dn1p + zlm_d0p + zlm_d1p + zlm_d2p;
    let neg_d_wave = zlm_dn2n + zlm_dn1n + zlm_d0n + zlm_d1n + zlm_d2n;
    let pos_real =
        zlm_s0p.real() * &pw_s_wave + &a2_1320 * &pos_d_wave.real() + &a2_1700 * &pos_d_wave.real();
    let pos_imag =
        zlm_s0p.imag() * &pw_s_wave + &a2_1320 * &pos_d_wave.imag() + &a2_1700 * &pos_d_wave.imag();
    let neg_real =
        zlm_s0n.real() * &pw_s_wave + &a2_1320 * &neg_d_wave.real() + &a2_1700 * &neg_d_wave.real();
    let neg_imag =
        zlm_s0n.imag() * &pw_s_wave + &a2_1320 * &neg_d_wave.imag() + &a2_1700 * &neg_d_wave.imag();
    let model = model!(pos_real, pos_imag, neg_real, neg_imag);
    let ds_data = Dataset::from_parquet("path/to/data.root", ReadMethod::EPSInBeam)?;
    let ds_accmc = Dataset::from_parquet("path/to/accmc.root", ReadMethod::EPSInBeam)?;
    let mut ell = ExtendedLogLikelihood::new(
        Manager::new(&model, &ds_data)?,
        Manager::new(&model, &ds_accmc)?,
    );
    ell.set_initial("a2_1320", "mass", 1.3182)?;
    ell.set_initial("a2_1320", "width", 0.1111)?;
    ell.fix("a2_1700", "mass", 1.698)?;
    ell.fix("a2_1700", "width", 0.265)?;
    ell.fix("pw_s_wave", "bin 10 im", 0.0)?;

    let mut nm = NelderMead::new(ell.clone(), &ell.get_initial(), None);
    minimize!(nm, 1000)?; // Run 1000 steps
    let (best_pars, best_fx) = nm.best();
    for (par_name, par_value) in ell.free_parameters().iter().zip(best_pars) {
        println!("{} -> {} (NLL = {})", par_name, par_value, best_fx);
    }
    Ok(())
}

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
  This module contains an all-encompassing error enum that almost every crate method will produce if it returns a Result.
• 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.
• utils
  This module holds some convenience methods for writing nice test functions for Amplitudes.

Macros

• assert_is_close
  A macro to assert if two floating point numbers are essentially equal. Similar to approx crate.
• convert
  Convenience macro for converting raw numeric values to a generic.
• convert_array
  Convenience macro for converting a raw numeric array to a generic array.
• convert_vec
  Convenience macro for converting a raw numeric Vec to a generic Vec.
• model
  Convenience macro for boxing up coherent sum terms into a Model.

Traits

• Field
  A trait representing a numeric field which can be used in calculating amplitudes.
• UnitVector
  A trait to normalize structs (mostly to use on nalgebra vectors without needing nalgebra::RealField)