ICODE-ADC/lammps
4
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

send eij to LAMMPS

Browse Source
This commit is contained in:
Linux-cpp-lisp
2023-05-25 17:49:02 -04:00
parent 1fd34ffac6
commit 3012426bd5
1 changed files with 6 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
examples/mliap/jax/mliap_unified_jax.py
View File

@ -15,11 +15,11 @@ def lj_potential(epsilon, sigma, rij):
lj1 = 4.0 * epsilon * sigma**12
lj2 = 4.0 * epsilon * sigma**6
eij = r6inv * (lj1 * r6inv - lj2)
return 0.5 * jnp.sum(eij)
eij = 0.5 * r6inv * (lj1 * r6inv - lj2)
return jnp.sum(eij), eij
# Construct a function computing _tot_e and its derivative
tot_e, fij = jax.value_and_grad(_tot_e)(rij)
return tot_e, fij
(_, eij), fij = jax.value_and_grad(_tot_e, has_aux=True)(rij)
return eij, fij
class MLIAPUnifiedJAX(MLIAPUnified):
@ -49,7 +49,7 @@ class MLIAPUnifiedJAX(MLIAPUnified):
# NOTE: Use data.rij_max with JAX.
rij = data.rij_max
e_tot, fij = lj_potential(self.epsilon, self.sigma, rij)
eij, fij = lj_potential(self.epsilon, self.sigma, rij)
data.energy = e_tot.item()
data.update_pair_energy(np.array(eij, dtype=np.float64))
data.update_pair_forces(np.array(fij, dtype=np.float64))