From ec7c1e03808fcf8020525a9a594507e89f20d5d2 Mon Sep 17 00:00:00 2001 From: denehoffman Date: Wed, 23 Oct 2024 12:04:29 -0400 Subject: [PATCH] feat: Add `LikelihoodTerm` trait and implement it for `NLL` Also changes the K-matrix benchmark to collect more samples --- benches/kmatrix_benchmark.rs | 10 +++- src/amplitudes/mod.rs | 100 +++++++++++++++++++---------------- src/lib.rs | 4 +- src/python.rs | 2 +- 4 files changed, 64 insertions(+), 52 deletions(-) diff --git a/benches/kmatrix_benchmark.rs b/benches/kmatrix_benchmark.rs index 74d5ab0..2297c98 100644 --- a/benches/kmatrix_benchmark.rs +++ b/benches/kmatrix_benchmark.rs @@ -1,3 +1,5 @@ +use std::time::Duration; + use criterion::{black_box, criterion_group, criterion_main, BatchSize, Criterion}; use laddu::{ amplitudes::{ @@ -5,7 +7,7 @@ use laddu::{ kmatrix::{KopfKMatrixA0, KopfKMatrixA2, KopfKMatrixF0, KopfKMatrixF2}, parameter, zlm::Zlm, - Manager, NLL, + LikelihoodTerm, Manager, NLL, }, data::open, utils::{ @@ -153,5 +155,9 @@ fn kmatrix_nll_benchmark(c: &mut Criterion) { }); } -criterion_group!(benches, kmatrix_nll_benchmark); +criterion_group! { + name = benches; + config = Criterion::default().measurement_time(Duration::from_secs(30)).sample_size(7000); + targets = kmatrix_nll_benchmark +} criterion_main!(benches); diff --git a/src/amplitudes/mod.rs b/src/amplitudes/mod.rs index c9f308b..fed6b23 100644 --- a/src/amplitudes/mod.rs +++ b/src/amplitudes/mod.rs @@ -395,6 +395,10 @@ impl Evaluator { } } +pub trait LikelihoodTerm { + fn evaluate(&self, parameters: &[Float]) -> Float; +} + /// An extended, unbinned negative log-likelihood evaluator. pub struct NLL { data_evaluator: Evaluator, @@ -467,6 +471,52 @@ impl NLL { self.mc_evaluator.resources.isolate_many(names); } + /// Project the stored [`Expression`] over the events in the [`Dataset`] stored by the + /// [`Evaluator`] with the given values for free parameters to obtain weights for each + /// Monte-Carlo event. This method takes the real part of the given expression (discarding + /// the imaginary part entirely, which does not matter if expressions are coherent sums + /// wrapped in [`Expression::norm_sqr`]). Event weights are determined by the following + /// formula: + /// + /// ```math + /// \text{weight}(\vec{p}; e) = \text{weight}(e) \mathcal{L}(e) \frac{N_{\text{Data}}}{N_{\text{MC}}} + /// ``` + #[cfg(feature = "rayon")] + pub fn project(&self, parameters: &[Float]) -> Vec { + let n_data = self.data_evaluator.dataset.weighted_len(); + let mc_result = self.mc_evaluator.evaluate(parameters); + let n_mc = self.mc_evaluator.dataset.weighted_len(); + mc_result + .par_iter() + .zip(self.mc_evaluator.dataset.par_iter()) + .map(|(l, e)| e.weight * l.re * (n_data / n_mc)) + .collect() + } + + /// Project the stored [`Expression`] over the events in the [`Dataset`] stored by the + /// [`Evaluator`] with the given values for free parameters to obtain weights for each + /// Monte-Carlo event. This method takes the real part of the given expression (discarding + /// the imaginary part entirely, which does not matter if expressions are coherent sums + /// wrapped in [`Expression::norm_sqr`]). Event weights are determined by the following + /// formula: + /// + /// ```math + /// \text{weight}(\vec{p}; e) = \text{weight}(e) \mathcal{L}(e) \frac{N_{\text{Data}}}{N_{\text{MC}}} + /// ``` + #[cfg(not(feature = "rayon"))] + pub fn project(&self, parameters: &[Float]) -> Vec { + let n_data = self.data_evaluator.dataset.weighted_len(); + let mc_result = self.mc_evaluator.evaluate(parameters); + let n_mc = self.mc_evaluator.dataset.weighted_len(); + mc_result + .iter() + .zip(self.mc_evaluator.dataset.iter()) + .map(|(l, e)| e.weight * l.re * (n_data / n_mc)) + .collect() + } +} + +impl LikelihoodTerm for NLL { /// Evaluate the stored [`Expression`] over the events in the [`Dataset`] stored by the /// [`Evaluator`] with the given values for free parameters. This method takes the /// real part of the given expression (discarding the imaginary part entirely, which @@ -477,7 +527,7 @@ impl NLL { /// NLL(\vec{p}) = -2 \left(\sum_{e \in \text{Data}} \text{weight}(e) \ln(\mathcal{L}(e)) - \frac{N_{\text{Data}}}{N_{\text{MC}}} \sum_{e \in \text{MC}} \text{weight}(e) \mathcal{L}(e) \right) /// ``` #[cfg(feature = "rayon")] - pub fn evaluate(&self, parameters: &[Float]) -> Float { + fn evaluate(&self, parameters: &[Float]) -> Float { let data_result = self.data_evaluator.evaluate(parameters); let n_data = self.data_evaluator.dataset.weighted_len(); let mc_result = self.mc_evaluator.evaluate(parameters); @@ -505,7 +555,7 @@ impl NLL { /// NLL(\vec{p}) = -2 \left(\sum_{e \in \text{Data}} \text{weight}(e) \ln(\mathcal{L}(e)) - \frac{N_{\text{Data}}}{N_{\text{MC}}} \sum_{e \in \text{MC}} \text{weight}(e) \mathcal{L}(e) \right) /// ``` #[cfg(not(feature = "rayon"))] - pub fn evaluate(&self, parameters: &[Float]) -> Float { + fn evaluate(&self, parameters: &[Float]) -> Float { let data_result = self.data_evaluator.evaluate(parameters); let n_data = self.data_evaluator.dataset.weighted_len(); let mc_result = self.mc_evaluator.evaluate(parameters); @@ -522,55 +572,11 @@ impl NLL { .sum(); -2.0 * (data_term - (n_data / n_mc) * mc_term) } - - /// Project the stored [`Expression`] over the events in the [`Dataset`] stored by the - /// [`Evaluator`] with the given values for free parameters to obtain weights for each - /// Monte-Carlo event. This method takes the real part of the given expression (discarding - /// the imaginary part entirely, which does not matter if expressions are coherent sums - /// wrapped in [`Expression::norm_sqr`]). Event weights are determined by the following - /// formula: - /// - /// ```math - /// \text{weight}(\vec{p}; e) = \text{weight}(e) \mathcal{L}(e) \frac{N_{\text{Data}}}{N_{\text{MC}}} - /// ``` - #[cfg(feature = "rayon")] - pub fn project(&self, parameters: &[Float]) -> Vec { - let n_data = self.data_evaluator.dataset.weighted_len(); - let mc_result = self.mc_evaluator.evaluate(parameters); - let n_mc = self.mc_evaluator.dataset.weighted_len(); - mc_result - .par_iter() - .zip(self.mc_evaluator.dataset.par_iter()) - .map(|(l, e)| e.weight * l.re * (n_data / n_mc)) - .collect() - } - - /// Project the stored [`Expression`] over the events in the [`Dataset`] stored by the - /// [`Evaluator`] with the given values for free parameters to obtain weights for each - /// Monte-Carlo event. This method takes the real part of the given expression (discarding - /// the imaginary part entirely, which does not matter if expressions are coherent sums - /// wrapped in [`Expression::norm_sqr`]). Event weights are determined by the following - /// formula: - /// - /// ```math - /// \text{weight}(\vec{p}; e) = \text{weight}(e) \mathcal{L}(e) \frac{N_{\text{Data}}}{N_{\text{MC}}} - /// ``` - #[cfg(not(feature = "rayon"))] - pub fn project(&self, parameters: &[Float]) -> Vec { - let n_data = self.data_evaluator.dataset.weighted_len(); - let mc_result = self.mc_evaluator.evaluate(parameters); - let n_mc = self.mc_evaluator.dataset.weighted_len(); - mc_result - .iter() - .zip(self.mc_evaluator.dataset.iter()) - .map(|(l, e)| e.weight * l.re * (n_data / n_mc)) - .collect() - } } impl Function for NLL { fn evaluate(&self, parameters: &[Float], _user_data: &mut ()) -> Result { - Ok(self.evaluate(parameters)) + Ok(LikelihoodTerm::evaluate(self, parameters)) } } diff --git a/src/lib.rs b/src/lib.rs index 385bc1b..1a3c6c8 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -283,8 +283,8 @@ pub mod prelude { constant, parameter, ylm::Ylm, zlm::Zlm, - Amplitude, AmplitudeID, Evaluator, Expression, Manager, MinimizerOptions, ParameterLike, - NLL, + Amplitude, AmplitudeID, Evaluator, Expression, LikelihoodTerm, Manager, MinimizerOptions, + ParameterLike, NLL, }; pub use crate::data::{open, open_binned, open_filtered, BinnedDataset, Dataset, Event}; pub use crate::resources::{ diff --git a/src/python.rs b/src/python.rs index 69965f8..ff85f34 100644 --- a/src/python.rs +++ b/src/python.rs @@ -14,7 +14,7 @@ pub(crate) mod laddu { use super::*; use crate as rust; - use crate::amplitudes::MinimizerOptions; + use crate::amplitudes::{LikelihoodTerm, MinimizerOptions}; use crate::utils::variables::Variable; use crate::utils::vectors::{FourMomentum, FourVector, ThreeMomentum, ThreeVector}; use crate::Float;