Learning a LJ potential

# Install konda
!pip -q install konda

import konda
konda.install()

# Accept Anaconda Terms of Service
!conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/main
!conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/r

# Create the environment
!konda create -q -n pinn_env python=3.10 -y

!konda activate pinn_env
!konda run "pip -q install 'tensorflow[and-cuda]==2.15.1' 'ase>=3.25' 'PyYAML~=6.0.1' 'numpy<2'"
!konda run "pip -q install git+https://github.com/Teoroo-CMC/PiNN"

#Install ase to runtime
import sys
!{sys.executable} -m pip -q install ase

%matplotlib inline

import numpy as np
import matplotlib.pyplot as plt
from ase import Atoms
from ase.calculators.lj import LennardJones

# Helper function: get the position given PES dimension(s)
def three_body_sample(atoms, a, r):
    x = a * np.pi / 180
    pos = [[0, 0, 0],
           [0, 2, 0],
           [0, r*np.cos(x), r*np.sin(x)]]
    atoms.set_positions(pos)
    return atoms

atoms = Atoms('H3', calculator=LennardJones())

na, nr = 50, 50
arange = np.linspace(30,180,na)
rrange = np.linspace(1,3,nr)

# Truth
agrid, rgrid = np.meshgrid(arange, rrange)
egrid = np.zeros([na, nr])
for i in range(na):
    for j in range(nr):
        atoms = three_body_sample(atoms, arange[i], rrange[j])
        egrid[i,j] = atoms.get_potential_energy()

# Samples
nsample = 100
asample, rsample = [], []
distsample = []
data = {'e_data':[], 'f_data':[], 'elems':[], 'coord':[]}
for i in range(nsample):
    a, r = np.random.choice(arange), np.random.choice(rrange)
    atoms = three_body_sample(atoms, a, r)
    dist = atoms.get_all_distances()
    dist = dist[np.nonzero(dist)]
    data['e_data'].append(atoms.get_potential_energy())
    data['f_data'].append(atoms.get_forces())
    data['coord'].append(atoms.get_positions())
    data['elems'].append(atoms.numbers)
    asample.append(a)
    rsample.append(r)
    distsample.append(dist)


#save data
! rm -rf "training_data" "validation_data"
! mkdir "training_data" "validation_data"
! mkdir "validation_data/PES_analysis" "validation_data/pairwise_potential_analysis"
np.savez_compressed("training_data/data.npz", **data)
np.save("training_data/na_nr.npy", [na, nr])

plt.pcolormesh(agrid, rgrid, egrid, shading='auto')
plt.plot(asample, rsample, 'rx')
plt.colorbar()

%%writefile train.py

import numpy as np
from pinn.io import sparse_batch, load_numpy

data = np.load("training_data/data.npz")


data = {k:np.array(v) for k,v in data.items()}
dataset = lambda: load_numpy(data, splits={'train':8, 'test':2})

train = lambda: dataset()['train'].shuffle(100).repeat().apply(sparse_batch(100))
test = lambda: dataset()['test'].repeat().apply(sparse_batch(100))

%%writefile model_specification.py

import pinn
import numpy as np
import os, warnings
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
import tensorflow as tf
from ase import Atoms
from ase.calculators.lj import LennardJones
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
index_warning = 'Converting sparse IndexedSlices'
warnings.filterwarnings('ignore', index_warning)

# specify the grid
na, nr = np.load("training_data/na_nr.npy")
arange = np.linspace(30,180,na)
rrange = np.linspace(1,3,nr)
agrid, rgrid = np.meshgrid(arange, rrange)


params={
    'model_dir': '/tmp/PiNet',
    'network': {
        'name': 'PiNet',
        'params': {
            'ii_nodes':[8,8],
            'pi_nodes':[8,8],
            'pp_nodes':[8,8],
            'out_nodes':[8,8],
            'depth': 4,
            'rc': 3.0,
            'atom_types':[1]}},
    'model':{
        'name': 'potential_model',
        'params': {
            'e_dress': {1:-0.3},  # element-specific energy dress
            'e_scale': 2, # energy scale for prediction
            'e_unit': 1.0,  # output unit of energy dur
            'log_e_per_atom': True, # log e_per_atom and its distribution
            'use_force': True}}}      # include force in Loss function
model = pinn.get_model(params)

