Demystifying Amplitude Analysis

Table of Contents

• Overview
• Installation
• Usage
• Implementing an Amplitude
  • Python Bindings
• TODOs

Overview

This crate contains all of the core mechanics

• rustitude will automatically parallelize amplitudes over the events in a dataset. There's no reason for a developer to ever write parallelized code themselves.
• Implementing Node on a struct is all that is needed to use it as an amplitude. This means developers need only write two to three total methods to describe the entire functionality of their amplitude, and one of these just gives the names and order of the amplitude's input parameters.
• A major goal of rustitude was to increase processing speed by sacrificing memory. This is done by precalculating parts of amplitudes which don't change when the free parameter inputs change. The simplest example of this is the Ylm amplitude (spherical harmonic), which can be entirely precalculated given the value of l and m. To calculate this amplitude at evaluation time, rustitude just needs to look up a value in an array!

Installation

Cargo provides the usual command for including this crate in a project:

cargo add rustitude-core

However, this project is best-used with some pre-made amplitudes (see rustitude-gluex, and these come bundled with the rustitude meta-crate, which also re-exports this crate's prelude. This can be installed via the same command:

cargo add rustitude

Implementing an Amplitude

While typical usage might be to use pre-made amplitudes in various combinations, it is important to know how to design an amplitude which will work seamlessly with this crate. Let's write down the rustitude version of the OmegaDalitz amplitude:

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(),
        ]
    }
}

Let's walk through this code. First, we need to define a struct which has all of the general information about the amplitude, and in this case some kind of Vec for storing precalculated data. We consider this precalculated data to correspond to a single dataset, and each dataset gets its own copy of the amplitude struct. Because this particular amplitude doesn't have any input parameters, we can #[derive(Default)] on it to make a default constructor, which allows the amplitude to be initialized with something like let amp = OmegaDalitz::default();. If we wanted a parameterized constructor, we have to define our own, and while Rust has no default name for constructors, pub fn new(...) -> rustitude_core::Amplitude_64 (or _32) is preferred.

Next, we implement the Node trait for the struct. Traits in Rust are kind of like abstract classes or interfaces in object-oriented languages, they provide some set of methods which a struct must implement. The first of these methods is fn precalculate(&mut self, dataset: &Dataset) -> Result<(), RustitudeError>. As the signature suggests, it takes a Dataset and mutates the struct in some way. It should raise a RustitudeError if anything goes wrong in the evaluation. The intended usage of this function is to precalculate some terms in the amplitude's mathematical expression, things which don't change when you update the free parameter inputs to the amplitude. In this case, the four input parameters, $\alpha$, $\beta$, $\gamma$, and $\delta$, are independent from dalitz_z, dalitz_sin3theta, and lambda, so we can safely calculate those ahead of time and just pull from their respective Vecs when needed later. I won't go too far into Rust's syntax here, but typical precalculation functions will start by iterating over the dataset's events in parallel (the line use rayon::prelude::*; is needed to use par_iter here) and collecting or unzipping that iterator into a Vec or group of Vecs.

The calculate step has the signature fn calculate(&self, parameters: &[f64], event: &Event) -> Result<Complex64, RustitudeError>. This means we need to take a list of parameters and a single event and turn them into a complex value. The Event struct contains an index field which can be used to access the precalculated storage arrays made in the previous step.

Finally, the parameters function just returns a list of the parameter names in the order they are expected to be input into calculate. In the event that an amplitude doesn't have any free parameters (like my implementation of the Ylm and Zlm amplitudes), we can omit this function entirely, as the default implementation returns vec![].

And that's it! However, it is important to be aware of the steps which need to be taken to allow this amplitude to be used through the Python interface.

Python Bindings

To make Python bindings, pyo3 needs to be included as a dependency:

use pyo3::prelude::*;
use rustitude_core::amplitude::Amplitude_64;

// OmegaDalitz code here

#[pyfunction(name = "OmegaDalitz")]
fn omega_dalitz(name: &str) -> PyResult<Amplitude_64> {
    Ok(Amplitude_64::new(name, OmegaDalitz::default()))
}

pub fn pyo3_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(omega_dalitz, m)?)?;
    Ok(())
}

Rather than bind the struct directly, we prefer to bind a function which returns a Amplitude_64 (or a PyResult of one), a wrapper struct that implements #[pyclass] and can be used in the Python interface. The pyo3_module function will then need to be added to the py-rustitude crate's lib.rs file. This step is a bit problematic, since it means new amplitudes cannot be added without modifying py-rustitude itself. For this reason, developers may want to work with their own fork of the repository rather than using the one installed by cargo if they wish to use Python with custom amplitudes. This is a limitation of pyo3, which doesn't recognize class bindings across crates. There is also a mechanism to implement Amplitudes via pure Python classes, but it is currently incompatible with multithreading due to the GIL (see the Python docs).

TODOs

In no particular order, here is a list of what (probably) needs to be done before I will stop making any breaking changes:

• Pure Rust parsing of ROOT files without the uproot backend (I have some moderate success with oxyroot, but there are still a few issues reading larger files)
• Add plotting methods
• A way to check if the number of parameters matches the input at compile time would be nice, not sure if it's possible though
• Amplitude serialization (I want to be able to send and share an amplitude configuration just like an AmpTools config file, but with the amplitude itself included)
• Lots of documentation
• Lots of tests