<jupyter_start><jupyter_text>Imports<jupyter_code>from ase_interface import ANIENS
from ase_interface import ensemblemolecule_multigpu
import hdnntools as hdt

import numpy as np
import ase

import time

import os
os.environ["OMP_NUM_THREADS"] = "2"

import  ase
from ase.md.langevin import Langevin
from ase.md.verlet import VelocityVerlet
from ase.io.trajectory import Trajectory
from ase import units
from ase.optimize import LBFGS
from ase.io import read, write
from ase.parallel import world
from ase import Atom, Atoms

import nglview as nv
from nglview.datafiles import PDB, XTC
import mdtraj as md
import pytraj as pt
from IPython.display import Image
from IPython.core.display import HTML 
from IPython.display import display, HTML<jupyter_output><empty_output><jupyter_text>Setup files to load (2LUF or TrpCage)<jupyter_code>wkdir1 = '/home/jsmith48/scratch/MD_TEST/2luf/'
wkdir = '/data/roitberg/example_md/simulations/hexane/heat/'

fname = '2luf_solv'
fname = 'hex_solv'

# Molecule file
molfile = wkdir + fname + '.pdb'

# Dynamics file
xyzfile = wkdir1 + 'mdcrd_ani1x.xyz'

# Trajectory file
trajfile = wkdir1 + 'traj.dat'

# Optimized structure out:
optfile = wkdir1 + 'opt_' +  fname +'.xyz'

dt = 0.25
C = 0.5 # Optimization convergence

ntdir = '/home/jsmith48/scratch/ANI-2x_retrain/train/ANI-2x_v62518/'
cns = ntdir + 'rHCNOSFCl-5.2R_16-3.8A_a4-8.params'
sae = ntdir + 'sae_linfit.dat'
nnf = ntdir + 'train'
Nn = 8<jupyter_output><empty_output><jupyter_text>Load ANI models<jupyter_code># Load molecule
mol = read(molfile)

# Setup NeuroChem
aens = ensemblemolecule_multigpu(cns, sae, nnf, Nn, [0])

# Set ANI calculator
mol.set_calculator(ANIENS(aens))

# Set PBC
L = 39.5
L = 40.75
mol.set_cell(([[L, 0, 0],
               [0, L, 0],
               [0, 0, L]]))

mol.set_pbc((True, True, True))<jupyter_output>Building...<jupyter_text>Calculate energy and force<jupyter_code>s_time = time.time()
E = mol.get_potential_energy()
F = mol.get_forces()
print('Total time:',time.time()-s_time)

print('Energy:',E)

spc = mol.get_chemical_symbols()
print('Number of atoms:',len(spc))<jupyter_output>Total time: 1.3641605377197266
Energy: -1941290.203032304
Number of atoms: 6020<jupyter_text>Optimize the system<jupyter_code># Optimize molecule
start_time = time.time()
dyn = LBFGS(mol)
dyn.run(fmax=C,steps=5000)
print('[ANI Total time:', time.time() - start_time, 'seconds]')

# Save optimized mol
spc = mol.get_chemical_symbols()
pos = mol.get_positions(wrap=False).reshape(1,len(spc),3)


hdt.writexyzfile(optfile, pos, spc)<jupyter_output>Step     Time          Energy         fmax
LBFGS:    0 21:36:50 -1941290.203042        8.5330
LBFGS:    1 21:36:50 -1941300.989833        6.8743
LBFGS:    2 21:36:50 -1941310.483095        5.2419
LBFGS:    3 21:36:50 -1941319.824452        3.7283
LBFGS:    4 21:36:51 -1941326.731189        2.8127
LBFGS:    5 21:36:51 -1941332.478459        2.4521
LBFGS:    6 21:36:51 -1941336.172151        2.0610
LBFGS:    7 21:36:52 -1941338.775514        1.7102
LBFGS:    8 21:36:52 -1941340.853215        1.1098
LBFGS:    9 21:36:52 -1941342.470461        1.0763
LBFGS:   10 21:36:53 -1941343.837420        1.1116
LBFGS:   11 21:36:53 -1941345.042754        1.0571
LBFGS:   12 21:36:53 -1941346.174576        0.9919
LBFGS:   13 21:36:53 -1941347.277518        0.8889
LBFGS:   14 21:36:54 -1941348.341134        0.7840
LBFGS:   15 21:36:54 -1941349.332973        0.8056
LBFGS:   16 21:36:54 -1941350.253131        0.8170
LBFGS:   17 21:36:55 -1941351.105492        0.8181
LBFGS:   18 21:36:55 -1941351.89[...]<jupyter_text>Setup thermostat<jupyter_code># We want to run MD with constant energy using the VelocityVerlet algorithm.
# coefficient to 0.04 atomic units.
dyn = Langevin(mol, dt * units.fs, 300.0 * units.kB, 0.04)
# Open MD output
mdcrd = open(xyzfile,'w')

