def generate_subgraphs(atomsobject, unwrap=True, get_edge=False):
from nff.io.ase import AtomsBatch
atoms = AtomsBatch(atomsobject)
z, adj, dmat, threshold = adjdistmat(atoms, unwrap=unwrap)
box_len = torch.Tensor(np.diag(atoms.get_cell()))
G = nx.from_numpy_matrix(adj)
for i, item in enumerate(z):
G.nodes[i]['z'] = item
sub_graphs = nx.connected_component_subgraphs(G)
edge_list = []
partitions = []
for i, sg in enumerate(sub_graphs):
partitions.append(list(sg.nodes))
if get_edge:
edge_list.append(list(sg.edges))
if len(edge_list) != 0:
return partitions, edge_list
else:
return partitions
def neural_neb_ase(reactantxyzfile,
productxyzfile,
nff_dir,
rxn_name,
steps=500,
n_images=24,
fmax=0.004,
isclimb=False):
#reactant and products as ase Atoms
initial = AtomsBatch(xyz_to_ase_atoms(reactantxyzfile),
cutoff=5.5,
nbr_torch=True,
directed=True)
final = AtomsBatch(xyz_to_ase_atoms(productxyzfile),
cutoff=5.5,
nbr_torch=True,
directed=True)
# Make a band consisting of n_images:
images = [initial]
images += [initial.copy() for i in range(n_images)]
images += [final]
neb = SingleCalculatorNEB(images, k=0.02, climb=isclimb)
neb.method = 'improvedtangent'
# Interpolate linearly the potisions of the n_images:
neb.interpolate()
neb.idpp_interpolate(optimizer=BFGS, steps=steps)
images = read('idpp.traj@-{}:'.format(str(n_images + 2)))
# # Set calculators:
nff_ase = NeuralFF.from_file(nff_dir, device='cuda:0')
neb.set_calculators(nff_ase)
# # Optimize:
optimizer = BFGS(neb, trajectory='{}/{}.traj'.format(nff_dir, rxn_name))
optimizer.run(fmax=fmax, steps=steps)
# Read NEB images from File
images = read('{}/{}.traj@-{}:'.format(nff_dir, rxn_name,
str(n_images + 2)))
return images
def _get_periodic_neighbor_list(self, cutoff, undirected=False):
from nff.io.ase import AtomsBatch
nbrlist = []
offsets = []
for nxyz, lattice in zip(self.props['nxyz'], self.props['lattice']):
atoms = AtomsBatch(nxyz[:, 0].long(),
positions=nxyz[:, 1:],
cell=lattice,
pbc=True,
cutoff=cutoff)
nbrs, offs = atoms.update_nbr_list()
nbrlist.append(nbrs)
offsets.append(offs)
self.props['nbr_list'] = nbrlist
self.props['offsets'] = offsets
return
def ev_run(ev_atoms,
nff_dir,
maxstepsize,
maxstep,
convergence,
device,
method='Powell'):
rmslist = []
maxlist = []
nff = NeuralFF.from_file(nff_dir, device=device)
for step in range(maxstep):
if step == 0:
xyz, grad, hessian, h = eigvec_following(ev_atoms, step,
maxstepsize, device,
method)
else:
xyz, grad, hessian, h = eigvec_following(ev_atoms, step,
maxstepsize, device,
method, hessian, grad, h)
if step == 0:
xyz_all = xyz
else:
xyz_all = torch.cat((xyz_all, xyz), dim=0)
rmslist.append(grad.pow(2).sqrt().mean())
maxlist.append(grad.pow(2).sqrt().max())
print("RMS: {}, MAX: {}".format(
grad.pow(2).sqrt().mean(),
grad.pow(2).sqrt().max()))
if grad.pow(2).sqrt().max() < convergence:
print("Optimization converged!")
break
ev_atoms = Atoms(numbers=ev_atoms.numbers,
positions=xyz.reshape(-1, 3).cpu().numpy(),
pbc=False)
ev_atoms = AtomsBatch(ev_atoms,
cutoff=5.5,
nbr_torch=True,
directed=True)
ev_atoms.set_calculator(nff)
return xyz, grad, xyz_all, rmslist, maxlist
def unwrap_xyz(self, mol_dic):
"""
Unwrap molecular coordinates by displacing atoms by box vectors
Args:
mol_dic (dict): dictionary of nodes of each disconnected subgraphs
"""
from nff.io.ase import AtomsBatch
for i in range(len(self.props['nxyz'])):
# makes atoms object
atoms = AtomsBatch(positions=self.props['nxyz'][i][:, 1:4],
numbers=self.props['nxyz'][i][:, 0],
cell=self.props["cell"][i],
pbc=True)
# recontruct coordinates based on subgraphs index
if self.props['smiles']:
mol_idx = mol_dic[self.props['smiles'][i]]
atoms.set_positions(reconstruct_atoms(atoms, mol_idx))
nxyz = atoms.get_nxyz()
self.props['nxyz'][i] = torch.Tensor(nxyz)
def gen_bond_prior(self, cutoff, bond_len_dict=None):
from nff.io.ase import AtomsBatch
from nff.io.ase import AtomsBatch
if not self.props:
raise TypeError("the dataset has no data yet")
bond_dict = {}
bond_count_dict = {}
mol_idx_dict = {}
#---------This part can be simplified---------#
for i in range(len(self.props['nxyz'])):
z = self.props['nxyz'][i][:, 0]
xyz = self.props['nxyz'][i][:, 1:4]
# generate arguments for ase Atoms boject
ase_param = {
"numbers":
z,
"positions":
xyz,
"pbc":
True,
"cell":
self.props['cell'][i] if 'cell' in self.props.keys() else None
}
atoms = Atoms(**ase_param)
sys_name = self.props['smiles'][i]
if sys_name not in bond_dict.keys():
print(sys_name)
i, j = neighbor_list("ij", atoms, DISTANCETHRESHOLDICT_Z)
bond_list = torch.LongTensor(np.stack((i, j), axis=1)).tolist()
bond_dict[sys_name] = bond_list
# generate molecular graph
# TODO: there is redundant code in generate_subgraphs
subgraph_index = generate_subgraphs(atoms)
mol_idx_dict[sys_name] = subgraph_index
# generate topologies
# TODO: include options to only generate bond topology
self.generate_topologies(bond_dic=bond_dict)
if 'cell' in self.props.keys():
self.unwrap_xyz(mol_idx_dict)
#---------This part can be simplified---------#
# generate bond length dictionary if not given
if not bond_len_dict:
bond_len_dict = self.gen_bond_stats()
# update bond len and offsets
all_bond_len = []
all_offsets = []
all_nbr_list = []
for i in range(len(self.props['nxyz'])):
z = self.props['nxyz'][i][:, 0]
xyz = self.props['nxyz'][i][:, 1:4]
bond_list = self.props['bonds'][i]
bond_type_list = torch.stack(
(z[bond_list[:, 0]], z[bond_list[:, 1]])).t()
bond_len_list = []
for bond in bond_type_list:
bond_type = tuple(torch.LongTensor(sorted(bond)).tolist())
bond_len_list.append(bond_len_dict[bond_type])
all_bond_len.append(torch.Tensor(bond_len_list).reshape(-1, 1))
# update offsets
ase_param = {
"numbers": z,
"positions": xyz,
"pbc": True,
"cutoff": cutoff,
"cell":
self.props['cell'][i] if 'cell' in self.props.keys() else None,
"nbr_torch": False
}
# the coordinates have been unwrapped and try to results offsets
atoms = AtomsBatch(**ase_param)
atoms.update_nbr_list()
all_offsets.append(atoms.offsets)
all_nbr_list.append(atoms.nbr_list)
# update
self.props['bond_len'] = all_bond_len
self.props['offsets'] = all_offsets
self.props['nbr_list'] = all_nbr_list
self._check_dictionary(deepcopy(self.props))