Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
reformated into SchNetPack input format.
"""
if output is None:
inputs = {}
else:
inputs = output
# Elemental composition
cell = np.array(atoms.cell.array, dtype=np.float32) # get cell array
inputs[Properties.Z] = torch.LongTensor(atoms.numbers.astype(np.int))
positions = atoms.positions.astype(np.float32)
if center_positions:
positions -= atoms.get_center_of_mass()
inputs[Properties.R] = torch.FloatTensor(positions)
inputs[Properties.cell] = torch.FloatTensor(cell)
# get atom environment
nbh_idx, offsets = environment_provider.get_environment(atoms)
# Get neighbors and neighbor mask
inputs[Properties.neighbors] = torch.LongTensor(nbh_idx.astype(np.int))
# Get cells
inputs[Properties.cell] = torch.FloatTensor(cell)
inputs[Properties.cell_offset] = torch.FloatTensor(offsets.astype(np.float32))
# If requested get neighbor lists for triples
if collect_triples:
nbh_idx_j, nbh_idx_k, offset_idx_j, offset_idx_k = collect_atom_triples(nbh_idx)
inputs[Properties.neighbor_pairs_j] = torch.LongTensor(nbh_idx_j.astype(np.int))
inputs[Properties.neighbor_pairs_k] = torch.LongTensor(nbh_idx_k.astype(np.int))
def __init__(
self,
model,
required_properties=[Properties.energy, Properties.forces],
force_handle=Properties.forces,
position_conversion=1.0 / MDUnits.angs2bohr,
force_conversion=1.0 / MDUnits.Ha2kcalpmol / MDUnits.angs2bohr,
property_conversion={},
detach=True,
):
super(SGDMLCalculator, self).__init__(
required_properties,
force_handle,
position_conversion,
force_conversion,
property_conversion,
detach,
)
self.model = model
from schnetpack.md.utils import MDUnits
from schnetpack.md.parsers.orca_parser import OrcaMainFileParser
from schnetpack import Properties
import torch
import os
import subprocess
from ase.data import chemical_symbols
from ase import Atoms
class OrcaCalculator(QMCalculator):
is_atomistic = [Properties.forces, Properties.shielding]
def __init__(
self,
required_properties,
force_handle,
compdir,
qm_executable,
orca_template,
orca_parser=OrcaMainFileParser,
position_conversion=1.0 / MDUnits.angs2bohr,
force_conversion=1.0,
property_conversion={},
queuer=None,
adaptive=False,
basename="input",
):
return self.formatters.format(self.parsed)
def reset(self):
"""
Reset state of parser
"""
self.read = False
self.parsed = None
class OrcaMainFileParser(OrcaOutputParser):
properties = [
"atoms",
Properties.forces,
Properties.energy,
Properties.dipole_moment,
Properties.polarizability,
Properties.shielding,
]
starts = {
"atoms": "CARTESIAN COORDINATES (ANGSTROEM)",
Properties.forces: "CARTESIAN GRADIENT",
Properties.energy: "FINAL SINGLE POINT ENERGY",
Properties.dipole_moment: "Total Dipole Moment",
Properties.polarizability: "The raw cartesian tensor (atomic units):",
Properties.shielding: "CHEMICAL SHIFTS",
}
stops = {
"atoms": "CARTESIAN COORDINATES (A.U.)",
Properties.forces: "Difference to translation invariance",
for i in index_list:
_i = i-n_tokens # atom index if ignoring tokens
while True:
partial_mol = mol.copy()
# get the type of the next atom
next_type = pred_types[_i+j]
# don't consider distance predictions at stop token prediction steps
if next_type == stop_token:
dist_mask = torch.zeros(i).float()
else:
dist_mask = torch.ones(i).float()
# always consider the type prediction
type_mask = torch.ones(i).float()
# assemble neighborhood mask and neighbor indices
neighbor_mask = torch.ones(i, i-1)
neighbors = partial_mol[Properties.neighbors][:i, :i-1]
# set position and labels of the current (focus) token (first atom)
cur = focus[_i + j]
pos = partial_mol[Properties.R][:i]
label_idx = i if i < n_atoms else 0
labels = partial_mol['_labels'][label_idx, :i]
pos = torch.cat((pos[cur:cur+1], pos[1:]), 0)
labels = torch.cat((labels[cur:cur+1], labels[1:]), 0)
partial_mol.update(
{Properties.R: pos,
Properties.Z: partial_mol[Properties.Z][:i],
'_labels': labels,
'_dist_mask': dist_mask,
'_type_mask': type_mask,
Properties.neighbor_mask: neighbor_mask,
Properties.neighbors: neighbors,
def build_batch(i):
amount = torch.sum(unfinished) # only get predictions for unfinished molecules
# build neighborhood and neighborhood mask
neighbors_i = neighbors[:i, :i-1].expand(amount, -1, -1).contiguous()
neighbor_mask = torch.ones_like(neighbors_i).float()
# set position of focus token (first entry of positions)
positions[unfinished, 0] = positions[unfinished, current_atoms[unfinished]]
# center positions on currently focused atom (for localized grid)
positions[unfinished, :i] -= \
positions[unfinished, current_atoms[unfinished]][:, None, :]
# build batch with data of the partial molecules
batch = {
Properties.R: positions[unfinished, :i],
Properties.Z: atom_numbers[unfinished, :i],
Properties.atom_mask: torch.zeros(amount, i, dtype=torch.float),
Properties.neighbors: neighbors_i,
Properties.neighbor_mask: neighbor_mask,
Properties.cell_offset: torch.zeros(amount, i, max(i-1, 1), n_dims),
Properties.cell: torch.zeros(amount, n_dims, n_dims),
'_next_types': atom_numbers[unfinished, i],
'_all_types': all_types.view(1, -1),
'_type_mask': torch.ones(amount, i, dtype=torch.float),
}
# put batch into torch variables and on gpu
batch = {
k: v.to(device)
for k, v in batch.items()
}
return batch
"""Compute atomic representations/embeddings.
Args:
inputs (dict of torch.Tensor): SchNetPack dictionary of input tensors.
Returns:
torch.Tensor: atom-wise representation.
list of torch.Tensor: intermediate atom-wise representations, if
return_intermediate=True was used.
"""
# get tensors from input dictionary
atomic_numbers = inputs[Properties.Z]
positions = inputs[Properties.R]
cell = inputs[Properties.cell]
cell_offset = inputs[Properties.cell_offset]
neighbors = inputs[Properties.neighbors]
neighbor_mask = inputs[Properties.neighbor_mask]
atom_mask = inputs[Properties.atom_mask]
# get atom embeddings for the input atomic numbers
x = self.embedding(atomic_numbers)
if False and self.charged_systems and Properties.charge in inputs.keys():
n_atoms = torch.sum(atom_mask, dim=1, keepdim=True)
charge = inputs[Properties.charge] / n_atoms # B
charge = charge[:, None] * self.charge # B x F
x = x + charge
# compute interatomic distance of every atom to its neighbors
r_ij = self.distances(
positions, neighbors, cell, cell_offset, neighbor_mask=neighbor_mask
idcs = np.where(d['valid'])[0]
if len(idcs) < 1:
i = end
continue
# shrink stats
idx_id = stat_heads.index('id')
idx_known = stat_heads.index('known')
new_stats = stats[:, start:end]
if 'new' in args.store and args.model_path is not None:
idcs = idcs[np.where(new_stats[idx_known, idcs] == 0)[0]]
new_stats = new_stats[:, idcs]
new_stats[idx_id] = np.arange(len(new_stats[idx_id])) # adjust ids
shrunk_stats = np.hstack((shrunk_stats, new_stats))
# shrink positions and atomic numbers
shrunk_res[key] = {Properties.R: d[Properties.R][idcs],
Properties.Z: d[Properties.Z][idcs]}
# store connectivity matrices if desired
if 'connectivity' in args.store:
shrunk_res[key].update(
{'connectivity': [d['connectivity'][k] for k in idcs]})
i = end
shrunk_res['stats'] = shrunk_stats
res = shrunk_res
# store results in new database
# get filename that is not yet taken for db
if os.path.isfile(target_db):
file_name, _ = os.path.splitext(target_db)
expand = 0
while True:
expand += 1