# Open MD output
traj = open(trajfile,'w')

# Define the printer
def storeenergy(a=mol, d=dyn, b=mdcrd, t=traj):  # store a reference to atoms in the definition.
    """Function to print the potential, kinetic and total energy."""
    epot = a.get_potential_energy() / len(a)
    ekin = a.get_kinetic_energy() / len(a)

    t.write(str(d.get_number_of_steps()) + ' ' + str(ekin / (1.5 * units.kB)) + ' ' + str(epot) + ' ' + str(ekin) + ' ' + str(epot+ekin) + '\n')
    b.write(str(len(a)) + '\n' + str(ekin / (1.5 * units.kB)) + ' Step: ' + str(d.get_number_of_steps()) + '\n')
    c = a.get_positions(wrap=False)
    for j, i in zip(a, c):
        b.write(str(j.symbol) + ' ' + str(i[0]) + ' ' + str(i[1]) + ' ' + str(i[2]) + '\n')

    print('Step: %d - %.1f ps Energy per atom: Epot = %.3f  Ekin = %.3f (T=%.3fK)  '
          'Etot = %.6f' % (d.get_number_of_steps(), (d.get_number_of_steps()*0.25)/1000, epot, ekin, ekin / (1.5 * units.kB), epot + ekin))
    
# Set the printer
dyn.attach(storeenergy, interval=100)<jupyter_output><empty_output><jupyter_text>Run MD<jupyter_code>start_time = time.time()
#dyn.set_temperature(300. * units.kB)
dyn.run(50000)
print('[ANI Total time:', time.time() - start_time, 'seconds]')<jupyter_output>Step: 100 - 0.0 ps Energy per atom: Epot = -322.482  Ekin = 0.005 (T=38.494K)  Etot = -322.477176
Step: 200 - 0.1 ps Energy per atom: Epot = -322.479  Ekin = 0.009 (T=66.300K)  Etot = -322.470721
Step: 300 - 0.1 ps Energy per atom: Epot = -322.477  Ekin = 0.012 (T=91.553K)  Etot = -322.464876
Step: 400 - 0.1 ps Energy per atom: Epot = -322.476  Ekin = 0.017 (T=128.040K)  Etot = -322.459845
Step: 500 - 0.1 ps Energy per atom: Epot = -322.478  Ekin = 0.022 (T=168.961K)  Etot = -322.456137
Step: 600 - 0.1 ps Energy per atom: Epot = -322.477  Ekin = 0.024 (T=184.807K)  Etot = -322.453187
Step: 700 - 0.2 ps Energy per atom: Epot = -322.475  Ekin = 0.025 (T=192.679K)  Etot = -322.449783
Step: 800 - 0.2 ps Energy per atom: Epot = -322.476  Ekin = 0.029 (T=222.736K)  Etot = -322.447495
Step: 900 - 0.2 ps Energy per atom: Epot = -322.475  Ekin = 0.030 (T=231.522K)  Etot = -322.445526
Step: 1000 - 0.2 ps Energy per atom: Epot = -322.473  Ekin = 0.029 (T=226.962K)  Etot = -322.443463
Step: 1100 -[...]<jupyter_text>Visualization<jupyter_code>m_traj = md.load_xyz(wkdir1+'mdcrd_ani1x.xyz',top=wkdir + fname + '.pdb')
m_view = nv.show_mdtraj(m_traj)
m_view.center()
m_view.add_representation( "ball+stick", selection='protein')
m_view.add_licorice('water and not hydrogen')
m_view<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1455259026687598, 'fraction_numerical': 0.14317111459968604, 'mean_word_length': 3.6, 'pattern_counts': {'":': 0, '<': 29, '<?xml version=': 0, '>': 29, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '8', 'dir': 'jupyter-structured-clean-dedup', 'id': '566654', 'license': 'MIT', 'n_tokens_mistral': 3293, 'n_tokens_neox': 2695, 'n_words': 618, 'path': 'notebooks/ASEANIMD_ANI-1x.ipynb', 'repo_name': 'vwcruzeiro/ANI-Tools', 'size': '6370'}