From 22f74030fe77388964532fa950b4022f2b560852 Mon Sep 17 00:00:00 2001 From: Felix Zinn Date: Fri, 15 Nov 2024 09:47:52 +0100 Subject: [PATCH] add bounds for parameter in toy example --- examples/toy_generation.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/examples/toy_generation.py b/examples/toy_generation.py index 5c5faa8..862e78d 100644 --- a/examples/toy_generation.py +++ b/examples/toy_generation.py @@ -1,5 +1,6 @@ import equinox as eqx import jax +import jax.numpy as jnp from jaxtyping import Array, PRNGKeyArray from model import hists, model, observation @@ -7,6 +8,10 @@ key = jax.random.PRNGKey(0) +# set lower and upper bounds for the mu parameter +model = eqx.tree_at(lambda t: t.mu.lower, model, jnp.array([0.0])) +model = eqx.tree_at(lambda t: t.mu.upper, model, jnp.array([10.0])) + # generate a new model with sampled parameters according to their constraint pdfs toymodel = evm.parameter.sample(model, key)