fix tests

sbijax/_src/nn/consistency_model.py

@@ -96,7 +96,7 @@ def vector_field(self, theta, time, context, **kwargs):
  return self._network(theta=theta, time=time, context=context, **kwargs)
 
 
-# pylint: disable=too-many-arguments
+# pylint: disable=too-many-arguments,too-many-instance-attributes
 class _CMResnet(hk.Module):
  """A simplified 1-d residual network."""
 

sbijax/_src/nn/continuous_normalizing_flow.py

@@ -47,7 +47,7 @@ def __call__(self, method, **kwargs):
  """
  return getattr(self, method)(**kwargs)
 
- def sample(self, context):
+ def sample(self, context, **kwargs):
  """Sample from the pushforward.
 
  Args:
@@ -61,7 +61,7 @@ def ode_func(time, theta_t):
  theta_t = theta_t.reshape(-1, self._n_dimension)
  time = jnp.full((theta_t.shape[0], 1), time)
  ret = self.vector_field(
- theta=theta_t, time=time, context=context, is_training=False
+ theta=theta_t, time=time, context=context, **kwargs
  )
  return ret.reshape(-1)
 

sbijax/_src/scmpe_test.py

@@ -0,0 +1,60 @@
+# pylint: skip-file
+
+import distrax
+import haiku as hk
+from jax import numpy as jnp
+
+from sbijax import SCMPE
+from sbijax.nn import make_consistency_model
+
+
+def prior_model_fns():
+ p = distrax.Independent(distrax.Normal(jnp.zeros(2), jnp.ones(2)), 1)
+ return p.sample, p.log_prob
+
+
+def simulator_fn(seed, theta):
+ p = distrax.MultivariateNormalDiag(theta, 0.1 * jnp.ones_like(theta))
+ y = p.sample(seed=seed)
+ return y
+
+
+def log_density_fn(theta, y):
+ prior = distrax.Uniform(jnp.full(2, -3.0), jnp.full(2, 3.0))
+ likelihood = distrax.MultivariateNormalDiag(
+ theta, 0.1 * jnp.ones_like(theta)
+ )
+
+ lp = jnp.sum(prior.log_prob(theta)) + jnp.sum(likelihood.log_prob(y))
+ return lp
+
+
+def test_scmpe():
+ rng_seq = hk.PRNGSequence(0)
+ y_observed = jnp.array([-1.0, 1.0])
+
+ prior_simulator_fn, prior_logdensity_fn = prior_model_fns()
+ fns = (prior_simulator_fn, prior_logdensity_fn), simulator_fn
+
+ estim = SCMPE(fns, make_consistency_model(2))
+ data, params = None, {}
+ for i in range(2):
+ data, _ = estim.simulate_data_and_possibly_append(
+ next(rng_seq),
+ params=params,
+ observable=y_observed,
+ data=data,
+ n_simulations=100,
+ n_chains=2,
+ n_samples=200,
+ n_warmup=100,
+ )
+ params, info = estim.fit(next(rng_seq), data=data, n_iter=2)
+ _ = estim.sample_posterior(
+ next(rng_seq),
+ params,
+ y_observed,
+ n_chains=2,
+ n_samples=200,
+ n_warmup=100,
+ )

sbijax/_src/sfmpe_test.py

@@ -0,0 +1,60 @@
+# pylint: skip-file
+
+import distrax
+import haiku as hk
+from jax import numpy as jnp
+
+from sbijax import SFMPE
+from sbijax.nn import make_ccnf
+
+
+def prior_model_fns():
+ p = distrax.Independent(distrax.Normal(jnp.zeros(2), jnp.ones(2)), 1)
+ return p.sample, p.log_prob
+
+
+def simulator_fn(seed, theta):
+ p = distrax.MultivariateNormalDiag(theta, 0.1 * jnp.ones_like(theta))
+ y = p.sample(seed=seed)
+ return y
+
+
+def log_density_fn(theta, y):
+ prior = distrax.Uniform(jnp.full(2, -3.0), jnp.full(2, 3.0))
+ likelihood = distrax.MultivariateNormalDiag(
+ theta, 0.1 * jnp.ones_like(theta)
+ )
+
+ lp = jnp.sum(prior.log_prob(theta)) + jnp.sum(likelihood.log_prob(y))
+ return lp
+
+
+def test_sfmpe():
+ rng_seq = hk.PRNGSequence(0)
+ y_observed = jnp.array([-1.0, 1.0])
+
+ prior_simulator_fn, prior_logdensity_fn = prior_model_fns()
+ fns = (prior_simulator_fn, prior_logdensity_fn), simulator_fn
+
+ estim = SFMPE(fns, make_ccnf(2))
+ data, params = None, {}
+ for i in range(2):
+ data, _ = estim.simulate_data_and_possibly_append(
+ next(rng_seq),
+ params=params,
+ observable=y_observed,
+ data=data,
+ n_simulations=100,
+ n_chains=2,
+ n_samples=200,
+ n_warmup=100,
+ )
+ params, info = estim.fit(next(rng_seq), data=data, n_iter=2)
+ _ = estim.sample_posterior(
+ next(rng_seq),
+ params,
+ y_observed,
+ n_chains=2,
+ n_samples=200,
+ n_warmup=100,
+ )