Merge pull request #15 from denehoffman/quality-of-life

Quality of life updates

.readthedocs.yaml

@@ -7,6 +7,7 @@ build:
   os: ubuntu-24.04
   tools:
     python: "3.12"
+    rust: "latest"
 
 sphinx:
   configuration: docs/source/conf.py
@@ -16,4 +17,7 @@ formats:
 
 python:
    install:
-   - requirements: docs/requirements.txt
+    - method: pip
+      path: .
+      extra_requirements:
+        - docs

Cargo.toml

@@ -11,6 +11,7 @@ license = "MIT OR Apache-2.0"
 keywords = ["PWA", "amplitude", "particle", "physics", "modeling"]
 categories = ["science", "mathematics"]
 exclude = ["/python_examples"]
+rust-version = "1.70.0"
 
 [lib]
 name = "laddu"
@@ -19,22 +20,22 @@ crate-type = ["cdylib", "rlib"]
 [dependencies]
 indexmap = "2.6.0"
 num = "0.4.3"
-nalgebra = "0.33.0"
-arrow = "53.1.0"
-parquet = "53.1.0"
+nalgebra = "0.33.2"
+arrow = "53.2.0"
+parquet = "53.2.0"
 factorial = "0.4.0"
 parking_lot = "0.12.3"
 dyn-clone = "1.0.17"
 auto_ops = "0.3.0"
 rand = "0.8.5"
 rand_chacha = "0.3.1"
 rayon = { version = "1.10.0", optional = true }
