Merge pull request #20 from denehoffman/acceptance-correction

Acceptance corrections and saving fits
denehoffman · Nov 20, 2024 · f7e5ae5 · f7e5ae5
2 parents f026533 + 644750b
commit f7e5ae5
Show file tree

Hide file tree

Showing 35 changed files with 7,188 additions and 1,782 deletions.
diff --git a/README.md b/README.md
@@ -197,23 +197,29 @@ def main():
 
     s0p = manager.register(ld.Scalar("s0p", parameter("s0p")))
     s0n = manager.register(ld.Scalar("s0n", parameter("s0n")))
-    d2p = manager.register(ld.ComplexScalar("d2", parameter("d2 re"), parameter("d2 im")))
+    d2p = manager.register(ld.ComplexScalar("d2p", parameter("d2 re"), parameter("d2 im")))
 
     pos_re = (s0p * z00p.real() + d2p * z22p.real()).norm_sqr()
     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
 
-    nll = ld.NLL(manager, ds_data, ds_mc)
-    status = nll.minimize(model, [1.0] * len(nll.parameters))
+    nll = ld.NLL(manager, model, ds_data, ds_mc)
+    status = nll.minimize([1.0] * len(nll.parameters))
     print(status)
-    fit_weights = nll.project(status.x, model)
+    fit_weights = nll.project(status.x)
+    s0p_weights = nll.project_with(status.x, ["z00p", "s0p"])
+    s0n_weights = nll.project_with(status.x, ["z00n", "s0n"])
+    d2p_weights = nll.project_with(status.x, ["z22p", "d2p"])
     masses_mc = res_mass.value_on(ds_mc)
     masses_data = res_mass.value_on(ds_data)
     weights_data = ds_data.weights
     plt.hist(masses_data, weights=weights_data, bins=80, range=(1.0, 2.0), label="Data", histtype="step")
     plt.hist(masses_mc, weights=fit_weights, bins=80, range=(1.0, 2.0), label="Fit", histtype="step")
+    plt.hist(masses_mc, weights=s0p_weights, bins=80, range=(1.0, 2.0), label="$S_0^+$", histtype="step")
+    plt.hist(masses_mc, weights=s0n_weights, bins=80, range=(1.0, 2.0), label="$S_0^-$", histtype="step")
+    plt.hist(masses_mc, weights=d2p_weights, bins=80, range=(1.0, 2.0), label="$D_2^+$", histtype="step")
     plt.legend()
     plt.savefig("demo.svg")
 

diff --git a/benches/kmatrix_benchmark.rs b/benches/kmatrix_benchmark.rs
@@ -139,7 +139,7 @@ fn kmatrix_nll_benchmark(c: &mut Criterion) {
     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, &ds_data, &ds_mc, &model);
