feat: Add LikelihoodTerm trait and implement it for NLL

Also changes the K-matrix benchmark to collect more samples

benches/kmatrix_benchmark.rs

@@ -1,11 +1,13 @@
+use std::time::Duration;
+
 use criterion::{black_box, criterion_group, criterion_main, BatchSize, Criterion};
 use laddu::{
     amplitudes::{
         constant,
         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);

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<Float> {
+        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<Float> {
+        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<Float> {
-        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<Float> {
-        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<Float, (), Infallible> for NLL {
     fn evaluate(&self, parameters: &[Float], _user_data: &mut ()) -> Result<Float, Infallible> {
-        Ok(self.evaluate(parameters))
+        Ok(LikelihoodTerm::evaluate(self, parameters))
     }
 }
 

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::{

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;