Python load_sae Examples

Example #1

File: torchani_update.py Project: ZiadFakhoury/TorchANI-Update

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


def overforward(self,
                species_aev: Tuple[Tensor, Tensor],
                cell: Optional[Tensor] = None,
                pbc: Optional[Tensor] = None) -> SpeciesEnergies:
    species, aev = species_aev

Example #2

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'

Example #3

File: model.py Project: libAtoms/testing-framework

import torchani
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.

model_dir = os.path.dirname(os.path.realpath(__file__))

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

species_order = ['H', 'C', 'N', 'O']
const_file = os.path.join(model_dir,'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,'sae_linfit.dat'))


nn = load_model(consts.species, model_dir)
model_trained = Sequential(aev_computer, nn, energy_shifter).to(device)

calculator = Calculator(species=species_order, model=model_trained)

name = 'ANI'