feat: add LadduError struct and work in proper error forwarding for…

… reading data and registering `Amplitude`s `Amplitude`s can no longer be registered with the same name by the same manager. This was causing some issues and there are better ways to handle it.
denehoffman · Oct 20, 2024 · edafd7f · edafd7f
1 parent 4023a65
commit edafd7f
Show file tree

Hide file tree

Showing 13 changed files with 212 additions and 135 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -30,6 +30,8 @@ rayon = { version = "1.10.0", optional = true }
 pyo3 = { version = "0.22.5", optional = true, features = ["num-complex"] }
 numpy = { version = "0.22.0", optional = true, features = ["nalgebra"] }
 ganesh = "0.11.2"
+thiserror = "1.0.64"
+shellexpand = "3.1.0"
 
 [dev-dependencies]
 approx = "0.5.1"

diff --git a/README.md b/README.md
@@ -121,7 +121,7 @@ impl MyBreitWigner {
 }
 
 impl Amplitude for MyBreitWigner {
-    fn register(&mut self, resources: &mut Resources) -> AmplitudeID {
+    fn register(&mut self, resources: &mut Resources) -> Result<AmplitudeID, LadduError> {
         self.pid_mass = resources.register_parameter(&self.mass);
         self.pid_width = resources.register_parameter(&self.width);
         resources.register_amplitude(&self.name)
@@ -162,8 +162,8 @@ let bw = manager.register(MyBreitWigner::new(
     &p1_mass,
     &p2_mass,
     &resonance_mass,
-));
-let mag = manager.register(Scalar::new("mag", parameter("magnitude")));
+)).unwrap();
+let mag = manager.register(Scalar::new("mag", parameter("magnitude"))).unwrap();
 let model = (mag * bw).norm_sqr();
 
 let nll = NLL::new(&manager, &ds_data, &ds_mc);

diff --git a/benches/kmatrix_benchmark.rs b/benches/kmatrix_benchmark.rs
@@ -24,92 +24,110 @@ fn kmatrix_nll_benchmark(c: &mut Criterion) {
     let polarization = Polarization::new(0, [1]);
     let resonance_mass = Mass::new([2, 3]);
     let mut manager = Manager::default();
-    let z00p = manager.register(Zlm::new(
-        "Z00+",
-        0,
-        0,
-        Sign::Positive,
-        &angles,
-        &polarization,
-    ));
-    let z00n = manager.register(Zlm::new(
-        "Z00-",
-        0,
-        0,
-        Sign::Negative,
-        &angles,
-        &polarization,
-    ));
-    let z22p = manager.register(Zlm::new(
-        "Z22+",
-        2,
-        2,
-        Sign::Positive,
-        &angles,
-        &polarization,
-    ));
-    let f0p = manager.register(KopfKMatrixF0::new(
-        "f0",
-        [
-            [constant(0.0), constant(0.0)],
-            [parameter("f0(980)+ re"), constant(0.0)],
-            [parameter("f0(1370)+ re"), parameter("f0(1370)+ im")],
-            [parameter("f0(1500)+ re"), parameter("f0(1500)+ im")],
-            [parameter("f0(1710)+ re"), parameter("f0(1710)+ im")],
-        ],
-        0,
-        &resonance_mass,
-    ));
-    let a0p = manager.register(KopfKMatrixA0::new(
-        "a0",
-        [
-            [parameter("a0(980)+ re"), parameter("a0(980)+ im")],
-            [parameter("a0(1450)+ re"), parameter("a0(1450)+ im")],
-        ],
-        0,
-        &resonance_mass,
-    ));
-    let f0n = manager.register(KopfKMatrixF0::new(
-        "f0",
-        [
-            [constant(0.0), constant(0.0)],
-            [parameter("f0(980)- re"), constant(0.0)],
-            [parameter("f0(1370)- re"), parameter("f0(1370)- im")],
-            [parameter("f0(1500)- re"), parameter("f0(1500)- im")],
-            [parameter("f0(1710)- re"), parameter("f0(1710)- im")],
-        ],
-        0,
-        &resonance_mass,
-    ));
-    let a0n = manager.register(KopfKMatrixA0::new(
-        "a0",
-        [
-            [parameter("a0(980)- re"), parameter("a0(980)- im")],
-            [parameter("a0(1450)- re"), parameter("a0(1450)- im")],
-        ],
-        0,
-        &resonance_mass,
-    ));
-    let f2 = manager.register(KopfKMatrixF2::new(
-        "f2",
-        [
-            [parameter("f2(1270) re"), parameter("f2(1270) im")],
-            [parameter("f2(1525) re"), parameter("f2(1525) im")],
-            [parameter("f2(1850) re"), parameter("f2(1850) im")],
-            [parameter("f2(1910) re"), parameter("f2(1910) im")],
-        ],
-        2,
-        &resonance_mass,
-    ));
-    let a2 = manager.register(KopfKMatrixA2::new(
-        "a2",
-        [
-            [parameter("a2(1320) re"), parameter("a2(1320) im")],
-            [parameter("a2(1700) re"), parameter("a2(1700) im")],
-        ],
-        2,
-        &resonance_mass,
-    ));
+    let z00p = manager
+        .register(Zlm::new(
+            "Z00+",
+            0,
+            0,
+            Sign::Positive,
+            &angles,
+            &polarization,
+        ))
+        .unwrap();
+    let z00n = manager
+        .register(Zlm::new(
+            "Z00-",
+            0,
+            0,
+            Sign::Negative,
+            &angles,
+            &polarization,
+        ))
+        .unwrap();
+    let z22p = manager
+        .register(Zlm::new(
+            "Z22+",
+            2,
+            2,
+            Sign::Positive,
+            &angles,
+            &polarization,
+        ))
+        .unwrap();
+    let f0p = manager
+        .register(KopfKMatrixF0::new(
+            "f0+",
+            [
+                [constant(0.0), constant(0.0)],
+                [parameter("f0(980)+ re"), constant(0.0)],
+                [parameter("f0(1370)+ re"), parameter("f0(1370)+ im")],
+                [parameter("f0(1500)+ re"), parameter("f0(1500)+ im")],
+                [parameter("f0(1710)+ re"), parameter("f0(1710)+ im")],
+            ],
+            0,
+            &resonance_mass,
+        ))
+        .unwrap();
+    let a0p = manager
+        .register(KopfKMatrixA0::new(
+            "a0+",
+            [
+                [parameter("a0(980)+ re"), parameter("a0(980)+ im")],
+                [parameter("a0(1450)+ re"), parameter("a0(1450)+ im")],
+            ],
+            0,
+            &resonance_mass,
+        ))
+        .unwrap();
+    let f0n = manager
+        .register(KopfKMatrixF0::new(
+            "f0-",
+            [
+                [constant(0.0), constant(0.0)],
+                [parameter("f0(980)- re"), constant(0.0)],
+                [parameter("f0(1370)- re"), parameter("f0(1370)- im")],
+                [parameter("f0(1500)- re"), parameter("f0(1500)- im")],
+                [parameter("f0(1710)- re"), parameter("f0(1710)- im")],
+            ],
+            0,
+            &resonance_mass,
+        ))
+        .unwrap();
+    let a0n = manager
+        .register(KopfKMatrixA0::new(
+            "a0-",
+            [
+                [parameter("a0(980)- re"), parameter("a0(980)- im")],
+                [parameter("a0(1450)- re"), parameter("a0(1450)- im")],
+            ],
+            0,
+            &resonance_mass,
+        ))
+        .unwrap();
+    let f2 = manager
+        .register(KopfKMatrixF2::new(
+            "f2",
+            [
+                [parameter("f2(1270) re"), parameter("f2(1270) im")],
+                [parameter("f2(1525) re"), parameter("f2(1525) im")],
+                [parameter("f2(1850) re"), parameter("f2(1850) im")],
+                [parameter("f2(1910) re"), parameter("f2(1910) im")],
+            ],
+            2,
+            &resonance_mass,
+        ))
+        .unwrap();
+    let a2 = manager
+        .register(KopfKMatrixA2::new(
+            "a2",
+            [
+                [parameter("a2(1320) re"), parameter("a2(1320) im")],
+                [parameter("a2(1700) re"), parameter("a2(1700) im")],
+            ],
+            2,
+            &resonance_mass,
+        ))
+        .unwrap();
     let s0p = f0p + a0p;
     let s0n = f0n + a0n;
     let d2p = f2 + a2;

diff --git a/src/amplitudes/breit_wigner.rs b/src/amplitudes/breit_wigner.rs
@@ -8,7 +8,7 @@ use crate::{
         functions::{blatt_weisskopf, breakup_momentum},
         variables::{Mass, Variable},
     },
-    Float, PI,
+    Float, LadduError, PI,
 };
 
 use super::{Amplitude, AmplitudeID};
@@ -63,7 +63,7 @@ impl BreitWigner {
 }
 
 impl Amplitude for BreitWigner {
-    fn register(&mut self, resources: &mut Resources) -> AmplitudeID {
+    fn register(&mut self, resources: &mut Resources) -> Result<AmplitudeID, LadduError> {
         self.pid_mass = resources.register_parameter(&self.mass);
         self.pid_width = resources.register_parameter(&self.width);
         resources.register_amplitude(&self.name)

diff --git a/src/amplitudes/common.rs b/src/amplitudes/common.rs
@@ -4,7 +4,7 @@ use crate::{
     amplitudes::{AmplitudeID, ParameterLike},
     data::Event,
     resources::{Cache, ParameterID, Parameters, Resources},
-    Float,
+    Float, LadduError,
 };
 
 use super::Amplitude;
@@ -30,7 +30,7 @@ impl Scalar {
 }
 
 impl Amplitude for Scalar {
-    fn register(&mut self, resources: &mut Resources) -> AmplitudeID {
+    fn register(&mut self, resources: &mut Resources) -> Result<AmplitudeID, LadduError> {
         self.pid = resources.register_parameter(&self.value);
         resources.register_amplitude(&self.name)
     }
@@ -66,7 +66,7 @@ impl ComplexScalar {
 }
 
 impl Amplitude for ComplexScalar {
-    fn register(&mut self, resources: &mut Resources) -> AmplitudeID {
+    fn register(&mut self, resources: &mut Resources) -> Result<AmplitudeID, LadduError> {
         self.pid_re = resources.register_parameter(&self.re);
         self.pid_im = resources.register_parameter(&self.im);
         resources.register_amplitude(&self.name)
@@ -103,7 +103,7 @@ impl PolarComplexScalar {
 }
 
 impl Amplitude for PolarComplexScalar {
-    fn register(&mut self, resources: &mut Resources) -> AmplitudeID {
+    fn register(&mut self, resources: &mut Resources) -> Result<AmplitudeID, LadduError> {
         self.pid_r = resources.register_parameter(&self.r);
         self.pid_theta = resources.register_parameter(&self.theta);
         resources.register_amplitude(&self.name)

diff --git a/src/amplitudes/kmatrix.rs b/src/amplitudes/kmatrix.rs
@@ -6,6 +6,7 @@ use num::traits::ConstOne;
 use num::traits::FloatConst;
 use num::Complex;
 
+use crate::LadduError;
 use crate::{
     amplitudes::{AmplitudeID, ParameterLike},
     data::Event,
@@ -214,7 +215,7 @@ impl KopfKMatrixF0 {
 }
 
 impl Amplitude for KopfKMatrixF0 {
-    fn register(&mut self, resources: &mut Resources) -> AmplitudeID {
+    fn register(&mut self, resources: &mut Resources) -> Result<AmplitudeID, LadduError> {
         for i in 0..self.couplings_indices_real.len() {
             self.couplings_indices_real[i] = resources.register_parameter(&self.couplings_real[i]);
             self.couplings_indices_imag[i] = resources.register_parameter(&self.couplings_imag[i]);
@@ -330,7 +331,7 @@ impl KopfKMatrixF2 {
 }
 
 impl Amplitude for KopfKMatrixF2 {
-    fn register(&mut self, resources: &mut Resources) -> AmplitudeID {
+    fn register(&mut self, resources: &mut Resources) -> Result<AmplitudeID, LadduError> {
         for i in 0..self.couplings_indices_real.len() {
             self.couplings_indices_real[i] = resources.register_parameter(&self.couplings_real[i]);
             self.couplings_indices_imag[i] = resources.register_parameter(&self.couplings_imag[i]);
@@ -438,7 +439,7 @@ impl KopfKMatrixA0 {
 }
 
 impl Amplitude for KopfKMatrixA0 {
-    fn register(&mut self, resources: &mut Resources) -> AmplitudeID {
+    fn register(&mut self, resources: &mut Resources) -> Result<AmplitudeID, LadduError> {
         for i in 0..self.couplings_indices_real.len() {
             self.couplings_indices_real[i] = resources.register_parameter(&self.couplings_real[i]);
             self.couplings_indices_imag[i] = resources.register_parameter(&self.couplings_imag[i]);
@@ -550,7 +551,7 @@ impl KopfKMatrixA2 {
 }
 
 impl Amplitude for KopfKMatrixA2 {
-    fn register(&mut self, resources: &mut Resources) -> AmplitudeID {
+    fn register(&mut self, resources: &mut Resources) -> Result<AmplitudeID, LadduError> {
         for i in 0..self.couplings_indices_real.len() {
             self.couplings_indices_real[i] = resources.register_parameter(&self.couplings_real[i]);
             self.couplings_indices_imag[i] = resources.register_parameter(&self.couplings_imag[i]);
@@ -661,7 +662,7 @@ impl KopfKMatrixRho {
 }
 
 impl Amplitude for KopfKMatrixRho {
-    fn register(&mut self, resources: &mut Resources) -> AmplitudeID {
+    fn register(&mut self, resources: &mut Resources) -> Result<AmplitudeID, LadduError> {
         for i in 0..self.couplings_indices_real.len() {
             self.couplings_indices_real[i] = resources.register_parameter(&self.couplings_real[i]);
             self.couplings_indices_imag[i] = resources.register_parameter(&self.couplings_imag[i]);
@@ -768,7 +769,7 @@ impl KopfKMatrixPi1 {
 }
 
 impl Amplitude for KopfKMatrixPi1 {
-    fn register(&mut self, resources: &mut Resources) -> AmplitudeID {
+    fn register(&mut self, resources: &mut Resources) -> Result<AmplitudeID, LadduError> {
         for i in 0..self.couplings_indices_real.len() {
             self.couplings_indices_real[i] = resources.register_parameter(&self.couplings_real[i]);
             self.couplings_indices_imag[i] = resources.register_parameter(&self.couplings_imag[i]);

diff --git a/src/amplitudes/mod.rs b/src/amplitudes/mod.rs
@@ -15,7 +15,7 @@ use rayon::prelude::*;
 use crate::{
     data::{Dataset, Event},
     resources::{Cache, Parameters, Resources},
-    Float,
+    Float, LadduError,
 };
 
 /// The Breit-Wigner amplitude.
@@ -66,7 +66,7 @@ pub trait Amplitude: DynClone + Send + Sync {
     /// the free parameters and cached values used by this [`Amplitude`]. It should end by
     /// returning an [`AmplitudeID`], which can be obtained from the
     /// [`Resources::register_amplitude`] method.
-    fn register(&mut self, resources: &mut Resources) -> AmplitudeID;
+    fn register(&mut self, resources: &mut Resources) -> Result<AmplitudeID, LadduError>;
     /// This method can be used to do some critical calculations ahead of time and
     /// store them in a [`Cache`]. These values can only depend on the data in an [`Event`],
     /// not on any free parameters in the fit. This method is opt-in since it is not required
@@ -254,11 +254,16 @@ pub struct Manager {
 impl Manager {
     /// Register the given [`Amplitude`] and return an [`AmplitudeID`] that can be used to build
     /// [`Expression`]s.
-    pub fn register(&mut self, amplitude: Box<dyn Amplitude>) -> AmplitudeID {
+    ///
+    /// # Errors
+    ///
+    /// The [`Amplitude`](crate::amplitudes::Amplitude)'s name must be unique and not already
+    /// registered, else this will return a [`RegistrationError`][LadduError::RegistrationError].
+    pub fn register(&mut self, amplitude: Box<dyn Amplitude>) -> Result<AmplitudeID, LadduError> {
         let mut amp = amplitude.clone();
-        let aid = amp.register(&mut self.resources);
+        let aid = amp.register(&mut self.resources)?;
         self.amplitudes.push(amp);
-        aid
+        Ok(aid)
     }
     /// Create an [`Evaluator`] which can compute the result of any [`Expression`] built on
     /// registered [`Amplitude`]s over the given [`Dataset`]. This method precomputes any relevant

diff --git a/src/amplitudes/ylm.rs b/src/amplitudes/ylm.rs
@@ -8,7 +8,7 @@ use crate::{
         functions::spherical_harmonic,
         variables::{Angles, Variable},
     },
-    Float,
+    Float, LadduError,
 };
 
 use super::Amplitude;
@@ -39,7 +39,7 @@ impl Ylm {
 }
 
 impl Amplitude for Ylm {
-    fn register(&mut self, resources: &mut Resources) -> AmplitudeID {
+    fn register(&mut self, resources: &mut Resources) -> Result<AmplitudeID, LadduError> {
         self.csid = resources.register_complex_scalar(None);
         resources.register_amplitude(&self.name)
     }