diff --git a/.gitignore b/.gitignore index a495fbb..680ee25 100644 --- a/.gitignore +++ b/.gitignore @@ -190,3 +190,5 @@ Cargo.lock # Others demos/ src/main.rs +*.perf.* +perf.* diff --git a/.justfile b/.justfile index e8f3033..fb68345 100644 --- a/.justfile +++ b/.justfile @@ -16,3 +16,17 @@ makedocs: odoc: firefox ./docs/build/html/index.html + +clean: + cargo clean + +profile: + RUSTFLAGS='-C force-frame-pointers=y' cargo build --profile perf + perf record -g target/perf/laddu + perf annotate -v --asm-raw --stdio + perf report -g graph,0.5,caller + +popen: + mv firefox.perf.data firefox.perf.data.old + perf script --input=perf.data -F +pid > firefox.perf.data + firefox https://profiler.firefox.com diff --git a/Cargo.toml b/Cargo.toml index 2109322..be226b9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -18,7 +18,7 @@ name = "laddu" crate-type = ["cdylib", "rlib"] [dependencies] -indexmap = "2.6.0" +indexmap = { version = "2.6.0", features = ["serde"] } num = "0.4.3" nalgebra = "0.33.2" arrow = "53.3.0" @@ -40,6 +40,8 @@ thiserror = "2.0.3" shellexpand = "3.1.0" accurate = "0.4.1" serde = "1.0.215" +serde_with = "3.11.0" +typetag = "0.2.18" serde-pickle = "1.2.0" bincode = "1.3.3" diff --git a/README.md b/README.md index 5ddb16b..bead637 100644 --- a/README.md +++ b/README.md @@ -84,10 +84,15 @@ is the relativistic width correction, $`q(m_a, m_b, m_c)`$ is the breakup moment Although this particular amplitude is already included in `laddu`, let's assume it isn't and imagine how we would write it from scratch: ```rust -use laddu::prelude::*; +use laddu::{ + ParameterLike, Event, Cache, Resources, Mass, + ParameterID, Parameters, Float, LadduError, PI, AmplitudeID, Complex, +}; +use laddu::traits::*; use laddu::utils::functions::{blatt_weisskopf, breakup_momentum}; +use laddu::{Deserialize, Serialize}; -#[derive(Clone)] +#[derive(Clone, Serialize, Deserialize)] pub struct MyBreitWigner { name: String, mass: ParameterLike, @@ -124,6 +129,7 @@ impl MyBreitWigner { } } +#[typetag::serde] impl Amplitude for MyBreitWigner { fn register(&mut self, resources: &mut Resources) -> Result { self.pid_mass = resources.register_parameter(&self.mass); @@ -151,6 +157,7 @@ impl Amplitude for MyBreitWigner { ### Calculating a Likelihood We could then write some code to use this amplitude. For demonstration purposes, let's just calculate an extended unbinned negative log-likelihood, assuming we have some data and Monte Carlo in the proper [parquet format](#data-format): ```rust +use laddu::{Scalar, Mass, Manager, NLL, parameter, open}; let ds_data = open("test_data/data.parquet").unwrap(); let ds_mc = open("test_data/mc.parquet").unwrap(); @@ -168,11 +175,12 @@ let bw = manager.register(MyBreitWigner::new( &resonance_mass, )).unwrap(); let mag = manager.register(Scalar::new("mag", parameter("magnitude"))).unwrap(); -let model = (mag * bw).norm_sqr(); +let expr = (mag * bw).norm_sqr(); +let model = manager.model(&expr); -let nll = NLL::new(&manager, &ds_data, &ds_mc); +let nll = NLL::new(&model, &ds_data, &ds_mc); println!("Parameters names and order: {:?}", nll.parameters()); -let result = nll.evaluate(&[1.27, 0.120, 100.0], &model); +let result = nll.evaluate(&[1.27, 0.120, 100.0]); println!("The extended negative log-likelihood is {}", result); ``` In practice, amplitudes can also be added together, their real and imaginary parts can be taken, and evaluators should mostly take the real part of whatever complex value comes out of the model. @@ -203,9 +211,10 @@ def main(): pos_im = (s0p * z00p.imag() + d2p * z22p.imag()).norm_sqr() neg_re = (s0n * z00n.real()).norm_sqr() neg_im = (s0n * z00n.imag()).norm_sqr() - model = pos_re + pos_im + neg_re + neg_im + expr = pos_re + pos_im + neg_re + neg_im + model = manager.model(expr) - nll = ld.NLL(manager, model, ds_data, ds_mc) + nll = ld.NLL(model, ds_data, ds_mc) status = nll.minimize([1.0] * len(nll.parameters)) print(status) fit_weights = nll.project(status.x) diff --git a/benches/kmatrix_benchmark.rs b/benches/kmatrix_benchmark.rs index 2e667c0..010a7e6 100644 --- a/benches/kmatrix_benchmark.rs +++ b/benches/kmatrix_benchmark.rs @@ -142,8 +142,9 @@ fn kmatrix_nll_benchmark(c: &mut Criterion) { let pos_im = (&s0p * z00p.imag() + &d2p * z22p.imag()).norm_sqr(); let neg_re = (&s0n * z00n.real()).norm_sqr(); 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 expr = pos_re + pos_im + neg_re + neg_im; + let model = manager.model(&expr); + let nll = NLL::new(&model, &ds_data, &ds_mc); let mut group = c.benchmark_group("K-Matrix NLL Performance"); for threads in 1..=num_cpus::get() { let pool = ThreadPoolBuilder::new() diff --git a/docs/source/tutorials/unbinned_fit.rst b/docs/source/tutorials/unbinned_fit.rst index 7a92aea..0c7b9ae 100644 --- a/docs/source/tutorials/unbinned_fit.rst +++ b/docs/source/tutorials/unbinned_fit.rst @@ -149,13 +149,13 @@ Next, we combine these together according to our model. For these amplitude poin positive_real_sum = (f0_1500 * bw0 * z00p.real() + f2_1525 * bw2 * z22p.real()).norm_sqr() positive_imag_sum = (f0_1500 * bw0 * z00p.imag() + f2_1525 * bw2 * z22p.imag()).norm_sqr() - model = positive_real_sum + positive_imag_sum + model = manager.model(positive_real_sum + positive_imag_sum) Now that we have the model, we want to fit the free parameters, which in this case are the complex photocouplings and the widths of each Breit-Wigner. We can do this by creating an ``NLL`` object which uses the data and accepted Monte-Carlo datasets to calculate the negative log-likelihood described earlier. .. code:: python - nll = ld.NLL(manager, model, data_ds, accmc_ds) + nll = ld.NLL(model, data_ds, accmc_ds) print(nll.parameters) # ['Re[f_0(1500)]', "Re[f_2'(1525)]", "Im[f_2'(1525)]", 'f_0 width', 'f_2 width'] @@ -265,7 +265,7 @@ To create an ``Evaluator`` object, we just need to load up the manager with the .. code:: python - gen_eval = manager.load(model, genmc_ds) + gen_eval = model.load(genmc_ds) tot_weights_acc = nll.project(status.x, mc_evaluator=gen_eval) f0_weights_acc = nll.project_with(status.x, ["[f_0(1500)]", "BW_0", "Z00+"], mc_evaluator=gen_eval) f2_weights_acc = nll.project_with(status.x, ["[f_2'(1525)]", "BW_2", "Z22+"], mc_evaluator=gen_eval) diff --git a/python/laddu/__init__.py b/python/laddu/__init__.py index aff22c0..bd3fcbb 100644 --- a/python/laddu/__init__.py +++ b/python/laddu/__init__.py @@ -5,7 +5,7 @@ import numpy as np -from laddu.amplitudes import Manager, constant, parameter +from laddu.amplitudes import Manager, constant, parameter, Model from laddu.amplitudes.breit_wigner import BreitWigner from laddu.amplitudes.common import ComplexScalar, PolarComplexScalar, Scalar from laddu.amplitudes.ylm import Ylm @@ -76,6 +76,7 @@ def open_amptools( 'Event', 'LikelihoodManager', 'Manager', + 'Model', 'Mandelstam', 'Mass', 'Observer', diff --git a/python/laddu/__init__.pyi b/python/laddu/__init__.pyi index 0db21af..eca267b 100644 --- a/python/laddu/__init__.pyi +++ b/python/laddu/__init__.pyi @@ -1,7 +1,7 @@ from abc import ABCMeta, abstractmethod from pathlib import Path -from laddu.amplitudes import Expression, Manager, constant, parameter +from laddu.amplitudes import Expression, Manager, constant, parameter, Model from laddu.amplitudes.breit_wigner import BreitWigner from laddu.amplitudes.common import ComplexScalar, PolarComplexScalar, Scalar from laddu.amplitudes.ylm import Ylm @@ -51,6 +51,7 @@ __all__ = [ 'Expression', 'LikelihoodManager', 'Manager', + 'Model', 'Mandelstam', 'Mass', 'Observer', diff --git a/python/laddu/amplitudes/__init__.py b/python/laddu/amplitudes/__init__.py index 3299190..10be39b 100644 --- a/python/laddu/amplitudes/__init__.py +++ b/python/laddu/amplitudes/__init__.py @@ -5,6 +5,7 @@ Evaluator, Expression, Manager, + Model, ParameterLike, constant, parameter, @@ -15,6 +16,7 @@ 'Expression', 'Amplitude', 'Manager', + 'Model', 'Evaluator', 'ParameterLike', 'parameter', diff --git a/python/laddu/amplitudes/__init__.pyi b/python/laddu/amplitudes/__init__.pyi index 0c981a3..c2147e5 100644 --- a/python/laddu/amplitudes/__init__.pyi +++ b/python/laddu/amplitudes/__init__.pyi @@ -33,9 +33,17 @@ class Manager: parameters: list[str] def __init__(self) -> None: ... def register(self, amplitude: Amplitude) -> AmplitudeID: ... - def load( - self, expression: Expression | AmplitudeID, dataset: Dataset - ) -> Evaluator: ... + def model(self, expression: Expression | AmplitudeID) -> Model: ... + +class Model: + parameters: list[str] + def __init__(self) -> None: ... + def load(self, dataset: Dataset) -> Evaluator: ... + def save_as(self, path: str) -> None: ... + @staticmethod + def load_from(path: str) -> Model: ... + def __getstate__(self) -> object: ... + def __setstate__(self, state: object) -> None: ... class Evaluator: parameters: list[str] @@ -56,6 +64,7 @@ __all__ = [ 'Expression', 'Amplitude', 'Manager', + 'Model', 'Evaluator', 'ParameterLike', 'parameter', diff --git a/python/laddu/likelihoods/__init__.pyi b/python/laddu/likelihoods/__init__.pyi index 678e1ce..12e25a1 100644 --- a/python/laddu/likelihoods/__init__.pyi +++ b/python/laddu/likelihoods/__init__.pyi @@ -120,7 +120,7 @@ class Status: def __init__(self) -> None: ... def save_as(self, path: str) -> None: ... @staticmethod - def load(path: str) -> Status: ... + def load_from(path: str) -> Status: ... def as_dict(self) -> dict[str, Any]: ... def __getstate__(self) -> object: ... def __setstate__(self, state: object) -> None: ... diff --git a/python/tests/test_amplitudes.py b/python/tests/test_amplitudes.py index 527270f..0826706 100644 --- a/python/tests/test_amplitudes.py +++ b/python/tests/test_amplitudes.py @@ -33,7 +33,8 @@ def test_constant_amplitude(): amp = Scalar('constant', constant(2.0)) aid = manager.register(amp) dataset = make_test_dataset() - evaluator = manager.load(aid, dataset) + model = manager.model(aid) + evaluator = model.load(dataset) result = evaluator.evaluate([]) assert result[0] == 2.0 + 0.0j @@ -43,7 +44,8 @@ def test_parametric_amplitude(): amp = Scalar('parametric', parameter('test_param')) aid = manager.register(amp) dataset = make_test_dataset() - evaluator = manager.load(aid, dataset) + model = manager.model(aid) + evaluator = model.load(dataset) result = evaluator.evaluate([3.0]) assert result[0] == 3.0 + 0.0j @@ -58,52 +60,62 @@ def test_expression_operations(): aid3 = manager.register(amp3) dataset = make_test_dataset() expr_add = aid1 + aid2 - eval_add = manager.load(expr_add, dataset) + model_add = manager.model(expr_add) + eval_add = model_add.load(dataset) result_add = eval_add.evaluate([]) assert result_add[0] == 2.0 + 1.0j expr_mul = aid1 * aid2 - eval_mul = manager.load(expr_mul, dataset) + model_mul = manager.model(expr_mul) + eval_mul = model_mul.load(dataset) result_mul = eval_mul.evaluate([]) assert result_mul[0] == 0.0 + 2.0j expr_add2 = expr_add + expr_mul - eval_add2 = manager.load(expr_add2, dataset) + model_add2 = manager.model(expr_add2) + eval_add2 = model_add2.load(dataset) result_add2 = eval_add2.evaluate([]) assert result_add2[0] == 2.0 + 3.0j expr_mul2 = expr_add * expr_mul - eval_mul2 = manager.load(expr_mul2, dataset) + model_mul2 = manager.model(expr_mul2) + eval_mul2 = model_mul2.load(dataset) result_mul2 = eval_mul2.evaluate([]) assert result_mul2[0] == -2.0 + 4.0j expr_real = aid3.real() - eval_real = manager.load(expr_real, dataset) + model_real = manager.model(expr_real) + eval_real = model_real.load(dataset) result_real = eval_real.evaluate([]) assert result_real[0] == 3.0 + 0.0j expr_mul2_real = expr_mul2.real() - eval_mul2_real = manager.load(expr_mul2_real, dataset) + model_mul2_real = manager.model(expr_mul2_real) + eval_mul2_real = model_mul2_real.load(dataset) result_mul2_real = eval_mul2_real.evaluate([]) assert result_mul2_real[0] == -2.0 + 0.0j expr_imag = aid3.imag() - eval_imag = manager.load(expr_imag, dataset) + model_imag = manager.model(expr_imag) + eval_imag = model_imag.load(dataset) result_imag = eval_imag.evaluate([]) assert result_imag[0] == 4.0 + 0.0j expr_mul2_imag = expr_mul2.imag() - eval_mul2_imag = manager.load(expr_mul2_imag, dataset) + model_mul2_imag = manager.model(expr_mul2_imag) + eval_mul2_imag = model_mul2_imag.load(dataset) result_mul2_imag = eval_mul2_imag.evaluate([]) assert result_mul2_imag[0] == 4.0 + 0.0j expr_norm = aid1.norm_sqr() - eval_norm = manager.load(expr_norm, dataset) + model_norm = manager.model(expr_norm) + eval_norm = model_norm.load(dataset) result_norm = eval_norm.evaluate([]) assert result_norm[0] == 4.0 + 0.0j expr_mul2_norm = expr_mul2.norm_sqr() - eval_mul2_norm = manager.load(expr_mul2_norm, dataset) + model_mul2_norm = manager.model(expr_mul2_norm) + eval_mul2_norm = model_mul2_norm.load(dataset) result_mul2_norm = eval_mul2_norm.evaluate([]) assert result_mul2_norm[0] == 20.0 + 0.0j @@ -117,7 +129,8 @@ def test_amplitude_activation(): dataset = make_test_dataset() expr = aid1 + aid2 - evaluator = manager.load(expr, dataset) + model = manager.model(expr) + evaluator = model.load(dataset) result = evaluator.evaluate([]) assert result[0] == 3.0 + 0.0j @@ -140,7 +153,8 @@ def test_gradient(): aid = manager.register(amp) dataset = make_test_dataset() expr = aid.norm_sqr() - evaluator = manager.load(expr, dataset) + model = manager.model(expr) + evaluator = model.load(dataset) params = [2.0] gradient = evaluator.evaluate_gradient(params) # For |f(x)|^2 where f(x) = x, the derivative should be 2x @@ -151,10 +165,14 @@ def test_gradient(): def test_parameter_registration(): manager = Manager() amp = Scalar('parametric', parameter('test_param')) - manager.register(amp) + aid = manager.register(amp) parameters = manager.parameters + model = manager.model(aid) + model_parameters = model.parameters assert len(parameters) == 1 assert parameters[0] == 'test_param' + assert len(model_parameters) == 1 + assert model_parameters[0] == 'test_param' def test_duplicate_amplitude_registration(): diff --git a/python/tests/test_breit_wigner.py b/python/tests/test_breit_wigner.py index f322628..fc02433 100644 --- a/python/tests/test_breit_wigner.py +++ b/python/tests/test_breit_wigner.py @@ -26,7 +26,8 @@ def test_bw_evaluation(): ) aid = manager.register(amp) dataset = make_test_dataset() - evaluator = manager.load(aid, dataset) + model = manager.model(aid) + evaluator = model.load(dataset) result = evaluator.evaluate([1.5, 0.3]) assert pytest.approx(result[0].real) == 1.4585691 assert pytest.approx(result[0].imag) == 1.4107341 @@ -39,7 +40,8 @@ def test_bw_gradient(): ) aid = manager.register(amp) dataset = make_test_dataset() - evaluator = manager.load(aid, dataset) + model = manager.model(aid) + evaluator = model.load(dataset) result = evaluator.evaluate_gradient([1.5, 0.3]) assert pytest.approx(result[0][0].real) == 1.3252039 assert pytest.approx(result[0][0].imag) == -11.6827505 diff --git a/python/tests/test_common.py b/python/tests/test_common.py index 60d0718..8a8391d 100644 --- a/python/tests/test_common.py +++ b/python/tests/test_common.py @@ -35,7 +35,8 @@ def test_scalar_creation_and_evaluation(): amp = Scalar('test_scalar', parameter('test_param')) aid = manager.register(amp) dataset = make_test_dataset() - evaluator = manager.load(aid, dataset) + model = manager.model(aid) + evaluator = model.load(dataset) result = evaluator.evaluate([2.5]) assert result[0].real == 2.5 assert result[0].imag == 0.0 @@ -47,7 +48,8 @@ def test_scalar_gradient(): aid = manager.register(amp) dataset = make_test_dataset() expr = aid.norm_sqr() - evaluator = manager.load(expr, dataset) + model = manager.model(expr) + evaluator = model.load(dataset) gradient = evaluator.evaluate_gradient([2.0]) assert gradient[0][0].real == 4.0 assert gradient[0][0].imag == 0.0 @@ -58,7 +60,8 @@ def test_complex_scalar_creation_and_evaluation(): amp = ComplexScalar('test_complex', parameter('re_param'), parameter('im_param')) aid = manager.register(amp) dataset = make_test_dataset() - evaluator = manager.load(aid, dataset) + model = manager.model(aid) + evaluator = model.load(dataset) result = evaluator.evaluate([1.5, 2.5]) assert result[0].real == 1.5 assert result[0].imag == 2.5 @@ -70,7 +73,8 @@ def test_complex_scalar_gradient(): aid = manager.register(amp) dataset = make_test_dataset() expr = aid.norm_sqr() - evaluator = manager.load(expr, dataset) + model = manager.model(expr) + evaluator = model.load(dataset) gradient = evaluator.evaluate_gradient([3.0, 4.0]) assert gradient[0][0].real == 6.0 assert gradient[0][0].imag == 0.0 @@ -83,7 +87,8 @@ def test_polar_complex_scalar_creation_and_evaluation(): amp = PolarComplexScalar('test_polar', parameter('r_param'), parameter('theta_param')) aid = manager.register(amp) dataset = make_test_dataset() - evaluator = manager.load(aid, dataset) + model = manager.model(aid) + evaluator = model.load(dataset) r = 2.0 theta = np.pi / 4.0 result = evaluator.evaluate([r, theta]) @@ -96,7 +101,8 @@ def test_polar_complex_scalar_gradient(): amp = PolarComplexScalar('test_polar', parameter('r_param'), parameter('theta_param')) aid = manager.register(amp) dataset = make_test_dataset() - evaluator = manager.load(aid, dataset) + model = manager.model(aid) + evaluator = model.load(dataset) r = 2.0 theta = np.pi / 4.0 gradient = evaluator.evaluate_gradient([r, theta]) diff --git a/python/tests/test_kmatrix.py b/python/tests/test_kmatrix.py index ead8ef5..cf0808c 100644 --- a/python/tests/test_kmatrix.py +++ b/python/tests/test_kmatrix.py @@ -44,7 +44,8 @@ def test_f0_evaluation(): ) aid = manager.register(amp) dataset = make_test_dataset() - evaluator = manager.load(aid, dataset) + model = manager.model(aid) + evaluator = model.load(dataset) result = evaluator.evaluate([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]) assert pytest.approx(result[0].real) == 0.2674945 assert pytest.approx(result[0].imag) == 0.7289451 @@ -67,7 +68,8 @@ def test_f0_gradient(): ) aid = manager.register(amp) dataset = make_test_dataset() - evaluator = manager.load(aid, dataset) + model = manager.model(aid) + evaluator = model.load(dataset) result = evaluator.evaluate_gradient( [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0] ) @@ -109,7 +111,8 @@ def test_f2_evaluation(): ) aid = manager.register(amp) dataset = make_test_dataset() - evaluator = manager.load(aid, dataset) + model = manager.model(aid) + evaluator = model.load(dataset) result = evaluator.evaluate([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8]) assert pytest.approx(result[0].real) == 0.02523304 assert pytest.approx(result[0].imag) == 0.3971239 @@ -131,7 +134,8 @@ def test_f2_gradient(): ) aid = manager.register(amp) dataset = make_test_dataset() - evaluator = manager.load(aid, dataset) + model = manager.model(aid) + evaluator = model.load(dataset) result = evaluator.evaluate_gradient([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8]) assert pytest.approx(result[0][0].real) == -0.3078948 assert pytest.approx(result[0][0].imag) == 0.3808689 @@ -165,7 +169,8 @@ def test_a0_evaluation(): ) aid = manager.register(amp) dataset = make_test_dataset() - evaluator = manager.load(aid, dataset) + model = manager.model(aid) + evaluator = model.load(dataset) result = evaluator.evaluate([0.1, 0.2, 0.3, 0.4]) assert pytest.approx(result[0].real) == -0.8002759 assert pytest.approx(result[0].imag) == -0.1359306 @@ -185,7 +190,8 @@ def test_a0_gradient(): ) aid = manager.register(amp) dataset = make_test_dataset() - evaluator = manager.load(aid, dataset) + model = manager.model(aid) + evaluator = model.load(dataset) result = evaluator.evaluate_gradient([0.1, 0.2, 0.3, 0.4]) assert pytest.approx(result[0][0].real) == 0.2906192 assert pytest.approx(result[0][0].imag) == -0.0998906 @@ -211,7 +217,8 @@ def test_a2_evaluation(): ) aid = manager.register(amp) dataset = make_test_dataset() - evaluator = manager.load(aid, dataset) + model = manager.model(aid) + evaluator = model.load(dataset) result = evaluator.evaluate([0.1, 0.2, 0.3, 0.4]) assert pytest.approx(result[0].real) == -0.2092661 assert pytest.approx(result[0].imag) == -0.0985062 @@ -231,7 +238,8 @@ def test_a2_gradient(): ) aid = manager.register(amp) dataset = make_test_dataset() - evaluator = manager.load(aid, dataset) + model = manager.model(aid) + evaluator = model.load(dataset) result = evaluator.evaluate_gradient([0.1, 0.2, 0.3, 0.4]) assert pytest.approx(result[0][0].real) == -0.5756896 assert pytest.approx(result[0][0].imag) == 0.9398863 @@ -257,7 +265,8 @@ def test_rho_evaluation(): ) aid = manager.register(amp) dataset = make_test_dataset() - evaluator = manager.load(aid, dataset) + model = manager.model(aid) + evaluator = model.load(dataset) result = evaluator.evaluate([0.1, 0.2, 0.3, 0.4]) assert pytest.approx(result[0].real) == 0.0948355 assert pytest.approx(result[0].imag) == 0.2609183 @@ -277,7 +286,8 @@ def test_rho_gradient(): ) aid = manager.register(amp) dataset = make_test_dataset() - evaluator = manager.load(aid, dataset) + model = manager.model(aid) + evaluator = model.load(dataset) result = evaluator.evaluate_gradient([0.1, 0.2, 0.3, 0.4]) assert pytest.approx(result[0][0].real) == 0.0265203 assert pytest.approx(result[0][0].imag) == -0.02660265 @@ -300,7 +310,8 @@ def test_pi1_evaluation(): ) aid = manager.register(amp) dataset = make_test_dataset() - evaluator = manager.load(aid, dataset) + model = manager.model(aid) + evaluator = model.load(dataset) result = evaluator.evaluate([0.1, 0.2]) assert pytest.approx(result[0].real) == -0.1101758 assert pytest.approx(result[0].imag) == 0.2638717 @@ -317,7 +328,8 @@ def test_pi1_gradient(): ) aid = manager.register(amp) dataset = make_test_dataset() - evaluator = manager.load(aid, dataset) + model = manager.model(aid) + evaluator = model.load(dataset) result = evaluator.evaluate_gradient([0.1, 0.2]) assert pytest.approx(result[0][0].real) == -14.7987174 assert pytest.approx(result[0][0].imag) == -5.8430094 diff --git a/python/tests/test_serde.py b/python/tests/test_serde.py new file mode 100644 index 0000000..989c74e --- /dev/null +++ b/python/tests/test_serde.py @@ -0,0 +1,89 @@ +from laddu import Event, Dataset, Vector3, Mass, Manager, parameter, Scalar +from laddu.amplitudes.kmatrix import ( + KopfKMatrixF0, + KopfKMatrixF2, +) +import pytest +import pickle + + +def make_test_event() -> Event: + return Event( + [ + Vector3(0.0, 0.0, 8.747).with_mass(0.0), + Vector3(0.119, 0.374, 0.222).with_mass(1.007), + Vector3(-0.112, 0.293, 3.081).with_mass(0.498), + Vector3(-0.007, -0.667, 5.446).with_mass(0.498), + ], + [Vector3(0.385, 0.022, 0.000)], + 0.48, + ) + + +def make_test_dataset() -> Dataset: + return Dataset([make_test_event()]) + + +def test_serde(): + manager = Manager() + res_mass = Mass([2, 3]) + f0 = KopfKMatrixF0( + 'f0', + ( + (parameter('p0'), parameter('p1')), + (parameter('p2'), parameter('p3')), + (parameter('p4'), parameter('p5')), + (parameter('p6'), parameter('p7')), + (parameter('p8'), parameter('p9')), + ), + 1, + res_mass, + ) + f2 = KopfKMatrixF2( + 'f2', + ( + (parameter('g0'), parameter('g1')), + (parameter('g2'), parameter('g3')), + (parameter('g4'), parameter('g5')), + (parameter('g6'), parameter('g7')), + ), + 1, + res_mass, + ) + s = Scalar('s', parameter('s')) + f0_aid = manager.register(f0) + f2_aid = manager.register(f2) + s_aid = manager.register(s) + expr = (f0_aid * s_aid + f2_aid).norm_sqr() + model = manager.model(expr) + dataset = make_test_dataset() + evaluator = model.load(dataset) + p = [ + 0.1, + 0.2, + 0.3, + 0.4, + 0.5, + 0.6, + 0.7, + 0.8, + 0.9, + 1.0, + 1.1, + 1.2, + 1.3, + 1.4, + 1.5, + 1.6, + 1.7, + 1.8, + 1.9, + ] + result = evaluator.evaluate(p) + pickled_model = pickle.dumps(model) + unpickled_model = pickle.loads(pickled_model) + unpickled_evaluator = unpickled_model.load(dataset) + unpickled_result = unpickled_evaluator.evaluate(p) + + assert pytest.approx(result[0].real) == pytest.approx(unpickled_result[0].real) + assert pytest.approx(result[0].imag) == pytest.approx(unpickled_result[0].imag) diff --git a/python/tests/test_ylm.py b/python/tests/test_ylm.py index 7ee3690..628b369 100644 --- a/python/tests/test_ylm.py +++ b/python/tests/test_ylm.py @@ -1,4 +1,4 @@ -from laddu import Event, Dataset, Vector3, Zlm, Angles, Manager, Polarization +from laddu import Event, Dataset, Vector3, Ylm, Angles, Manager import pytest @@ -19,26 +19,26 @@ def make_test_dataset() -> Dataset: return Dataset([make_test_event()]) -def test_zlm_evaluation(): +def test_ylm_evaluation(): manager = Manager() angles = Angles(0, [1], [2], [2, 3], 'Helicity') - polarization = Polarization(0, [1]) - amp = Zlm('zlm', 1, 1, '+', angles, polarization) + amp = Ylm('ylm', 1, 1, angles) aid = manager.register(amp) dataset = make_test_dataset() - evaluator = manager.load(aid, dataset) + model = manager.model(aid) + evaluator = model.load(dataset) result = evaluator.evaluate([]) - assert pytest.approx(result[0].real) == 0.04284127 - assert pytest.approx(result[0].imag) == -0.2385963 + assert pytest.approx(result[0].real) == 0.2713394 + assert pytest.approx(result[0].imag) == 0.1426897 def test_ylm_gradient(): manager = Manager() angles = Angles(0, [1], [2], [2, 3], 'Helicity') - polarization = Polarization(0, [1]) - amp = Zlm('zlm', 1, 1, '+', angles, polarization) + amp = Ylm('ylm', 1, 1, angles) aid = manager.register(amp) dataset = make_test_dataset() - evaluator = manager.load(aid, dataset) + model = manager.model(aid) + evaluator = model.load(dataset) result = evaluator.evaluate_gradient([]) assert len(result[0]) == 0 # amplitude has no parameters diff --git a/python/tests/test_zlm.py b/python/tests/test_zlm.py index 01fe59d..3dcf662 100644 --- a/python/tests/test_zlm.py +++ b/python/tests/test_zlm.py @@ -1,4 +1,4 @@ -from laddu import Event, Dataset, Vector3, Ylm, Angles, Manager +from laddu import Event, Dataset, Vector3, Zlm, Angles, Manager, Polarization import pytest @@ -19,24 +19,28 @@ def make_test_dataset() -> Dataset: return Dataset([make_test_event()]) -def test_ylm_evaluation(): +def test_zlm_evaluation(): manager = Manager() angles = Angles(0, [1], [2], [2, 3], 'Helicity') - amp = Ylm('ylm', 1, 1, angles) + polarization = Polarization(0, [1]) + amp = Zlm('zlm', 1, 1, '+', angles, polarization) aid = manager.register(amp) dataset = make_test_dataset() - evaluator = manager.load(aid, dataset) + model = manager.model(aid) + evaluator = model.load(dataset) result = evaluator.evaluate([]) - assert pytest.approx(result[0].real) == 0.2713394 - assert pytest.approx(result[0].imag) == 0.1426897 + assert pytest.approx(result[0].real) == 0.04284127 + assert pytest.approx(result[0].imag) == -0.2385963 -def test_ylm_gradient(): +def test_zlm_gradient(): manager = Manager() angles = Angles(0, [1], [2], [2, 3], 'Helicity') - amp = Ylm('ylm', 1, 1, angles) + polarization = Polarization(0, [1]) + amp = Zlm('zlm', 1, 1, '+', angles, polarization) aid = manager.register(amp) dataset = make_test_dataset() - evaluator = manager.load(aid, dataset) + model = manager.model(aid) + evaluator = model.load(dataset) result = evaluator.evaluate_gradient([]) assert len(result[0]) == 0 # amplitude has no parameters diff --git a/src/amplitudes/breit_wigner.rs b/src/amplitudes/breit_wigner.rs index 2040d53..734da10 100644 --- a/src/amplitudes/breit_wigner.rs +++ b/src/amplitudes/breit_wigner.rs @@ -1,5 +1,6 @@ use nalgebra::DVector; use num::Complex; +use serde::{Deserialize, Serialize}; use crate::{ amplitudes::ParameterLike, @@ -23,7 +24,7 @@ use super::{Amplitude, AmplitudeID}; /// \Gamma = \Gamma_0 \frac{m_0}{m} \frac{q(m, m_1, m_2)}{q(m_0, m_1, m_2)} \left(\frac{B_{\ell}(m, m_1, m_2)}{B_{\ell}(m_0, m_1, m_2)}\right)^2 /// ``` /// is the relativistic width correction, $`q(m_a, m_b, m_c)`$ is the breakup momentum of a particle with mass $`m_a`$ decaying into two particles with masses $`m_b`$ and $`m_c`$, $`B_{\ell}(m_a, m_b, m_c)`$ is the Blatt-Weisskopf barrier factor for the same decay assuming particle $`a`$ has angular momentum $`\ell`$, $`m_0`$ is the mass of the resonance, $`\Gamma_0`$ is the nominal width of the resonance, $`m_1`$ and $`m_2`$ are the masses of the decay products, and $`m`$ is the "input" mass. -#[derive(Clone)] +#[derive(Clone, Serialize, Deserialize)] pub struct BreitWigner { name: String, mass: ParameterLike, @@ -63,6 +64,7 @@ impl BreitWigner { } } +#[typetag::serde] impl Amplitude for BreitWigner { fn register(&mut self, resources: &mut Resources) -> Result { self.pid_mass = resources.register_parameter(&self.mass); @@ -130,7 +132,8 @@ mod tests { let dataset = Arc::new(test_dataset()); let expr = aid.into(); - let evaluator = manager.load(&expr, &dataset); + let model = manager.model(&expr); + let evaluator = model.load(&dataset); let result = evaluator.evaluate(&[1.5, 0.3]); @@ -154,7 +157,8 @@ mod tests { let dataset = Arc::new(test_dataset()); let expr = aid.into(); - let evaluator = manager.load(&expr, &dataset); + let model = manager.model(&expr); + let evaluator = model.load(&dataset); let result = evaluator.evaluate_gradient(&[1.5, 0.3]); assert_relative_eq!(result[0][0].re, 1.3252039, epsilon = 1e-7); diff --git a/src/amplitudes/common.rs b/src/amplitudes/common.rs index cd1b837..5b587b7 100644 --- a/src/amplitudes/common.rs +++ b/src/amplitudes/common.rs @@ -1,5 +1,6 @@ use nalgebra::DVector; use num::Complex; +use serde::{Deserialize, Serialize}; use crate::{ amplitudes::{AmplitudeID, ParameterLike}, @@ -11,7 +12,7 @@ use crate::{ use super::Amplitude; /// A scalar-valued [`Amplitude`] which just contains a single parameter as its value. -#[derive(Clone)] +#[derive(Clone, Serialize, Deserialize)] pub struct Scalar { name: String, value: ParameterLike, @@ -30,6 +31,7 @@ impl Scalar { } } +#[typetag::serde] impl Amplitude for Scalar { fn register(&mut self, resources: &mut Resources) -> Result { self.pid = resources.register_parameter(&self.value); @@ -55,7 +57,7 @@ impl Amplitude for Scalar { /// A complex-valued [`Amplitude`] which just contains two parameters representing its real and /// imaginary parts. -#[derive(Clone)] +#[derive(Clone, Serialize, Deserialize)] pub struct ComplexScalar { name: String, re: ParameterLike, @@ -78,6 +80,7 @@ impl ComplexScalar { } } +#[typetag::serde] impl Amplitude for ComplexScalar { fn register(&mut self, resources: &mut Resources) -> Result { self.pid_re = resources.register_parameter(&self.re); @@ -107,7 +110,7 @@ impl Amplitude for ComplexScalar { /// A complex-valued [`Amplitude`] which just contains two parameters representing its magnitude and /// phase. -#[derive(Clone)] +#[derive(Clone, Serialize, Deserialize)] pub struct PolarComplexScalar { name: String, r: ParameterLike, @@ -130,6 +133,7 @@ impl PolarComplexScalar { } } +#[typetag::serde] impl Amplitude for PolarComplexScalar { fn register(&mut self, resources: &mut Resources) -> Result { self.pid_r = resources.register_parameter(&self.r); @@ -175,7 +179,8 @@ mod tests { let dataset = Arc::new(test_dataset()); let expr = aid.into(); // Direct amplitude evaluation - let evaluator = manager.load(&expr, &dataset); + let model = manager.model(&expr); + let evaluator = model.load(&dataset); let params = vec![2.5]; let result = evaluator.evaluate(¶ms); @@ -192,7 +197,8 @@ mod tests { let dataset = Arc::new(test_dataset()); let expr = aid.norm_sqr(); // |f(x)|^2 - let evaluator = manager.load(&expr, &dataset); + let model = manager.model(&expr); + let evaluator = model.load(&dataset); let params = vec![2.0]; let gradient = evaluator.evaluate_gradient(¶ms); @@ -210,7 +216,8 @@ mod tests { let dataset = Arc::new(test_dataset()); let expr = aid.into(); - let evaluator = manager.load(&expr, &dataset); + let model = manager.model(&expr); + let evaluator = model.load(&dataset); let params = vec![1.5, 2.5]; // Real and imaginary parts let result = evaluator.evaluate(¶ms); @@ -227,7 +234,8 @@ mod tests { let dataset = Arc::new(test_dataset()); let expr = aid.norm_sqr(); // |f(x + iy)|^2 - let evaluator = manager.load(&expr, &dataset); + let model = manager.model(&expr); + let evaluator = model.load(&dataset); let params = vec![3.0, 4.0]; // Real and imaginary parts let gradient = evaluator.evaluate_gradient(¶ms); @@ -248,7 +256,8 @@ mod tests { let dataset = Arc::new(test_dataset()); let expr = aid.into(); - let evaluator = manager.load(&expr, &dataset); + let model = manager.model(&expr); + let evaluator = model.load(&dataset); let r = 2.0; let theta = PI / 4.0; @@ -269,7 +278,8 @@ mod tests { let dataset = Arc::new(test_dataset()); let expr = aid.into(); // f(r,θ) = re^(iθ) - let evaluator = manager.load(&expr, &dataset); + let model = manager.model(&expr); + let evaluator = model.load(&dataset); let r = 2.0; let theta = PI / 4.0; diff --git a/src/amplitudes/kmatrix.rs b/src/amplitudes/kmatrix.rs index 546c311..f3d4530 100644 --- a/src/amplitudes/kmatrix.rs +++ b/src/amplitudes/kmatrix.rs @@ -5,6 +5,7 @@ use nalgebra::{SMatrix, SVector}; use num::traits::ConstOne; use num::traits::FloatConst; use num::Complex; +use serde::{Deserialize, Serialize}; use crate::LadduError; use crate::{ @@ -21,7 +22,7 @@ use crate::{ use super::Amplitude; /// An Adler zero term used in a K-matrix. -#[derive(Copy, Clone, Debug)] +#[derive(Copy, Clone, Debug, Serialize, Deserialize)] pub struct AdlerZero { /// The zero position $`s_0`$. pub s_0: Float, @@ -30,7 +31,7 @@ pub struct AdlerZero { } /// Methods for computing various parts of a K-matrix with fixed couplings and mass poles. -#[derive(Clone, Debug)] +#[derive(Clone, Debug, Serialize, Deserialize)] pub struct FixedKMatrix { g: SMatrix, c: SMatrix, @@ -140,7 +141,7 @@ impl FixedKMatrix Result { for i in 0..self.couplings_indices_real.len() { @@ -284,7 +286,7 @@ impl Amplitude for KopfKMatrixF0 { /// (free production couplings only). /// /// [^1]: Kopf, B., Albrecht, M., Koch, H., Küßner, M., Pychy, J., Qin, X., & Wiedner, U. (2021). Investigation of the lightest hybrid meson candidate with a coupled-channel analysis of $`\bar{p}p`$-, $`\pi^- p`$- and $`\pi \pi`$-Data. The European Physical Journal C, 81(12). [doi:10.1140/epjc/s10052-021-09821-2](https://doi.org/10.1140/epjc/s10052-021-09821-2) -#[derive(Clone)] +#[derive(Clone, Serialize, Deserialize)] pub struct KopfKMatrixF2 { name: String, channel: usize, @@ -361,6 +363,7 @@ impl KopfKMatrixF2 { } } +#[typetag::serde] impl Amplitude for KopfKMatrixF2 { fn register(&mut self, resources: &mut Resources) -> Result { for i in 0..self.couplings_indices_real.len() { @@ -420,7 +423,7 @@ impl Amplitude for KopfKMatrixF2 { /// (free production couplings only). /// /// [^1]: Kopf, B., Albrecht, M., Koch, H., Küßner, M., Pychy, J., Qin, X., & Wiedner, U. (2021). Investigation of the lightest hybrid meson candidate with a coupled-channel analysis of $`\bar{p}p`$-, $`\pi^- p`$- and $`\pi \pi`$-Data. The European Physical Journal C, 81(12). [doi:10.1140/epjc/s10052-021-09821-2](https://doi.org/10.1140/epjc/s10052-021-09821-2) -#[derive(Clone)] +#[derive(Clone, Serialize, Deserialize)] pub struct KopfKMatrixA0 { name: String, channel: usize, @@ -489,6 +492,7 @@ impl KopfKMatrixA0 { } } +#[typetag::serde] impl Amplitude for KopfKMatrixA0 { fn register(&mut self, resources: &mut Resources) -> Result { for i in 0..self.couplings_indices_real.len() { @@ -548,7 +552,7 @@ impl Amplitude for KopfKMatrixA0 { /// (free production couplings only). /// /// [^1]: Kopf, B., Albrecht, M., Koch, H., Küßner, M., Pychy, J., Qin, X., & Wiedner, U. (2021). Investigation of the lightest hybrid meson candidate with a coupled-channel analysis of $`\bar{p}p`$-, $`\pi^- p`$- and $`\pi \pi`$-Data. The European Physical Journal C, 81(12). [doi:10.1140/epjc/s10052-021-09821-2](https://doi.org/10.1140/epjc/s10052-021-09821-2) -#[derive(Clone)] +#[derive(Clone, Serialize, Deserialize)] pub struct KopfKMatrixA2 { name: String, channel: usize, @@ -621,6 +625,7 @@ impl KopfKMatrixA2 { } } +#[typetag::serde] impl Amplitude for KopfKMatrixA2 { fn register(&mut self, resources: &mut Resources) -> Result { for i in 0..self.couplings_indices_real.len() { @@ -680,7 +685,7 @@ impl Amplitude for KopfKMatrixA2 { /// (free production couplings only). /// /// [^1]: Kopf, B., Albrecht, M., Koch, H., Küßner, M., Pychy, J., Qin, X., & Wiedner, U. (2021). Investigation of the lightest hybrid meson candidate with a coupled-channel analysis of $`\bar{p}p`$-, $`\pi^- p`$- and $`\pi \pi`$-Data. The European Physical Journal C, 81(12). [doi:10.1140/epjc/s10052-021-09821-2](https://doi.org/10.1140/epjc/s10052-021-09821-2) -#[derive(Clone)] +#[derive(Clone, Serialize, Deserialize)] pub struct KopfKMatrixRho { name: String, channel: usize, @@ -752,6 +757,7 @@ impl KopfKMatrixRho { } } +#[typetag::serde] impl Amplitude for KopfKMatrixRho { fn register(&mut self, resources: &mut Resources) -> Result { for i in 0..self.couplings_indices_real.len() { @@ -811,7 +817,7 @@ impl Amplitude for KopfKMatrixRho { /// (free production couplings only). /// /// [^1]: Kopf, B., Albrecht, M., Koch, H., Küßner, M., Pychy, J., Qin, X., & Wiedner, U. (2021). Investigation of the lightest hybrid meson candidate with a coupled-channel analysis of $`\bar{p}p`$-, $`\pi^- p`$- and $`\pi \pi`$-Data. The European Physical Journal C, 81(12). [doi:10.1140/epjc/s10052-021-09821-2](https://doi.org/10.1140/epjc/s10052-021-09821-2) -#[derive(Clone)] +#[derive(Clone, Serialize, Deserialize)] pub struct KopfKMatrixPi1 { name: String, channel: usize, @@ -879,6 +885,7 @@ impl KopfKMatrixPi1 { } } +#[typetag::serde] impl Amplitude for KopfKMatrixPi1 { fn register(&mut self, resources: &mut Resources) -> Result { for i in 0..self.couplings_indices_real.len() { @@ -963,7 +970,8 @@ mod tests { let dataset = Arc::new(test_dataset()); let expr = aid.into(); - let evaluator = manager.load(&expr, &dataset); + let model = manager.model(&expr); + let evaluator = model.load(&dataset); let result = evaluator.evaluate(&[0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]); @@ -991,7 +999,8 @@ mod tests { let dataset = Arc::new(test_dataset()); let expr = aid.into(); - let evaluator = manager.load(&expr, &dataset); + let model = manager.model(&expr); + let evaluator = model.load(&dataset); let result = evaluator.evaluate_gradient(&[0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]); @@ -1037,7 +1046,8 @@ mod tests { let dataset = Arc::new(test_dataset()); let expr = aid.into(); - let evaluator = manager.load(&expr, &dataset); + let model = manager.model(&expr); + let evaluator = model.load(&dataset); let result = evaluator.evaluate(&[0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8]); @@ -1064,7 +1074,8 @@ mod tests { let dataset = Arc::new(test_dataset()); let expr = aid.into(); - let evaluator = manager.load(&expr, &dataset); + let model = manager.model(&expr); + let evaluator = model.load(&dataset); let result = evaluator.evaluate_gradient(&[0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8]); @@ -1103,7 +1114,8 @@ mod tests { let dataset = Arc::new(test_dataset()); let expr = aid.into(); - let evaluator = manager.load(&expr, &dataset); + let model = manager.model(&expr); + let evaluator = model.load(&dataset); let result = evaluator.evaluate(&[0.1, 0.2, 0.3, 0.4]); @@ -1128,7 +1140,8 @@ mod tests { let dataset = Arc::new(test_dataset()); let expr = aid.into(); - let evaluator = manager.load(&expr, &dataset); + let model = manager.model(&expr); + let evaluator = model.load(&dataset); let result = evaluator.evaluate_gradient(&[0.1, 0.2, 0.3, 0.4]); @@ -1159,7 +1172,8 @@ mod tests { let dataset = Arc::new(test_dataset()); let expr = aid.into(); - let evaluator = manager.load(&expr, &dataset); + let model = manager.model(&expr); + let evaluator = model.load(&dataset); let result = evaluator.evaluate(&[0.1, 0.2, 0.3, 0.4]); @@ -1184,7 +1198,8 @@ mod tests { let dataset = Arc::new(test_dataset()); let expr = aid.into(); - let evaluator = manager.load(&expr, &dataset); + let model = manager.model(&expr); + let evaluator = model.load(&dataset); let result = evaluator.evaluate_gradient(&[0.1, 0.2, 0.3, 0.4]); @@ -1215,7 +1230,8 @@ mod tests { let dataset = Arc::new(test_dataset()); let expr = aid.into(); - let evaluator = manager.load(&expr, &dataset); + let model = manager.model(&expr); + let evaluator = model.load(&dataset); let result = evaluator.evaluate(&[0.1, 0.2, 0.3, 0.4]); @@ -1240,7 +1256,8 @@ mod tests { let dataset = Arc::new(test_dataset()); let expr = aid.into(); - let evaluator = manager.load(&expr, &dataset); + let model = manager.model(&expr); + let evaluator = model.load(&dataset); let result = evaluator.evaluate_gradient(&[0.1, 0.2, 0.3, 0.4]); @@ -1263,7 +1280,8 @@ mod tests { let dataset = Arc::new(test_dataset()); let expr = aid.into(); - let evaluator = manager.load(&expr, &dataset); + let model = manager.model(&expr); + let evaluator = model.load(&dataset); let result = evaluator.evaluate(&[0.1, 0.2]); @@ -1280,7 +1298,8 @@ mod tests { let dataset = Arc::new(test_dataset()); let expr = aid.into(); - let evaluator = manager.load(&expr, &dataset); + let model = manager.model(&expr); + let evaluator = model.load(&dataset); let result = evaluator.evaluate_gradient(&[0.1, 0.2]); diff --git a/src/amplitudes/mod.rs b/src/amplitudes/mod.rs index 7a0df65..c86004a 100644 --- a/src/amplitudes/mod.rs +++ b/src/amplitudes/mod.rs @@ -11,6 +11,7 @@ use num::Complex; use parking_lot::RwLock; #[cfg(feature = "rayon")] use rayon::prelude::*; +use serde::{Deserialize, Serialize}; use crate::{ data::{Dataset, Event}, @@ -30,7 +31,7 @@ pub mod ylm; pub mod zlm; /// An enum containing either a named free parameter or a constant value. -#[derive(Clone, Default)] +#[derive(Clone, Default, Serialize, Deserialize)] pub enum ParameterLike { /// A named free parameter. Parameter(String), @@ -61,6 +62,7 @@ pub fn constant(value: Float) -> ParameterLike { /// cache values which do not depend on free parameters. /// /// See [`BreitWigner`](breit_wigner::BreitWigner), [`Ylm`](ylm::Ylm), and [`Zlm`](zlm::Zlm) for examples which use all of these features. +#[typetag::serde(tag = "type")] pub trait Amplitude: DynClone + Send + Sync { /// This method should be used to tell the [`Resources`] manager about all of /// the free parameters and cached values used by this [`Amplitude`]. It should end by @@ -176,7 +178,7 @@ pub(crate) struct GradientValues(pub(crate) Vec>>); /// A tag which refers to a registered [`Amplitude`]. This is the base object which can be used to /// build [`Expression`]s and should be obtained from the [`Manager::register`] method. -#[derive(Clone, Default, Debug)] +#[derive(Clone, Default, Debug, Serialize, Deserialize)] pub struct AmplitudeID(pub(crate) String, pub(crate) usize); impl Display for AmplitudeID { @@ -192,8 +194,11 @@ impl From for Expression { } /// An expression tree which contains [`AmplitudeID`]s and operators over them. -#[derive(Clone)] +#[derive(Clone, Serialize, Deserialize, Default)] pub enum Expression { + #[default] + /// A expression equal to zero. + Zero, /// A registered [`Amplitude`] referenced by an [`AmplitudeID`]. Amp(AmplitudeID), /// The sum of two [`Expression`]s. @@ -260,6 +265,7 @@ impl Expression { Expression::Real(a) => Complex::new(a.evaluate(amplitude_values).re, 0.0), Expression::Imag(a) => Complex::new(a.evaluate(amplitude_values).im, 0.0), Expression::NormSqr(a) => Complex::new(a.evaluate(amplitude_values).norm_sqr(), 0.0), + Expression::Zero => Complex::ZERO, } } pub(crate) fn evaluate_gradient( @@ -292,6 +298,7 @@ impl Expression { a.evaluate_gradient(amplitude_values, gradient_values) .map(|g| Complex::new(2.0 * (g * conj_f_a).re, 0.0)) } + Expression::Zero => DVector::zeros(0), } } /// Takes the real part of the given [`Expression`]. @@ -322,10 +329,11 @@ impl Expression { Self::Real(_) => "Re".to_string(), Self::Imag(_) => "Im".to_string(), Self::NormSqr(_) => "NormSqr".to_string(), + Self::Zero => "0".to_string(), }; writeln!(f, "{}{}{}", parent_prefix, immediate_prefix, display_string)?; match self { - Self::Amp(_) => {} + Self::Amp(_) | Self::Zero => {} Self::Add(a, b) | Self::Mul(a, b) => { let terms = [a, b]; let mut it = terms.iter().peekable(); @@ -348,7 +356,7 @@ impl Expression { /// A manager which can be used to register [`Amplitude`]s with [`Resources`]. This structure is /// essential to any analysis and should be constructed using the [`Manager::default()`] method. -#[derive(Default, Clone)] +#[derive(Default, Clone, Serialize, Deserialize)] pub struct Manager { amplitudes: Vec>, resources: Resources, @@ -372,20 +380,45 @@ impl Manager { self.amplitudes.push(amp); Ok(aid) } - /// Create an [`Evaluator`] which can compute the result of the given [`Expression`] built on + /// Turns an [`Expression`] made from registered [`Amplitude`]s into a [`Model`]. + pub fn model(&self, expression: &Expression) -> Model { + Model { + manager: self.clone(), + expression: expression.clone(), + } + } +} + +/// A struct which contains a set of registerd [`Amplitude`]s (inside a [`Manager`]) +/// and an [`Expression`]. +/// +/// This struct implements [`serde::Serialize`] and [`serde::Deserialize`] and is intended +/// to be used to store models to disk. +#[derive(Clone, Serialize, Deserialize)] +pub struct Model { + pub(crate) manager: Manager, + pub(crate) expression: Expression, +} + +impl Model { + /// Get the list of parameter names in the order they appear in the [`Model`]'s [`Manager`] field. + pub fn parameters(&self) -> Vec { + self.manager.parameters() + } + /// Create an [`Evaluator`] which can compute the result of the internal [`Expression`] built on /// registered [`Amplitude`]s over the given [`Dataset`]. This method precomputes any relevant /// information over the [`Event`]s in the [`Dataset`]. - pub fn load(&self, expression: &Expression, dataset: &Arc) -> Evaluator { - let loaded_resources = Arc::new(RwLock::new(self.resources.clone())); + pub fn load(&self, dataset: &Arc) -> Evaluator { + let loaded_resources = Arc::new(RwLock::new(self.manager.resources.clone())); loaded_resources.write().reserve_cache(dataset.len()); - for amplitude in &self.amplitudes { + for amplitude in &self.manager.amplitudes { amplitude.precompute_all(dataset, &mut loaded_resources.write()); } Evaluator { - amplitudes: self.amplitudes.clone(), + amplitudes: self.manager.amplitudes.clone(), resources: loaded_resources.clone(), dataset: dataset.clone(), - expression: expression.clone(), + expression: self.expression.clone(), } } } @@ -620,7 +653,8 @@ mod tests { events: vec![Arc::new(test_event())], }); let expr = Expression::Amp(aid); - let evaluator = manager.load(&expr, &dataset); + let model = manager.model(&expr); + let evaluator = model.load(&dataset); let result = evaluator.evaluate(&[]); assert_eq!(result[0], Complex::from(2.0)); } @@ -632,7 +666,8 @@ mod tests { let aid = manager.register(amp).unwrap(); let dataset = Arc::new(test_dataset()); let expr = Expression::Amp(aid); - let evaluator = manager.load(&expr, &dataset); + let model = manager.model(&expr); + let evaluator = model.load(&dataset); let result = evaluator.evaluate(&[3.0]); assert_eq!(result[0], Complex::from(3.0)); } @@ -652,61 +687,71 @@ mod tests { // Test (amp) addition let expr_add = &aid1 + &aid2; - let eval_add = manager.load(&expr_add, &dataset); + let model_add = manager.model(&expr_add); + let eval_add = model_add.load(&dataset); let result_add = eval_add.evaluate(&[]); assert_eq!(result_add[0], Complex::new(2.0, 1.0)); // Test (amp) multiplication let expr_mul = &aid1 * &aid2; - let eval_mul = manager.load(&expr_mul, &dataset); + let model_mul = manager.model(&expr_mul); + let eval_mul = model_mul.load(&dataset); let result_mul = eval_mul.evaluate(&[]); assert_eq!(result_mul[0], Complex::new(0.0, 2.0)); // Test (expr) addition let expr_add2 = &expr_add + &expr_mul; - let eval_add2 = manager.load(&expr_add2, &dataset); + let model_add2 = manager.model(&expr_add2); + let eval_add2 = model_add2.load(&dataset); let result_add2 = eval_add2.evaluate(&[]); assert_eq!(result_add2[0], Complex::new(2.0, 3.0)); // Test (expr) multiplication let expr_mul2 = &expr_add * &expr_mul; - let eval_mul2 = manager.load(&expr_mul2, &dataset); + let model_mul2 = manager.model(&expr_mul2); + let eval_mul2 = model_mul2.load(&dataset); let result_mul2 = eval_mul2.evaluate(&[]); assert_eq!(result_mul2[0], Complex::new(-2.0, 4.0)); // Test (amp) real let expr_real = aid3.real(); - let eval_real = manager.load(&expr_real, &dataset); + let model_real = manager.model(&expr_real); + let eval_real = model_real.load(&dataset); let result_real = eval_real.evaluate(&[]); assert_eq!(result_real[0], Complex::new(3.0, 0.0)); // Test (expr) real let expr_mul2_real = expr_mul2.real(); - let eval_mul2_real = manager.load(&expr_mul2_real, &dataset); + let model_mul2_real = manager.model(&expr_mul2_real); + let eval_mul2_real = model_mul2_real.load(&dataset); let result_mul2_real = eval_mul2_real.evaluate(&[]); assert_eq!(result_mul2_real[0], Complex::new(-2.0, 0.0)); // Test (expr) imag let expr_mul2_imag = expr_mul2.imag(); - let eval_mul2_imag = manager.load(&expr_mul2_imag, &dataset); + let model_mul2_imag = manager.model(&expr_mul2_imag); + let eval_mul2_imag = model_mul2_imag.load(&dataset); let result_mul2_imag = eval_mul2_imag.evaluate(&[]); assert_eq!(result_mul2_imag[0], Complex::new(4.0, 0.0)); // Test (amp) imag let expr_imag = aid3.imag(); - let eval_imag = manager.load(&expr_imag, &dataset); + let model_imag = manager.model(&expr_imag); + let eval_imag = model_imag.load(&dataset); let result_imag = eval_imag.evaluate(&[]); assert_eq!(result_imag[0], Complex::new(4.0, 0.0)); // Test (amp) norm_sqr let expr_norm = aid1.norm_sqr(); - let eval_norm = manager.load(&expr_norm, &dataset); + let model_norm = manager.model(&expr_norm); + let eval_norm = model_norm.load(&dataset); let result_norm = eval_norm.evaluate(&[]); assert_eq!(result_norm[0], Complex::new(4.0, 0.0)); // Test (expr) norm_sqr let expr_mul2_norm = expr_mul2.norm_sqr(); - let eval_mul2_norm = manager.load(&expr_mul2_norm, &dataset); + let model_mul2_norm = manager.model(&expr_mul2_norm); + let eval_mul2_norm = model_mul2_norm.load(&dataset); let result_mul2_norm = eval_mul2_norm.evaluate(&[]); assert_eq!(result_mul2_norm[0], Complex::new(20.0, 0.0)); } @@ -722,7 +767,8 @@ mod tests { let dataset = Arc::new(test_dataset()); let expr = &aid1 + &aid2; - let evaluator = manager.load(&expr, &dataset); + let model = manager.model(&expr); + let evaluator = model.load(&dataset); // Test initial state (all active) let result = evaluator.evaluate(&[]); @@ -752,7 +798,8 @@ mod tests { let aid = manager.register(amp).unwrap(); let dataset = Arc::new(test_dataset()); let expr = aid.norm_sqr(); - let evaluator = manager.load(&expr, &dataset); + let model = manager.model(&expr); + let evaluator = model.load(&dataset); let params = vec![2.0]; let gradient = evaluator.evaluate_gradient(¶ms); @@ -767,10 +814,14 @@ mod tests { let mut manager = Manager::default(); let amp = Scalar::new("parametric", parameter("test_param")); - manager.register(amp).unwrap(); + let aid = manager.register(amp).unwrap(); let parameters = manager.parameters(); + let model = manager.model(&aid.into()); + let model_parameters = model.parameters(); assert_eq!(parameters.len(), 1); assert_eq!(parameters[0], "test_param"); + assert_eq!(model_parameters.len(), 1); + assert_eq!(model_parameters[0], "test_param"); } #[test] diff --git a/src/amplitudes/ylm.rs b/src/amplitudes/ylm.rs index bbed1eb..66bb2d7 100644 --- a/src/amplitudes/ylm.rs +++ b/src/amplitudes/ylm.rs @@ -1,5 +1,6 @@ use nalgebra::DVector; use num::Complex; +use serde::{Deserialize, Serialize}; use crate::{ amplitudes::AmplitudeID, @@ -15,7 +16,7 @@ use crate::{ use super::Amplitude; /// An [`Amplitude`] for the spherical harmonic function $`Y_\ell^m(\theta, \phi)`$. -#[derive(Clone)] +#[derive(Clone, Serialize, Deserialize)] pub struct Ylm { name: String, l: usize, @@ -39,6 +40,7 @@ impl Ylm { } } +#[typetag::serde] impl Amplitude for Ylm { fn register(&mut self, resources: &mut Resources) -> Result { self.csid = resources.register_complex_scalar(None); @@ -91,7 +93,8 @@ mod tests { let dataset = Arc::new(test_dataset()); let expr = aid.into(); - let evaluator = manager.load(&expr, &dataset); + let model = manager.model(&expr); + let evaluator = model.load(&dataset); let result = evaluator.evaluate(&[]); @@ -108,7 +111,8 @@ mod tests { let dataset = Arc::new(test_dataset()); let expr = aid.into(); - let evaluator = manager.load(&expr, &dataset); + let model = manager.model(&expr); + let evaluator = model.load(&dataset); let result = evaluator.evaluate_gradient(&[]); assert_eq!(result[0].len(), 0); // amplitude has no parameters diff --git a/src/amplitudes/zlm.rs b/src/amplitudes/zlm.rs index 39f994c..bfca351 100644 --- a/src/amplitudes/zlm.rs +++ b/src/amplitudes/zlm.rs @@ -1,5 +1,6 @@ use nalgebra::DVector; use num::Complex; +use serde::{Deserialize, Serialize}; use crate::{ amplitudes::AmplitudeID, @@ -20,7 +21,7 @@ use super::Amplitude; /// [here](https://arxiv.org/abs/1906.04841)[^1] /// /// [^1]: Mathieu, V., Albaladejo, M., Fernández-Ramírez, C., Jackura, A. W., Mikhasenko, M., Pilloni, A., & Szczepaniak, A. P. (2019). Moments of angular distribution and beam asymmetries in $`\eta\pi^0`$ photoproduction at GlueX. _Physical Review D_, **100**(5). [doi:10.1103/physrevd.100.054017](https://doi.org/10.1103/PhysRevD.100.054017) -#[derive(Clone)] +#[derive(Clone, Serialize, Deserialize)] pub struct Zlm { name: String, l: usize, @@ -55,6 +56,7 @@ impl Zlm { } } +#[typetag::serde] impl Amplitude for Zlm { fn register(&mut self, resources: &mut Resources) -> Result { self.csid = resources.register_complex_scalar(None); @@ -122,7 +124,8 @@ mod tests { let dataset = Arc::new(test_dataset()); let expr = aid.into(); - let evaluator = manager.load(&expr, &dataset); + let model = manager.model(&expr); + let evaluator = model.load(&dataset); let result = evaluator.evaluate(&[]); @@ -140,7 +143,8 @@ mod tests { let dataset = Arc::new(test_dataset()); let expr = aid.into(); - let evaluator = manager.load(&expr, &dataset); + let model = manager.model(&expr); + let evaluator = model.load(&dataset); let result = evaluator.evaluate_gradient(&[]); assert_eq!(result[0].len(), 0); // amplitude has no parameters diff --git a/src/data.rs b/src/data.rs index 08aaa49..ad6869e 100644 --- a/src/data.rs +++ b/src/data.rs @@ -519,7 +519,7 @@ impl BinnedDataset { #[cfg(test)] mod tests { - use crate::prelude::ThreeMomentum; + use crate::traits::ThreeMomentum; use super::*; use approx::{assert_relative_eq, assert_relative_ne}; diff --git a/src/lib.rs b/src/lib.rs index db03160..ec60cdd 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -50,10 +50,15 @@ //! Although this particular amplitude is already included in `laddu`, let's assume it isn't and imagine how we would write it from scratch: //! //! ```rust -//! use laddu::prelude::*; +//! use laddu::{ +//! ParameterLike, Event, Cache, Resources, Mass, +//! ParameterID, Parameters, Float, LadduError, PI, AmplitudeID, Complex, +//! }; +//! use laddu::traits::*; //! use laddu::utils::functions::{blatt_weisskopf, breakup_momentum}; +//! use laddu::{Deserialize, Serialize}; //! -//! #[derive(Clone)] +//! #[derive(Clone, Serialize, Deserialize)] //! pub struct MyBreitWigner { //! name: String, //! mass: ParameterLike, @@ -90,6 +95,7 @@ //! } //! } //! +//! #[typetag::serde] //! impl Amplitude for MyBreitWigner { //! fn register(&mut self, resources: &mut Resources) -> Result { //! self.pid_mass = resources.register_parameter(&self.mass); @@ -123,10 +129,15 @@ //! ### Calculating a Likelihood //! We could then write some code to use this amplitude. For demonstration purposes, let's just calculate an extended unbinned negative log-likelihood, assuming we have some data and Monte Carlo in the proper [parquet format](#data-format): //! ```rust -//! # use laddu::prelude::*; +//! # use laddu::{ +//! # ParameterLike, Event, Cache, Resources, +//! # ParameterID, Parameters, Float, LadduError, PI, AmplitudeID, Complex, +//! # }; +//! # use laddu::traits::*; //! # use laddu::utils::functions::{blatt_weisskopf, breakup_momentum}; -//! # -//! # #[derive(Clone)] +//! # use laddu::{Deserialize, Serialize}; +//! +//! # #[derive(Clone, Serialize, Deserialize)] //! # pub struct MyBreitWigner { //! # name: String, //! # mass: ParameterLike, @@ -163,6 +174,7 @@ //! # } //! # } //! # +//! # #[typetag::serde] //! # impl Amplitude for MyBreitWigner { //! # fn register(&mut self, resources: &mut Resources) -> Result { //! # self.pid_mass = resources.register_parameter(&self.mass); @@ -186,6 +198,7 @@ //! # Complex::from(f * n) / d //! # } //! # } +//! use laddu::{Scalar, Mass, Manager, NLL, parameter, open}; //! let ds_data = open("test_data/data.parquet").unwrap(); //! let ds_mc = open("test_data/mc.parquet").unwrap(); //! @@ -203,9 +216,10 @@ //! &resonance_mass, //! )).unwrap(); //! let mag = manager.register(Scalar::new("mag", parameter("magnitude"))).unwrap(); -//! let model = (mag * bw).norm_sqr(); +//! let expr = (mag * bw).norm_sqr(); +//! let model = manager.model(&expr); //! -//! let nll = NLL::new(&manager, &model, &ds_data, &ds_mc); +//! let nll = NLL::new(&model, &ds_data, &ds_mc); //! println!("Parameters names and order: {:?}", nll.parameters()); //! let result = nll.evaluate(&[1.27, 0.120, 100.0]); //! println!("The extended negative log-likelihood is {}", result); @@ -278,37 +292,43 @@ pub mod likelihoods; pub mod resources; /// Utility functions, enums, and traits pub mod utils; - -/// All of the typical things you might want in the namespace for a typical `laddu` analysis. -pub mod prelude { - pub use crate::amplitudes::{ - breit_wigner::BreitWigner, - common::{ComplexScalar, PolarComplexScalar, Scalar}, - constant, parameter, - ylm::Ylm, - zlm::Zlm, - Amplitude, AmplitudeID, Evaluator, Expression, Manager, ParameterLike, - }; - pub use crate::data::{open, BinnedDataset, Dataset, Event}; - pub use crate::likelihoods::{ - LikelihoodEvaluator, LikelihoodExpression, LikelihoodID, LikelihoodManager, LikelihoodTerm, - MinimizerOptions, ReadWrite, NLL, - }; - pub use crate::resources::{ - Cache, ComplexMatrixID, ComplexScalarID, ComplexVectorID, MatrixID, ParameterID, - Parameters, Resources, ScalarID, VectorID, - }; - pub use crate::utils::enums::{Channel, Frame, Sign}; - pub use crate::utils::variables::{ - Angles, CosTheta, Mandelstam, Mass, Phi, PolAngle, PolMagnitude, Polarization, Variable, - }; +/// Useful traits for all crate structs +pub mod traits { + pub use crate::amplitudes::Amplitude; + pub use crate::likelihoods::LikelihoodTerm; + pub use crate::utils::variables::Variable; pub use crate::utils::vectors::{FourMomentum, FourVector, ThreeMomentum, ThreeVector}; - pub use crate::{Float, LadduError, PI}; - pub use ganesh::Status; - pub use nalgebra::{DVector, Vector3, Vector4}; - pub use num::Complex; + pub use crate::ReadWrite; } +pub use crate::amplitudes::{ + breit_wigner::BreitWigner, + common::{ComplexScalar, PolarComplexScalar, Scalar}, + constant, parameter, + ylm::Ylm, + zlm::Zlm, + AmplitudeID, Evaluator, Expression, Manager, Model, ParameterLike, +}; +pub use crate::data::{open, BinnedDataset, Dataset, Event}; +pub use crate::likelihoods::{ + LikelihoodEvaluator, LikelihoodExpression, LikelihoodID, LikelihoodManager, MinimizerOptions, + NLL, +}; +pub use crate::resources::{ + Cache, ComplexMatrixID, ComplexScalarID, ComplexVectorID, MatrixID, ParameterID, Parameters, + Resources, ScalarID, VectorID, +}; +pub use crate::utils::enums::{Channel, Frame, Sign}; +pub use crate::utils::variables::{ + Angles, CosTheta, Mandelstam, Mass, Phi, PolAngle, PolMagnitude, Polarization, +}; + +// Re-exports +pub use ganesh::Status; +pub use nalgebra::{DVector, Vector3, Vector4}; +pub use num::Complex; +pub use serde::{Deserialize, Serialize}; + /// A module containing Python bindings with PyO3. #[cfg(feature = "python")] pub mod python; @@ -377,3 +397,73 @@ impl From for PyErr { } } } + +use serde::de::DeserializeOwned; +use std::{ + fmt::Debug, + fs::File, + io::{BufReader, BufWriter}, + path::Path, +}; +/// A trait which allows structs with [`Serialize`] and [`Deserialize`](`serde::Deserialize`) to be +/// written and read from files with a certain set of types/extensions. +/// +/// Currently, Python's pickle format is supported supported, since it's an easy-to-parse standard +/// that supports floating point values better that JSON or TOML +pub trait ReadWrite: Serialize + DeserializeOwned { + /// Create a null version of the object which acts as a shell into which Python's `pickle` module + /// can load data. This generally shouldn't be used to construct the struct in regular code. + fn create_null() -> Self; + /// Save a [`serde`]-object to a file path, using the extension to determine the file format + fn save_as>(&self, file_path: T) -> Result<(), LadduError> { + let expanded_path = shellexpand::full(file_path.as_ref())?; + let file_path = Path::new(expanded_path.as_ref()); + let extension = file_path + .extension() + .and_then(|ext| ext.to_str().map(|ext| ext.to_string())) + .unwrap_or("".to_string()); + let file = File::create(file_path)?; + let mut writer = BufWriter::new(file); + match extension.as_str() { + "pkl" | "pickle" => serde_pickle::to_writer(&mut writer, self, Default::default())?, + _ => { + return Err(LadduError::Custom(format!( + "Unsupported file extension: {}\nValid options are \".pkl\" or \".pickle\"", + extension + ))) + } + }; + Ok(()) + } + /// Load a [`serde`]-object from a file path, using the extension to determine the file format + fn load_from>(file_path: T) -> Result { + let file_path = Path::new(&*shellexpand::full(file_path.as_ref())?).canonicalize()?; + let extension = file_path + .extension() + .and_then(|ext| ext.to_str().map(|ext| ext.to_string())) + .unwrap_or("".to_string()); + let file = File::open(file_path)?; + let reader = BufReader::new(file); + match extension.as_str() { + "pkl" | "pickle" => Ok(serde_pickle::from_reader(reader, Default::default())?), + _ => Err(LadduError::Custom(format!( + "Unsupported file extension: {}\nValid options are \".pkl\" or \".pickle\"", + extension + ))), + } + } +} + +impl ReadWrite for Status { + fn create_null() -> Self { + Status::default() + } +} +impl ReadWrite for Model { + fn create_null() -> Self { + Model { + manager: Manager::default(), + expression: Expression::default(), + } + } +} diff --git a/src/likelihoods.rs b/src/likelihoods.rs index d643c70..84aad5f 100644 --- a/src/likelihoods.rs +++ b/src/likelihoods.rs @@ -2,17 +2,14 @@ use std::{ collections::HashMap, convert::Infallible, fmt::{Debug, Display}, - fs::File, - io::{BufReader, BufWriter}, - path::Path, sync::Arc, }; use crate::{ - amplitudes::{AmplitudeValues, Evaluator, Expression, GradientValues, Manager}, + amplitudes::{AmplitudeValues, Evaluator, GradientValues, Model}, data::Dataset, resources::Parameters, - Float, LadduError, + Float, }; use accurate::{sum::Klein, traits::*}; use auto_ops::*; @@ -25,55 +22,6 @@ use num::Complex; #[cfg(feature = "rayon")] use rayon::prelude::*; -use serde::{de::DeserializeOwned, Serialize}; - -/// A trait which allows structs with [`Serialize`] and [`Deserialize`](`serde::Deserialize`) to be -/// written and read from files with a certain set of types/extensions. -/// -/// Currently, Python's pickle format is supported supported, since it's an easy-to-parse standard -/// that supports floating point values better that JSON or TOML -pub trait ReadWrite: Serialize + DeserializeOwned { - /// Save a [`serde`]-object to a file path, using the extension to determine the file format - fn save_as>(&self, file_path: T) -> Result<(), LadduError> { - let expanded_path = shellexpand::full(file_path.as_ref())?; - let file_path = Path::new(expanded_path.as_ref()); - let extension = file_path - .extension() - .and_then(|ext| ext.to_str().map(|ext| ext.to_string())) - .unwrap_or("".to_string()); - let file = File::create(file_path)?; - let mut writer = BufWriter::new(file); - match extension.as_str() { - "pkl" | "pickle" => serde_pickle::to_writer(&mut writer, self, Default::default())?, - _ => { - return Err(LadduError::Custom(format!( - "Unsupported file extension: {}\nValid options are \".pkl\" or \".pickle\"", - extension - ))) - } - }; - Ok(()) - } - /// Load a [`serde`]-object from a file path, using the extension to determine the file format - fn load>(file_path: T) -> Result { - let file_path = Path::new(&*shellexpand::full(file_path.as_ref())?).canonicalize()?; - let extension = file_path - .extension() - .and_then(|ext| ext.to_str().map(|ext| ext.to_string())) - .unwrap_or("".to_string()); - let file = File::open(file_path.clone())?; - let reader = BufReader::new(file); - match extension.as_str() { - "pkl" | "pickle" => Ok(serde_pickle::from_reader(reader, Default::default())?), - _ => Err(LadduError::Custom(format!( - "Unsupported file extension: {}\nValid options are \".pkl\" or \".pickle\"", - extension - ))), - } - } -} - -impl ReadWrite for Status {} /// A trait which describes a term that can be used like a likelihood (more correctly, a negative /// log-likelihood) in a minimization. @@ -99,15 +47,10 @@ impl NLL { /// Construct an [`NLL`] from a [`Manager`] and two [`Dataset`]s (data and Monte Carlo), as /// well as an [`Expression`]. This is the equivalent of the [`Manager::load`] method, /// but for two [`Dataset`]s and a different method of evaluation. - pub fn new( - manager: &Manager, - expression: &Expression, - ds_data: &Arc, - ds_accmc: &Arc, - ) -> Box { + pub fn new(model: &Model, ds_data: &Arc, ds_accmc: &Arc) -> Box { Self { - data_evaluator: manager.clone().load(expression, ds_data), - accmc_evaluator: manager.clone().load(expression, ds_accmc), + data_evaluator: model.load(ds_data), + accmc_evaluator: model.load(ds_accmc), } .into() } @@ -437,7 +380,7 @@ impl LikelihoodTerm for NLL { if *active { amp.compute(&data_parameters, event, cache) } else { - Complex::new(0.0, 0.0) + Complex::ZERO } }) .collect(), @@ -488,7 +431,7 @@ impl LikelihoodTerm for NLL { if *active { amp.compute(&mc_parameters, event, cache) } else { - Complex::new(0.0, 0.0) + Complex::ZERO } }) .collect(), @@ -552,7 +495,7 @@ impl LikelihoodTerm for NLL { if *active { amp.compute(&data_parameters, event, cache) } else { - Complex::new(0.0, 0.0) + Complex::ZERO } }) .collect(), @@ -601,7 +544,7 @@ impl LikelihoodTerm for NLL { if *active { amp.compute(&mc_parameters, event, cache) } else { - Complex::new(0.0, 0.0) + Complex::ZERO } }) .collect(), diff --git a/src/python.rs b/src/python.rs index 172dd3b..b6c1dd6 100644 --- a/src/python.rs +++ b/src/python.rs @@ -16,7 +16,7 @@ pub(crate) mod laddu { use crate as rust; use crate::likelihoods::LikelihoodTerm as RustLikelihoodTerm; use crate::likelihoods::MinimizerOptions; - use crate::prelude::ReadWrite; + use crate::traits::ReadWrite; use crate::utils::variables::Variable; use crate::utils::vectors::{FourMomentum, FourVector, ThreeMomentum, ThreeVector}; use crate::Float; @@ -1887,15 +1887,6 @@ pub(crate) mod laddu { #[pyclass] struct Manager(rust::amplitudes::Manager); - /// An Amplitude which can be registered by a Manager - /// - /// See Also - /// -------- - /// laddu.Manager - /// - #[pyclass] - struct Amplitude(Box); - #[pymethods] impl Manager { #[new] @@ -1934,20 +1925,18 @@ pub(crate) mod laddu { fn register(&mut self, amplitude: &Amplitude) -> PyResult { Ok(AmplitudeID(self.0.register(amplitude.0.clone())?)) } - /// Load an Expression by precalculating each term over the given Dataset + /// Generate a Model from the given expression made of registered Amplitudes /// /// Parameters /// ---------- /// expression : Expression or AmplitudeID /// The expression to use in precalculation - /// dataset : Dataset - /// The Dataset to use in precalculation /// /// Returns /// ------- - /// Evaluator - /// An object that can be used to evaluate the `expression` over each event in the - /// `dataset` + /// Model + /// An object which represents the underlying mathematical model and can be loaded with + /// a Dataset /// /// Notes /// ----- @@ -1956,7 +1945,7 @@ pub(crate) mod laddu { /// expression. These parameters will have no effect on evaluation, but they must be /// included in function calls. /// - fn load(&self, expression: &Bound<'_, PyAny>, dataset: &Dataset) -> PyResult { + fn model(&self, expression: &Bound<'_, PyAny>) -> PyResult { let expression = if let Ok(expression) = expression.extract::() { Ok(expression.0) } else if let Ok(aid) = expression.extract::() { @@ -1966,10 +1955,121 @@ pub(crate) mod laddu { "'expression' must either by an Expression or AmplitudeID", )) }?; - Ok(Evaluator(self.0.load(&expression, &dataset.0))) + Ok(Model(self.0.model(&expression))) + } + } + + /// A class which represents a model composed of registered Amplitudes + /// + #[pyclass] + struct Model(rust::amplitudes::Model); + + #[pymethods] + impl Model { + /// The free parameters used by the Manager + /// + /// Returns + /// ------- + /// parameters : list of str + /// The list of parameter names + /// + #[getter] + fn parameters(&self) -> Vec { + self.0.parameters() + } + /// Load a Model by precalculating each term over the given Dataset + /// + /// Parameters + /// ---------- + /// dataset : Dataset + /// The Dataset to use in precalculation + /// + /// Returns + /// ------- + /// Evaluator + /// An object that can be used to evaluate the `expression` over each event in the + /// `dataset` + /// + /// Notes + /// ----- + /// While the given `expression` will be the one evaluated in the end, all registered + /// Amplitudes will be loaded, and all of their parameters will be included in the final + /// expression. These parameters will have no effect on evaluation, but they must be + /// included in function calls. + /// + fn load(&self, dataset: &Dataset) -> PyResult { + Ok(Evaluator(self.0.load(&dataset.0))) + } + /// Save the Model to a file + /// + /// Parameters + /// ---------- + /// path : str + /// The path of the new file (overwrites if the file exists!) + /// + /// Raises + /// ------ + /// IOError + /// If anything fails when trying to write the file + /// + /// Notes + /// ----- + /// Valid file path names must have either the ".pickle" or ".pkl" extension + /// + fn save_as(&self, path: &str) -> PyResult<()> { + self.0.save_as(path)?; + Ok(()) + } + /// Load a Model from a file + /// + /// Parameters + /// ---------- + /// path : str + /// The path of the existing fit file + /// + /// Returns + /// ------- + /// Model + /// The model contained in the file + /// + /// Raises + /// ------ + /// IOError + /// If anything fails when trying to read the file + /// + /// Notes + /// ----- + /// Valid file path names must have either the ".pickle" or ".pkl" extension + /// + #[staticmethod] + fn load_from(path: &str) -> PyResult { + Ok(Model(crate::amplitudes::Model::load_from(path)?)) + } + #[new] + fn new() -> Self { + Model(crate::amplitudes::Model::create_null()) + } + fn __getstate__<'py>(&self, py: Python<'py>) -> PyResult> { + Ok(PyBytes::new_bound( + py, + serialize(&self.0).unwrap().as_slice(), + )) + } + fn __setstate__(&mut self, state: Bound<'_, PyBytes>) -> PyResult<()> { + *self = Model(deserialize(state.as_bytes()).unwrap()); + Ok(()) } } + /// An Amplitude which can be registered by a Manager + /// + /// See Also + /// -------- + /// laddu.Manager + /// + #[pyclass] + struct Amplitude(Box); + /// A class which can be used to evaluate a stored Expression /// /// See Also @@ -2300,10 +2400,8 @@ pub(crate) mod laddu { /// /// Parameters /// ---------- - /// manager : Manager - /// The Manager to use for precalculation - /// expression : Expression or AmplitudeID - /// The Expression to evaluate + /// model: Model + /// The Model to evaluate /// ds_data : Dataset /// A Dataset representing true signal data /// ds_accmc : Dataset @@ -2316,25 +2414,10 @@ pub(crate) mod laddu { #[pymethods] impl NLL { #[new] - #[pyo3(signature = (manager, expression, ds_data, ds_accmc))] - fn new( - manager: &Manager, - expression: &Bound<'_, PyAny>, - ds_data: &Dataset, - ds_accmc: &Dataset, - ) -> PyResult { - let expression = if let Ok(expression) = expression.extract::() { - Ok(expression.0) - } else if let Ok(aid) = expression.extract::() { - Ok(rust::amplitudes::Expression::Amp(aid.0)) - } else { - Err(PyTypeError::new_err( - "'expression' must either by an Expression or AmplitudeID", - )) - }?; + #[pyo3(signature = (model, ds_data, ds_accmc))] + fn new(model: &Model, ds_data: &Dataset, ds_accmc: &Dataset) -> PyResult { Ok(Self(rust::likelihoods::NLL::new( - &manager.0, - &expression, + &model.0, &ds_data.0, &ds_accmc.0, ))) @@ -3251,12 +3334,12 @@ pub(crate) mod laddu { /// Valid file path names must have either the ".pickle" or ".pkl" extension /// #[staticmethod] - fn load(path: &str) -> PyResult { - Ok(Status(ganesh::Status::load(path)?)) + fn load_from(path: &str) -> PyResult { + Ok(Status(ganesh::Status::load_from(path)?)) } #[new] fn new() -> Self { - Status(ganesh::Status::default()) + Status(ganesh::Status::create_null()) } fn __getstate__<'py>(&self, py: Python<'py>) -> PyResult> { Ok(PyBytes::new_bound( diff --git a/src/resources.rs b/src/resources.rs index b3a7d29..8554629 100644 --- a/src/resources.rs +++ b/src/resources.rs @@ -3,6 +3,8 @@ use std::{array, collections::HashMap}; use indexmap::IndexSet; use nalgebra::{SMatrix, SVector}; use num::Complex; +use serde::{Deserialize, Serialize}; +use serde_with::serde_as; use crate::{ amplitudes::{AmplitudeID, ParameterLike}, @@ -42,7 +44,7 @@ impl<'a> Parameters<'a> { } /// The main resource manager for cached values, amplitudes, parameters, and constants. -#[derive(Default, Debug, Clone)] +#[derive(Default, Debug, Clone, Serialize, Deserialize)] pub struct Resources { amplitudes: HashMap, pub(crate) active: Vec, @@ -60,7 +62,7 @@ pub struct Resources { /// A single cache entry corresponding to precomputed data for a particular /// [`Event`](crate::data::Event) in a [`Dataset`](crate::data::Dataset). -#[derive(Clone, Debug)] +#[derive(Clone, Debug, Serialize, Deserialize)] pub struct Cache(Vec); impl Cache { fn new(cache_size: usize) -> Self { @@ -162,7 +164,7 @@ impl Cache { } /// An object which acts as a tag to refer to either a free parameter or a constant value. -#[derive(Default, Copy, Clone, Debug)] +#[derive(Default, Copy, Clone, Debug, Serialize, Deserialize)] pub enum ParameterID { /// A free parameter. Parameter(usize), @@ -174,16 +176,17 @@ pub enum ParameterID { } /// A tag for retrieving or storing a scalar value in a [`Cache`]. -#[derive(Copy, Clone, Default, Debug)] +#[derive(Copy, Clone, Default, Debug, Serialize, Deserialize)] pub struct ScalarID(usize); /// A tag for retrieving or storing a complex scalar value in a [`Cache`]. -#[derive(Copy, Clone, Default, Debug)] +#[derive(Copy, Clone, Default, Debug, Serialize, Deserialize)] pub struct ComplexScalarID(usize, usize); /// A tag for retrieving or storing a vector in a [`Cache`]. -#[derive(Copy, Clone, Debug)] -pub struct VectorID([usize; R]); +#[serde_as] +#[derive(Copy, Clone, Debug, Serialize, Deserialize)] +pub struct VectorID(#[serde_as(as = "[_; R]")] [usize; R]); impl Default for VectorID { fn default() -> Self { @@ -192,8 +195,12 @@ impl Default for VectorID { } /// A tag for retrieving or storing a complex-valued vector in a [`Cache`]. -#[derive(Copy, Clone, Debug)] -pub struct ComplexVectorID([usize; R], [usize; R]); +#[serde_as] +#[derive(Copy, Clone, Debug, Serialize, Deserialize)] +pub struct ComplexVectorID( + #[serde_as(as = "[_; R]")] [usize; R], + #[serde_as(as = "[_; R]")] [usize; R], +); impl Default for ComplexVectorID { fn default() -> Self { @@ -202,8 +209,11 @@ impl Default for ComplexVectorID { } /// A tag for retrieving or storing a matrix in a [`Cache`]. -#[derive(Copy, Clone, Debug)] -pub struct MatrixID([[usize; C]; R]); +#[serde_as] +#[derive(Copy, Clone, Debug, Serialize, Deserialize)] +pub struct MatrixID( + #[serde_as(as = "[[_; C]; R]")] [[usize; C]; R], +); impl Default for MatrixID { fn default() -> Self { @@ -212,8 +222,12 @@ impl Default for MatrixID { } /// A tag for retrieving or storing a complex-valued matrix in a [`Cache`]. -#[derive(Copy, Clone, Debug)] -pub struct ComplexMatrixID([[usize; C]; R], [[usize; C]; R]); +#[serde_as] +#[derive(Copy, Clone, Debug, Serialize, Deserialize)] +pub struct ComplexMatrixID( + #[serde_as(as = "[[_; C]; R]")] [[usize; C]; R], + #[serde_as(as = "[[_; C]; R]")] [[usize; C]; R], +); impl Default for ComplexMatrixID { fn default() -> Self { diff --git a/src/utils/variables.rs b/src/utils/variables.rs index 5dac0db..cc57bd2 100644 --- a/src/utils/variables.rs +++ b/src/utils/variables.rs @@ -1,4 +1,5 @@ use nalgebra::Vector3; +use serde::{Deserialize, Serialize}; use std::sync::Arc; #[cfg(feature = "rayon")] @@ -33,7 +34,7 @@ pub trait Variable: Clone + Send + Sync { /// A struct for obtaining the mass of a particle by indexing the four-momenta of an event, adding /// together multiple four-momenta if more than one index is given. -#[derive(Clone, Debug)] +#[derive(Clone, Debug, Serialize, Deserialize)] pub struct Mass(Vec); impl Mass { /// Create a new [`Mass`] from the sum of the four-momenta at the given indices in the @@ -50,7 +51,7 @@ impl Variable for Mass { /// A struct for obtaining the $`\cos\theta`$ (cosine of the polar angle) of a decay product in /// a given reference frame of its parent resonance. -#[derive(Clone, Debug)] +#[derive(Clone, Debug, Serialize, Deserialize)] pub struct CosTheta { beam: usize, recoil: Vec, @@ -128,7 +129,7 @@ impl Variable for CosTheta { /// A struct for obtaining the $`\phi`$ angle (azimuthal angle) of a decay product in a given /// reference frame of its parent resonance. -#[derive(Clone, Debug)] +#[derive(Clone, Debug, Serialize, Deserialize)] pub struct Phi { beam: usize, recoil: Vec, @@ -205,7 +206,7 @@ impl Variable for Phi { } /// A struct for obtaining both spherical angles at the same time. -#[derive(Clone, Debug)] +#[derive(Clone, Debug, Serialize, Deserialize)] pub struct Angles { /// See [`CosTheta`]. pub costheta: CosTheta, @@ -239,7 +240,7 @@ impl Angles { } /// A struct defining the polarization angle for a beam relative to the production plane. -#[derive(Clone, Debug)] +#[derive(Clone, Debug, Serialize, Deserialize)] pub struct PolAngle { beam: usize, recoil: Vec, @@ -268,7 +269,7 @@ impl Variable for PolAngle { } /// A struct defining the polarization magnitude for a beam relative to the production plane. -#[derive(Copy, Clone, Default, Debug)] +#[derive(Copy, Clone, Default, Debug, Serialize, Deserialize)] pub struct PolMagnitude { beam: usize, } @@ -286,7 +287,7 @@ impl Variable for PolMagnitude { } /// A struct for obtaining both the polarization angle and magnitude at the same time. -#[derive(Clone, Debug)] +#[derive(Clone, Debug, Serialize, Deserialize)] pub struct Polarization { /// See [`PolMagnitude`]. pub pol_magnitude: PolMagnitude, @@ -316,7 +317,7 @@ impl Polarization { /// $`t = (p_1 - p_3)^2 = (p_4 - p_2)^2`$ /// /// $`u = (p_1 - p_4)^2 = (p_3 - p_2)^2`$ -#[derive(Clone, Debug)] +#[derive(Clone, Debug, Serialize, Deserialize)] pub struct Mandelstam { p1: Vec, p2: Vec,