예제 #1
0
 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)
예제 #2
0
 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)
예제 #3
0
 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
예제 #4
0
 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)
예제 #5
0
 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')