simulation.py

import argparse, pickle, os, multiprocessing
from pprint import pprint
from datetime import datetime
import sys
from copy import deepcopy
from math import floor
from prettytable import PrettyTable
import time

from simtk.openmm.app import *
from simtk.openmm import *
from simtk.unit import *
from openmmtools.integrators import NoseHooverChainVelocityVerletIntegrator

from sys import stdout
import numpy as np
import os, pickle, progressbar
from termcolor import colored

from trajectory_class import Trajectory, interval_divide, get_max_processes
from custom_units import *
from convert_xyz_to_npz import convert_xyz_to_npz

# Using cython to make specific parts involving
# lots of floating point operations faster
sys.path.append('cython_files/pbc_adjustment')
import pbc_adjustment

def timing(f):
    # measure time taken by a function call
    def wrap(*args):
        time1 = time.time()
        ret = f(*args)
        time2 = time.time()
        print('{:s} function took {:.3f} ms'.format(f.__name__, (time2-time1)*1000.0))

        return ret
    return wrap

def get_force_magnitude(r, eps=0.238 * kilocalories_per_mole, sig=3.4 * angstrom ):
    if r == 0:
        raise ValueError
    return 4 * eps * ( -(12 * (sig**12) )/(r**13)  +  (6.0*(sig**6))/(r**7) )


def get_my_forces_helper(args):
    return get_my_forces(*args)

#@timing
def get_my_forces(dist_matrix, st_atom_ind, en_atom_ind, b_size):
    force_unit = kilocalorie/(nanometer*mole)
    ret = []

    dist_matrix_values_only = np.array([i._value for i in dist_matrix])
    
    for atom_ind in range(st_atom_ind, en_atom_ind+1):
        
        # Getting the dx array using numpy
        #del dx_array[atom_ind] # above computed dx_array is not used anymore. So deleting just for comparision with numpy's dx_array_2
        temp_b_size = np.array([i._value for i in b_size]) * b_size[0].unit
        
        dx_array_2 = dist_matrix_values_only - dist_matrix_values_only[atom_ind] # Added May 26

        temp_b_size_2 = np.array([i._value for i in temp_b_size])
        
        pbc_adjustment.pbc_adjust(dx_array_2, temp_b_size_2)
        
        dx_array_2 = np.delete(dx_array_2, atom_ind, 0) * dist_matrix.unit

        total_distances = np.linalg.norm(dx_array_2, axis=1) * dx_array_2.unit # Get total distances
        indices_within_cutoff = np.argwhere(total_distances <= lj_cutoff)
        indices_within_cutoff = np.reshape(indices_within_cutoff, (-1))
        if len(indices_within_cutoff) == 0:
            raise ValueError("No atoms within cutoff range. Maybe somethong wrong with the system.")
        total_distances = total_distances[indices_within_cutoff]
        dx_array_2 = dx_array_2[indices_within_cutoff]
        
        force_magnitudes = [get_force_magnitude(i) for i in total_distances]
        force_magnitudes = np.array([i._value for i in force_magnitudes]) * force_magnitudes[0].unit

        forces_due_to_each_atom = (dx_array_2 * force_magnitudes[:, np.newaxis]) / total_distances[:, np.newaxis]
        final_force = np.sum(forces_due_to_each_atom, axis=0)
        
        ret.append(final_force)
    return ret

def append_traj(traj_filename, traj_array, st_ind, en_ind, skip_every_n_frames):
    with open(traj_filename, 'a') as out_file:
        for traj_array_frame in traj_array[st_ind:en_ind+1:skip_every_n_frames]:
            out_file.write(str(len(traj_array_frame)) + '\n')
            #out_file.write(" generated by psp" + '\n')
            out_file.write(str(" Distance_unit = angstroms, Vel unit = nm/ps, frc unit = kJ/nm mol" + "\n"))

            if '_pos.xyz' in traj_filename:
                    for c in traj_array_frame: # openmm uses nanometre, multiply by 10 to dump angstroms
                        out_file.write(str("   Ar" + "\t" + str(10*c[0]._value) + "\t" + str(10*c[1]._value) + "\t" + str(10*c[2]._value) + "\n"))
            else:
                if '_frc.xyz' in traj_filename:
                    assert traj_array[st_ind][0].unit == kilojoule / (nanometer * mole)
                if '_vel.xyz' in traj_filename:
                    assert traj_array[st_ind][0].unit == nanometer / picosecond                
                for c in traj_array_frame:
                    out_file.write(str("   Ar" + "\t" + str(c[0]._value) + "\t" + str(c[1]._value) + "\t" + str(c[2]._value) + "\n"))
    return

