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