def get_mol_with_targets(record, entry) -> MolWithTargets: # offmol offmol = Molecule.from_qcschema(entry) # trajectory containing xyz, energies, and gradients trajectory = record.get_trajectory() # xyz molecules = [snapshot.get_molecule() for snapshot in trajectory] xyz = np.array([mol.geometry for mol in molecules]) # energies and gradients energies_and_gradients = list(map(get_energy_and_gradient, trajectory)) energies = np.array([e for (e, _) in energies_and_gradients]) gradients = np.array([g for (_, g) in energies_and_gradients]) return MolWithTargets(offmol, xyz, energies, gradients)
def get_graph(collection, record_name): # get record and trajectory record = collection.get_record(record_name, specification="default") entry = collection.get_entry(record_name) from openforcefield.topology import Molecule mol = Molecule.from_qcschema(entry) try: trajectory = record.get_trajectory() except: return None if trajectory is None: return None g = esp.Graph(mol) # energy is already hartree g.nodes['g'].data['u_ref'] = torch.tensor( [ Quantity( snapshot.properties.scf_total_energy, esp.units.HARTREE_PER_PARTICLE ).value_in_unit( esp.units.ENERGY_UNIT ) for snapshot in trajectory ], dtype=torch.get_default_dtype(), )[None, :] g.nodes['n1'].data['xyz'] = torch.tensor( np.stack( [ Quantity( snapshot.get_molecule().geometry, unit.bohr, ).value_in_unit( esp.units.DISTANCE_UNIT ) for snapshot in trajectory ], axis=1 ), requires_grad=True, dtype=torch.get_default_dtype(), ) g.nodes['n1'].data['u_ref_prime'] = torch.stack( [ torch.tensor( Quantity( snapshot.dict()['return_result'], esp.units.HARTREE_PER_PARTICLE / unit.bohr, ).value_in_unit( esp.units.FORCE_UNIT ), dtype=torch.get_default_dtype(), ) for snapshot in trajectory ], dim=1 ) return g