#!/usr/bin/env python
# coding: utf-8
# In[ ]:
"""Performs processing of Molecular Dynamics Trajectories to extract cluster models centered on the
solute molecule and including all solvent molecules within a given range"""
# # Main Routine
# In[ ]:
[docs]class ClustersTask:
def __init__(self,**kwargs):
self.wrapper = None
self.script_settings = None
self.task_command = 'clusters'
args = self.make_parser().parse_args("")
for arg in vars(args):
setattr(self,arg,getattr(args,arg))
[docs] def run(self):
"""
Main routine for processing clusters and running them for excitations.
The steps involved are:
1. Load the trajectory file of MD snapshots, which for a given solvent solute pair
it expects to find at '{solute}_{solvent}_solv.traj' in the current directory
2. 'Carve' spheres out of the trajectory, that is to say:
a. Delete counterions (checking they are not within the sphere first)
b. Delete any whole solvent molecules for which no atoms of the solvent molecule
are within ``self.radius`` of any atom in the solute molecule.
c. Label solvent molecules with a tag if they are within ``self.kernel`` by the
criteria above.
d. Rotate and center the cluster in the simulation cell, using the two most-distant
pairs of carbon atoms in the solute to find a common 'plane' for the solute snapshots.
3. Calculate electronic excitations for each cluster. This can currently be done with the
ONETEP wrapper functions or with NWChem but other wrappers can be added.
*Arguments*
self: Argument list for the task, with attributes including:
``solute``, ``solvent``, ``radius``, ``output``, ``task_id``, ``counterions``, ``charges``
``kernel``, ``basis``, ``func``, ``boxsize``, ``impsolv``,``nroots``, ``cleanup``, ``continuation``
wrapper: class
Wrapper to use in the Clusters task
*Output*:
If ``self.task_id`` is None, excited state calculations for the whole trajectory.
If has a int value, then an excited state calculation for just that frame in the trajectory.
"""
from ase.io import Trajectory
import os
if self.kernel == 0.0:
self.kernel = None
if (self.output=='dat' or self.output=='onetep') and self.basis=='6-31g*':
# set default basis for ONETEP calculations:
# 10 a0 NGWFs, 800 eV cutoff, 6 eV cond NGWF energy range
calc_params = {'ngwf_rad': 10.0,
'cutoff': '800 eV',
'energy_range': '6 eV',
'target': 0,
'func': self.func}
if self.calc_params == {}:
self.calc_params = self
if (self.output=='nwo' or self.output=='nwchem' or self.output=='orca'):
# set default basis for NWChem calculations:
calc_params = {'basis': self.basis,
'func': self.func,
'disp': self.disp,
'target': None}
if self.calc_params == {}:
self.calc_params = calc_params
# implicit solvent name defaults to solvent name but should usually be specified separately
if self.impsolv is None:
self.impsolv = self.solvent
# Determine charge on molecule (TODO: spin)
charge = 0
spin = 0
if isinstance(self.charges,dict):
if self.solute in self.charges:
charge = self.charges[self.solute]
if isinstance(self.charges,int) or isinstance(self.charges,str):
charge = int(self.charges)
if True:
traj_carved, traj_carved_name = self.carve_spheres(self.solute,self.solvent,self.counterions,self.radius,
self.kernel,self.max_solvent_mols,self.boxsize,self.task_id)
else: # hack to deactivate carving and use existing trajectory
traj_carved = list(enumerate(Trajectory(f'{self.solute}_{self.solvent}_gs_A_carved.traj')))
traj_carved_name = None
writeonly = (self.output=="nw") or (self.output=="dat") or (self.output=="xyz")
print(f'{len(list(traj_carved))} clusters carved from full snapshots.')
if writeonly:
print(f'Writing in {self.output} format.')
else:
print(f'Calculating excitations using {self.output}.')
if self.output=="nwi" or self.output=="nwchem" or self.output=="orca":
from esteem.trajectories import recalculate_trajectory
seed = f'{self.solute}_{self.solvent}'
target = list(range(0,self.nroots+1))
traj_label = 'A'
trajname_suffix = self.output
input_target = 0
# If we are processing the whole trajectory, set the full range and no offset
if self.task_id is None:
input_traj_range = None
output_traj_offset = 0
input_suffix = 'carved'
# If we are processing just one frame, set range to 1 and offset to task_id
else:
input_traj_range = range(0,1)
output_traj_offset = self.task_id
input_suffix = f'carved_{self.task_id:04}'
# Now run through the trajectory, calculating singlepoint energy for each frame
recalculate_trajectory(seed,target,traj_label,trajname_suffix,input_target,input_suffix,
self.wrapper.singlepoint,calc_params=calc_params,
input_traj_range=input_traj_range,output_traj_offset=output_traj_offset,
solvent=self.impsolv)
elif self.output=="dat" or self.output=="onetep":
self.calc_all_excitations(self.solute,self.solvent,traj_carved,self.wrapper.excitations,charge,
self.calc_params,self.impsolv,self.nroots,writeonly,
self.continuation,self.cleanup)
elif self.output=="xyz":
self.write_traj_xyz(self.solute+"_"+self.solvent+"_solv",traj_carved)
elif self.output=="amber":
self.write_amber_minimised(self.solute,self.solvent,traj_carved)
else:
raise Exception('Error: Unrecognised output format. The value of self.output was: {}'.format(self.output))
# Delete any temporary trajectory files made during a task-parallel run
if self.task_id is not None and traj_carved_name is not None:
os.remove(traj_carved_name)
[docs] def carve_spheres(self,soluseed,solvseed,counterions=None,solvent_radius=0.0,kernel_radius=0.0,
nmol_solvent_targ=None,boxsize=None,task_id=None):
"""Carves out spheres from a periodic solvent model, centered on the solute"""
from ase.io import Trajectory, read
# Determine trajectory names
traj = Trajectory(f'{soluseed}_{solvseed}_solv.traj')
traj_carved_name = f'{soluseed}_{solvseed}_gs_A_carved.traj'
if task_id is not None:
if task_id >= len(traj) or task_id < 0:
raise Exception(f"Invalid task_id: {task_id} outside range 0:{len(traj)}")
traj = [traj[task_id]]
traj_carved_name = f'{soluseed}_{solvseed}_gs_A_carved_{task_id:04}.traj'
else:
task_id = 0
print(f'Writing carved snapshots to {traj_carved_name}')
traj_carved = Trajectory(traj_carved_name,'w')
# Find counterions, if set
counterion_atoms = Atoms()
if isinstance(counterions,dict):
if soluseed in counterions:
counterion_atoms = Atoms(counterions[soluseed])
if isinstance(counterions,str):
counterion_atoms = Atoms(counterions)
nat_solute = len(read(soluseed+".xyz"))
nat_solvent = len(read(solvseed+".xyz"))
nat_counterions = len(counterion_atoms)
nat_tot = len(traj[0])
nmol_solvent = int((nat_tot-nat_solute-nat_counterions)/nat_solvent)
for t in traj:
carve_sphere(t,solvent_radius,kernel_radius,nat_solute,nat_solvent,
nat_counterions,nmol_solvent,nmol_solvent_targ,boxsize)
# Write model to trajectory
traj_carved.write(t)
traj_carved.close()
traj_carved = Trajectory(traj_carved_name)
return list(enumerate(traj_carved,start=task_id)),traj_carved_name
[docs] def calc_all_excitations(self,soluseed,solvseed,traj_carved,excit_func,
charge=0,calc_params={},impsolv=None,nroots=1,
writeonly=False,continuation=False,cleanup=False):
"""Loop over trajectory frames and do an excited states calculation for each one"""
for i,t in traj_carved:
frameseed = f'{soluseed}_{solvseed}_solv{i:04}'
excitations, energies = excit_func(t,frameseed,nroots=nroots,
calc_params=calc_params,solvent=impsolv,charge=charge,
writeonly=writeonly,continuation=continuation,
cleanup=cleanup)
print(f'Excitations for {frameseed}: {excitations}')
[docs] def calc_all_vibrations(self,soluseed,solvseed,traj_carved,geom_opt_func,freq_func,
charge=0,calc_params={},impsolv=None,nroots=1,
writeonly=False,continuation=False,cleanup=False):
"""Loop over trajectory frames and do a vibrational frequency calculation for each one"""
for i,t in traj_carved:
frameseed = f'{soluseed}_{solvseed}_solv{i:04}'
from ase.constraints import FixAtoms
c = FixAtoms(indices=[atom.index for atom in t if atom.tag == 0])
t.set_constraint(c)
energy,force,positions = geom_opt_func(t,frameseed,
calc_params=calc_params,solvent=impsolv,charge=charge,
writeonly=writeonly,continuation=continuation,
cleanup=cleanup)
freq_func(t,frameseed,calc_params=calc_params,solvent=impsolv,charge=charge,
writeonly=writeonly,continuation=continuation,
cleanup=cleanup)
print(f'Vibrational Frequencies for {frameseed} completed')
def write_traj_xyz(self,seed,traj_carved):
for i,t in traj_carved:
frameseed = f'{seed}{i:04}'
write(frameseed+".xyz",t)
# Load routines from Amber.py
def write_amber_minimised(self,soluseed,solvseed,traj_carved):
from os import path, makedirs, getcwd, chdir
from shutil import copyfile
from ase.io import read, write
seed = f'{self.solute}_{self.solvent}'
target = None
traj_label = 'A'
trajname_suffix = self.output
input_target = 0
if self.task_id is None:
input_suffix = 'carved'
else:
input_suffix = f'carved_{self.task_id:04}'
# Create directory for output files and change to it
wdir = trajname_suffix
origdir = getcwd()
if not path.exists(wdir):
makedirs(wdir)
copyfile(f'{soluseed}.xyz',f'{wdir}/{soluseed}.xyz')
copyfile(f'{solvseed}.xyz',f'{wdir}/{solvseed}.xyz')
chdir(wdir)
nat_solute = len(read(f'{soluseed}.xyz'))
nat_solvent = len(read(f'{solvseed}.xyz'))
# Set large Ewald cutoff radius
self.wrapper.cut = 40
# Prepare Solute inputs
if not path.exists(f'{soluseed}.prmtop'):
print(f'Preparing Amber inputs for {soluseed} solute')
self.wrapper.prepare_input(soluseed,netcharge=0,offset=0)
else:
print(f'Using pre-existing Amber inputs for {soluseed} solute')
solute = read(f'{soluseed}.pdb')
e0_am_solu = self.wrapper.singlepoint(solute,soluseed)
print(f'Amber solute ground state energy: {e0_am_solu}')
# Prepare Solvent inputs
if not path.exists(f'{solvseed}.prmtop'):
print(f'Preparing Amber inputs for {solvseed} solvent')
if solvseed == soluseed:
offset = 0
else:
offset = 99
self.wrapper.prepare_input(solvseed,netcharge=0,offset=offset)
else:
print(f'Using pre-existing Amber inputs for {solvseed} solvent')
solvent = read(f'{solvseed}.pdb')
e0_am_solv = self.wrapper.singlepoint(solvent,solvseed)
print(f'Amber solvent ground state energy: {e0_am_solv}')
# Write parameter and topology file for this frame via Amber wrapper
print(f'Creating frame prmtop files for {soluseed} in {solvseed} clusters')
for i,t in traj_carved:
nat_tot = len(t)
nmol_solvent = int((nat_tot-nat_solute)/nat_solvent)
label = f'{seed}_gs_{traj_label}_{trajname_suffix}{i:04}'
write(f'{label}.pdb',t)
if True: # set to false to deactivate writing prmtop and use existing versions
self.wrapper.make_frame_prmtop(label,t,soluseed,solvseed,nat_solute,nat_solvent,nmol_solvent)
chdir(origdir) # go back to parent directory as recalculate_trajectory expects to start there
# Recalculate trajectory energies with Amber wrapper
from esteem.trajectories import recalculate_trajectory
calc_params = {}
recalculate_trajectory(seed,target,traj_label,trajname_suffix,input_target,input_suffix,
self.wrapper.singlepoint,calc_params=calc_params)
def make_parser(self):
# Parse command line values
main_help = ('Given a trajectory containing MD snapshots of a solvated \n'+
'molecule, carve out spheres of a given radius to produce \n'+
'inputs for explicit solvent calculations.')
epi_help = ('Note: Expects the input trajectory to be named in the format \n'+
'<solute>_<solvent>_solv.traj\n'+
'Writes output to trajectory <solute>_<solvent>_carved.traj \n'+
'and to <solute>_<solvent>_solvXXX.<ext> where XXX is the \n'+
'index of the snapshot and <ext> is the file extension of the \n'+
'input file for the code being used (specified by -o).')
from argparse import ArgumentParser, SUPPRESS
parser = ArgumentParser(description=main_help,epilog=epi_help)
parser.add_argument('--solute','-u',required=False,type=str,help='Solute name')
parser.add_argument('--solvent','-v',required=False,type=str,help='Solvent name')
parser.add_argument('--radius','-r',default=5.0,type=float,help='Maximum distance from any atom of a solvent molecule to atom of a solute molecule for it to be included in cluster')
parser.add_argument('--kernel','-k',default=0.0,type=float,help='Maximum distance from any atom of a solvent molecule to atom of a solute molecule for it to be included in tddft kernel')
parser.add_argument('--max_solvent_mols','-M',default=None,type=int,help='Maximum number of solvent molecules to include (chosen randomly)')
parser.add_argument('--disp','-d',default=None,type=str,help='Dispersion correction (set to True to activate)')
parser.add_argument('--boxsize','-B',default='50.0',type=float,help='Size of box in which electronic excitation calculation is performed')
parser.add_argument('--output','-o',default='nw',type=str,help='Format of output: takes values nw, dat, nwchem and onetep (former 2 write input files only, latter 2 perform calculation)')
parser.add_argument('--nroots','-n',default=5,type=int,help='Number of excitations to calculate')
parser.add_argument('--task_id','-t',default=None,type=int,help='Task ID of the current job, inherited from driver')
parser.add_argument('--cleanup','-c',default=False,action='store_true',help='Whether to delete all temporary files after the job')
parser.add_argument('--continuation','-N',default=False,action='store_true',help='Whether to continue from a previous run of this file')
parser.add_argument('--counterions','-C',default={},type=str,help='Counterion(s) to add, eg Na')
parser.add_argument('--charges','-Q',default={},nargs='?',type=dict,help='Charges on molecular species. Not for command-line use')
parser.add_argument('--md_suffix','-m',default="md",nargs='?',type=str,help=SUPPRESS)
parser.add_argument('--exc_suffix','-e',default="exc",nargs='?',type=str,help=SUPPRESS)
parser.add_argument('--calc_params','-Z',default={},nargs='?',type=str,help=SUPPRESS)
# Wrapper Dependent
parser.add_argument('--basis','-b',default='6-31g*',nargs='*',type=str,help='Basis string or tuple')
parser.add_argument('--func','-f',default='PBE',type=str,help='Functional for electronic excitation calculation')
parser.add_argument('--impsolv','-i',default=None,type=str,help='Implicit solvent string (may differ from solvent name)')
return parser
def validate_args(args):
default_args = make_parser().parse_args(['--solute','a','--solvent','b'])
for arg in vars(args):
if arg not in default_args:
raise Exception(f"Unrecognised argument '{arg}'")
# # Helper Functions to delete/label molecules and rotate clusters
# In[1]:
import numpy as np
from ase import Atoms
def carve_sphere(t,solvent_radius,kernel_radius,nat_solute,nat_solvent,
nat_counterions,nmol_solvent,nmol_solvent_targ,boxsize=None):
nat_tot = len(t)
# Delete counterions further away than solvent_radius from
# any solute atom
if (nat_counterions>0):
t = delete_counterions(t,solvent_radius,nat_tot,nat_solute,nat_counterions)
# Delete solvent molecules further away than solvent_radius from
# any solute atom
t = delete_distant_molecules(t,solvent_radius,nat_tot-nat_counterions,
nat_solvent,nmol_solvent,nat_solute)
if nmol_solvent_targ is not None:
t = delete_excess_molecules(t,nat_tot-nat_counterions,nat_solvent,nat_solute,
nmol_solvent_targ)
# Set a tag of 2 on the nearby solvent molecules within kernel_radius
# from a solute molecule. Also gives tag 1 to solute molecule atoms.
t = label_nearby_molecules(t,kernel_radius,nat_tot-nat_counterions,
nat_solvent,nmol_solvent,nat_solute)
# Rotates model so longest axes of solute are in common plane for
# ease of viewing, and adds a box
t = rotate_and_center_solute(t,nat_solute,boxsize)
[docs]def delete_counterions(t,solvent_radius,nat_tot,nat_solute,nat_counterions):
"""
Deletes counterions from an Atoms model.
Assumes they appear in the range [nat_solute:nat_solute+nat_counterions]
"""
Rij = np.zeros((nat_counterions,nat_solute))
# Get entries in matrix of distances between all solute atoms and counterions
for i in range(nat_solute):
Rij[:,i] = np.array(t.get_distances(i,range(nat_solute,nat_solute+nat_counterions),mic=True))
# Reshape to (number of counterions) x (atoms in solute)
Rij = Rij.reshape((nat_counterions,nat_solute))
# Turn into Boolean values of whether distance is less than threshold
flags = (Rij<solvent_radius)
# Reduce over solute atoms
flags = np.any(flags,axis=1)
# Delete all the atoms we don't want
del_idx = [i for i in range(nat_solute,nat_solute+nat_counterions) if flags[i-nat_solute]==0]
if len(del_idx)!=nat_counterions:
print('WARNING: counterion(s) not deleted! Too close to solute')
del t[del_idx]
return t
[docs]def delete_excess_molecules(t,nat_tot,nat_solvent,nat_solute,nmol_solvent_targ):
"""
Deletes solvent molecules
"""
import random
nat_tot_cluster = len(t)
nmol_solvent = int((nat_tot_cluster-nat_solute)/nat_solvent)
solvent_mols_to_keep = random.sample(range(nmol_solvent),nmol_solvent_targ)
del_idx = [i for i in range(nat_solute,nat_tot_cluster)
if (i-nat_solute)//nat_solvent not in solvent_mols_to_keep]
del t[del_idx]
return t
[docs]def delete_distant_molecules(t,solvent_radius,nat_tot,nat_solvent,nmol_solvent,nat_solute):
"""
Deletes solvent molecules beyond a certain radius from the solute from an Atoms model
Assumes that the first nat_solute atoms are the solute, and after that solvent
atoms are arranged in nmol_solvent groups of size nat_solvent
"""
Rij = np.zeros((nat_solvent*nmol_solvent,nat_solute))
# Get entries in matrix of distances between all solute and solvent atoms
for i in range(nat_solute):
Rij[:,i] = np.array(t.get_distances(i,range(nat_solute,nat_tot),mic=True))
# Reshape to (number of solvent molecules) x (atoms in each solvent molecule) x (atoms in solute)
Rij = Rij.reshape((nmol_solvent,nat_solvent,nat_solute))
# Turn into Boolean values of whether distance is less than threshold
flags = (Rij<solvent_radius)
# Reduce over solute atoms each solvent atom could be in range of
flags = np.any(flags,axis=2)
# Reduce over solvent atoms in the same molecule as a solvent atom which is in range of a solute atom
flags = np.any(flags,axis=1)
# Expand out so there is the same entry for each atom in each molecule
flags = np.repeat(flags,nat_solvent)
# Delete all the atoms we don't want
del_idx = [i for i in range(nat_solute,nat_tot) if flags[i-nat_solute]==0]
del t[del_idx]
return t
[docs]def label_nearby_molecules(t,kernel_radius,nat_tot,nat_solvent,nmol_solvent,nat_solute):
"""
Adds a tag to solvent molecules within a certain radius from the solute to an Atoms model
Assumes that the first nat_solute atoms are the solute, and after that solvent
atoms are arranged in nmol_solvent groups of size nat_solvent
"""
if kernel_radius is None:
return t
nat_tot_cluster = len(t)
nmol_solvent_cluster = int((nat_tot_cluster-nat_solute)/nat_solvent)
Rij = np.zeros((nat_solvent*nmol_solvent_cluster,nat_solute))
# Get entries in matrix of distances between all solute and solvent atoms
if nmol_solvent_cluster>0:
for i in range(nat_solute):
Rij[:,i] = np.array(t.get_distances(i,range(nat_solute,nat_tot_cluster),mic=True))
# Reshape to (number of solvent molecules) x (atoms in each solvent molecule) x (atoms in solute)
Rij = Rij.reshape((nmol_solvent_cluster,nat_solvent,nat_solute))
# Turn into Boolean values of whether distance is less than threshold
flags = (Rij<kernel_radius)
# Reduce over solute atoms each solvent atom could be in range of
flags = np.any(flags,axis=2)
# Reduce over solvent atoms in the same molecule as a solvent atom which is in range of a solute atom
flags = np.any(flags,axis=1)
# Expand out so there is the same entry for each atom in each molecule
flags = np.repeat(flags,nat_solvent)
# Set all tags to zero
tags = np.zeros(nat_tot_cluster)
# Tag the solute molecule as "1"
tags[0:nat_solute] = 1
# Tag the rest as "2"
tags[nat_solute:nat_tot_cluster] = 2
# Now find those in range kernel_radius and tag them as "0"
flag_idx = [i for i in range(nat_solute,nat_tot_cluster) if flags[i-nat_solute]==0]
for i in flag_idx:
tags[i] = 0
t.set_tags(tags)
return t
[docs]def rotate_and_center_solute(t,nat_solute=None,boxsize=None):
"""Rotates a cluster model so the solute is centered and lies in the xy plane"""
# Extract solute atoms as separate atoms object
if nat_solute is None:
nat_solute = len(t)
solute = t[0:nat_solute]
# any molecule of less than 10 atoms is not rotated
rotate_threshold = 10
if nat_solute>rotate_threshold:
# Choose atoms along long and short axes
d=solute.get_all_distances()
for i in range(len(solute)):
for j in range(len(solute)):
if not (solute.get_chemical_symbols()[i]=='C' and
solute.get_chemical_symbols()[j]=='C'):
d[i,j]=0
# Find furthest-apart pair of carbon atoms, these define '1' and '4' atoms
ind = np.unravel_index(np.argmax(d, axis=None), d.shape)
p1=ind[0]; p4=ind[1]
# Find next-furthest-apart pair of carbon atoms, defining '2' and '3' atoms
d[p1,p4]=0; d[p4,p1]=0
ind2 = np.unravel_index(np.argmax(d, axis=None), d.shape)
# Ensure consistent orientation by choosing lowest overall index as '1'
# and set others based on ensuring |p1-p2| < |p1-p3|
p1 = min(min(ind,ind2)); p4 = max(min(ind,ind2));
p2 = min(max(ind,ind2)); p3 = max(max(ind,ind2));
if solute.get_distance(p1,p2)<solute.get_distance(p1,p3):
p3 = min(max(ind,ind2)); p2 = max(max(ind,ind2))
# Rotate so average of p1-p2 and p3-p4 points along x
vector1 = (t.positions[p1] - t.positions[p2] +
t.positions[p3] - t.positions[p4])
t.rotate(vector1,(1,0,0))
# Rotate so average of p1-p3 and p2-p4 points along y
vector2 = (t.positions[p3] - t.positions[p1] +
t.positions[p4] - t.positions[p2])
t.rotate(vector2,(0,1,0))
# Rotate so cross product points along z
vector1 = (t.positions[p1] - t.positions[p2] +
t.positions[p3] - t.positions[p4])
vector2 = (t.positions[p3] - t.positions[p1] +
t.positions[p4] - t.positions[p2])
t.rotate(np.cross(vector1,vector2),(0,0,1))
# Translate so center is at middle of box and set cell
if boxsize is not None:
# use average of the positions of the carbon atoms defining the longest
# lengths in the molecule to define the "center"
if nat_solute>rotate_threshold:
vector = ((boxsize*0.5,boxsize*0.5,boxsize*0.5) -
(t.positions[p1]+t.positions[p2]+t.positions[p3]+t.positions[p4])/4)
else:
vector = (boxsize*0.5,boxsize*0.5,boxsize*0.5) - solute.get_center_of_mass()
t.translate(vector)
t.set_cell(boxsize * np.identity(3))
else:
# center molecule on origin
if nat_solute>rotate_threshold:
vector = -(t.positions[p1]+t.positions[p2]+t.positions[p3]+t.positions[p4])/4
else:
vector = -solute.get_center_of_mass()
t.translate(vector)
t.pbc = False
t.set_cell(None)
return t
def dist_test(mol):
all_dist = mol.get_all_distances()
if (all_dist>5).any() or np.logical_and((all_dist<0.82),(all_dist>0)).any():
print(all_dist)
return True
else:
return False
def debug_cluster_trajectory(traj,nat_solu,nat_solv,start=0,end=-1,
before=20,after=20,badness_test_func=None,radius=None):
found_one = False
recent_frames = []
for j,a in enumerate(traj[start:end]):
print(j)
nat = len(a)
solu = a[0:nat_solu]
nmol = int((nat - nat_solu)/nat_solv)
for i in range(nmol):
solv = a[nat_solu + i*nat_solv:nat_solu+(i+1)*nat_solv]
if badness_test_func(solv):
print(f'Bad Molecule found at frame {j} molecule {i}')
recent_frames = []
for p in range(-before,after):
if radius is not None:
t = mdl[start+j+p].copy()
solu = t[0:nat_solu].copy()
solv = t[nat_solu + i*nat_solv:nat_solu+(i+1)*nat_solv].copy()
del t[list(range(0,nat_solu))]
del t[list(range(i*nat_solv,(i+1)*nat_solv))]
t = solv + t
delete_distant_molecules(t,4.0,nat-nat_solu,nat_solv,nmol-1,nat_solv)
t = t + solu
recent_frames.append(t)
else:
recent_frames.append(traj[start+j+p][nat_solu + i*nat_solv:nat_solu+(i+1)*nat_solv])
#recent_frames.append(traj[start+j+p])
found_one = True
break
if found_one:
break
return recent_frames
# # Handle inputs
# In[ ]:
if __name__ == '__main__':
from esteem.wrappers import nwchem
clus = ClustersTask()
args = clus.make_parser().parse_args()
clus.wrapper = nwchem.NWChemWrapper()
clus.run()