From bc13a35d28a1935489d616171388d6e35fcbe15e Mon Sep 17 00:00:00 2001 From: denehoffman Date: Fri, 22 Nov 2024 17:51:27 -0500 Subject: [PATCH 1/8] chore: bump dependencies --- Cargo.toml | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 599bd61..8123c31 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -21,8 +21,8 @@ crate-type = ["cdylib", "rlib"] indexmap = "2.6.0" num = "0.4.3" nalgebra = "0.33.2" -arrow = "53.2.0" -parquet = "53.2.0" +arrow = "53.3.0" +parquet = "53.3.0" factorial = "0.4.0" parking_lot = "0.12.3" dyn-clone = "1.0.17" @@ -36,11 +36,11 @@ pyo3 = { version = "0.22.6", optional = true, features = [ ] } numpy = { version = "0.22.1", optional = true, features = ["nalgebra"] } ganesh = "0.13.1" -thiserror = "2.0.0" +thiserror = "2.0.3" shellexpand = "3.1.0" accurate = "0.4.1" -serde = "1.0.214" -serde-pickle = "1.1.1" +serde = "1.0.215" +serde-pickle = "1.2.0" bincode = "1.3.3" [dev-dependencies] From a88808575462b3f0e19fc520baf5bd8252fafd99 Mon Sep 17 00:00:00 2001 From: denehoffman Date: Fri, 22 Nov 2024 17:51:27 -0500 Subject: [PATCH 2/8] chore: bump dependencies --- Cargo.toml | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 599bd61..8123c31 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -21,8 +21,8 @@ crate-type = ["cdylib", "rlib"] indexmap = "2.6.0" num = "0.4.3" nalgebra = "0.33.2" -arrow = "53.2.0" -parquet = "53.2.0" +arrow = "53.3.0" +parquet = "53.3.0" factorial = "0.4.0" parking_lot = "0.12.3" dyn-clone = "1.0.17" @@ -36,11 +36,11 @@ pyo3 = { version = "0.22.6", optional = true, features = [ ] } numpy = { version = "0.22.1", optional = true, features = ["nalgebra"] } ganesh = "0.13.1" -thiserror = "2.0.0" +thiserror = "2.0.3" shellexpand = "3.1.0" accurate = "0.4.1" -serde = "1.0.214" -serde-pickle = "1.1.1" +serde = "1.0.215" +serde-pickle = "1.2.0" bincode = "1.3.3" [dev-dependencies] From 9a8f23ad1a3944149428a08ec408022cf8b4b1ec Mon Sep 17 00:00:00 2001 From: denehoffman Date: Mon, 2 Dec 2024 22:56:57 -0500 Subject: [PATCH 3/8] fix: minor fixes for building without rayon/pyo3 --- src/lib.rs | 2 ++ src/likelihoods.rs | 14 ++++++-------- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index c8cd751..db03160 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -263,7 +263,9 @@ #![warn(clippy::perf, clippy::style, missing_docs)] +#[cfg(feature = "python")] use pyo3::PyErr; + use thiserror::Error; /// [`Amplitude`](crate::amplitudes::Amplitude)s and methods for making and evaluating them. diff --git a/src/likelihoods.rs b/src/likelihoods.rs index a92eb32..d643c70 100644 --- a/src/likelihoods.rs +++ b/src/likelihoods.rs @@ -382,7 +382,6 @@ impl LikelihoodTerm for NLL { #[cfg(not(feature = "rayon"))] fn evaluate(&self, parameters: &[Float]) -> Float { let data_result = self.data_evaluator.evaluate(parameters); - let n_data = self.data_evaluator.dataset.len() as Float; let mc_result = self.accmc_evaluator.evaluate(parameters); let n_mc = self.accmc_evaluator.dataset.len() as Float; let data_term: Float = data_result @@ -392,7 +391,7 @@ impl LikelihoodTerm for NLL { .sum_with_accumulator::>(); let mc_term: Float = mc_result .iter() - .zip(self.mc_evaluator.dataset.iter()) + .zip(self.accmc_evaluator.dataset.iter()) .map(|(l, e)| e.weight * l.re) .sum_with_accumulator::>(); -2.0 * (data_term - mc_term / n_mc) @@ -521,7 +520,6 @@ impl LikelihoodTerm for NLL { fn evaluate_gradient(&self, parameters: &[Float]) -> DVector { let data_resources = self.data_evaluator.resources.read(); let data_parameters = Parameters::new(parameters, &data_resources.constants); - let n_data = self.data_evaluator.dataset.weighted_len(); let mc_resources = self.accmc_evaluator.resources.read(); let mc_parameters = Parameters::new(parameters, &mc_resources.constants); let n_mc = self.accmc_evaluator.dataset.len() as Float; @@ -575,14 +573,14 @@ impl LikelihoodTerm for NLL { .sum(); let mc_term: DVector = self - .mc_evaluator + .accmc_evaluator .dataset .iter() .zip(mc_resources.caches.iter()) .map(|(event, cache)| { let mut gradient_values = - vec![DVector::zeros(parameters.len()); self.mc_evaluator.amplitudes.len()]; - self.mc_evaluator + vec![DVector::zeros(parameters.len()); self.accmc_evaluator.amplitudes.len()]; + self.accmc_evaluator .amplitudes .iter() .zip(mc_resources.active.iter()) @@ -595,7 +593,7 @@ impl LikelihoodTerm for NLL { ( event.weight, AmplitudeValues( - self.mc_evaluator + self.accmc_evaluator .amplitudes .iter() .zip(mc_resources.active.iter()) @@ -614,7 +612,7 @@ impl LikelihoodTerm for NLL { .map(|(weight, amp_vals, grad_vals)| { ( weight, - self.mc_evaluator + self.accmc_evaluator .expression .evaluate_gradient(&_vals, &grad_vals), ) From 0b6218054f80043b5bce7fc7968893dd2223cab4 Mon Sep 17 00:00:00 2001 From: denehoffman Date: Mon, 2 Dec 2024 22:57:20 -0500 Subject: [PATCH 4/8] feat(bench): updated benchmark to run over available parallelism --- Cargo.toml | 1 + benches/kmatrix_benchmark.rs | 46 +++++++++++++++++++++++++----------- 2 files changed, 33 insertions(+), 14 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 8123c31..6371356 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -48,6 +48,7 @@ approx = "0.5.1" criterion = { version = "2.7.2", package = "codspeed-criterion-compat", features = [ "html_reports", ] } +num_cpus = "1.16.0" [[bench]] name = "kmatrix_benchmark" diff --git a/benches/kmatrix_benchmark.rs b/benches/kmatrix_benchmark.rs index acfa16d..afc97b9 100644 --- a/benches/kmatrix_benchmark.rs +++ b/benches/kmatrix_benchmark.rs @@ -1,6 +1,6 @@ use std::time::Duration; -use criterion::{black_box, criterion_group, criterion_main, BatchSize, Criterion}; +use criterion::{black_box, criterion_group, criterion_main, BatchSize, BenchmarkId, Criterion}; use laddu::{ amplitudes::{ constant, @@ -19,6 +19,10 @@ use laddu::{ }; use rand::{distributions::Uniform, prelude::*}; +#[cfg(feature = "rayon")] +use rayon::ThreadPoolBuilder; + +#[cfg(feature = "rayon")] fn kmatrix_nll_benchmark(c: &mut Criterion) { let ds_data = open("benches/bench.parquet").unwrap(); let ds_mc = open("benches/bench.parquet").unwrap(); @@ -140,20 +144,34 @@ fn kmatrix_nll_benchmark(c: &mut Criterion) { let neg_im = (&s0n * z00n.imag()).norm_sqr(); let model = pos_re + pos_im + neg_re + neg_im; let nll = NLL::new(&manager, &model, &ds_data, &ds_mc); - let mut rng = rand::thread_rng(); - let range = Uniform::new(-100.0, 100.0); - c.bench_function("kmatrix benchmark (nll)", |b| { - b.iter_batched( - || { - let p: Vec = (0..nll.parameters().len()) - .map(|_| rng.sample(range)) - .collect(); - p + let mut group = c.benchmark_group("K-Matrix NLL Performance"); + for threads in 1..=num_cpus::get() { + let pool = ThreadPoolBuilder::new() + .num_threads(threads) + .build() + .unwrap(); + group.bench_with_input( + BenchmarkId::from_parameter(threads), + &threads, + |b, &_threads| { + pool.install(|| { + let mut rng = rand::thread_rng(); + let range = Uniform::new(-100.0, 100.0); + b.iter_batched( + || { + let p: Vec = (0..nll.parameters().len()) + .map(|_| rng.sample(range)) + .collect(); + p + }, + |p| black_box(nll.evaluate(&p)), + BatchSize::SmallInput, + ) + }) }, - |p| black_box(nll.evaluate(&p)), - BatchSize::SmallInput, - ) - }); + ); + } + group.finish(); } criterion_group! { From 81273d1c06d5cb9c822048abf243180d675d9931 Mon Sep 17 00:00:00 2001 From: denehoffman Date: Tue, 3 Dec 2024 00:25:44 -0500 Subject: [PATCH 5/8] refactor: change order of four-vector components and modify operation of `boost` This is a *breaking change* and effects every application that uses four-vectors. The `boost_along` method was removed because its usage can be confusing when boosting in a negative direction. Negating a four-vector negates the energy component because these are just regular vectors with a thin wrapper. To avoid confusion and potential mistakes, the method was removed entirely. Only the `boost` method should be used. The behavior of the `boost` method has also changed to match the Python `vector` and ROOT `TLorentzVector` conventions. Boosting a vector like `p.boost(&-p.beta())` will result in a vector with no three-momentum. Previously, this same operation was done without the `-` sign, but this is confusing since it seems to imply that doubling the velocity of a vector moves it to its center-of-momentum frame. The order of four-vectors was changed such that energy now comes last: `[px, py, pz, e]`. This was done because the `nalgebra::Vector4` struct has methods like `p.xy()` that get the first two components, which previously would have given `[e, px]` while the user might expect `[px, py]`. This behavior is now the default. Two fortunate results of writing this in Rust and consolidating everything to one crate. First, the test suite made it easy to refactor. Second, most Python users will be unaffected by this breaking change, even though it is breaking, simply because it does not change any operation outside of raw vector math. Since all the amplitudes are coded in Rust and included in the crate, the modified versions will just work in Python, and no Python amplitudes could exist (yet). If anyone implemented an amplitude in Rust, however, it might need modification, but only if it performs vector operations. Many amplitudes use `Variable`s defined in this crate as a shortcut around such math, so they might also be entirely unaffected. --- python/laddu/convert.py | 35 ++++++------ python/laddu/utils/vectors/__init__.pyi | 3 +- src/data.rs | 13 ++--- src/python.rs | 71 +++++++++++++------------ src/utils/variables.rs | 12 ++--- src/utils/vectors.rs | 46 +++++++--------- 6 files changed, 90 insertions(+), 90 deletions(-) diff --git a/python/laddu/convert.py b/python/laddu/convert.py index dee9cde..584555e 100644 --- a/python/laddu/convert.py +++ b/python/laddu/convert.py @@ -20,7 +20,9 @@ from loguru import logger -def read_root_file(input_file, tree_name, pol_in_beam, pol_angle, pol_magnitude, num_entries=None): +def read_root_file( + input_file, tree_name, pol_in_beam, pol_angle, pol_magnitude, num_entries=None +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: """Read ROOT file and extract data with optional polarization from the beam.""" logger.info(f"Reading ROOT file: {input_file}") tfile = uproot.open(input_file) @@ -40,10 +42,10 @@ def read_root_file(input_file, tree_name, pol_in_beam, pol_angle, pol_magnitude, Pz_final = np.array(list(tree["Pz_FinalState"].array(library="np", entry_stop=num_entries))) # pyright: ignore # Handle beam four-vector: (nevents, 4) - p4_beam = np.stack([E_beam, Px_beam, Py_beam, Pz_beam], axis=-1) + p4_beam = np.stack([Px_beam, Py_beam, Pz_beam, E_beam], axis=-1) # Handle final state four-vectors: (nevents, nparticles, 4) - p4_final = np.stack([E_final, Px_final, Py_final, Pz_final], axis=-1) + p4_final = np.stack([Px_final, Py_final, Pz_final, E_final], axis=-1) # Check if EPS branch exists and update eps if needed if "EPS" in tree: @@ -58,8 +60,9 @@ def read_root_file(input_file, tree_name, pol_in_beam, pol_angle, pol_magnitude, logger.info("Using beam's momentum for polarization (eps).") eps = np.stack([Px_beam, Py_beam, Pz_beam], axis=-1)[:, np.newaxis] # Reset beam momentum - p4_beam[:, 1:] = 0 # Set Px, Py to 0 - p4_beam[:, 3] = E_beam # Set Pz = E for beam + p4_beam[:, 0] = 0 # Set Px to 0 + p4_beam[:, 1] = 0 # Set Py to 0 + p4_beam[:, 2] = E_beam # Set Pz = E for beam elif pol_angle is not None and pol_magnitude is not None: logger.info(f"Using input polarization angle ({pol_angle}) and magnitude ({pol_magnitude}).") eps_x = pol_magnitude * np.cos(pol_angle) * np.ones_like(E_beam) @@ -74,10 +77,10 @@ def read_root_file(input_file, tree_name, pol_in_beam, pol_angle, pol_magnitude, logger.info("Concatenating beam and final state particles.") p4s = np.concatenate([p4_beam[:, np.newaxis, :], p4_final], axis=1) - return p4s.astype(np.float32), weight, eps.astype(np.float32) + return p4s.astype(np.float32), eps.astype(np.float32), weight -def save_as_parquet(p4s, weight, eps, output_file): +def save_as_parquet(p4s, eps, weight, output_file): """Save the processed data into Parquet format.""" logger.info("Saving data to Parquet format.") @@ -85,10 +88,10 @@ def save_as_parquet(p4s, weight, eps, output_file): columns = {} n_particles = p4s.shape[1] for i in range(n_particles): - columns[f"p4_{i}_E"] = p4s[:, i, 0] - columns[f"p4_{i}_Px"] = p4s[:, i, 1] - columns[f"p4_{i}_Py"] = p4s[:, i, 2] - columns[f"p4_{i}_Pz"] = p4s[:, i, 3] + columns[f"p4_{i}_Px"] = p4s[:, i, 0] + columns[f"p4_{i}_Py"] = p4s[:, i, 1] + columns[f"p4_{i}_Pz"] = p4s[:, i, 2] + columns[f"p4_{i}_E"] = p4s[:, i, 3] n_eps = eps.shape[1] for i in range(n_eps): @@ -110,12 +113,12 @@ def convert_from_amptools( output_path: Path, tree_name: str = "kin", pol_in_beam: bool = False, # noqa: FBT001, FBT002 - pol_angle_deg: float | None = None, + pol_angle: float | None = None, pol_magnitude: float | None = None, num_entries: int | None = None, ): - p4s, weight, eps = read_root_file(input_path, tree_name, pol_in_beam, pol_angle_deg, pol_magnitude, num_entries) - save_as_parquet(p4s, weight, eps, output_path) + p4s, eps, weight = read_root_file(input_path, tree_name, pol_in_beam, pol_angle, pol_magnitude, num_entries) + save_as_parquet(p4s, eps, weight, output_path) def run(): @@ -136,6 +139,4 @@ def run(): df_read = pd.read_parquet(output_file) print("Output Parquet File (head):") # noqa: T201 print(df_read.head()) # noqa: T201 - print("Output Columns:") # noqa: T201 - for column in df_read.columns: - print(column) # noqa: T201 + print(f"Total Entries: {len(df_read)}") # noqa: T201 diff --git a/python/laddu/utils/vectors/__init__.pyi b/python/laddu/utils/vectors/__init__.pyi index ec65178..f9da5fe 100644 --- a/python/laddu/utils/vectors/__init__.pyi +++ b/python/laddu/utils/vectors/__init__.pyi @@ -31,10 +31,9 @@ class Vector4: beta: Vector3 m: float m2: float - def __init__(self, e: float, px: float, py: float, pz: float): ... + def __init__(self, px: float, py: float, pz: float, e: float): ... def __add__(self, other: Vector4) -> Vector4: ... def boost(self, beta: Vector3) -> Vector4: ... - def boost_along(self, other: Vector4) -> Vector4: ... def to_numpy(self) -> npt.NDArray[np.float64]: ... @staticmethod def from_momentum(momentum: Vector3, mass: float) -> Vector4: ... diff --git a/src/data.rs b/src/data.rs index e95ca35..200e6d1 100644 --- a/src/data.rs +++ b/src/data.rs @@ -358,7 +358,7 @@ fn batch_to_event(batch: &RecordBatch, row: usize) -> Event { .downcast_ref::() .unwrap() .value(row) as Float; - p4s.push(Vector4::new(e, px, py, pz)); + p4s.push(Vector4::new(px, py, pz, e)); } // TODO: insert empty vectors if not provided @@ -532,11 +532,12 @@ mod tests { #[test] fn test_event_p4_sum() { let event = test_event(); + println!("{}", event); let sum = event.get_p4_sum([2, 3]); - assert_relative_eq!(sum[0], event.p4s[2].e() + event.p4s[3].e()); - assert_relative_eq!(sum[1], event.p4s[2].px() + event.p4s[3].px()); - assert_relative_eq!(sum[2], event.p4s[2].py() + event.p4s[3].py()); - assert_relative_eq!(sum[3], event.p4s[2].pz() + event.p4s[3].pz()); + assert_relative_eq!(sum[0], event.p4s[2].px() + event.p4s[3].px()); + assert_relative_eq!(sum[1], event.p4s[2].py() + event.p4s[3].py()); + assert_relative_eq!(sum[2], event.p4s[2].pz() + event.p4s[3].pz()); + assert_relative_eq!(sum[3], event.p4s[2].e() + event.p4s[3].e()); } #[test] @@ -600,7 +601,7 @@ mod tests { struct BeamEnergy; impl Variable for BeamEnergy { fn value(&self, event: &Event) -> Float { - event.p4s[0][0] + event.p4s[0].e() } } diff --git a/src/python.rs b/src/python.rs index e0704ee..0929c7b 100644 --- a/src/python.rs +++ b/src/python.rs @@ -264,6 +264,22 @@ pub(crate) mod laddu { fn to_numpy<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray1> { PyArray1::from_slice_bound(py, self.0.as_slice()) } + /// Convert an array into a 3-vector + /// + /// Parameters + /// ---------- + /// array_like + /// An array containing the components of this ``Vector3`` + /// + /// Returns + /// ------- + /// laddu_vec: Vector3 + /// A copy of the input array as a ``laddu`` vector + /// + #[staticmethod] + fn from_array(array: Vec) -> Self { + Self::new(array[0], array[1], array[2]) + } fn __repr__(&self) -> String { format!("{:?}", self.0) } @@ -274,15 +290,15 @@ pub(crate) mod laddu { /// A 4-momentum vector formed from energy and Cartesian 3-momentum components /// - /// This vector is ordered with energy as the zeroeth component (:math:`[E, p_x, p_y, p_z]`) and assumes a :math:`(+---)` + /// This vector is ordered with energy as the fourth component (:math:`[p_x, p_y, p_z, E]`) and assumes a :math:`(---+)` /// signature /// /// Parameters /// ---------- - /// e : float - /// The energy component /// px, py, pz : float /// The Cartesian components of the 3-vector + /// e : float + /// The energy component /// /// #[pyclass] @@ -291,8 +307,8 @@ pub(crate) mod laddu { #[pymethods] impl Vector4 { #[new] - fn new(e: Float, px: Float, py: Float, pz: Float) -> Self { - Self(nalgebra::Vector4::new(e, px, py, pz)) + fn new(px: Float, py: Float, pz: Float, e: Float) -> Self { + Self(nalgebra::Vector4::new(px, py, pz, e)) } fn __add__(&self, other: &Bound<'_, PyAny>) -> PyResult { if let Ok(other_vec) = other.extract::>() { @@ -382,7 +398,7 @@ pub(crate) mod laddu { /// The resulting 4-momentum is equal to the original boosted to an inertial frame with /// relative velocity :math:`\beta`: /// - /// .. math:: \left[E'; \vec{p}'\right] = \left[ \gamma E - \vec{\beta}\cdot\vec{p}; \vec{p} + \left(\frac{(\gamma - 1) \vec{p}\cdot\vec{\beta}}{\beta^2} - \gamma E\right)\vec{\beta}\right] + /// .. math:: \left[\vec{p}'; E'\right] = \left[ \vec{p} + \left(\frac{(\gamma - 1) \vec{p}\cdot\vec{\beta}}{\beta^2} + \gamma E\right)\vec{\beta}; \gamma E + \vec{\beta}\cdot\vec{p} \right] /// /// Parameters /// ---------- @@ -398,7 +414,6 @@ pub(crate) mod laddu { /// -------- /// Vector4.beta /// Vector4.gamma - /// Vector4.boost_along /// fn boost(&self, beta: &Vector3) -> Self { Self(self.0.boost(&beta.0)) @@ -477,7 +492,6 @@ pub(crate) mod laddu { /// -------- /// Vector4.beta /// Vector4.boost - /// Vector4.boost_along /// #[getter] fn gamma(&self) -> Float { @@ -498,7 +512,6 @@ pub(crate) mod laddu { /// -------- /// Vector4.gamma /// Vector4.boost - /// Vector4.boost_along /// #[getter] fn beta(&self) -> Vector3 { @@ -542,30 +555,6 @@ pub(crate) mod laddu { fn m2(&self) -> Float { self.0.m2() } - /// Boost the given 4-momentum into the rest frame of another - /// - /// The resulting 4-momentum is equal to the original boosted into the rest frame - /// of `other` - /// - /// Boosting a vector by itself should yield that vector in its own rest frame - /// - /// Parameters - /// ---------- - /// other : Vector4 - /// The 4-momentum whose rest-frame is the boost target - /// - /// Returns - /// ------- - /// Vector4 - /// The boosted 4-momentum - /// - /// See Also - /// -------- - /// Vector4.boost - /// - fn boost_along(&self, other: &Self) -> Self { - Self(self.0.boost_along(&other.0)) - } /// Construct a 4-momentum from a 3-momentum and corresponding `mass` /// /// The mass-energy equivalence is used to compute the energy of the 4-momentum: @@ -598,6 +587,22 @@ pub(crate) mod laddu { fn to_numpy<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray1> { PyArray1::from_slice_bound(py, self.0.as_slice()) } + /// Convert an array into a 4-vector + /// + /// Parameters + /// ---------- + /// array_like + /// An array containing the components of this ``Vector4`` + /// + /// Returns + /// ------- + /// laddu_vec: Vector4 + /// A copy of the input array as a ``laddu`` vector + /// + #[staticmethod] + fn from_array(array: Vec) -> Self { + Self::new(array[0], array[1], array[2], array[3]) + } fn __str__(&self) -> String { format!("{}", self.0) } diff --git a/src/utils/variables.rs b/src/utils/variables.rs index 7898cd3..5dac0db 100644 --- a/src/utils/variables.rs +++ b/src/utils/variables.rs @@ -96,10 +96,10 @@ impl Variable for CosTheta { let recoil = event.get_p4_sum(&self.recoil); let daughter = event.get_p4_sum(&self.daughter); let resonance = event.get_p4_sum(&self.resonance); - let daughter_res = daughter.boost_along(&resonance); + let daughter_res = daughter.boost(&-resonance.beta()); match self.frame { Frame::Helicity => { - let recoil_res = recoil.boost_along(&resonance); + let recoil_res = recoil.boost(&-resonance.beta()); let z = -recoil_res.vec3().unit(); let y = beam.vec3().cross(&-recoil.vec3()).unit(); let x = y.cross(&z); @@ -111,7 +111,7 @@ impl Variable for CosTheta { angles.costheta() } Frame::GottfriedJackson => { - let beam_res = beam.boost_along(&resonance); + let beam_res = beam.boost(&-resonance.beta()); let z = beam_res.vec3().unit(); let y = beam.vec3().cross(&-recoil.vec3()).unit(); let x = y.cross(&z); @@ -174,10 +174,10 @@ impl Variable for Phi { let recoil = event.get_p4_sum(&self.recoil); let daughter = event.get_p4_sum(&self.daughter); let resonance = event.get_p4_sum(&self.resonance); - let daughter_res = daughter.boost_along(&resonance); + let daughter_res = daughter.boost(&-resonance.beta()); match self.frame { Frame::Helicity => { - let recoil_res = recoil.boost_along(&resonance); + let recoil_res = recoil.boost(&-resonance.beta()); let z = -recoil_res.vec3().unit(); let y = beam.vec3().cross(&-recoil.vec3()).unit(); let x = y.cross(&z); @@ -189,7 +189,7 @@ impl Variable for Phi { angles.phi() } Frame::GottfriedJackson => { - let beam_res = beam.boost_along(&resonance); + let beam_res = beam.boost(&-resonance.beta()); let z = beam_res.vec3().unit(); let y = beam.vec3().cross(&-recoil.vec3()).unit(); let x = y.cross(&z); diff --git a/src/utils/vectors.rs b/src/utils/vectors.rs index f1447a1..5a28b89 100644 --- a/src/utils/vectors.rs +++ b/src/utils/vectors.rs @@ -4,9 +4,9 @@ use crate::Float; /// An object which behaves as a four-vector. pub trait FourVector { - /// The magnitude of the vector (with $`+---`$ signature). + /// The magnitude of the vector (with $`---+`$ signature). fn mag(&self) -> Float; - /// The squared magnitude of the vector (with $`+---`$ signature). + /// The squared magnitude of the vector (with $`---+`$ signature). fn mag2(&self) -> Float; /// Yields the three-vector part. fn vec3(&self) -> VectorView; @@ -34,8 +34,6 @@ pub trait FourMomentum: FourVector { fn m(&self) -> Float; /// The squared mass of the corresponding object. fn m2(&self) -> Float; - /// Yields the vector boosted along the direction of another [`FourMomentum`]. - fn boost_along(&self, other: &Self) -> Self; /// Creates a [`FourMomentum`] from a [`ThreeMomentum`] and a mass using $`E^2 = m^2 + p^2`$. fn from_momentum(momentum: &dyn ThreeMomentum, mass: Float) -> Self; /// Pretty-prints the four-momentum. @@ -83,36 +81,36 @@ impl FourVector for Vector4 { } fn mag2(&self) -> Float { - self[0] * self[0] - (self[1] * self[1] + self[2] * self[2] + self[3] * self[3]) + self[3] * self[3] - (self[0] * self[0] + self[1] * self[1] + self[2] * self[2]) } fn boost(&self, beta: &Vector3) -> Self { let b2 = beta.dot(beta); let gamma = 1.0 / Float::sqrt(1.0 - b2); let p3 = - self.vec3() + beta * ((gamma - 1.0) * self.vec3().dot(beta) / b2 - gamma * self[0]); - Vector4::new(gamma * (self[0] - beta.dot(&self.vec3())), p3.x, p3.y, p3.z) + self.vec3() + beta * ((gamma - 1.0) * self.vec3().dot(beta) / b2 + gamma * self[3]); + Vector4::new(p3.x, p3.y, p3.z, gamma * (self[3] + beta.dot(&self.vec3()))) } fn vec3(&self) -> VectorView { - self.fixed_rows::<3>(1) + self.fixed_rows::<3>(0) } } impl FourMomentum for Vector4 { - fn e(&self) -> Float { + fn px(&self) -> Float { self[0] } - fn px(&self) -> Float { + fn py(&self) -> Float { self[1] } - fn py(&self) -> Float { + fn pz(&self) -> Float { self[2] } - fn pz(&self) -> Float { + fn e(&self) -> Float { self[3] } @@ -138,13 +136,9 @@ impl FourMomentum for Vector4 { self.mag2() } - fn boost_along(&self, other: &Self) -> Self { - self.boost(&other.beta()) - } - fn from_momentum(momentum: &dyn ThreeMomentum, mass: Float) -> Self { let e = Float::sqrt(mass.powi(2) + momentum.mag2()); - Self::new(e, momentum.px(), momentum.py(), momentum.pz()) + Self::new(momentum.px(), momentum.py(), momentum.pz(), e) } } @@ -237,7 +231,7 @@ mod tests { #[test] fn test_four_momentum_basics() { - let p = vector![10.0, 3.0, 4.0, 5.0]; + let p = vector![3.0, 4.0, 5.0, 10.0]; assert_eq!(p.e(), 10.0); assert_eq!(p.px(), 3.0); assert_eq!(p.py(), 4.0); @@ -258,8 +252,8 @@ mod tests { #[test] fn test_three_momentum_basics() { - let p = vector![10.0, 3.0, 4.0, 5.0]; - let q = vector![0.0, 1.2, -3.4, 7.6]; + let p = vector![3.0, 4.0, 5.0, 10.0]; + let q = vector![1.2, -3.4, 7.6, 0.0]; let p3_view = p.momentum(); let q3_view = q.momentum(); assert_eq!(p3_view.px(), 3.0); @@ -291,18 +285,18 @@ mod tests { #[test] fn test_boost_com() { - let p = vector![10.0, 3.0, 4.0, 5.0]; - let zero = p.boost_along(&(-p)); + let p = vector![3.0, 4.0, 5.0, 10.0]; + let zero = p.boost(&-p.beta()); + assert_relative_eq!(zero[0], 0.0); assert_relative_eq!(zero[1], 0.0); assert_relative_eq!(zero[2], 0.0); - assert_relative_eq!(zero[3], 0.0); } #[test] fn test_boost() { - let p1 = vector![10.0, 3.0, 4.0, 5.0]; - let p2 = vector![9.0, 3.4, 2.3, 1.2]; - let p1_boosted = p1.boost_along(&p2); + let p1 = vector![3.0, 4.0, 5.0, 10.0]; + let p2 = vector![3.4, 2.3, 1.2, 9.0]; + let p1_boosted = p1.boost(&-p2.beta()); assert_relative_eq!(p1_boosted.e(), 8.157632144622882); assert_relative_eq!(p1_boosted.px(), -0.6489200627053444); assert_relative_eq!(p1_boosted.py(), 1.5316128987581492); From 9cc913a9419bbfd09a6520269acace8099884e62 Mon Sep 17 00:00:00 2001 From: denehoffman Date: Tue, 3 Dec 2024 17:19:52 -0500 Subject: [PATCH 6/8] refactor: get rid of `from_momentum` method and replace with methods coming from 3-vectors Instead of `Vector4::from_momentum`, use `Vector3::with_mass` or `Vector3::with_energy` (and similar for python). Also adds `from_array` methods for Python to make for easier conversion from `numpy` arrays or lists. --- python/laddu/utils/vectors/__init__.pyi | 8 +++- src/data.rs | 18 ++++---- src/python.rs | 56 +++++++++++++++---------- src/utils/vectors.rs | 45 ++++++++++++++++---- 4 files changed, 89 insertions(+), 38 deletions(-) diff --git a/python/laddu/utils/vectors/__init__.pyi b/python/laddu/utils/vectors/__init__.pyi index f9da5fe..17e197f 100644 --- a/python/laddu/utils/vectors/__init__.pyi +++ b/python/laddu/utils/vectors/__init__.pyi @@ -1,3 +1,5 @@ +from collections.abc import Sequence + import numpy as np import numpy.typing as npt @@ -17,6 +19,10 @@ class Vector3: def dot(self, other: Vector3) -> float: ... def cross(self, other: Vector3) -> Vector3: ... def to_numpy(self) -> npt.NDArray[np.float64]: ... + @staticmethod + def from_array(array: Sequence) -> Vector3: ... + def with_mass(self, mass: float) -> Vector4: ... + def with_energy(self, mass: float) -> Vector4: ... class Vector4: mag: float @@ -36,4 +42,4 @@ class Vector4: def boost(self, beta: Vector3) -> Vector4: ... def to_numpy(self) -> npt.NDArray[np.float64]: ... @staticmethod - def from_momentum(momentum: Vector3, mass: float) -> Vector4: ... + def from_array(array: Sequence) -> Vector4: ... diff --git a/src/data.rs b/src/data.rs index 200e6d1..45f4d59 100644 --- a/src/data.rs +++ b/src/data.rs @@ -24,10 +24,10 @@ pub fn test_event() -> Event { use crate::utils::vectors::*; Event { p4s: vec![ - Vector4::from_momentum(&vector![0.0, 0.0, 8.747], 0.0), // beam - Vector4::from_momentum(&vector![0.119, 0.374, 0.222], 1.007), // "proton" - Vector4::from_momentum(&vector![-0.112, 0.293, 3.081], 0.498), // "kaon" - Vector4::from_momentum(&vector![-0.007, -0.667, 5.446], 0.498), // "kaon" + vector![0.0, 0.0, 8.747].with_mass(0.0), // beam + vector![0.119, 0.374, 0.222].with_mass(1.007), // "proton" + vector![-0.112, 0.293, 3.081].with_mass(0.498), // "kaon" + vector![-0.007, -0.667, 5.446].with_mass(0.498), // "kaon" ], eps: vec![vector![0.385, 0.022, 0.000]], weight: 0.48, @@ -519,6 +519,8 @@ impl BinnedDataset { #[cfg(test)] mod tests { + use crate::prelude::ThreeMomentum; + use super::*; use approx::{assert_relative_eq, assert_relative_ne}; #[test] @@ -571,8 +573,8 @@ mod tests { let mut dataset = test_dataset(); dataset.events.push(Arc::new(Event { p4s: vec![ - Vector4::from_momentum(&vector![0.0, 0.0, 5.0], 0.0), - Vector4::from_momentum(&vector![0.0, 0.0, 1.0], 1.0), + vector![0.0, 0.0, 5.0].with_mass(0.0), + vector![0.0, 0.0, 1.0].with_mass(1.0), ], eps: vec![], weight: 1.0, @@ -587,12 +589,12 @@ mod tests { fn test_binned_dataset() { let mut dataset = Dataset::default(); dataset.events.push(Arc::new(Event { - p4s: vec![Vector4::from_momentum(&vector![0.0, 0.0, 1.0], 1.0)], + p4s: vec![vector![0.0, 0.0, 1.0].with_mass(1.0)], eps: vec![], weight: 1.0, })); dataset.events.push(Arc::new(Event { - p4s: vec![Vector4::from_momentum(&vector![0.0, 0.0, 2.0], 2.0)], + p4s: vec![vector![0.0, 0.0, 2.0].with_mass(2.0)], eps: vec![], weight: 2.0, })); diff --git a/src/python.rs b/src/python.rs index 0929c7b..8791ea3 100644 --- a/src/python.rs +++ b/src/python.rs @@ -254,6 +254,40 @@ pub(crate) mod laddu { fn pz(&self) -> Float { self.0.pz() } + /// Convert a 3-vector momentum to a 4-momentum with the given mass + /// + /// The mass-energy equivalence is used to compute the energy of the 4-momentum: + /// + /// .. math:: E = \sqrt{m^2 + p^2} + /// + /// Parameters + /// ---------- + /// mass: float + /// The mass of the new 4-momentum + /// + /// Returns + /// ------- + /// Vector4 + /// A new 4-momentum with the given mass + /// + fn with_mass(&self, mass: Float) -> Vector4 { + Vector4(self.0.with_mass(mass)) + } + /// Convert a 3-vector momentum to a 4-momentum with the given energy + /// + /// Parameters + /// ---------- + /// energy: float + /// The mass of the new 4-momentum + /// + /// Returns + /// ------- + /// Vector4 + /// A new 4-momentum with the given energy + /// + fn with_energy(&self, mass: Float) -> Vector4 { + Vector4(self.0.with_energy(mass)) + } /// Convert the 3-vector to a ``numpy`` array /// /// Returns @@ -555,28 +589,6 @@ pub(crate) mod laddu { fn m2(&self) -> Float { self.0.m2() } - /// Construct a 4-momentum from a 3-momentum and corresponding `mass` - /// - /// The mass-energy equivalence is used to compute the energy of the 4-momentum: - /// - /// .. math:: E = \sqrt{m^2 + p^2} - /// - /// Parameters - /// ---------- - /// momentum : Vector3 - /// The spatial 3-momentum - /// mass : float - /// The associated rest mass - /// - /// Returns - /// ------- - /// Vector4 - /// A new 4-momentum with the given spatial `momentum` and `mass` - /// - #[staticmethod] - fn from_momentum(momentum: &Vector3, mass: Float) -> Self { - Self(nalgebra::Vector4::from_momentum(&momentum.0, mass)) - } /// Convert the 4-vector to a `numpy` array /// /// Returns diff --git a/src/utils/vectors.rs b/src/utils/vectors.rs index 5a28b89..dcd17c1 100644 --- a/src/utils/vectors.rs +++ b/src/utils/vectors.rs @@ -34,8 +34,6 @@ pub trait FourMomentum: FourVector { fn m(&self) -> Float; /// The squared mass of the corresponding object. fn m2(&self) -> Float; - /// Creates a [`FourMomentum`] from a [`ThreeMomentum`] and a mass using $`E^2 = m^2 + p^2`$. - fn from_momentum(momentum: &dyn ThreeMomentum, mass: Float) -> Self; /// Pretty-prints the four-momentum. fn to_p4_string(&self) -> String { format!( @@ -73,6 +71,10 @@ pub trait ThreeMomentum: ThreeVector { fn py(&self) -> Float; /// Momentum in the $`z`$-direction fn pz(&self) -> Float; + /// Converts this three-momentum to a four-momentum with the given mass. + fn with_mass(&self, mass: Float) -> Vector4; + /// Converts this three-momentum to a four-momentum with the given energy. + fn with_energy(&self, energy: Float) -> Vector4; } impl FourVector for Vector4 { @@ -135,11 +137,6 @@ impl FourMomentum for Vector4 { fn m2(&self) -> Float { self.mag2() } - - fn from_momentum(momentum: &dyn ThreeMomentum, mass: Float) -> Self { - let e = Float::sqrt(mass.powi(2) + momentum.mag2()); - Self::new(momentum.px(), momentum.py(), momentum.pz(), e) - } } impl ThreeMomentum for Vector3 { @@ -154,6 +151,15 @@ impl ThreeMomentum for Vector3 { fn pz(&self) -> Float { self[2] } + + fn with_mass(&self, mass: Float) -> Vector4 { + let e = Float::sqrt(mass.powi(2) + self.mag2()); + Vector4::new(self.px(), self.py(), self.pz(), e) + } + + fn with_energy(&self, energy: Float) -> Vector4 { + Vector4::new(self.px(), self.py(), self.pz(), energy) + } } impl ThreeVector for Vector3 { @@ -194,6 +200,15 @@ impl<'a> ThreeMomentum for VectorView<'a, Float, U3, U1, U4> { fn pz(&self) -> Float { self[2] } + + fn with_mass(&self, mass: Float) -> Vector4 { + let e = Float::sqrt(mass.powi(2) + self.mag2()); + Vector4::new(self.px(), self.py(), self.pz(), e) + } + + fn with_energy(&self, energy: Float) -> Vector4 { + Vector4::new(self.px(), self.py(), self.pz(), energy) + } } impl<'a> ThreeVector for VectorView<'a, Float, U3, U1, U4> { @@ -229,6 +244,22 @@ mod tests { use super::*; + #[test] + fn test_three_to_four_momentum_conversion() { + let p3 = vector![1.0, 2.0, 3.0]; + let target_p4 = vector![1.0, 2.0, 3.0, 10.0]; + let p4_from_mass = p3.with_mass(target_p4.m()); + assert_eq!(target_p4.e(), p4_from_mass.e()); + assert_eq!(target_p4.px(), p4_from_mass.px()); + assert_eq!(target_p4.py(), p4_from_mass.py()); + assert_eq!(target_p4.pz(), p4_from_mass.pz()); + let p4_from_energy = p3.with_energy(target_p4.e()); + assert_eq!(target_p4.e(), p4_from_energy.e()); + assert_eq!(target_p4.px(), p4_from_energy.px()); + assert_eq!(target_p4.py(), p4_from_energy.py()); + assert_eq!(target_p4.pz(), p4_from_energy.pz()); + } + #[test] fn test_four_momentum_basics() { let p = vector![3.0, 4.0, 5.0, 10.0]; From 0e3a4a6df73646db4775d58fe8216a104459b332 Mon Sep 17 00:00:00 2001 From: denehoffman Date: Tue, 3 Dec 2024 17:22:33 -0500 Subject: [PATCH 7/8] feat: add basic implementation to read directly from ROOT files This method is much slower than reading a parquet file, and I don't know if there really is a way to speed that up without using a Rust-side crate. I'm going to put that on the back-burner, but this works for now if people are desperate. --- python/laddu/__init__.py | 83 ++++++++++++++++++++++++++------------- python/laddu/__init__.pyi | 67 ++++++++++++++++++------------- 2 files changed, 96 insertions(+), 54 deletions(-) diff --git a/python/laddu/__init__.py b/python/laddu/__init__.py index e667424..9ee7706 100644 --- a/python/laddu/__init__.py +++ b/python/laddu/__init__.py @@ -1,12 +1,17 @@ +from __future__ import annotations + from abc import ABCMeta, abstractmethod +from typing import TYPE_CHECKING + +import numpy as np from laddu.amplitudes import Manager, constant, parameter from laddu.amplitudes.breit_wigner import BreitWigner from laddu.amplitudes.common import ComplexScalar, PolarComplexScalar, Scalar from laddu.amplitudes.ylm import Ylm from laddu.amplitudes.zlm import Zlm -from laddu.convert import convert_from_amptools -from laddu.data import BinnedDataset, Dataset, open +from laddu.convert import convert_from_amptools, read_root_file +from laddu.data import BinnedDataset, Dataset, Event, open # noqa: A004 from laddu.likelihoods import NLL, LikelihoodManager, Status from laddu.utils.variables import Angles, CosTheta, Mandelstam, Mass, Phi, PolAngle, Polarization, PolMagnitude from laddu.utils.vectors import Vector3, Vector4 @@ -14,6 +19,9 @@ from . import amplitudes, convert, data, likelihoods, utils from .laddu import version +if TYPE_CHECKING: + from pathlib import Path + __version__ = version() @@ -23,38 +31,59 @@ def callback(self, step: int, status: Status) -> tuple[Status, bool]: pass +def open_amptools( + path: str | Path, + tree: str = "kin", + *, + pol_in_beam: bool = False, + pol_angle: float | None = None, + pol_magnitude: float | None = None, + num_entries: int | None = None, +) -> Dataset: + pol_angle_rad = pol_angle * np.pi / 180 if pol_angle else None + p4s_list, eps_list, weight_list = read_root_file(path, tree, pol_in_beam, pol_angle_rad, pol_magnitude, num_entries) + return Dataset( + [ + Event([Vector4.from_array(p4) for p4 in p4s], [Vector3.from_array(eps_vec) for eps_vec in eps], weight) + for p4s, eps, weight in zip(p4s_list, eps_list, weight_list) + ] + ) + + __all__ = [ - "__version__", - "convert", - "convert_from_amptools", - "Dataset", - "open", + "NLL", + "Angles", "BinnedDataset", - "utils", - "data", - "amplitudes", - "likelihoods", - "Vector3", - "Vector4", + "BreitWigner", + "ComplexScalar", "CosTheta", - "Phi", - "Angles", - "PolMagnitude", - "PolAngle", - "Polarization", + "Dataset", + "Event", + "LikelihoodManager", + "Manager", "Mandelstam", "Mass", - "Manager", - "LikelihoodManager", - "NLL", - "Status", "Observer", - "parameter", - "constant", - "Scalar", - "ComplexScalar", + "Phi", + "PolAngle", + "PolMagnitude", "PolarComplexScalar", + "Polarization", + "Scalar", + "Status", + "Vector3", + "Vector4", "Ylm", "Zlm", - "BreitWigner", + "__version__", + "amplitudes", + "constant", + "convert", + "convert_from_amptools", + "data", + "likelihoods", + "open", + "open_amptools", + "parameter", + "utils", ] diff --git a/python/laddu/__init__.pyi b/python/laddu/__init__.pyi index 45a539b..d266b0f 100644 --- a/python/laddu/__init__.pyi +++ b/python/laddu/__init__.pyi @@ -1,4 +1,5 @@ from abc import ABCMeta, abstractmethod +from pathlib import Path from laddu.amplitudes import Expression, Manager, constant, parameter from laddu.amplitudes.breit_wigner import BreitWigner @@ -6,51 +7,63 @@ from laddu.amplitudes.common import ComplexScalar, PolarComplexScalar, Scalar from laddu.amplitudes.ylm import Ylm from laddu.amplitudes.zlm import Zlm from laddu.convert import convert_from_amptools -from laddu.data import BinnedDataset, Dataset, open +from laddu.data import BinnedDataset, Dataset, Event, open # noqa: A004 from laddu.likelihoods import NLL, LikelihoodManager, Status from laddu.utils.variables import Angles, CosTheta, Mandelstam, Mass, Phi, PolAngle, Polarization, PolMagnitude from laddu.utils.vectors import Vector3, Vector4 from . import amplitudes, convert, data, utils +__version__: str + class Observer(metaclass=ABCMeta): @abstractmethod def callback(self, step: int, status: Status) -> tuple[Status, bool]: ... -__version__: str +def open_amptools( + path: str | Path, + tree: str = "kin", + *, + pol_in_beam: bool = False, + pol_angle: float | None = None, + pol_magnitude: float | None = None, + num_entries: int | None = None, +) -> Dataset: ... __all__ = [ - "__version__", - "convert", - "convert_from_amptools", - "Dataset", - "open", + "NLL", + "Angles", "BinnedDataset", - "utils", - "data", - "amplitudes", - "Vector3", - "Vector4", + "BreitWigner", + "ComplexScalar", "CosTheta", - "Phi", - "Angles", - "PolMagnitude", - "PolAngle", - "Polarization", + "Dataset", + "Event", + "Expression", + "LikelihoodManager", + "Manager", "Mandelstam", "Mass", - "Manager", - "LikelihoodManager", - "NLL", - "Expression", - "Status", "Observer", - "parameter", - "constant", - "Scalar", - "ComplexScalar", + "Phi", + "PolAngle", + "PolMagnitude", "PolarComplexScalar", + "Polarization", + "Scalar", + "Status", + "Vector3", + "Vector4", "Ylm", "Zlm", - "BreitWigner", + "__version__", + "amplitudes", + "constant", + "convert", + "convert_from_amptools", + "data", + "open", + "open_amptools", + "parameter", + "utils", ] From 27a61a172c2df2e3e041be8b21d96886ab2ac8c1 Mon Sep 17 00:00:00 2001 From: denehoffman Date: Tue, 3 Dec 2024 17:52:48 -0500 Subject: [PATCH 8/8] fix: correct parallelism to allow for proper codspeed benchmarking --- benches/kmatrix_benchmark.rs | 26 ++++++++++++-------------- 1 file changed, 12 insertions(+), 14 deletions(-) diff --git a/benches/kmatrix_benchmark.rs b/benches/kmatrix_benchmark.rs index afc97b9..2e667c0 100644 --- a/benches/kmatrix_benchmark.rs +++ b/benches/kmatrix_benchmark.rs @@ -154,20 +154,18 @@ fn kmatrix_nll_benchmark(c: &mut Criterion) { BenchmarkId::from_parameter(threads), &threads, |b, &_threads| { - pool.install(|| { - let mut rng = rand::thread_rng(); - let range = Uniform::new(-100.0, 100.0); - b.iter_batched( - || { - let p: Vec = (0..nll.parameters().len()) - .map(|_| rng.sample(range)) - .collect(); - p - }, - |p| black_box(nll.evaluate(&p)), - BatchSize::SmallInput, - ) - }) + let mut rng = rand::thread_rng(); + let range = Uniform::new(-100.0, 100.0); + b.iter_batched( + || { + let p: Vec = (0..nll.parameters().len()) + .map(|_| rng.sample(range)) + .collect(); + p + }, + |p| pool.install(|| black_box(nll.evaluate(&p))), + BatchSize::SmallInput, + ) }, ); }