Extending

Theories and likelihoods are Equinox modules: plug in any differentiable code or data.

From a function

For toy models and quick fits, wrap functions written with jax.numpy:

import jax.numpy as jnp
from cosmonuts import Model, Uniform
from cosmonuts.external import FunctionLikelihood, FunctionTheory

x = jnp.linspace(0.0, 1.0, 20)
y = 1.0 + 2.0 * x
s = 0.1

line = FunctionTheory(lambda a, b: {"y": lambda x: a + b * x}, provides={"y"})
fit = FunctionLikelihood(
    lambda provider: -0.5 * jnp.sum((provider.y(x) - y) ** 2 / s**2),
    requires={"y"}, name="points",
)
model = Model(line, [fit], {"a": Uniform(-5, 5), "b": Uniform(-5, 5)})

FunctionTheory takes its parameters from the function's argument names. FunctionLikelihood returns \(\log L\) and reads the theory's output through an argument named provider. NoTheory covers likelihoods that only read parameters.

In a run file, the same lambdas are written as strings (see Run files and CLI).

Subclass

• A Theory sets params and implements __call__(params), returning a provider with the quantities of Provider it serves and an ok flag.
• A Likelihood sets requires and implements log_likelihood(provider, params).
• GaussianLikelihood needs only data, a covariance and theory_vector(provider, params).

Keep the code differentiable and jit-able; return ok=False instead of raising when the theory fails at a parameter point.

Optional backends

If your theory needs a heavy dependency, import it inside the methods that use it, as the bundled ClassETheory and EFT4PyTheory do, and expose it as an extra.

To feed a bundled likelihood, return the quantities it requires with their standard signatures. Run-file names resolve to your classes by import path:

theory:
  mypackage.mytheory:MyTheory: {option: 1}