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)