def test_strat_int_vjp():
D, ts, y0, args, f, g = make_sde()
flat_args, _ = ravel_pytree(args)
bm = make_brownian_motion(ts[0], np.zeros(y0.shape), ts[1],
random.PRNGKey(0))
dt = 1e-4
eps = 1e-6
def strat_int(argv):
y0, args = argv
ys = stratonovich_integrate(f, g, y0, ts, bm, dt, args)
return np.sum(ys[1])
ys = stratonovich_integrate(f, g, y0, ts, bm, dt, args)
v_yt = np.ones_like(y0)
v_argst = np.zeros_like(flat_args)
y0_rec, exact_grad = vjp_strat_integrate(v_yt, v_argst, ys[-1], f, g, ts,
bm, dt, args)
numerical_grad = numerical_gradient(strat_int, (y0, args), eps=eps)
print("states:", y0, y0_rec)
assert np.allclose(y0, y0_rec, rtol=1e-2, atol=1e-02)
flat_grads, unravel = ravel_pytree(exact_grad)
print("numerical grad: ", unravel(numerical_grad))
print(" exact grad: ", exact_grad)
assert np.allclose(numerical_grad, flat_grads, rtol=1e-2, atol=1e-2)
def test_init_dataloader(self):
class Model(elegy.model.model_base.ModelBase):
def init_step(self, x, y_true, states: elegy.States):
return states.update(a=x.shape, b=y_true.shape)
def pred_step(self, x, states):
states = states.update(c=3)
return x + 1, states
model = Model()
x = np.random.uniform(size=(10, 1))
y = np.random.uniform(size=(10, 3))
dataset = TensorDataset(torch.from_numpy(x), torch.from_numpy(y))
dataloader = DataLoader(dataset, batch_size=2)
model.init(x=dataloader, batch_size=2)
assert model.states.a == (2, 1)
assert model.states.b == (2, 3)
assert model.initialized
y_pred = model.predict(x=dataloader)
assert jnp.allclose(y_pred, x + 1)
y_pred = model.predict(x=dataloader)
assert jnp.allclose(y_pred, x + 1)
y_pred
def test_train_controller(self, pip_feed, mode):
_, train_loss, score = train_controller(
self.controller,
self.sim,
pip_feed=pip_feed,
mode=mode,
duration=0.87,
dt=0.03,
epochs=50,
optimizer_params=self.optimizer_params,
scheduler='Cosine',
tensorboard_dir=None,
model_parameters=self.model_parameters,
print_loss=1)
print('pip_feed:' + str(pip_feed))
print('mode:' + str(mode))
print('train_loss:' + str(train_loss))
print('score:' + str(score))
# TODO(dsuo) remove one of parallel/sequential if results same
if (pip_feed, mode) == ('parallel', 'multipip'):
self.assertTrue(jnp.allclose(float(score), 20.96012))
elif (pip_feed, mode) == ('parallel', 'singular'):
self.assertTrue(jnp.allclose(float(score), 33.595024))
elif (pip_feed, mode) == ('sequential', 'multipip'):
self.assertTrue(jnp.allclose(score, 21.095024))
elif (pip_feed, mode) == ('sequential', 'singular'):
self.assertTrue(jnp.allclose(score, 33.380775))
def test_loop_over_tt(self, variant):
dt = 0.03
horizon = 29
waveform = BreathWaveform.create()
expiratory = Expiratory.create(waveform=waveform)
controller_state = self.controller.init()
expiratory_state = expiratory.init()
state, obs = self.sim.reset()
"""jit_loop_over_tt = jax.jit(
functools.partial(
loop_over_tt,
controller=self.controller,
expiratory=expiratory,
env=self.sim,
dt=dt))
"""
variant_loop_over_tt = variant(
functools.partial(loop_over_tt,
controller=self.controller,
expiratory=expiratory,
env=self.sim,
dt=dt))
_, (_, u_ins, _, pressures, _) = jax.lax.scan(
variant_loop_over_tt,
(state, obs, controller_state, expiratory_state, 0),
jnp.arange(horizon))
print('test_loop_over_tt u_ins:' + str(list(u_ins)))
print('test_loop_over_tt pressure:' + str(list(pressures)))
self.assertTrue(jnp.allclose(u_ins, self.expected_u_ins))
self.assertTrue(jnp.allclose(pressures, self.expected_pressures))
def test_train_simulator(self, scheduler):
r = 5
c = 10
batch_size = 2
epochs = 50
activation_fn_name = 'relu'
self.sim, test_loss = train_simulator(self.dataset,
self.sim,
self.sim.num_boundary_models,
activation_fn_name,
r,
c,
batch_size=batch_size,
epochs=epochs,
patience=10,
scheduler=scheduler,
print_loss=1,
use_tensorboard=False,
mode='train')
print('test_loss:' + str(test_loss))
print('scheduler: ' + str(scheduler))
if scheduler == 'Cosine':
self.assertTrue(
jnp.allclose(float(test_loss), 0.60894716, atol=1e-3))
elif scheduler == 'ReduceLROnPlateau':
self.assertTrue(
jnp.allclose(float(test_loss), 0.60894716, atol=1e-3))
def test_dist():
a, b = randn(100, 3), randn(200, 3)
print("Testing distance computation in input space, i.e. standard arrays",
file=sys.stderr)
for (name, pow) in [("sqeuclidean", 2), ("euclidean", 1)]:
for dist_args in [(a, b), (a, a)]:
ground_truth = cdist(*dist_args, name)
d = dist(*dist_args, power=pow)
print("*", name, pow, np.abs(ground_truth - d).max())
assert (allclose(ground_truth, d, atol=1e-05))
print("Testing distance computation between vectors of RKHS elements",
file=sys.stderr)
k = LinearKernel()
va, vb = FiniteVec(k, a), FiniteVec(k, b)
for pow in [2, 1]:
for insp_args, dist_args in [[(a, b), (va, vb)], [(a, a), (va, va)]]:
lin_rkhsdist_sq = (
np.diag(insp_args[0] @ insp_args[0].T)[:, np.newaxis] -
2 * insp_args[0] @ insp_args[1].T +
np.diag(insp_args[1] @ insp_args[1].T)[np.newaxis, :])
ground_truth = np.power(lin_rkhsdist_sq, pow / 2.)
d = dist(*dist_args, power=pow)
print("*", pow, np.abs(ground_truth - d).max())
assert (allclose(ground_truth, d, atol=1e-05))
def test_temperature_dist_warper(self):
input_ids = None
length = 20
scores = self._get_uniform_logits(batch_size=2, length=length)
# tweak scores to not be uniform anymore
scores[1, 5] = (1 / length) + 0.1 # peak, 1st batch
scores[1, 10] = (1 / length) - 0.4 # valley, 1st batch
# compute softmax
probs = jax.nn.softmax(scores, axis=-1)
temp_dist_warper_sharper = FlaxTemperatureLogitsWarper(temperature=0.5)
temp_dist_warper_smoother = FlaxTemperatureLogitsWarper(temperature=1.3)
warped_prob_sharp = jax.nn.softmax(temp_dist_warper_sharper(input_ids, scores.copy(), cur_len=None), axis=-1)
warped_prob_smooth = jax.nn.softmax(temp_dist_warper_smoother(input_ids, scores.copy(), cur_len=None), axis=-1)
# uniform distribution stays uniform
self.assertTrue(jnp.allclose(probs[0, :], warped_prob_sharp[0, :], atol=1e-3))
self.assertTrue(jnp.allclose(probs[0, :], warped_prob_smooth[0, :], atol=1e-3))
# sharp peaks get higher, valleys get lower
self.assertLess(probs[1, :].max(), warped_prob_sharp[1, :].max())
self.assertGreater(probs[1, :].min(), warped_prob_sharp[1, :].min())
# smooth peaks get lower, valleys get higher
self.assertGreater(probs[1, :].max(), warped_prob_smooth[1, :].max())
self.assertLess(probs[1, :].min(), warped_prob_smooth[1, :].min())
def test_jacobian():
fdm_jac0 = fdm.jacobian(hh)(x_input)
jax_jac0 = jax.jacrev(hh)(x_input)
with check:
assert np.allclose(fdm_jac0, jax_jac0)
rngkey = jax.random.PRNGKey(0)
v = jax.random.normal(rngkey, shape=(V.dim(), ), dtype="float64")
fdm_vjp0 = v @ fdm_jac0
jax_vjp0 = jax.vjp(hh, x_input)[1](v)
with check:
assert np.allclose(fdm_vjp0, jax_vjp0)
fdm_jac1 = fdm.jacobian(gg)(y_input)
jax_jac1 = jax.jacrev(gg)(y_input)
with check:
assert np.allclose(fdm_jac1, jax_jac1)
rngkey = jax.random.PRNGKey(1)
v = jax.random.normal(rngkey, shape=(V.dim(), ), dtype="float64")
fdm_vjp1 = v @ fdm_jac1
jax_vjp1 = jax.vjp(gg, y_input)[1](v)
with check:
assert np.allclose(fdm_vjp1, jax_vjp1)
def test_build_bijector_from_info():
x = np.array([[1, 2, 3, 4], [5, 6, 7, 8]])
init_fun, info1 = Chain(
Reverse(),
Chain(ColorTransform(0, 1, 1), Roll(-1)),
Scale(-0.5),
Softplus(0, 1),
Chain(Roll(), Scale(-4)),
)
params, forward_fun, inverse_fun = init_fun(random.PRNGKey(0), 4)
xfwd1, log_det1 = forward_fun(params, x)
init_fun, info2 = build_bijector_from_info(info1)
assert info1 == info2
params, forward_fun, inverse_fun = init_fun(random.PRNGKey(0), 4)
xfwd2, log_det2 = forward_fun(params, x)
assert np.allclose(xfwd1, xfwd2)
assert np.allclose(log_det1, log_det2)
invx, inv_log_det = inverse_fun(params, xfwd2)
assert np.allclose(x, invx)
assert np.allclose(log_det2, -inv_log_det)
def test_get_breaths(self):
print(self.analyzer.get_breaths())
u_in = jnp.array([1 for i in range(29)])
pressure = jnp.array([1 for i in range(29)])
for (u, p) in self.analyzer.get_breaths():
assert jnp.allclose(u, u_in)
assert jnp.allclose(p, pressure)
def test_optimizer_epoch(self):
optax_op = optax.adam(1e-3)
lr_schedule = lambda step, epoch: epoch
optimizer = elegy.Optimizer(optax_op,
lr_schedule=lr_schedule,
steps_per_epoch=2)
params = np.random.uniform((3, 4))
grads = np.random.uniform((3, 4))
rng = elegy.RNGSeq(42)
optimizer_states = optimizer.init(
rng=rng,
net_params=params,
)
assert jnp.allclose(optimizer.current_lr(optimizer_states), 0)
params, optimizer_states = optimizer.apply(params, grads,
optimizer_states, rng)
assert jnp.allclose(optimizer.current_lr(optimizer_states), 0)
params, optimizer_states = optimizer.apply(params, grads,
optimizer_states, rng)
assert jnp.allclose(optimizer.current_lr(optimizer_states), 1)
params, optimizer_states = optimizer.apply(params, grads,
optimizer_states, rng)
assert jnp.allclose(optimizer.current_lr(optimizer_states), 1)
params, optimizer_states = optimizer.apply(params, grads,
optimizer_states, rng)
def check_riemann_sym(g: Array,
R: Array,
rtol: float = 1e-5,
atol: float = 1e-8):
r"""
Checks four symmetries/identities of a Riemann curvature tensor. Batched.
Args:
g: the metric tensor :math:`g_{ab}(x)`
R: the Riemann tensor at the same point, :math:`R^r_{\ smn}(x)`
rtol: relative tolerance
atol: absolute tolerance
Raises:
ValueError: if one of the symmetries/identities does not hold
"""
R_lower = jnp.einsum("...as,...sbcd->...abcd", g, R)
if not jnp.allclose(R_lower, -jnp.einsum("...abdc->...abcd", R_lower),
rtol, atol):
raise ValueError("R_rsmn is not antisymmetric in the last two indices")
if not jnp.allclose(R_lower, -jnp.einsum("...bacd->...abcd", R_lower),
rtol, atol):
raise ValueError(
"R_rsmn is not antisymmetric in the first two indices")
if not jnp.allclose(
R_lower + jnp.einsum("...acdb->...abcd", R_lower) +
jnp.einsum("...adbc->...abcd", R_lower),
jnp.zeros_like(R_lower),
rtol,
atol,
):
raise ValueError("R_rsmn does not satisfy the first Bianchi identity")
if not jnp.allclose(R_lower, jnp.einsum("...cdab->...abcd", R_lower), rtol,
atol):
raise ValueError("R_rsmn does not satisfy interchange symmetry")
def test_multiply():
x = np.linspace(-2.5, 15, 5)[:, np.newaxis].astype(np.float32)
y = randn(x.size)[:, np.newaxis].astype(np.float32)
gk_x = GaussianKernel(0.1)
x_e1 = FiniteVec.construct_RKHS_Elem(gk_x, x)
x_e2 = FiniteVec.construct_RKHS_Elem(gk_x, y)
x_fv = FiniteVec(gk_x,
np.vstack([x, y]),
prefactors=np.hstack([x_e1.prefactors] * 2),
points_per_split=x.size)
oper_feat_vec = FiniteVec(gk_x, x)
oper = FiniteOp(oper_feat_vec, oper_feat_vec, np.eye(len(x)))
res_e1 = multiply(oper, x_e1)
res_e2 = multiply(oper, x_e2)
res_v = multiply(oper, x_fv)
assert np.allclose(
res_e1.prefactors,
(oper.matr @ oper.inp_feat.inner(x_e1)
).flatten()), "Application of operator to RKHS element failed."
assert np.allclose(
res_v.inspace_points,
np.vstack([res_e1.inspace_points, res_e2.inspace_points])
), "Application of operator to all vectors in RKHS vector failed at inspace points."
assert np.allclose(
res_v.prefactors, np.hstack([
res_e1.prefactors, res_e2.prefactors
])), "Application of operator to all vectors in RKHS vector failed."
assert np.allclose(
multiply(oper, oper).matr, oper.inp_feat.inner(
oper.outp_feat)), "Application of operator to operator failed."
def test_make_D_mat():
noiseterms = 2
p = 3
eta = 0 * jnp.arange(noiseterms * p).reshape(noiseterms, p)
zeta = 2 * jnp.arange(noiseterms * p).reshape(noiseterms, p)
D_mat = pychastic.vectorized_I_generation.make_D_mat(eta, zeta)
D_mat_loopy = pychastic.vectorized_I_generation.make_D_mat_loopy(eta, zeta)
assert jnp.allclose(D_mat, D_mat_loopy)
eta = 2 * jnp.arange(noiseterms * p).reshape(noiseterms, p)
zeta = 0 * jnp.arange(noiseterms * p).reshape(noiseterms, p)
D_mat = pychastic.vectorized_I_generation.make_D_mat(eta, zeta)
D_mat_loopy = pychastic.vectorized_I_generation.make_D_mat_loopy(eta, zeta)
assert jnp.allclose(D_mat, D_mat_loopy)
eta = 3 * jnp.arange(noiseterms * p).reshape(noiseterms, p)
zeta = 5 * jnp.arange(noiseterms * p).reshape(noiseterms, p)
D_mat = pychastic.vectorized_I_generation.make_D_mat(eta, zeta)
D_mat_loopy = pychastic.vectorized_I_generation.make_D_mat_loopy(eta, zeta)
assert jnp.allclose(D_mat, D_mat_loopy)
def isdm(mat):
"""Checks whether a given matrix is a valid density matrix.
Args:
mat (:obj:`jnp.ndarray`): Input matrix
Returns:
bool: ``True`` if input matrix is a valid density matrix;
``False`` otherwise
"""
isdensity = True
if (
isket(mat) == True
or isbra(mat) == True
or isherm(mat) == False
or jnp.allclose(jnp.real(jnp.trace(mat)), 1, atol=1e-09) == False
):
isdensity = False
else:
evals, _ = jnp.linalg.eig(mat)
for eig in evals:
if eig < 0 and jnp.allclose(eig, 0, atol=1e-06) == False:
isdensity = False
break
return isdensity
def test_dp_noise_perturbation(self):
svi_state = DPSVIState(None, self.rng, .3)
grads_list = [
jnp.mean(px_grads, axis=0) for px_grads in self.px_grads_list
]
new_svi_state, grads = \
self.svi._perturb_and_reassemble_gradients(
svi_state, grads_list, self.batch_size, self.tree_def
)
self.assertIs(svi_state.optim_state, new_svi_state.optim_state)
self.assertFalse(jnp.allclose(svi_state.rng_key,
new_svi_state.rng_key))
self.assertEqual(self.tree_def, jax.tree_structure(grads))
expected_std = (self.dp_scale * self.clipping_threshold /
self.batch_size) * svi_state.observation_scale
for site, px_site in zip(jax.tree_leaves(grads),
jax.tree_leaves(self.px_grads_list)):
self.assertEqual(px_site.shape[1:], site.shape)
self.assertTrue(
jnp.allclose(expected_std, jnp.std(site), atol=1e-2),
f"expected stdev {expected_std} but was {jnp.std(site)}")
self.assertTrue(jnp.allclose(0., jnp.mean(site), atol=1e-2))
def test_slice(self):
class MAE(elegy.metrics.Mean):
def call(self, y_true, y_pred):
return super().call(jnp.abs(y_true - y_pred))
y_true = dict(a=jnp.array([1.0, 2.0, 3.0]))
y_pred = dict(a=jnp.array([2.0, 3.0, 4.0]))
mae = MAE(on="a")
# raises because it doesn't use kwargs
with pytest.raises(BaseException):
sample_loss = mae.call_with_defaults()(y_true, y_pred)
# raises because it doesn't use __call__ which filters
with pytest.raises(BaseException):
sample_loss = mae.call(y_true=y_true, y_pred=y_pred)
loss = mae.call_with_defaults()(y_true=y_true, y_pred=y_pred)
assert jnp.allclose(loss, 1)
y_pred = dict(a=jnp.array([3.0, 4.0, 5.0]))
loss = mae.call_with_defaults()(y_true=y_true, y_pred=y_pred)
assert jnp.allclose(loss, 1.5)
def test_parametrized_jit(jitted_fun):
net = parametrized(jitted_fun)
inputs = random_inputs((2, ))
params = net.init_parameters(inputs, key=PRNGKey(0))
assert 'fun' == type(params).__name__
assert (3, ) == params.dense.bias.shape
params_ = net.init_parameters(inputs, key=PRNGKey(0))
assert_parameters_equal(params, params_)
out = net.apply(params, inputs)
assert out.shape == (3, )
assert np.allclose([0.84194356, -1.5927866, -1.7411114], out)
# run twice to cover cached jit call
out_ = net.apply(params, inputs)
assert np.allclose(out, out_)
out = net.apply(params, inputs, jit=True)
assert np.allclose(out, out_)
out_ = net.apply(params, inputs, jit=True)
assert np.allclose(out, out_)
out_ = net.apply_from({net: params}, inputs, jit=True)
assert np.allclose(out, out_)
def test_sample_with_intermediates(self):
locs = jnp.array([[-5., -5.], [0., 0.], [5., 5.]])
scales = jnp.ones_like(locs) * 0.1
pis = jnp.array([.5, .3, .2])
mix = GaussianMixture(locs, scales, pis)
rng = jax.random.PRNGKey(2963)
n_total = 1000
vals, interm = mix.sample_with_intermediates(rng, sample_shape=(10, n_total//10))
zs = interm[0]
self.assertEqual((10, n_total//10), jnp.shape(zs))
self.assertEqual((10, n_total//10, 2), jnp.shape(vals))
self.assertTrue(jnp.alltrue(zs >= 0) and jnp.alltrue(zs < 3))
unq_vals, unq_counts = np.unique(zs, return_counts=True)
unq_counts = unq_counts / n_total
# crude test that samples are from mixture (ratio of samples per component is plausible)
pis_stddev = jnp.sqrt(pis*(1-pis)/n_total)
self.assertTrue(jnp.allclose(unq_counts, pis, atol=3*pis_stddev))
for i in range(3):
# crude test that samples are from mixture (mean is within 3 times stddev per component)
self.assertTrue(jnp.allclose(locs[i], jnp.mean(vals[zs == i], axis=0), atol=3*scales[i]/jnp.sqrt(n_total)))
def test_adjoint_g_dynamics():
# Check that the function that computes the product of the augmented
# diffusion dynamics against a vector actually does the same thing as
# computing the diffusion matrix explicitly.
D, ts, y0, args, f, g = make_sde()
flat_args, unravel = ravel_pytree(args)
def flat_f(y, t, flat_args):
return f(y, t, unravel(flat_args))
def flat_g(y, t, flat_args):
return g(y, t, unravel(flat_args))
aug_y, unpack = ravel_pytree((y0, y0, np.zeros(np.size(flat_args))))
f_aug, g_aug, aug_gdg = make_ito_adjoint_dynamics(flat_f, flat_g, unpack)
# Test g_aug
sigma = make_explicit_sigma(flat_g, unpack)
explicit = sigma(aug_y, ts[0], flat_args)
implicit = jacobian(g_aug, argnums=3)(aug_y, ts[0], flat_args,
np.ones(y0.shape))
assert np.allclose(explicit, implicit)
# Test aug_gdg (Milstein correction factor)
explicit_milstein = make_explicit_milstein(sigma, aug_y, ts[0], flat_args)
implicit_milstein = jacobian(aug_gdg, argnums=3)(aug_y, ts[0], flat_args,
np.ones(y0.shape))
print(explicit_milstein)
print(implicit_milstein)
assert np.allclose(explicit_milstein, implicit_milstein)
def test_mvn(shape=(1000, 5)):
key = jr.PRNGKey(time.time_ns())
data = 5 * jr.normal(key, shape=shape)
mvn = dists.MultivariateNormalFullCovariance.fit(data)
assert np.allclose(data.mean(axis=0), mvn.loc, atol=1e-6)
assert np.allclose(np.cov(data, rowvar=False, bias=True),
mvn.covariance(),
atol=1e-6)
def test_ricci_scalar(n=10):
"""
See https://math.stackexchange.com/questions/3057630/calculation-of-christoffel-symbol-for-unit-sphere.
"""
xs = get_xs(n)
Rss = vmap(lambda x: get_ricci_scalar(x, get_g))(xs)
jnp.allclose(Rss, 2.0)
def test_boundary_conditions(self):
a = jnp.array(0.)
b = jnp.array(1.)
# Unconstrained boundary slopes.
bijector = _make_bijector(params_shape=(),
range_min=float(a),
range_max=float(b),
boundary_slopes='unconstrained')
log_slope_a = self.variant(bijector.forward_log_det_jacobian)(a)
log_slope_b = self.variant(bijector.forward_log_det_jacobian)(b)
self.assertEqual(self.variant(bijector.forward)(a), a)
self.assertEqual(self.variant(bijector.forward)(b), b)
self.assertFalse(jnp.allclose(log_slope_a, 0.))
self.assertFalse(jnp.allclose(log_slope_b, 0.))
# Lower boundary slope equal to 1.
bijector = _make_bijector(params_shape=(),
range_min=float(a),
range_max=float(b),
boundary_slopes='lower_identity')
log_slope_a = self.variant(bijector.forward_log_det_jacobian)(a)
log_slope_b = self.variant(bijector.forward_log_det_jacobian)(b)
self.assertEqual(self.variant(bijector.forward)(a), a)
self.assertEqual(self.variant(bijector.forward)(b), b)
self.assertEqual(log_slope_a, 0.)
self.assertFalse(jnp.allclose(log_slope_b, 0.))
# Upper boundary slope equal to 1.
bijector = _make_bijector(params_shape=(),
range_min=float(a),
range_max=float(b),
boundary_slopes='upper_identity')
log_slope_a = self.variant(bijector.forward_log_det_jacobian)(a)
log_slope_b = self.variant(bijector.forward_log_det_jacobian)(b)
self.assertEqual(self.variant(bijector.forward)(a), a)
self.assertEqual(self.variant(bijector.forward)(b), b)
self.assertFalse(jnp.allclose(log_slope_a, 0.))
self.assertEqual(log_slope_b, 0.)
# Both boundary slopes equal to 1.
bijector = _make_bijector(params_shape=(),
range_min=float(a),
range_max=float(b),
boundary_slopes='identity')
log_slope_a = self.variant(bijector.forward_log_det_jacobian)(a)
log_slope_b = self.variant(bijector.forward_log_det_jacobian)(b)
self.assertEqual(self.variant(bijector.forward)(a), a)
self.assertEqual(self.variant(bijector.forward)(b), b)
self.assertEqual(log_slope_a, 0.)
self.assertEqual(log_slope_b, 0.)
# Circular spline (periodic slope).
bijector = _make_bijector(params_shape=(),
range_min=float(a),
range_max=float(b),
boundary_slopes='circular')
log_slope_a = self.variant(bijector.forward_log_det_jacobian)(a)
log_slope_b = self.variant(bijector.forward_log_det_jacobian)(b)
self.assertEqual(self.variant(bijector.forward)(a), a)
self.assertEqual(self.variant(bijector.forward)(b), b)
self.assertEqual(log_slope_a, log_slope_b)
self.assertFalse(jnp.allclose(log_slope_b, 0.))
def test_matvec_vjp_jit():
def f(x, v):
res, vjp_fun = jax.vjp(mv, x)
(vjp, ) = vjp_fun(v)
return res, vjp
res, vjp = jax.jit(f)(x_local, vprime)
assert jnp.allclose(Ax, res)
assert jnp.allclose(ATvprime_local, vjp)
def test_matvecT_vjp_jit():
def f(y, v):
res, vjp_fun = jax.vjp(mvT, y)
(jvp, ) = vjp_fun(v)
return res, jvp
res, jvp = jax.jit(f)(y, v_local)
assert jnp.allclose(ATy_local, res)
assert jnp.allclose(Av, jvp)
def test_nat_to_std(nat_to_std, std_to_nat, dist):
std = nat_to_std(*dist)
nat = std_to_nat(*std)
std2 = nat_to_std(*nat)
for s1, s2 in zip(std, std2):
assert np.allclose(s1, s2)
for n1, n2 in zip(nat, dist):
assert np.allclose(n1, n2)
print('Passed the nat to std test!')
def test_is_bijective(self, bijector, args):
init_fun, bijector_info = bijector(*args)
params, forward_fun, inverse_fun = init_fun(random.PRNGKey(0),
x.shape[-1])
fwd_outputs, fwd_log_det = forward_fun(params, x)
inv_outputs, inv_log_det = inverse_fun(params, fwd_outputs)
assert np.allclose(inv_outputs, x, atol=1e-6)
assert np.allclose(fwd_log_det, -inv_log_det, atol=1e-6)
def test_cluster_from_ndarray():
""" """
N = 100
energy = rjax.uniform(rng, (N,), minval=0., maxval=3.)
costh = rjax.uniform(rng, (N,), minval=-1., maxval=1.)
phi = rjax.uniform(rng, (N,), minval=-np.pi, maxval=np.pi)
clu = Cluster.from_ndarray(np.column_stack([energy, costh, phi]))
assert np.allclose(clu.energy, energy)
assert np.allclose(clu.costh, costh)
assert np.allclose(clu.phi, phi)
assert clu.as_array.shape == (N, 3)
def test_cluster_ctor_arrays():
""" """
N = 100
energy = rjax.uniform(rng, (N,), minval=0., maxval=3.)
costh = rjax.uniform(rng, (N,), minval=-1., maxval=1.)
phi = rjax.uniform(rng, (N,), minval=-np.pi, maxval=np.pi)
clu = Cluster(energy, costh, phi)
assert np.allclose(clu.energy, energy)
assert np.allclose(clu.costh, costh)
assert np.allclose(clu.phi, phi)
assert clu.as_array.shape == (N, 3)
def test_helix_ctor_ndarray():
""" """
N = 100
d0, phi0, omega, z0, tanl = [rjax.uniform(rng, (N, )) for i in range(5)]
hel = Helix(d0, phi0, omega, z0, tanl)
assert np.allclose(d0, hel.d0)
assert np.allclose(phi0, hel.phi0)
assert np.allclose(omega, hel.omega)
assert np.allclose(z0, hel.z0)
assert np.allclose(tanl, hel.tanl)