-pyo3 = { version = "0.22.5", optional = true, features = [
+pyo3 = { version = "0.22.6", optional = true, features = [
   "num-complex",
   "abi3-py37",
 ] }
-numpy = { version = "0.22.0", optional = true, features = ["nalgebra"] }
-ganesh = "0.13.0"
+numpy = { version = "0.22.1", optional = true, features = ["nalgebra"] }
+ganesh = "0.13.1"
 thiserror = "2.0.0"
 shellexpand = "3.1.0"
 accurate = "0.4.1"
@@ -60,7 +61,9 @@ default = ["rayon", "python"]
 extension-module = ["pyo3/extension-module"]
 rayon = ["dep:rayon"]
 f32 = []
-python = ["dep:pyo3", "dep:numpy", "extension-module"]
+python = ["pyo3", "numpy", "extension-module"]
+pyo3 = ["dep:pyo3"]
+numpy = ["dep:numpy"]
 
 [profile.release]
 lto = true

python/laddu/likelihoods/__init__.pyi

@@ -49,6 +49,9 @@ class NLL:
     def isolate(self, name: str | list[str]) -> None: ...
     def evaluate(self, parameters: list[float] | npt.NDArray[np.float64]) -> float: ...
     def project(self, parameters: list[float] | npt.NDArray[np.float64]) -> npt.NDArray[np.float64]: ...
+    def project_with(
+        self, parameters: list[float] | npt.NDArray[np.float64], name: str | list[str]
+    ) -> npt.NDArray[np.float64]: ...
     def minimize(
         self,
         p0: list[float] | npt.NDArray[np.float64],

src/likelihoods.rs

@@ -193,6 +193,68 @@ impl NLL {
             .map(|(l, e)| e.weight * l.re / 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 differs from the standard [`NLL::project`] in that it first
+    /// isolates the selected [`Amplitude`](`crate::amplitudes::Amplitude`)s by name, but returns
+    /// the [`NLL`] to its prior state after calculation.
+    ///
+    /// 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) / N_{\text{MC}}
+    /// ```
+    #[cfg(feature = "rayon")]
+    pub fn project_with<T: AsRef<str>>(&self, parameters: &[Float], names: &[T]) -> Vec<Float> {
+        let current_active_data = self.data_evaluator.resources.read().active.clone();
+        let current_active_mc = self.mc_evaluator.resources.read().active.clone();
+        self.isolate_many(names);
+        let mc_result = self.mc_evaluator.evaluate(parameters);
+        let n_mc = self.mc_evaluator.dataset.len() as Float;
+        let res = mc_result
+            .par_iter()
+            .zip(self.mc_evaluator.dataset.par_iter())
+            .map(|(l, e)| e.weight * l.re / n_mc)
+            .collect();
+        self.data_evaluator.resources.write().active = current_active_data;
+        self.mc_evaluator.resources.write().active = current_active_mc;
+        res
+    }
+
+    /// 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 differs from the standard [`NLL::project`] in that it first
+    /// isolates the selected [`Amplitude`](`crate::amplitudes::Amplitude`)s by name, but returns
+    /// the [`NLL`] to its prior state after calculation.
+    ///
+    /// 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) / N_{\text{MC}}
+    /// ```
+    #[cfg(not(feature = "rayon"))]
+    pub fn project_with<T: AsRef<str>>(&self, parameters: &[Float], names: &[T]) -> Vec<Float> {
+        let current_active_data = self.data_evaluator.resources.read().active.clone();
+        let current_active_mc = self.mc_evaluator.resources.read().active.clone();
+        self.isolate_many(names);
+        let mc_result = self.mc_evaluator.evaluate(parameters);
+        let n_mc = self.mc_evaluator.dataset.len() as Float;
+        let res = mc_result
+            .iter()
+            .zip(self.mc_evaluator.dataset.iter())
+            .map(|(l, e)| e.weight * l.re / n_mc)
+            .collect();
+        self.data_evaluator.resources.write().active = current_active_data;
+        self.mc_evaluator.resources.write().active = current_active_mc;
+        res
+    }
 }
 
 impl LikelihoodTerm for NLL {

src/python.rs

@@ -54,6 +54,9 @@ pub(crate) mod laddu {
         fn __add__(&self, other: Self) -> Self {
             Self(self.0 + other.0)
         }
+        fn __radd__(&self, other: Self) -> Self {
+            other.__add__(self.clone())
+        }
         /// The dot product
         ///
         /// Calculates the dot product of two Vector3s.
@@ -268,6 +271,9 @@ pub(crate) mod laddu {
         fn __add__(&self, other: Self) -> Self {
             Self(self.0 + other.0)
         }
+        fn __radd__(&self, other: Self) -> Self {
+            other.__add__(self.clone())
+        }
         /// The magnitude of the 4-vector
         ///
         /// This is calculated as:
@@ -1344,6 +1350,15 @@ pub(crate) mod laddu {
                 Err(PyTypeError::new_err("Unsupported operand type for +"))
             }
         }
+        fn __radd__(&self, other: &Bound<'_, PyAny>) -> PyResult<Expression> {
+            if let Ok(other_aid) = other.extract::<PyRef<AmplitudeID>>() {
+                Ok(Expression(other_aid.0.clone() + self.0.clone()))
+            } else if let Ok(other_expr) = other.extract::<Expression>() {
+                Ok(Expression(other_expr.0.clone() + self.0.clone()))
+            } else {
+                Err(PyTypeError::new_err("Unsupported operand type for +"))
+            }
+        }
         fn __mul__(&self, other: &Bound<'_, PyAny>) -> PyResult<Expression> {
             if let Ok(other_aid) = other.extract::<PyRef<AmplitudeID>>() {
                 Ok(Expression(self.0.clone() * other_aid.0.clone()))
@@ -1404,6 +1419,15 @@ pub(crate) mod laddu {
                 Err(PyTypeError::new_err("Unsupported operand type for +"))
             }
         }
+        fn __radd__(&self, other: &Bound<'_, PyAny>) -> PyResult<Expression> {
+            if let Ok(other_aid) = other.extract::<PyRef<AmplitudeID>>() {
+                Ok(Expression(other_aid.0.clone() + self.0.clone()))
+            } else if let Ok(other_expr) = other.extract::<Expression>() {
+                Ok(Expression(other_expr.0.clone() + self.0.clone()))
+            } else {
+                Err(PyTypeError::new_err("Unsupported operand type for +"))
+            }
+        }
         fn __mul__(&self, other: &Bound<'_, PyAny>) -> PyResult<Expression> {
             if let Ok(other_aid) = other.extract::<PyRef<AmplitudeID>>() {
                 Ok(Expression(self.0.clone() * other_aid.0.clone()))
@@ -1992,6 +2016,53 @@ pub(crate) mod laddu {
             PyArray1::from_slice_bound(py, &self.0.project(&parameters))
         }
 
+        /// Project the model over the Monte Carlo dataset with the given parameter values, first
+        /// isolating the given terms by name. The NLL is then reset to its previous state of
+        /// activation.
+        ///
+        /// This is defined as
+        ///
+        /// .. math:: e_w(\vec{p}) = \frac{e_w}{N_{MC}} \mathcal{L}(e)
+        ///
+        /// Parameters
+        /// ----------
+        /// parameters : list of float
+        ///     The values to use for the free parameters
+        /// arg : str or list of str
+        ///     Names of Amplitudes to be isolated
+        ///
+        /// Returns
+        /// -------
+        /// result : array_like
+        ///     Weights for every Monte Carlo event which represent the fit to data
+        ///
+        /// Raises
+        /// ------
+        /// TypeError
+        ///     If `arg` is not a str or list of str
+        ///
+        fn project_with<'py>(
+            &self,
+            py: Python<'py>,
+            parameters: Vec<Float>,
+            arg: &Bound<'_, PyAny>,
+        ) -> PyResult<Bound<'py, PyArray1<Float>>> {
+            let names = if let Ok(string_arg) = arg.extract::<String>() {
+                vec![string_arg]
+            } else if let Ok(list_arg) = arg.downcast::<PyList>() {
+                let vec: Vec<String> = list_arg.extract()?;
+                vec
+            } else {
+                return Err(PyTypeError::new_err(
+                    "Argument must be either a string or a list of strings",
+                ));
+            };
+            Ok(PyArray1::from_slice_bound(
+                py,
+                &self.0.project_with(&parameters, &names),
+            ))
+        }
+
         /// Minimize the NLL with respect to the free parameters in the model
         ///
         /// This method "runs the fit". Given an initial position `p0` and optional `bounds`, this
@@ -2122,6 +2193,15 @@ pub(crate) mod laddu {
                 Err(PyTypeError::new_err("Unsupported operand type for +"))
             }
         }
+        fn __radd__(&self, other: &Bound<'_, PyAny>) -> PyResult<LikelihoodExpression> {
+            if let Ok(other_aid) = other.extract::<PyRef<LikelihoodID>>() {
+                Ok(LikelihoodExpression(other_aid.0.clone() + self.0.clone()))
+            } else if let Ok(other_expr) = other.extract::<LikelihoodExpression>() {
+                Ok(LikelihoodExpression(other_expr.0.clone() + self.0.clone()))
+            } else {
+                Err(PyTypeError::new_err("Unsupported operand type for +"))
+            }
+        }
         fn __mul__(&self, other: &Bound<'_, PyAny>) -> PyResult<LikelihoodExpression> {
             if let Ok(other_aid) = other.extract::<PyRef<LikelihoodID>>() {
                 Ok(LikelihoodExpression(self.0.clone() * other_aid.0.clone()))
@@ -2150,6 +2230,15 @@ pub(crate) mod laddu {
                 Err(PyTypeError::new_err("Unsupported operand type for +"))
             }
         }
+        fn __radd__(&self, other: &Bound<'_, PyAny>) -> PyResult<LikelihoodExpression> {
+            if let Ok(other_aid) = other.extract::<PyRef<LikelihoodID>>() {
+                Ok(LikelihoodExpression(other_aid.0.clone() + self.0.clone()))
+            } else if let Ok(other_expr) = other.extract::<LikelihoodExpression>() {
+                Ok(LikelihoodExpression(other_expr.0.clone() + self.0.clone()))
+            } else {
+                Err(PyTypeError::new_err("Unsupported operand type for +"))
+            }
+        }
         fn __mul__(&self, other: &Bound<'_, PyAny>) -> PyResult<LikelihoodExpression> {
             if let Ok(other_aid) = other.extract::<PyRef<LikelihoodID>>() {
                 Ok(LikelihoodExpression(self.0.clone() * other_aid.0.clone()))