pfackeldey · pfackeldey · May 6, 2025 · May 6, 2025 · May 6, 2025 · May 6, 2025
diff --git a/.gitignore b/.gitignore
@@ -71,6 +71,7 @@ instance/
 # Sphinx documentation
 docs/_build/
 docs/composition.html
+docs/jupyter_execute
 
 # PyBuilder
 .pybuilder/

diff --git a/docs/api/effect.md b/docs/api/effect.md
@@ -1,7 +1,7 @@
-# evermore.effect
+# Effects
 
 ```{eval-rst}
-.. automodule:: evermore.effect
+.. automodule:: evermore.binned.effect
     :show-inheritance:
     :members:
 ```
diff --git a/docs/api/index.md b/docs/api/index.md
@@ -1,10 +1,12 @@
 # API Reference
 
 ```{toctree}
+parameter.md
+transform.md
+sample.md
 effect.md
-loss.md
 modifier.md
-parameter.md
+loss.md
 pdf.md
 util.md
 ```
diff --git a/docs/api/loss.md b/docs/api/loss.md
@@ -1,4 +1,4 @@
-# evermore.loss
+# Loss
 
 ```{eval-rst}
 .. automodule:: evermore.loss

diff --git a/docs/api/modifier.md b/docs/api/modifier.md
@@ -1,7 +1,7 @@
-# evermore.modifier
+# evermore.binned.modifier
 
 ```{eval-rst}
-.. automodule:: evermore.modifier
+.. automodule:: evermore.binned.modifier
     :show-inheritance:
     :members:
 ```
diff --git a/docs/api/parameter.md b/docs/api/parameter.md
@@ -1,7 +1,7 @@
-# evermore.parameter
+# Parameters
 
 ```{eval-rst}
-.. automodule:: evermore.parameter
+.. automodule:: evermore.parameters.parameter
     :show-inheritance:
     :members:
 ```
diff --git a/docs/api/pdf.md b/docs/api/pdf.md
@@ -1,4 +1,4 @@
-# evermore.pdf
+# PDFs
 
 ```{eval-rst}
 .. automodule:: evermore.pdf

diff --git a/docs/api/sample.md b/docs/api/sample.md
@@ -0,0 +1,7 @@
+# Sampling Parameters
+
+```{eval-rst}
+.. automodule:: evermore.parameters.sample
+    :show-inheritance:
+    :members:
+```
diff --git a/docs/api/transform.md b/docs/api/transform.md
@@ -0,0 +1,7 @@
+# Parameter Transformations
+
+```{eval-rst}
+.. automodule:: evermore.parameters.transform
+    :show-inheritance:
+    :members:
+```
diff --git a/docs/api/util.md b/docs/api/util.md
@@ -1,4 +1,4 @@
-# evermore.util
+# Utilities
 
 ```{eval-rst}
 .. automodule:: evermore.util

diff --git a/docs/conf.py b/docs/conf.py
@@ -49,7 +49,7 @@
     "sphinx.ext.viewcode",
     "sphinx.ext.mathjax",
     "sphinx.ext.napoleon",
-    "sphinx.ext.autosectionlabel",
+    # "sphinx.ext.autosectionlabel",
     "sphinx_copybutton",
     "sphinx_design",
     "sphinx_togglebutton",

diff --git a/docs/tips_and_tricks.md b/docs/tips_and_tricks.md
@@ -68,35 +68,35 @@ with open("composition.html", "w") as f:
 
 evermore provides a set of parameter transformations that can be used to modify the parameter values.
 This can be useful for example to ensure that the parameter values are within a certain range or that they are always positive.
-evermore provides two predefined transformations: [`evm.parameter.MinuitTransform`](#evermore.parameter.MinuitTransform) (for bounded parameters) and [`evm.parameter.SoftPlusTransform`](#evermore.parameter.SoftPlusTransform) (for positive parameters).
+evermore provides two predefined transformations: [`evm.transform.MinuitTransform`](#evermore.parameters.transform.MinuitTransform) (for bounded parameters) and [`evm.transform.SoftPlusTransform`](#evermore.parameters.transform.SoftPlusTransform) (for positive parameters).
 
 
 ```{code-cell} ipython3
 import evermore as evm
 import wadler_lindig as wl
 
 
-enforce_positivity = evm.parameter.SoftPlusTransform()
+enforce_positivity = evm.transform.SoftPlusTransform()
 pytree = {
     "a": evm.Parameter(2.0, transform=enforce_positivity),
     "b": evm.Parameter(0.1, transform=enforce_positivity),
 }
 
 # unwrap (or "transform")
-pytree_t = evm.parameter.unwrap(pytree)
+pytree_t = evm.transform.unwrap(pytree)
 # wrap back (or "inverse transform")