+    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| {

diff --git a/docs/requirements.txt b/docs/requirements.txt
diff --git a/docs/source/conf.py b/docs/source/conf.py
@@ -17,7 +17,7 @@
 # -- General configuration ---------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration
 
-extensions = ["sphinx.ext.autodoc", "sphinx.ext.autosummary", "sphinx.ext.napoleon"]
+extensions = ["sphinx.ext.autodoc", "sphinx.ext.autosummary", "sphinx.ext.napoleon", "sphinx_copybutton"]
 
 autodoc_default_options = {
     "members": True,

diff --git a/docs/source/index.rst b/docs/source/index.rst
@@ -1,13 +1,13 @@
 laddu
 =====
 
-``laddu`` is an library for analyzing particle physics data. It is written in Rust with bindings to Python which are documented here.
+`laddu <https://github.com/denehoffman/laddu>`_ is an library for analyzing particle physics data. It is written in Rust with bindings to Python which are documented here.
 
-Right now, this page just documents the basic Python API of the library, but the plan is to eventually write a tutorial here for general use. See the `GitHub Repository <https://github.com/denehoffman/laddu>`_ for more details.
 
 .. toctree::
-   :caption: API Reference
-   :maxdepth: 2
+   :caption: Contents
+   :maxdepth: 1
    :name: sec-api
-
+
+   tutorials/index
    modules/index
diff --git a/docs/source/modules/index.rst b/docs/source/modules/index.rst
@@ -1,5 +1,7 @@
-Submodules
-==========
+API Reference
+=============
+
+This is the main API documentation for ``laddu``.
 
 .. toctree::
    :maxdepth: 1

diff --git a/docs/source/modules/laddu/amplitudes/index.rst b/docs/source/modules/laddu/amplitudes/index.rst
@@ -1,4 +1,4 @@
-Amplitudes
+amplitudes
 ==========
 
 .. automodule:: laddu.amplitudes

diff --git a/docs/source/modules/laddu/data.rst b/docs/source/modules/laddu/data.rst
@@ -1,4 +1,4 @@
-Data
+data
 ====
 
 .. automodule:: laddu.data

diff --git a/docs/source/modules/laddu/likelihoods.rst b/docs/source/modules/laddu/likelihoods.rst
@@ -1,4 +1,4 @@
-Likelihoods
+likelihoods
 ===========
 
 .. automodule:: laddu.likelihoods

diff --git a/docs/source/modules/laddu/utils/index.rst b/docs/source/modules/laddu/utils/index.rst
@@ -1,4 +1,4 @@
-Utilities
+utils
 =========
 
 .. toctree::

diff --git a/docs/source/tutorials/binned_fit.rst b/docs/source/tutorials/binned_fit.rst
@@ -0,0 +1,4 @@
+Binned Fitting Tutorial
+=======================
+
+.. note:: Tutorial under construction!
diff --git a/docs/source/tutorials/index.rst b/docs/source/tutorials/index.rst
@@ -0,0 +1,11 @@
+Tutorials
+=========
+
+To get started with ``laddu``, please refer to the following tutorials for the basic operation of the library and the API reference for more details.
+
+.. toctree::
+   :maxdepth: 1
+
+   unbinned_fit
+   binned_fit
+   live_plotting
diff --git a/docs/source/tutorials/live_plotting.rst b/docs/source/tutorials/live_plotting.rst
@@ -0,0 +1,4 @@
+Live Plotting
+=============
+
+.. note:: Tutorial under construction!
diff --git a/docs/source/tutorials/unbinned_fit.rst b/docs/source/tutorials/unbinned_fit.rst
diff --git a/docs/source/tutorials/unbinned_fit_efficiency.png b/docs/source/tutorials/unbinned_fit_efficiency.png
diff --git a/docs/source/tutorials/unbinned_fit_mass_angle_plot.png b/docs/source/tutorials/unbinned_fit_mass_angle_plot.png
diff --git a/docs/source/tutorials/unbinned_fit_result.png b/docs/source/tutorials/unbinned_fit_result.png
diff --git a/docs/source/tutorials/unbinned_fit_result_acceptance_corrected.png b/docs/source/tutorials/unbinned_fit_result_acceptance_corrected.png
diff --git a/pyproject.toml b/pyproject.toml
@@ -25,7 +25,7 @@ amptools-to-laddu = "laddu:convert.run"
 
 [project.optional-dependencies]
 tests = ["pytest"]
-docs = ["sphinx", "sphinx-rtd-theme"]
+docs = ["sphinx", "sphinx-rtd-theme", "sphinx-copybutton"]
 
 [tool.maturin]
 python-source = "python"

diff --git a/python/laddu/amplitudes/__init__.pyi b/python/laddu/amplitudes/__init__.pyi
@@ -32,7 +32,7 @@ class Amplitude: ...
 class Manager:
     def __init__(self) -> None: ...
     def register(self, amplitude: Amplitude) -> AmplitudeID: ...
-    def load(self, dataset: Dataset, expression: Expression) -> Evaluator: ...
+    def load(self, expression: Expression, dataset: Dataset) -> Evaluator: ...
 
 class Evaluator:
     parameters: list[str]

diff --git a/python/laddu/likelihoods/__init__.pyi b/python/laddu/likelihoods/__init__.pyi
@@ -4,7 +4,7 @@ from typing import Any, Literal
 import numpy as np
 import numpy.typing as npt
 
-from laddu.amplitudes import Expression, Manager
+from laddu.amplitudes import Evaluator, Expression, Manager
 from laddu.data import Dataset
 
 class LikelihoodID:
@@ -43,18 +43,30 @@ class LikelihoodEvaluator:
 class NLL:
     parameters: list[str]
     data: Dataset
-    mc: Dataset
-    def __init__(self, manager: Manager, ds_data: Dataset, ds_mc: Dataset, expression: Expression) -> None: ...
+    accmc: Dataset
+    def __init__(
+        self,
+        manager: Manager,
+        expression: Expression,
+        ds_data: Dataset,
+        ds_accmc: Dataset,
+    ) -> None: ...
     def as_term(self) -> LikelihoodTerm: ...
     def activate(self, name: str | list[str]) -> None: ...
     def activate_all(self) -> None: ...
     def deactivate(self, name: str | list[str]) -> None: ...
     def deactivate_all(self) -> None: ...
     def isolate(self, name: str | list[str]) -> None: ...
     def evaluate(self, parameters: list[float] | npt.NDArray[np.float64]) -> float: ...
-    def project(self, parameters: list[float] | npt.NDArray[np.float64]) -> npt.NDArray[np.float64]: ...
+    def project(
+        self, parameters: list[float] | npt.NDArray[np.float64], *, mc_evaluator: Evaluator | None = None
+    ) -> npt.NDArray[np.float64]: ...
     def project_with(
-        self, parameters: list[float] | npt.NDArray[np.float64], name: str | list[str]
+        self,
+        parameters: list[float] | npt.NDArray[np.float64],
+        name: str | list[str],
+        *,
+        mc_evaluator: Evaluator | None = None,
     ) -> npt.NDArray[np.float64]: ...
     def minimize(
         self,

diff --git a/python_examples/example_1/accmc_1.parquet b/python_examples/example_1/accmc_1.parquet
-Original file line number
+Diff line change
@@ -1,4 +1,4 @@
-    Data
+    data
     ====
     .. automodule:: laddu.data
@@ Expand Down @@