def testSAE(self):
shifter = torchani.EnergyShifter(None)
torchani.data.load(dataset_path).subtract_self_energies(shifter)
true_self_energies = torch.tensor([
-19.354171758844188, -19.354171758844046, -54.712238523648587,
-75.162829556770987
],
dtype=torch.float64)
self.assertEqual(true_self_energies, shifter.self_energies)
def testShapeInference(self):
shifter = torchani.EnergyShifter(None)
ds = torchani.data.load(dataset_path).subtract_self_energies(shifter)
len(ds)
ds = ds.species_to_indices()
len(ds)
ds = ds.shuffle()
len(ds)
ds = ds.collate(batch_size)
len(ds)
def testDataloader(self):
shifter = torchani.EnergyShifter(None)
dataset = list(
torchani.data.load(dataset_path).subtract_self_energies(
shifter).species_to_indices().shuffle())
loader = torch.utils.data.DataLoader(
dataset,
batch_size=batch_size,
collate_fn=torchani.data.collate_fn,
num_workers=64)
for _ in loader:
pass
def testSAE(self):
tolerance = 1e-5
shifter = torchani.EnergyShifter(None)
torchani.data.load(dataset_path).subtract_self_energies(shifter)
true_self_energies = torch.tensor([
-19.354171758844188, -19.354171758844046, -54.712238523648587,
-75.162829556770987
],
dtype=torch.float64)
diff = torch.abs(true_self_energies - shifter.self_energies)
for e in diff:
self.assertLess(e, tolerance)
def __init__(self,
label="ase",
atoms=None,
coordinates='tmp_ase.pdb',
**kwargs):
self.label = label
coordinates = os.path.dirname(label) + "/" + coordinates
self.coordinates = coordinates
self.atoms = atoms
self.energy_free = None
self.forces = []
self.aev_computer = torchani.SortedAEV(const_file=const_file,
device=device)
self.nn = torchani.ModelOnAEV(self.aev_computer,
derivative=True,
from_nc=network_dir,
ensemble=8)
self.shift_energy = torchani.EnergyShifter(sae_file)
device=parser.device)
EtaA = torch.tensor([8.0000000e+00], device=parser.device)
ShfA = torch.tensor(
[9.0000000e-01, 1.5500000e+00, 2.2000000e+00, 2.8500000e+00],
device=parser.device)
num_species = 4
aev_computer = torchani.AEVComputer(Rcr, Rca, EtaR, ShfR, EtaA, Zeta, ShfA,
ShfZ, num_species)
nn = torchani.ANIModel([atomic() for _ in range(4)])
model = torch.nn.Sequential(aev_computer, nn).to(parser.device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.000001)
mse = torch.nn.MSELoss(reduction='none')
print('=> loading dataset...')
shifter = torchani.EnergyShifter(None)
dataset = list(
torchani.data.load(parser.dataset_path).subtract_self_energies(
shifter).species_to_indices().shuffle().collate(parser.batch_size))
print('=> start warming up')
total_batch_counter = 0
for epoch in range(0, WARM_UP_BATCHES + 1):
print('Epoch: %d/inf' % (epoch + 1, ))
progbar = pkbar.Kbar(target=len(dataset) - 1, width=8)
for i, properties in enumerate(dataset):
if not parser.dry_run and total_batch_counter == WARM_UP_BATCHES:
print('=> warm up finished, start profiling')