-pytree_tt = evm.parameter.wrap(pytree_t)
+pytree_tt = evm.transform.wrap(pytree_t)
 
 wl.pprint(("Original", pytree), width=150, short_arrays=False)
 wl.pprint(("Transformed", pytree_t), width=150, short_arrays=False)
 wl.pprint(("Transformed back / Original", pytree_tt), width=150, short_arrays=False)
 ```
 
-Transformations always transform into the unconstrained real space (using [`evm.parameter.unwrap`](#evermore.parameter.unwrap)) and back to the constrained space (using [`evm.parameter.wrap`](#evermore.parameter.wrap)).
+Transformations always transform into the unconstrained real space (using [`evm.transform.unwrap`](#evermore.parameters.transform.unwrap)) and back to the constrained space (using [`evm.transform.wrap`](#evermore.parameters.transform.wrap)).
 Typically, you would transform your parameters as a first step inside your loss (or model) function.
 Then, a minimizer can optimize the transformed parameters in the unconstrained space. Finally, you can transform them back to the constrained space for further processing.
 
-Custom transformations can be defined by subclassing [`evm.parameter.ParameterTransformation`](#evermore.parameter.ParameterTransformation) and implementing the [`wrap`](#evermore.parameter.ParameterTransformation.wrap) and [`unwrap`](#evermore.parameter.ParameterTransformation.unwrap) methods.
+Custom transformations can be defined by subclassing [`evm.transform.ParameterTransformation`](#evermore.parameters.transform.ParameterTransformation) and implementing the [`wrap`](#evermore.parameters.transform.ParameterTransformation.wrap) and [`unwrap`](#evermore.parameters.transform.ParameterTransformation.unwrap) methods.
 
 
 ## Parameter Partitioning
@@ -154,7 +154,7 @@ params = {"a": evm.NormalParameter(), "b": evm.NormalParameter()}
 rng_key = jax.random.key(0)
 rng_keys = jax.random.split(rng_key, 100)
 
-vec_sample = jax.vmap(evm.parameter.sample, in_axes=(None, 0))
+vec_sample = jax.vmap(evm.sample.sample_uncorrelated, in_axes=(None, 0))
 
 with treescope.active_autovisualizer.set_scoped(treescope.ArrayAutovisualizer()):
     tree = vec_sample(params, rng_keys)

diff --git a/src/evermore/__init__.py b/src/evermore/__init__.py
@@ -5,7 +5,7 @@
 from __future__ import annotations
 
 __author__ = "Peter Fackeldey"
-__email__ = "peter.fackeldey@rwth-aachen.de"
+__email__ = "peter.fackeldey@princeton.edu"
 __copyright__ = "Copyright 2023, Peter Fackeldey"
 __credits__ = ["Peter Fackeldey"]
 __contact__ = "https://github.com/pfackeldey/evermore"
@@ -27,7 +27,9 @@
     "modifier",
     "parameter",
     "pdf",
+    "sample",
     "staterror",
+    "transform",
     "util",
     "visualization",
 ]
@@ -38,17 +40,23 @@ def __dir__():
 
 
 from evermore import (  # noqa: E402
-    effect,
     loss,
-    modifier,
-    parameter,
     pdf,
-    staterror,
     util,
     visualization,
 )
-from evermore.modifier import Modifier  # noqa: E402
-from evermore.parameter import (  # noqa: E402,
+from evermore.binned import (  # noqa: E402
+    effect,
+    modifier,
+    staterror,
+)
+from evermore.binned.modifier import Modifier  # noqa: E402
+from evermore.parameters import (  # noqa: E402
+    parameter,
+    sample,
+    transform,
+)
+from evermore.parameters.parameter import (  # noqa: E402
     NormalParameter,
     Parameter,
 )
diff --git a/src/evermore/binned/__init__.py b/src/evermore/binned/__init__.py
diff --git a/src/evermore/effect.py → src/evermore/binned/effect.py b/src/evermore/effect.py → src/evermore/binned/effect.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import abc
 from collections.abc import Callable
 from typing import Literal
@@ -7,8 +9,7 @@
 import jax.numpy as jnp
 from jaxtyping import Array, PyTree
 
-from evermore.custom_types import OffsetAndScale
-from evermore.parameter import Parameter
+from evermore.parameters.parameter import Parameter
 from evermore.util import atleast_1d_float_array
 from evermore.visualization import SupportsTreescope
 
@@ -25,6 +26,18 @@ def __dir__():
     return __all__
 
 
+class OffsetAndScale(eqx.Module):
+    offset: Array = eqx.field(converter=atleast_1d_float_array, default=0.0)
+    scale: Array = eqx.field(converter=atleast_1d_float_array, default=1.0)
+
+    def broadcast(self) -> OffsetAndScale:
+        shape = jnp.broadcast_shapes(self.offset.shape, self.scale.shape)
+        return type(self)(
+            offset=jnp.broadcast_to(self.offset, shape),
+            scale=jnp.broadcast_to(self.scale, shape),
+        )
+
+
 class Effect(eqx.Module, SupportsTreescope):
     @abc.abstractmethod
     def __call__(self, parameter: PyTree[Parameter], hist: Array) -> OffsetAndScale: ...

diff --git a/src/evermore/modifier.py → src/evermore/binned/modifier.py b/src/evermore/modifier.py → src/evermore/binned/modifier.py
@@ -3,21 +3,20 @@
 import abc
 from collections.abc import Callable
 from functools import partial
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Protocol, runtime_checkable
 
 import equinox as eqx
 import jax
 import jax.numpy as jnp
 from jaxtyping import Array, ArrayLike, PyTree
 
-from evermore.custom_types import ModifierLike, OffsetAndScale
-from evermore.effect import DEFAULT_EFFECT
-from evermore.parameter import Parameter
+from evermore.binned.effect import DEFAULT_EFFECT, OffsetAndScale
+from evermore.parameters.parameter import Parameter
 from evermore.util import tree_stack
 from evermore.visualization import SupportsTreescope
 
 if TYPE_CHECKING:
-    from evermore.effect import Effect
+    from evermore.binned.effect import Effect
 
 __all__ = [
     "BooleanMask",
@@ -35,6 +34,13 @@ def __dir__():
     return __all__
 
 
+@runtime_checkable
+class ModifierLike(Protocol):
+    def offset_and_scale(self, hist: Array) -> OffsetAndScale: ...
+    def __call__(self, hist: Array) -> Array: ...
+    def __matmul__(self, other: ModifierLike) -> Compose: ...
+
+
 class AbstractModifier(eqx.Module):
     @abc.abstractmethod
     def offset_and_scale(self: ModifierLike, hist: Array) -> OffsetAndScale: ...
@@ -46,19 +52,19 @@ def __call__(self: ModifierLike, hist: Array) -> Array: ...
     def __matmul__(self: ModifierLike, other: ModifierLike) -> Compose: ...
 
 
-class ApplyFn(eqx.Module):
+class ApplyFn(AbstractModifier):
     @jax.named_scope("evm.modifier.ApplyFn")
     def __call__(self: ModifierLike, hist: Array) -> Array:
         os = self.offset_and_scale(hist=hist)
         return os.scale * (hist + os.offset)
 
 
-class MatMulCompose(eqx.Module):
+class MatMulCompose(AbstractModifier):
     def __matmul__(self: ModifierLike, other: ModifierLike) -> Compose:
         return Compose(self, other)
 
 
-class ModifierBase(ApplyFn, MatMulCompose, AbstractModifier, SupportsTreescope):
+class ModifierBase(ApplyFn, MatMulCompose, SupportsTreescope):
     """
     This serves as a base class for all modifiers.
     It automatically implements the __call__ method to apply the scale factors to the hist array

diff --git a/src/evermore/staterror.py → src/evermore/binned/staterror.py b/src/evermore/staterror.py → src/evermore/binned/staterror.py
@@ -7,10 +7,9 @@
 import jax.numpy as jnp
 from jaxtyping import Array, PyTree
 
-from evermore.custom_types import ModifierLike
-from evermore.effect import Identity
-from evermore.modifier import Modifier, Where
-from evermore.parameter import NormalParameter
+from evermore.binned.effect import Identity
+from evermore.binned.modifier import Modifier, ModifierLike, Where
+from evermore.parameters.parameter import NormalParameter
 from evermore.util import sum_over_leaves
 from evermore.visualization import SupportsTreescope
 

diff --git a/src/evermore/custom_types.py b/src/evermore/custom_types.py
diff --git a/src/evermore/loss.py b/src/evermore/loss.py
@@ -3,8 +3,8 @@
 import jax.numpy as jnp
 from jaxtyping import Array, PyTree
 
-from evermore.custom_types import PDFLike
-from evermore.parameter import Parameter, _params_map
+from evermore.parameters.parameter import Parameter, _params_map
+from evermore.pdf import PDFLike
 
 __all__ = [
     "get_log_probs",
@@ -15,7 +15,23 @@ def __dir__():
     return __all__
 
 
-def get_log_probs(module: PyTree) -> PyTree:
+def get_log_probs(params: PyTree) -> PyTree:
+    """
+    Compute the log probabilities for all parameters.
+
+    This function iterates over all parameters in the provided PyTree params,
+    applies their associated prior distributions (if any), and computes the
+    log probability for each parameter. If a parameter does not have a prior
+    distribution, a default log probability of 0.0 is returned.
+
+    Args:
+        params (PyTree): A PyTree containing parameters to compute log probabilities for.
+
+    Returns:
+        PyTree: A PyTree with the same structure as the input, where each parameter
+        is replaced by its corresponding log probability.
+    """
+
     def _constraint(param: Parameter) -> Array:
         prior = param.prior
         if isinstance(prior, PDFLike):
@@ -25,4 +41,4 @@ def _constraint(param: Parameter) -> Array:
         return jnp.array([0.0])
 
     # constraints from pdfs
-    return _params_map(_constraint, module)
+    return _params_map(_constraint, params)