#add the model to train.py
!cat model_specification.py >> train.py

%%writefile -a train.py

train_spec = tf.estimator.TrainSpec(input_fn=train, max_steps=5e3)
eval_spec = tf.estimator.EvalSpec(input_fn=test, steps=10)
tf.estimator.train_and_evaluate(model, train_spec, eval_spec)

print("Training done.")

%rm -rf /tmp/PiNet #removes previously trained directory to rewrite it
!konda run "python train.py"

%%writefile PES_analysis.py

def three_body_sample(atoms, a, r):
    x = a * np.pi / 180
    pos = [[0, 0, 0],
           [0, 2, 0],
           [0, r*np.cos(x), r*np.sin(x)]]
    atoms.set_positions(pos)
    return atoms

#add the model
!cat model_specification.py >> PES_analysis.py

%%writefile -a PES_analysis.py

atoms = Atoms('H3', calculator=pinn.get_calc(model))
epred = np.zeros([na, nr])
for i in range(na):
    for j in range(nr):
        a, r = arange[i], rrange[j]
        atoms = three_body_sample(atoms, a, r)
        epred[i,j] = atoms.get_potential_energy()
np.save("validation_data/PES_analysis/epred.npy", epred)

!konda run "python PES_analysis.py"

import numpy as np
epred = np.load("validation_data/PES_analysis/epred.npy")

plt.pcolormesh(agrid, rgrid, epred, shading='auto')
plt.colorbar()
plt.title('NN predicted PES')
plt.figure()
plt.pcolormesh(agrid, rgrid, np.abs(egrid-epred), shading='auto')
plt.plot(asample, rsample, 'rx')
plt.title('NN Prediction error and sampled points')
plt.colorbar()

%%writefile pairwise_potential_analysis.py

#add the model
!cat model_specification.py >> pairwise_potential_analysis.py

%%writefile -a pairwise_potential_analysis.py

atoms1 = Atoms('H2', calculator=pinn.get_calc(model))
atoms2 = Atoms('H2', calculator=LennardJones())

nr2 = 100
rrange2 = np.linspace(1,1.9,nr2)
epred = np.zeros(nr2)
etrue = np.zeros(nr2)

for i in range(nr2):
    pos = [[0, 0, 0],
           [rrange2[i], 0, 0]]
    atoms1.set_positions(pos)
    atoms2.set_positions(pos)
    epred[i] = atoms1.get_potential_energy()
    etrue[i] = atoms2.get_potential_energy()

#save data
np.save("validation_data/pairwise_potential_analysis/epred.npy", epred)
np.save("validation_data/pairwise_potential_analysis/etrue.npy", epred)
np.save("validation_data/pairwise_potential_analysis/nr2.npy", nr2)

!konda run "python pairwise_potential_analysis.py"

nr2 = np.load("validation_data/pairwise_potential_analysis/nr2.npy")
epred = np.load("validation_data/pairwise_potential_analysis/epred.npy")
epred = np.load("validation_data/pairwise_potential_analysis/epred.npy")
etrue = np.load("validation_data/pairwise_potential_analysis/etrue.npy")

rrange2 = np.linspace(1,1.9,nr2)

f, (ax1, ax2) = plt.subplots(2,1, gridspec_kw = {'height_ratios':[3, 1]})
ax1.plot(rrange2, epred)
ax1.plot(rrange2, etrue,'--')
ax1.legend(['Prediction', 'Truth'], loc=4)
_=ax2.hist(np.concatenate(distsample,0), 20, range=(1,1.9))

%%writefile MD.py

from ase import units
from ase.io import Trajectory
from ase.md.nvtberendsen import NVTBerendsen
from ase.md.velocitydistribution import thermalize_momenta

#add the model
!cat model_specification.py >> MD.py

%%writefile -a MD.py

atoms = Atoms('H', cell=[2, 2, 2], pbc=True)
atoms = atoms.repeat([5,5,5])
atoms.rattle()
atoms.calc = pinn.get_calc(model)
thermalize_momenta(atoms, temperature_K=300)
dyn = NVTBerendsen(atoms, 0.5 * units.fs, 300, taut=0.5*100*units.fs)
dyn.attach(Trajectory('ase_nvt.traj', 'w', atoms).write, interval=10)
dyn.run(5000)

!konda run "python MD.py"