def test_neutromeratio_energy_calculations_with_ANI_in_vacuum():
from ..tautomers import Tautomer
from ..constants import kT
from ..analysis import setup_alchemical_system_and_energy_function
import numpy as np
from ..ani import AlchemicalANI1ccx, ANI1ccx, ANI2x
# read in exp_results.pickle
with open("data/test_data/exp_results.pickle", "rb") as f:
exp_results = pickle.load(f)
# vacuum system
name = "molDWRow_298"
energy_function, tautomer, flipped = setup_alchemical_system_and_energy_function(
name=name, ANImodel=AlchemicalANI1ccx, env="vacuum")
# read in pregenerated traj
traj_path = (
"data/test_data/vacuum/molDWRow_298/molDWRow_298_lambda_0.0000_in_vacuum.dcd"
)
top_path = "data/test_data/vacuum/molDWRow_298/molDWRow_298.pdb"
traj, _ = _get_traj(traj_path, top_path,
[tautomer.hybrid_hydrogen_idx_at_lambda_0])
coordinates = [x.xyz[0] for x in traj[:10]] * unit.nanometer
# test energies with neutromeratio ANI objects
# first with ANI1ccx
model = neutromeratio.ani.ANI1ccx()
torch.set_num_threads(1)
energy_function = neutromeratio.ANI_force_and_energy(
model=model, atoms=tautomer.initial_state_ligand_atoms, mol=None)
energy = energy_function.calculate_energy(coordinates)
assert is_quantity_close(
energy.energy[0].in_units_of(unit.kilojoule_per_mole),
(-906555.29945346 * unit.kilojoule_per_mole),
rtol=1e-5,
)
# with ANI2x
model = neutromeratio.ani.ANI2x()
torch.set_num_threads(1)
energy_function = neutromeratio.ANI_force_and_energy(
model=model, atoms=tautomer.initial_state_ligand_atoms, mol=None)
energy = energy_function.calculate_energy(coordinates)
assert is_quantity_close(
energy.energy[0].in_units_of(unit.kilojoule_per_mole),
(-907243.8987177598 * unit.kilojoule_per_mole),
rtol=1e-5,
)
def test_restraint_with_AlchemicalANI1ccx():
import numpy as np
from ..constants import kT
from ..ani import AlchemicalANI1ccx
from ..analysis import setup_alchemical_system_and_energy_function
# generate tautomer object
name = "molDWRow_298"
energy_function, tautomer, flipped = setup_alchemical_system_and_energy_function(
name=name, env="vacuum", ANImodel=AlchemicalANI1ccx)
# read in pregenerated traj
traj_path = (
"data/test_data/vacuum/molDWRow_298/molDWRow_298_lambda_0.0000_in_vacuum.dcd"
)
top_path = "data/test_data/vacuum/molDWRow_298/molDWRow_298.pdb"
traj, _ = _get_traj(traj_path, top_path)
x0 = [x.xyz[0] for x in traj[0]] * unit.nanometer
energy = energy_function.calculate_energy(x0, lambda_value=0.0)
e_ = energy.energy.value_in_unit(unit.kilojoule_per_mole)
print(e_)
print(energy)
assert is_quantity_close(
energy.energy.in_units_of(unit.kilojoule_per_mole),
(-906912.01647632 * unit.kilojoule_per_mole),
rtol=1e-9,
)
def test_break_up_AlchemicalANI2x_timings():
from ..constants import kT
from ..analysis import setup_alchemical_system_and_energy_function
from ..ani import AlchemicalANI2x, CompartimentedAlchemicalANI2x
# vacuum system
name = "molDWRow_298"
# read in pregenerated traj
traj_path = (
"data/test_data/vacuum/molDWRow_298/molDWRow_298_lambda_0.0000_in_vacuum.dcd"
)
top_path = "data/test_data/vacuum/molDWRow_298/molDWRow_298.pdb"
# test energies with neutromeratio AlchemicalANI objects
# with ANI2x
traj, _ = _get_traj(traj_path, top_path)
coordinates = [x.xyz[0] for x in traj[:100]] * unit.nanometer
energy_function, _, _ = setup_alchemical_system_and_energy_function(
name=name, ANImodel=AlchemicalANI2x, env="vacuum"
)
AlchemicalANI2x_energy_lambda_1 = energy_function.calculate_energy(
coordinates, lambda_value=1.0
)
energy_function, tautomer, flipped = setup_alchemical_system_and_energy_function(
name=name, ANImodel=CompartimentedAlchemicalANI2x, env="vacuum"
)
CompartimentedAlchemicalANI2x_1 = energy_function.calculate_energy(
coordinates, lambda_value=1.0
)
for e1, e2 in zip(
AlchemicalANI2x_energy_lambda_1.energy, CompartimentedAlchemicalANI2x_1.energy
):
assert is_quantity_close(
e1.in_units_of(unit.kilojoule_per_mole),
e2.in_units_of(unit.kilojoule_per_mole),
rtol=1e-5,
)
def test_neutromeratio_energy_calculations_with_AlchemicalANI1ccx():
from ..tautomers import Tautomer
import numpy as np
from ..constants import kT
from ..analysis import setup_alchemical_system_and_energy_function
from ..ani import AlchemicalANI1ccx, AlchemicalANI2x, ANI1ccx
# read in exp_results.pickle
with open("data/test_data/exp_results.pickle", "rb") as f:
exp_results = pickle.load(f)
######################################################################
# vacuum
######################################################################
name = "molDWRow_298"
energy_function, tautomer, flipped = setup_alchemical_system_and_energy_function(
name=name, env="vacuum", ANImodel=AlchemicalANI1ccx)
# read in pregenerated traj
traj_path = (
"data/test_data/vacuum/molDWRow_298/molDWRow_298_lambda_0.0000_in_vacuum.dcd"
)
top_path = "data/test_data/vacuum/molDWRow_298/molDWRow_298.pdb"
traj, _ = _get_traj(traj_path, top_path)
x0 = [x.xyz[0] for x in traj[:10]] * unit.nanometer
energy_1 = energy_function.calculate_energy(x0, lambda_value=1.0)
for e1, e2 in zip(
energy_1.energy,
[
-906555.29945346,
-905750.20471091,
-906317.24952004,
-906545.17543265,
-906581.65215098,
-906618.2832786,
-906565.05631782,
-905981.82167316,
-904681.20632002,
-904296.8214631,
] * unit.kilojoule_per_mole,
):
assert is_quantity_close(e1, e2, rtol=1e-2)
base_path=f"../data/test_data/{env}/{name}",
)
with profiler.profile(record_shapes=True, profile_memory=True) as prof:
with profiler.record_function("model_inference"):
CompartimentedAlchemicalANI2x_1 = energy_function.calculate_energy(
coordinates, lambda_value=1.0, include_restraint_energy_contribution=False
)
s = prof.self_cpu_time_total / 1000000
print(f"First time to calculate energies with CompartimentedAlchemicalANI2x: {s:.3f} s")
with profiler.profile(record_shapes=True, profile_memory=True) as prof:
with profiler.record_function("model_inference"):
CompartimentedAlchemicalANI2x_1 = energy_function.calculate_energy(
coordinates, lambda_value=1.0, include_restraint_energy_contribution=False
)
s = prof.self_cpu_time_total / 1000000
print(
f"Second time to calculate energies with CompartimentedAlchemicalANI2x: {s:.3f} s"
)
for e1, e2 in zip(
AlchemicalANI2x_energy_lambda_1.energy, CompartimentedAlchemicalANI2x_1.energy
):
assert is_quantity_close(
e1.in_units_of(unit.kilojoule_per_mole),
e2.in_units_of(unit.kilojoule_per_mole),
rtol=1e-5,
)
def test_generating_droplet():
from ..analysis import setup_alchemical_system_and_energy_function
from ..utils import generate_tautomer_class_stereobond_aware
from ..ani import AlchemicalANI1ccx
import numpy as np
exp_results = pickle.load(open("data/test_data/exp_results.pickle", "rb"))
name = "molDWRow_298"
diameter = 10
t1_smiles = exp_results[name]["t1-smiles"]
t2_smiles = exp_results[name]["t2-smiles"]
# generate both rdkit mol
t_type, tautomers, flipped = generate_tautomer_class_stereobond_aware(
name, t1_smiles, t2_smiles, nr_of_conformations=5
)
tautomer = tautomers[0]
tautomer.perform_tautomer_transformation()
m = tautomer.add_droplet(
tautomer.hybrid_topology,
tautomer.get_hybrid_coordinates(),
diameter=diameter * unit.angstrom,
restrain_hydrogen_bonds=True,
restrain_hydrogen_angles=False,
top_file=f"data/test_data/{name}_in_droplet.pdb",
)
# define the alchemical atoms
alchemical_atoms = [
tautomer.hybrid_hydrogen_idx_at_lambda_1,
tautomer.hybrid_hydrogen_idx_at_lambda_0,
]
# extract hydrogen donor idx and hydrogen idx for from_mol
model = AlchemicalANI1ccx(alchemical_atoms=alchemical_atoms)
model = model.to(device)
# perform initial sampling
energy_function = neutromeratio.ANI_force_and_energy(
model=model,
atoms=tautomer.ligand_in_water_atoms,
mol=None,
)
for r in tautomer.ligand_restraints:
energy_function.add_restraint_to_lambda_protocol(r)
for r in tautomer.hybrid_ligand_restraints:
energy_function.add_restraint_to_lambda_protocol(r)
x0 = tautomer.get_ligand_in_water_coordinates()
energy = energy_function.calculate_energy(x0)
print(energy.energy[0])
assert is_quantity_close(
energy.energy[0], (-15146778.81228019 * unit.kilojoule_per_mole), rtol=1e-5
)
tautomer.add_COM_for_hybrid_ligand(
np.array([diameter / 2, diameter / 2, diameter / 2]) * unit.angstrom
)
for r in tautomer.solvent_restraints:
energy_function.add_restraint_to_lambda_protocol(r)
for r in tautomer.com_restraints:
energy_function.add_restraint_to_lambda_protocol(r)
energy = energy_function.calculate_energy(x0)
assert is_quantity_close(
energy.energy[0], (-15018039.455067404 * unit.kilojoule_per_mole), rtol=1e-5
)
energy_function, tautomer, _ = setup_alchemical_system_and_energy_function(
name=name,
env="droplet",
ANImodel=AlchemicalANI1ccx,
base_path="data/test_data/",
diameter=diameter,
)
energy = energy_function.calculate_energy(x0)
assert is_quantity_close(
energy.energy[0], (-15018040.86806798 * unit.kilojoule_per_mole), rtol=1e-7
)
del model
def test_break_up_AlchemicalANI2x_energies():
from ..constants import kT
from ..analysis import setup_alchemical_system_and_energy_function
from ..ani import AlchemicalANI2x, CompartimentedAlchemicalANI2x
# vacuum system
name = "molDWRow_298"
# read in pregenerated traj
traj_path = (
"data/test_data/vacuum/molDWRow_298/molDWRow_298_lambda_0.0000_in_vacuum.dcd"
)
top_path = "data/test_data/vacuum/molDWRow_298/molDWRow_298.pdb"
# test energies with neutromeratio AlchemicalANI objects
# with ANI2x
traj, _ = _get_traj(traj_path, top_path)
coordinates = [x.xyz[0] for x in traj[:10]] * unit.nanometer
energy_function, _, _ = setup_alchemical_system_and_energy_function(
name=name, ANImodel=AlchemicalANI2x, env="vacuum"
)
AlchemicalANI2x_energy_lambda_0 = energy_function.calculate_energy(
coordinates, lambda_value=0.0
)
AlchemicalANI2x_energy_lambda_1 = energy_function.calculate_energy(
coordinates, lambda_value=1.0
)
# making sure that the AlchemicalANI and the ANI results are the same on the endstates (with removed dummy atoms)
assert is_quantity_close(
AlchemicalANI2x_energy_lambda_1.energy[0].in_units_of(unit.kilojoule_per_mole),
(-907243.8987177598 * unit.kilojoule_per_mole),
rtol=1e-5,
)
energy_function, _, _ = setup_alchemical_system_and_energy_function(
name=name, ANImodel=CompartimentedAlchemicalANI2x, env="vacuum"
)
CompartimentedAlchemicalANI2x_1 = energy_function.calculate_energy(
coordinates, lambda_value=1.0
)
# making sure that the AlchemicalANI and the CompartimentedAlchemicalANI2x results are the same
assert is_quantity_close(
CompartimentedAlchemicalANI2x_1.energy[0].in_units_of(unit.kilojoule_per_mole),
(-907243.8987177598 * unit.kilojoule_per_mole),
rtol=1e-5,
)
CompartimentedAlchemicalANI2x_0 = energy_function.calculate_energy(
coordinates, lambda_value=0.0
)
print(CompartimentedAlchemicalANI2x_0.energy[0])
# making sure that the AlchemicalANI and the CompartimentedAlchemicalANI2x results are the same
assert is_quantity_close(
CompartimentedAlchemicalANI2x_0.energy[0].in_units_of(unit.kilojoule_per_mole),
AlchemicalANI2x_energy_lambda_0.energy[0].in_units_of(unit.kilojoule_per_mole),
rtol=1e-5,
)
def test_neutromeratio_energy_calculations_with_AlchemicalANI1ccx_in_droplet():
from ..analysis import setup_alchemical_system_and_energy_function
from ..ani import AlchemicalANI1ccx, ANI1ccx
# read in exp_results.pickle
with open("data/test_data/exp_results.pickle", "rb") as f:
exp_results = pickle.load(f)
######################################################################
# droplet
######################################################################
name = "molDWRow_298"
energy_function, tautomer, flipped = setup_alchemical_system_and_energy_function(
name=name,
ANImodel=AlchemicalANI1ccx,
env="droplet",
base_path="data/test_data/droplet/molDWRow_298/",
diameter=10,
)
# read in pregenerated traj
traj_path = (
"data/test_data/droplet/molDWRow_298/molDWRow_298_lambda_0.0000_in_droplet.dcd"
)
top_path = "data/test_data/droplet/molDWRow_298/molDWRow_298_in_droplet.pdb"
traj, _ = _get_traj(traj_path, top_path)
# overwrite the coordinates that rdkit generated with the first frame in the traj
x0 = [x.xyz[0] for x in traj[:10]] * unit.nanometer
energy_1 = energy_function.calculate_energy(x0, lambda_value=1.0)
energy_0 = energy_function.calculate_energy(x0, lambda_value=0.0)
assert len(tautomer.ligand_in_water_atoms) == len(x0[0])
print(energy_1.energy.in_units_of(unit.kilojoule_per_mole))
print(energy_1.energy)
for e1, e2 in zip(
energy_1.energy,
[
-3514015.25593752,
-3513765.22323459,
-3512655.23621176,
-3512175.08728504,
-3512434.17610022,
-3513923.68325093,
-3513997.76968092,
-3513949.58322023,
-3513957.74678051,
-3514045.84769267,
]
* unit.kilojoule_per_mole,
):
assert is_quantity_close(e1, e2, rtol=1e-5)
for e1, e2 in zip(
energy_0.energy,
[
-3514410.82062014,
-3514352.85161146,
-3514328.06891661,
-3514324.15896465,
-3514323.94519662,
-3514326.6575155,
-3514320.69798817,
-3514326.79413299,
-3514309.49535377,
-3514308.0598529,
]
* unit.kilojoule_per_mole,
):
assert is_quantity_close(e1, e2, rtol=1e-5)
######################################################################
# compare with ANI1ccx -- test1
# this test removes all restraints
name = "molDWRow_298"
energy_function, tautomer, flipped = setup_alchemical_system_and_energy_function(
name=name,
ANImodel=AlchemicalANI1ccx,
env="droplet",
base_path="data/test_data/droplet/molDWRow_298/",
diameter=10,
)
# remove restraints
energy_function.list_of_lambda_restraints = []
# overwrite the coordinates that rdkit generated with the first frame in the traj
x0 = [x.xyz[0] for x in traj[:10]] * unit.nanometer
energy_1 = energy_function.calculate_energy(x0, lambda_value=1.0)
energy_0 = energy_function.calculate_energy(x0, lambda_value=0.0)
model = ANI1ccx()
traj, _ = _get_traj(traj_path, top_path, [tautomer.hybrid_hydrogen_idx_at_lambda_1])
atoms = (
tautomer.ligand_in_water_atoms[: tautomer.hybrid_hydrogen_idx_at_lambda_1]
+ tautomer.ligand_in_water_atoms[tautomer.hybrid_hydrogen_idx_at_lambda_1 + 1 :]
)
energy_function = neutromeratio.ANI_force_and_energy(
model=model, atoms=atoms, mol=None
)
coordinates = [x.xyz[0] for x in traj[:10]] * unit.nanometer
assert len(atoms) == len(coordinates[0])
energies_ani1ccx_0 = energy_function.calculate_energy(coordinates)
assert is_quantity_close(
energy_0.energy[0].in_units_of(unit.kilojoule_per_mole),
energies_ani1ccx_0.energy[0].in_units_of(unit.kilojoule_per_mole),
)
######################################################################
# compare with ANI1ccx -- test2
# includes restraint energy
name = "molDWRow_298"
energy_function, tautomer, flipped = setup_alchemical_system_and_energy_function(
name=name,
ANImodel=AlchemicalANI1ccx,
env="droplet",
base_path="data/test_data/droplet/molDWRow_298/",
diameter=10,
)
# read in pregenerated traj
traj_path = (
"data/test_data/droplet/molDWRow_298/molDWRow_298_lambda_0.0000_in_droplet.dcd"
)
top_path = "data/test_data/droplet/molDWRow_298/molDWRow_298_in_droplet.pdb"
traj, _ = _get_traj(traj_path, top_path)
x0 = [x.xyz[0] for x in traj[:10]] * unit.nanometer
energy_1 = energy_function.calculate_energy(x0, lambda_value=1.0)
energy_0 = energy_function.calculate_energy(x0, lambda_value=0.0)
model = ANI1ccx()
traj, _ = _get_traj(traj_path, top_path, [tautomer.hybrid_hydrogen_idx_at_lambda_1])
atoms = (
tautomer.ligand_in_water_atoms[: tautomer.hybrid_hydrogen_idx_at_lambda_1]
+ tautomer.ligand_in_water_atoms[tautomer.hybrid_hydrogen_idx_at_lambda_1 + 1 :]
)
energy_function = neutromeratio.ANI_force_and_energy(
model=model, atoms=atoms, mol=None
)
coordinates = [x.xyz[0] for x in traj[:10]] * unit.nanometer
assert len(atoms) == len(coordinates[0])
energies_ani1ccx_0 = energy_function.calculate_energy(coordinates)
# get restraint energy
energy_0_restraint = energy_0.restraint_energy_contribution.in_units_of(
unit.kilojoule_per_mole
)
print(energy_0_restraint[0])
print(energy_0.energy[0])
print(energy_0)
# subtracting restraint energies
energy_0_minus_restraint = energy_0.energy[0] - energy_0_restraint[0]
assert is_quantity_close(
energy_0_minus_restraint,
energies_ani1ccx_0.energy[0].in_units_of(unit.kilojoule_per_mole),
)
def test_neutromeratio_energy_calculations_with_AlchemicalANI1ccx_in_vacuum():
from ..constants import kT
from ..analysis import setup_alchemical_system_and_energy_function
from ..ani import AlchemicalANI1ccx, ANI1ccx
######################################################################
# vacuum
######################################################################
name = "molDWRow_298"
energy_function, tautomer, _ = setup_alchemical_system_and_energy_function(
name=name, env="vacuum", ANImodel=AlchemicalANI1ccx
)
# read in pregenerated traj
traj_path = (
"data/test_data/vacuum/molDWRow_298/molDWRow_298_lambda_0.0000_in_vacuum.dcd"
)
top_path = "data/test_data/vacuum/molDWRow_298/molDWRow_298.pdb"
traj, _ = _get_traj(traj_path, top_path)
x0 = [x.xyz[0] for x in traj[:10]] * unit.nanometer
energy_1 = energy_function.calculate_energy(x0, lambda_value=1.0)
for e1, e2 in zip(
energy_1.energy,
[
-906555.29945346,
-905750.20471091,
-906317.24952004,
-906545.17543265,
-906581.65215098,
-906618.2832786,
-906565.05631782,
-905981.82167316,
-904681.20632002,
-904296.8214631,
]
* unit.kilojoule_per_mole,
):
assert is_quantity_close(e1, e2, rtol=1e-2)
energy_0 = energy_function.calculate_energy(x0, lambda_value=0.0)
assert is_quantity_close(
energy_0.energy[0].in_units_of(unit.kilojoule_per_mole),
(-906912.01647632 * unit.kilojoule_per_mole),
rtol=1e-9,
)
######################################################################
# compare with ANI1ccx
traj, _ = _get_traj(traj_path, top_path, [tautomer.hybrid_hydrogen_idx_at_lambda_1])
model = ANI1ccx()
energy_function = neutromeratio.ANI_force_and_energy(
model=model, atoms=tautomer.final_state_ligand_atoms, mol=None
)
coordinates = [x.xyz[0] for x in traj[0]] * unit.nanometer
assert len(tautomer.initial_state_ligand_atoms) == len(coordinates[0])
assert is_quantity_close(
energy_0.energy[0], energy_function.calculate_energy(coordinates).energy
)
traj, _ = _get_traj(traj_path, top_path, [tautomer.hybrid_hydrogen_idx_at_lambda_0])
energy_function = neutromeratio.ANI_force_and_energy(
model=model, atoms=tautomer.initial_state_ligand_atoms, mol=None
)
coordinates = [x.xyz[0] for x in traj[0]] * unit.nanometer
assert len(tautomer.final_state_ligand_atoms) == len(coordinates[0])
assert is_quantity_close(
energy_1.energy[0], energy_function.calculate_energy(coordinates).energy
)
def test_neutromeratio_energy_calculations_with_AlchemicalANI2x_in_vacuum():
from ..constants import kT
from ..analysis import setup_alchemical_system_and_energy_function
from ..ani import AlchemicalANI2x
# vacuum system
name = "molDWRow_298"
energy_function, tautomer, _ = setup_alchemical_system_and_energy_function(
name=name, ANImodel=AlchemicalANI2x, env="vacuum"
)
# read in pregenerated traj
traj_path = (
"data/test_data/vacuum/molDWRow_298/molDWRow_298_lambda_0.0000_in_vacuum.dcd"
)
top_path = "data/test_data/vacuum/molDWRow_298/molDWRow_298.pdb"
# test energies with neutromeratio ANI objects
# with ANI2x
model = neutromeratio.ani.ANI2x()
# final state
traj, _ = _get_traj(traj_path, top_path, [tautomer.hybrid_hydrogen_idx_at_lambda_1])
coordinates = [x.xyz[0] for x in traj[:10]] * unit.nanometer
assert len(tautomer.initial_state_ligand_atoms) == len(coordinates[0])
energy_function = neutromeratio.ANI_force_and_energy(
model=model, atoms=tautomer.final_state_ligand_atoms, mol=None
)
ANI2x_energy_final_state = energy_function.calculate_energy(
coordinates,
)
# initial state
traj, _ = _get_traj(traj_path, top_path, [tautomer.hybrid_hydrogen_idx_at_lambda_0])
coordinates = [x.xyz[0] for x in traj[:10]] * unit.nanometer
energy_function = neutromeratio.ANI_force_and_energy(
model=model, atoms=tautomer.initial_state_ligand_atoms, mol=None
)
ANI2x_energy_initial_state = energy_function.calculate_energy(
coordinates,
)
assert is_quantity_close(
ANI2x_energy_initial_state.energy[0].in_units_of(unit.kilojoule_per_mole),
(-907243.8987177598 * unit.kilojoule_per_mole),
rtol=1e-5,
)
# compare with energies for AlchemicalANI object
# vacuum system
traj, _ = _get_traj(traj_path, top_path)
coordinates = [x.xyz[0] for x in traj[:10]] * unit.nanometer
energy_function, tautomer, flipped = setup_alchemical_system_and_energy_function(
name=name, ANImodel=AlchemicalANI2x, env="vacuum"
)
AlchemicalANI2x_energy_lambda_0 = energy_function.calculate_energy(
coordinates, lambda_value=0.0
)
AlchemicalANI2x_energy_lambda_1 = energy_function.calculate_energy(
coordinates, lambda_value=1.0
)
# making sure that the AlchemicalANI and the ANI results are the same on the endstates (with removed dummy atoms)
assert is_quantity_close(
AlchemicalANI2x_energy_lambda_1.energy[0].in_units_of(unit.kilojoule_per_mole),
(-907243.8987177598 * unit.kilojoule_per_mole),
rtol=1e-5,
)
assert is_quantity_close(
AlchemicalANI2x_energy_lambda_0.energy[0].in_units_of(unit.kilojoule_per_mole),
ANI2x_energy_final_state.energy[0].in_units_of(unit.kilojoule_per_mole),
rtol=1e-5,
)
def test_neutromeratio_energy_calculations_with_ANI_in_droplet():
from ..tautomers import Tautomer
from ..constants import kT
from ..analysis import setup_alchemical_system_and_energy_function
import numpy as np
from ..ani import AlchemicalANI1ccx
# read in exp_results.pickle
with open("data/test_data/exp_results.pickle", "rb") as f:
exp_results = pickle.load(f)
# droplet system
name = "molDWRow_298"
energy_function, tautomer, flipped = setup_alchemical_system_and_energy_function(
name=name,
ANImodel=AlchemicalANI1ccx,
env="droplet",
diameter=10,
base_path="data/test_data/droplet/molDWRow_298/",
)
# read in pregenerated traj
traj_path = (
"data/test_data/droplet/molDWRow_298/molDWRow_298_lambda_0.0000_in_droplet.dcd"
)
top_path = "data/test_data/droplet/molDWRow_298/molDWRow_298_in_droplet.pdb"
# _get_traj remove the dummy atom
traj, _ = _get_traj(traj_path, top_path,
[tautomer.hybrid_hydrogen_idx_at_lambda_0])
# overwrite the coordinates that rdkit generated with the first frame in the traj
torch.set_num_threads(1)
coordinates = [x.xyz[0] for x in traj[:10]] * unit.nanometer
# removes the dummy atom
atoms = (
tautomer.ligand_in_water_atoms[:tautomer.
hybrid_hydrogen_idx_at_lambda_0] +
tautomer.
ligand_in_water_atoms[tautomer.hybrid_hydrogen_idx_at_lambda_0 + 1:])
assert len(tautomer.ligand_in_water_atoms) == len(atoms) + 1
# ANI1ccx
model = neutromeratio.ani.ANI1ccx()
energy_function = neutromeratio.ANI_force_and_energy(model=model,
atoms=atoms,
mol=None)
energy = energy_function.calculate_energy(coordinates)
assert is_quantity_close(
energy.energy[0].in_units_of(unit.kilojoule_per_mole),
(-3514015.2561722626 * unit.kilojoule_per_mole),
rtol=1e-5,
)
# ANI2x
model = neutromeratio.ani.ANI2x()
atoms = (
tautomer.ligand_in_water_atoms[:tautomer.
hybrid_hydrogen_idx_at_lambda_0] +
tautomer.
ligand_in_water_atoms[tautomer.hybrid_hydrogen_idx_at_lambda_0 + 1:])
assert len(tautomer.ligand_in_water_atoms) == len(atoms) + 1
energy_function = neutromeratio.ANI_force_and_energy(model=model,
atoms=atoms,
mol=None)
energy = energy_function.calculate_energy(coordinates)
assert is_quantity_close(
energy.energy[0].in_units_of(unit.kilojoule_per_mole),
(-3515114.528875586 * unit.kilojoule_per_mole),
rtol=1e-5,
)