def test_small_inverse(self):
key = random.PRNGKey(0)
for _ in range(STOCHASTIC_SAMPLES):
key, split = random.split(key)
mat = random.normal(split, (2, 2))
inv_mat = space._small_inverse(mat)
inv_mat_real_onp = onp.linalg.inv(mat)
inv_mat_real_np = np.linalg.inv(mat)
self.assertAllClose(inv_mat, inv_mat_real_onp, True)
self.assertAllClose(inv_mat, inv_mat_real_np, True)
def test_transform_inverse(self, spatial_dimension, dtype):
key = random.PRNGKey(0)
for _ in range(STOCHASTIC_SAMPLES):
key, split1, split2 = random.split(key, 3)
R = random.normal(split1, (PARTICLE_COUNT, spatial_dimension),
dtype=dtype)
T = random.normal(split2, (spatial_dimension, spatial_dimension),
dtype=dtype)
T_inv = space._small_inverse(T)
R_test = space.transform(T_inv, space.transform(T, R))
self.assertAllClose(R, R_test, True)
