mcase_npt

#!/usr/bin/env python
import numpy as np
import mcase.particle_moves as mcp
import mcase.cell_moves as mcc
from ase.io import read, write
from ase.units import kB, GPa
from mcase.reader import read_input
from sys import argv

settings = read_input(argv[1])
rng = np.random.default_rng()

# setup system
walker = read(settings['xyz'])

# determine what the calculator is
if settings['calculator'] == 'nep':
    from pynep.calculate import NEP
    walker.calc = NEP(settings['model'])
elif settings['calculator'] == 'dp':
    from deepmd.calculator import DP
    walker.calc = DP(model=settings['model'])
elif settings['calculator'] == 'nequip':
    from nequip.ase import NequIPCalculator
    walker.calc = NequIPCalculator.from_deployed_model(settings['model'],\
            device='cuda')
else:
    print('calculator not recognized')
    exit()

# setup simulation
temperature = float(settings['temperature'])
pressure = float(settings['pressure'])*GPa
timestep = float(settings['timestep'])
cellstep = float(settings['cellstep'])
anglestep = float(settings['anglestep'])
number_samples = int(settings['samples'])
beta = 1.0/(kB*temperature)
scale = np.sqrt(2*timestep)
save_interval = int(settings['save_every'])

# successful read, print out stuff
print('read inputs from ' + argv[1])
print('particles read from ' + settings['xyz'])
print(settings['calculator'] + ' model read from ' + settings['model'])
print('temperature: %s K (beta = %s eV^-1)' % (temperature, beta))
print('pressure: %s eV/A^3 (%s GPa)' % (pressure, settings['pressure']))
print('timestep: %s (length scale = %s)' % (timestep, scale))
print('cell step: %s A' % cellstep)
print('angle step: %s degree' % anglestep)
print('number of samples to calculate: %d' % number_samples)
print('writing samples to ' + settings['save_xyz'])
print('will write every %d steps' % save_interval)

# start sampling
p_accepted = 0 # for particle moves
c_accepted = 0 # for cell moves
r = walker.get_positions()
energy = walker.get_potential_energy()
force = walker.get_forces()
cell = walker.cell.cellpar()
volume = walker.get_volume()
N = len(walker)

for i in range(number_samples):
    # particle move
    dr = rng.normal(size=r.shape, scale=scale) + timestep*beta*force
    r_new = r + dr
    walker.set_positions(r_new)
    energy_new = walker.get_potential_energy()
    force_new = walker.get_forces()
    # accept or reject
    if rng.random() < mcp.acceptance_ratio(r, r_new, energy, energy_new,\
            force, force_new, beta, timestep):
        energy = energy_new
        force = np.copy(force_new)
        r = walker.get_positions(wrap=True) # wrap to avoid blowing up
        p_accepted += 1
    else:
        walker.set_positions(r)
    # cell move
    dc = rng.random(size=6) - 0.5
    dc[:3] *= cellstep
    dc[3:] *= anglestep
    cell_new = cell + dc
    walker.set_cell(cell_new, scale_atoms=True)
    volume_new = walker.get_volume()
    energy_new = walker.get_potential_energy()
    if rng.random() < mcc.old_acceptance_ratio(energy, energy_new,\
            volume, volume_new, N, pressure, beta):
        energy = energy_new
        force = walker.get_forces()
        r = walker.get_positions(wrap=True)
        cell = np.copy(cell_new)
        volume = volume_new
        c_accepted += 1
    else:
        walker.set_cell(cell, scale_atoms=True)
    print(energy, cell[0], cell[1], cell[2], cell[3], cell[4], cell[5])
    if i % save_interval == (save_interval - 1):
        write(settings['save_xyz'], walker, format='extxyz', append=True,
                write_info=False, write_results=False)

# finished
print('particle move acceptance rate: %.2f' % (p_accepted / number_samples))
print('cell move acceptance rate: %.2f' % (c_accepted / number_samples))