def test_EMT_from_db_nbrlist(self, spatial_dimension, dtype, low_pressure): if spatial_dimension == 2: N = 64 else: N = 128 if dtype == jnp.float32: max_grad_thresh = 1e-5 atol = 1e-4 rtol = 1e-3 else: max_grad_thresh = 1e-10 atol = 1e-8 rtol = 1e-5 for index in range(NUM_SAMPLES): cijkl, R, sigma, box = test_util.load_elasticity_test_data( spatial_dimension, low_pressure, dtype, index) displacement, shift = space.periodic_general( box, fractional_coordinates=True) neighbor_fn, energy_fn = energy.soft_sphere_neighbor_list( displacement, box, sigma=sigma, fractional_coordinates=True) nbrs = neighbor_fn.allocate(R) assert (jnp.max(jnp.abs(grad(energy_fn)(R, nbrs))) < max_grad_thresh) EMT_fn = jit( elasticity.athermal_moduli(energy_fn, check_convergence=True)) C, converged = EMT_fn(R, box, neighbor=nbrs) assert (C.dtype == dtype) assert (C.shape == (spatial_dimension, spatial_dimension, spatial_dimension, spatial_dimension)) if converged: self.assertAllClose(cijkl, elasticity._extract_elements(C, False), atol=atol, rtol=rtol) #make sure the symmetries are there self.assertAllClose(C, jnp.einsum("ijkl->jikl", C)) self.assertAllClose(C, jnp.einsum("ijkl->ijlk", C)) self.assertAllClose(C, jnp.einsum("ijkl->lkij", C))
def test_EMT_from_db_dynamic(self, spatial_dimension, dtype, low_pressure): if spatial_dimension == 2: N = 64 else: N = 128 if dtype == jnp.float32: max_grad_thresh = 1e-5 atol = 1e-4 rtol = 1e-3 else: max_grad_thresh = 1e-10 atol = 1e-8 rtol = 1e-5 for index in range(NUM_SAMPLES): cijkl, R, sigma, box = test_util.load_elasticity_test_data( spatial_dimension, low_pressure, dtype, index) R = space.transform(box, R) box = box[0, 0] displacement, shift = space.periodic(box) #Below we use the wrong sigma, so we must pass it dynamically energy_fn = energy.soft_sphere_pair(displacement, sigma=1.0) maxgrad = jnp.max(jnp.abs(grad(energy_fn)(R, sigma=sigma))) assert (maxgrad < max_grad_thresh) EMT_fn = jit( elasticity.athermal_moduli(energy_fn, check_convergence=True)) C, converged = EMT_fn(R, box, sigma=sigma) assert (C.dtype == dtype) assert (C.shape == (spatial_dimension, spatial_dimension, spatial_dimension, spatial_dimension)) if converged: self.assertAllClose(cijkl, elasticity._extract_elements(C, False), atol=atol, rtol=rtol) #make sure the symmetries are there self.assertAllClose(C, jnp.einsum("ijkl->jikl", C)) self.assertAllClose(C, jnp.einsum("ijkl->ijlk", C)) self.assertAllClose(C, jnp.einsum("ijkl->lkij", C))