def saveAEVC(AEVC: torchani.AEVComputer, n_species: int, path) -> None:
"""
Save AEVComputer.
Parameters
----------
AEVC: torchani.AEVComputer
AEVComputer
n_species: int
Number of species
path:
Save path
"""
Rcr, EtaR, ShfR, Rca, ShfZ, EtaA, Zeta, ShfA = AEVC.constants()
torch.save(
{
"args": {
"Rcr": Rcr,
"EtaR": EtaR,
"ShfR": ShfR,
"Rca": Rca,
"ShfZ": ShfZ,
"EtaA": EtaA,
"Zeta": Zeta,
"ShfA": ShfA,
"num_species": n_species,
},
"state_dict": AEVC.state_dict(),
},
path,
)
mlflow.log_artifact(path)
from torchani.neurochem import load_model
from torchani.ase import Calculator
from torchani import AEVComputer
from torchani.neurochem import load_sae
from torchani.nn import Sequential
import os
# A module defining a module needs to define only one variable,
# named `calculator`, which should be an instance of the ase.calculator.Calculator,
# a subclass of this, or a compatible class implementing the calculator interface.
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model_dir = os.path.dirname(os.path.realpath(__file__))
species_order = ['H', 'C', 'N', 'O']
const_file = os.path.join(model_dir,
'../ANI_common/rHCNO-5.2R_16-3.5A_a4-8.params')
consts = torchani.neurochem.Constants(const_file)
aev_computer = AEVComputer(**consts)
energy_shifter = load_sae(
os.path.join(model_dir, '../ANI_common/sae_linfit.dat'))
nn = load_model(consts.species, os.path.join(model_dir, '../ANI_common'))
nn.load_state_dict(
torch.load(os.path.join(model_dir, 'model1200EF_pre_1_best.pt')))
model_trained = Sequential(aev_computer, nn, energy_shifter).to(device)
calculator = Calculator(species=species_order, model=model_trained)
name = 'ANI_1200K'
def __init__(self, *args, **kwargs):
super(BFSGraphDataset, self).__init__(*args, **kwargs)
self.random_order = kwargs["random_order"]
self.max_prev_nodes = kwargs["max_prev_nodes"]
self.num_edge_classes = kwargs
self.max_num_nodes = max(self.num_atoms_all)
assert self.max_num_nodes == self.restrict_max_atoms or \
self.restrict_max_atoms < 0, \
"restrict_max_atoms number is too high: " + \
"maximum number of nodes in molecules is {:d}".format(
self.max_num_nodes
)
original_start_node_label = kwargs.get("original_start_node_label",
None)
if "node_relabel_map" not in kwargs:
# define relabelling from Periodic Table numbers to {0, 1, ...}
unique_labels = set()
for index in range(len(self)):
sample = super(BFSGraphDataset, self).__getitem__(index)
node_feature_matrix = sample['node_feature_matrix']
adj_matrix = sample['adj_matrix']
labels = set(node_feature_matrix.flatten().tolist())
unique_labels.update(labels)
# discard 0 padding
unique_labels.discard(0)
self.node_relabel_map = {
v: i
for i, v in enumerate(sorted(unique_labels))
}
else:
self.node_relabel_map = kwargs["node_relabel_map"]
self.inverse_node_relabel_map = {
i: v
for v, i in self.node_relabel_map.items()
}
if original_start_node_label is not None:
self.start_node_label = \
self.node_relabel_map[original_start_node_label]
else:
self.start_node_label = None
if "edge_relabel_map" not in kwargs:
raise NotImplementedError()
else:
self.edge_relabel_map = kwargs["edge_relabel_map"]
self.inverse_edge_relabel_map = {
i: v
for v, i in sorted(self.edge_relabel_map.items(), reverse=True)
}
self.num_node_classes = len(self.inverse_node_relabel_map)
self.num_edge_classes = len(self.inverse_edge_relabel_map)
if self.has_3D:
self.const_file = kwargs["const_file"]
consts = torchani.neurochem.Constants(self.const_file)
self.aev_computer = AEVComputer(**consts)
self.species_converter = SpeciesConverter(consts.species)
from torchani import ANIModel, AEVComputer
from torchani import nn
from torchani import neurochem
import os
from typing import NamedTuple, Tuple, Optional
local = os.path.dirname(os.path.realpath(__file__))
model_path = os.path.join(local, r"ani-model-zoo-master\resources\ani-2x_8x")
const = neurochem.Constants(
os.path.join(model_path, 'rHCNOSFCl-5.1R_16-3.5A_a8-4.params'))
aev_computer = AEVComputer(Rcr=const.Rcr,
Rca=const.Rca,
EtaR=const.EtaR,
ShfR=const.ShfR,
EtaA=const.EtaA,
Zeta=const.Zeta,
ShfA=const.ShfA,
ShfZ=const.ShfZ,
num_species=const.num_species)
EShifter = neurochem.load_sae(os.path.join(model_path, 'sae_linfit.dat'))
ANI2_Network = neurochem.load_model(
const.species, os.path.join(model_path, r'train7\networks'))
class SpeciesEnergies(NamedTuple):
species: Tensor
energies: Tensor