def convert_forces_to_atomic_units(f):
    return f.in_units_of(hartree / bohr / mole) / constants.AVOGADRO_CONSTANT_NA

def str2bool(v):
    if isinstance(v, bool):
       return v
    if v.lower() in ('yes', 'true', 't', 'y', '1'):
        return True 
    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
        return False
    else:
        raise argparse.ArgumentTypeError('Boolean value expected.')

if __name__ == "__main__":
    now_time = datetime.now().strftime('%b-%d-%Y--%H-%M-%S')

    pdb_filename = 'argon_random_poisson_translated_0.pdb'
    #pdb_filename = 'argon_500_atoms_equilibrated.pdb'

    box_size = 1.67 * nanometer # box for 96 atoms system
    #box_size = 2.8947833204138995 * nanometer # box for  500 atoms system with same density as 96 atoms and 1.67 nanometer system

    assert pdb_filename is not None
    pdb = PDBFile(pdb_filename)

    # If not a cuboid, then give jut 1 side cube
    if not isinstance(box_size, list):
        box_size = [box_size for _ in range(3)]

    pdb.topology.setUnitCellDimensions([box_size_val._value for box_size_val in box_size]*box_size[0].unit)

    print(pdb.topology.getPeriodicBoxVectors())
    print(pdb.topology.getUnitCellDimensions())
    assert tuple([box_size_val._value for box_size_val in box_size])*box_size[0].unit == pdb.topology.getUnitCellDimensions() == tuple([pdb.topology.getPeriodicBoxVectors()[box_vector_ind][box_vector_ind]._value for box_vector_ind in range(3)])*pdb.topology.getPeriodicBoxVectors().unit

    steps = 10000
    TEMPERATURE = 90.0*kelvin
    PRESSURE = 1.8*1.01325*bar # 1 atm = 1.01325 bar
    skip_every_n_frames = 1

    simulation_timestep = 0.002 * picoseconds

    parser = argparse.ArgumentParser(description="Sim with custom NN force")
    parser.add_argument('--model_folder' , help='folder containing the saved model')
    parser.add_argument('--mode' , help='Run in classical force field mode or NN predictor mode')
    parser.add_argument('--save_traj_at_end' , type=str2bool, help='True if whole traj will be saved at the end from RAM. False if traj is saved every step. Set False if steps is a high value coz whole traj wont fit in RAM.')
    parser.add_argument('--write_pos_only' , type=str2bool, help='If True, just records the pos trajectory. If False, records pos , vel and frc traj.')
    parser.add_argument('--also_save_npz' , type=str2bool, help='If True, also saves all traj files as npz at the end(converts the .xyz files to npz). If false, just xyz files are saved. xyz files are still saved. note that npz      files are written at the end, so they wont be written partially if the simulation crashes. xyz files will have the partial trajectory.')

    args = parser.parse_args()
    model_folder = args.model_folder
    mode = args.mode
    save_traj_at_end = args.save_traj_at_end
    write_pos_only = args.write_pos_only
    also_save_npz = args.also_save_npz

    if not save_traj_at_end:
        clear_traj_every_n_frames = 5000 # clear traj from RAM every clear_traj_every_n_frames steps
        assert clear_traj_every_n_frames % skip_every_n_frames == 0 # skip_every_n_frames must be a multiple of clear_traj_every_n_frames

    if mode not in ['classical', 'classical-and-predictor']:
        raise ValueError('Invalid mode supplied')

    # creating system
    forcefield = ForceField("empty_ff.xml")
    system = forcefield.createSystem(
                                     pdb.topology,
                                     removeCMMotion = True,
                                     )
    num_of_particles = system.getNumParticles()
    print(colored("Argon system initialized with "+str(num_of_particles)+" particles", 'green'))

    output_file_prefix = 'sim__' + str(num_of_particles) + '_atoms__' + str(steps) + '_timesteps__' + 'skip_every_n_frames_' + str(skip_every_n_frames) + '__' + str(TEMPERATURE._value) + '_kelvin__' + str(PRESSURE._value) + '_bar__'  + '_'.join(["{0:.4f}".format(box_size_val._value) for box_size_val in box_size]) + '_box_size_nanometers__' + mode + '_mode__' + now_time
    print("output_file_prefix = ", output_file_prefix)

    lj_cutoff = 0.80 * nanometer

    # Process pool initialization
    if mode in ['classical-and-predictor']:
        max_processes = get_max_processes(num_of_particles)
        if num_of_particles < 2000:  # a bad atoms-per-process ratio will slow down the computation because of communication overheads
            max_processes = 10
        process_pool = multiprocessing.Pool(max_processes)

    if mode in ['classical-and-predictor']:
        sys.path.append('./saved_weights/')
        #from force_predictor import NeuralNetwork
        from force_predictor_multiple_densities import NeuralNetwork
        nn_model = NeuralNetwork(dummy=True).load_all(model_folder)
        nn_model.predict_forces(np.random.random((num_of_particles, 50, 3)), process_pool) # Sample prediction for the model to "warm up". Without this, a BLAS error is thrown for some reason

        predictor_force = CustomExternalForce("- cx*x - cy*y - cz*z")
        predictor_force.addPerParticleParameter('cx')
        predictor_force.addPerParticleParameter('cy')
        predictor_force.addPerParticleParameter('cz')
        predictor_force.setForceGroup(1)

        params = [0, 0, 0]
        for i in range(num_of_particles):
            predictor_force.addParticle(i, params)

        system.addForce(predictor_force)
    
    # Adding the classical force field
    if mode in ['classical', 'classical-and-predictor']:
        argon_charge = 0.0 * elementary_charge
        cutoff_type = openmm.NonbondedForce.CutoffPeriodic

        mass = 39.9 * amu  # argon
        sigma = 3.4 * angstrom  # argon

        cutoff = 8.0 * angstrom
        
        epsilon = 0.238 * kilocalories_per_mole  # argon

        dispersion_correction = True
        nb = NonbondedForce()
        nb.setNonbondedMethod(cutoff_type)
        nb.setCutoffDistance(cutoff)
        print("Cutoff dist = ", nb.getCutoffDistance())
        nb.setUseDispersionCorrection(dispersion_correction)
        for i in range(num_of_particles):
            nb.addParticle(argon_charge, sigma, epsilon)
        system.addForce(nb)
        print("switching dist = ", nb.getSwitchingDistance())
        print("Nonbonded method = ", nb.getNonbondedMethod())

    # Add thermostat and barostat
    thermostat = AndersenThermostat(TEMPERATURE*kelvin, 1/picosecond)
    system.addForce(thermostat)

    #print(colored("WARNING: Barostat is on", 'magenta'))
    #barostat = MonteCarloBarostat(PRESSURE*bar, TEMPERATURE*kelvin)
    #system.addForce(barostat)
    
    # setting up integrator
    #integrator = LangevinIntegrator(TEMPERATURE, 1/picosecond, simulation_timestep) # no thermostat needed if using LangevinIntegrator
    #integrator = VerletIntegrator(simulation_timestep) # if using verlet, use AndersenThermostat
    #integrator = AndersenVelocityVerletIntegrator(TEMPERATURE, 1/picosecond, simulation_timestep)
    integrator = NoseHooverChainVelocityVerletIntegrator(system=system, temperature=TEMPERATURE, timestep=simulation_timestep)


    #creating simulation context
    simulation = Simulation(pdb.topology, system, integrator)
    simulation.context.setPositions(pdb.positions)

     # adding reporters
    reporter_step_ct = 1  # After how many timesteps each detail will be logged to reporter
    simulation.reporters.append(StateDataReporter(output_file_prefix+'_state_data.dat',
                                                    reporter_step_ct,
                                                    step=True,
                                                    potentialEnergy = True,
                                                    kineticEnergy = True,
                                                    totalEnergy = True,
                                                    temperature = True,
                                                    volume = True,
                                                    density = True,
                                                    separator='\t'))



    #minimizing system
    print("Executing simulation.minimizeEnergy()")
    simulation.minimizeEnergy() # do a classical minimization first
   
    traj = Trajectory(box_size=box_size, time_step=simulation_timestep)

    print_flag = False
    #running simulation with steps broken up
    bar = progressbar.ProgressBar(max_value=steps-1, 
                                    #redirect_stderr=True
                                    )
    pos_traj_filename = output_file_prefix+'_pos.xyz'
    vel_traj_filename = output_file_prefix+'_vel.xyz'
    frc_traj_filename = output_file_prefix+'_frc.xyz'
    
    # Running simulation step by step
    for step in range(steps):

        state = simulation.context.getState(getPositions=True,
                                            getForces=True,
                                            getVelocities=True,
                                            )

        traj.append_frame(state.getPositions(asNumpy=True))
        if not write_pos_only:
            traj.append_vel_frame(state.getVelocities(asNumpy=True))
            traj.append_frc_frame(state.getForces(asNumpy=True))

        if mode in ['classical-and-predictor']:
            envs = traj.get_envs_for_frame(step, 50, process_pool)
            envs = envs.in_units_of(angstroms)
            assert envs.unit == angstroms
            envs, _ = traj.separate_values_and_unit(envs)
            forces = nn_model.predict_forces(envs, process_pool) # these are in Hartree / Bohr
            forces = forces * (hartree / bohr)
            
            for i in range(num_of_particles):
                predictor_force.setParticleParameters(i, i, forces[i].in_units_of(kilojoule/nanometer)*AVOGADRO_CONSTANT_NA) # THESE FORCES GOING IN MUST BE IN kilojoule/(nm mole)

            predictor_force.updateParametersInContext(simulation.context)
        
        if print_flag:
            # just getting information for what is going on with the forces and velocities
            print(colored('Forces that are fed in in step '+ str(step + 1), 'magenta'))
            temp_forces = state.getForces(asNumpy=False)
            print(colored(temp_forces, 'magenta'))
        
        simulation.step(1)

        if print_flag:
            # getting information for what is going on with the positions
            print(colored('Positions given after step ' + str(step + 1), 'red'))
            temp_pos = state.getPositions(asNumpy=True)
            print(colored(temp_pos, 'red'))
        
        bar.update(step)

        if not save_traj_at_end:

            if step > 0 and (step) % clear_traj_every_n_frames == 0:
                append_traj(pos_traj_filename, traj.frames, 0, clear_traj_every_n_frames-1, skip_every_n_frames)
                if not write_pos_only:
                    append_traj(vel_traj_filename, traj.vel_frames, 0, clear_traj_every_n_frames-1, skip_every_n_frames)
                    append_traj(frc_traj_filename, traj.frc_frames, 0, clear_traj_every_n_frames-1, skip_every_n_frames)

                traj.delete_frames(0, clear_traj_every_n_frames-1)
                if not write_pos_only:
                    traj.delete_vel_frames(0, clear_traj_every_n_frames-1)
                    traj.delete_frc_frames(0, clear_traj_every_n_frames-1)

    if not save_traj_at_end:
        append_traj(pos_traj_filename, traj.frames, 0, traj.no_of_frames-1, skip_every_n_frames)
        if not write_pos_only:
            append_traj(vel_traj_filename, traj.vel_frames, 0, traj.vel_no_of_frames-1, skip_every_n_frames)
            append_traj(frc_traj_filename, traj.frc_frames, 0, traj.frc_no_of_frames-1, skip_every_n_frames)

    print()

    if mode in ['classical-and-predictor']:
        process_pool.close()

    if save_traj_at_end:
        if not write_pos_only:
            raise NotImplementedError("Not implemented for vel and frc trajectories")
        print("Writing to Trajectory file, skipping every ", skip_every_n_frames, " frames.")
        traj.save_as_xyz(pos_traj_filename, skip_every_n_frames=skip_every_n_frames)

    if also_save_npz:
        args_list_for_npz_write = [(pos_traj_filename, pdb_filename)]
        if not write_pos_only:
            args_list_for_npz_write.append((vel_traj_filename, pdb_filename))
            args_list_for_npz_write.append((frc_traj_filename, pdb_filename))
        npz_write_process_pool = multiprocessing.Pool(len(args_list_for_npz_write))
        npz_write_process_pool.starmap(convert_xyz_to_npz, args_list_for_npz_write)
        npz_write_process_pool.close()