def test_restraint_with_AlchemicalANI1ccx_for_batches_in_vacuum():
import numpy as np
from ..constants import kT
from ..ani import AlchemicalANI1ccx
from ..analysis import setup_alchemical_system_and_energy_function
# generate smiles
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 = energy_function.calculate_energy(x0, lambda_value=0.0)
e_ = energy.energy.value_in_unit(unit.kilojoule_per_mole)
comp_ = [
-906912.01647632,
-906874.76383595,
-906873.86032954,
-906880.6874288,
-906873.53381254,
-906867.20345442,
-906862.08888594,
-906865.00135805,
-906867.07548311,
-906871.27971175,
]
for e, comp_e in zip(e_, comp_):
np.isclose(e, comp_e, rtol=1e-6)
def test_timing_for_single_energy_calculation_with_AlchemicalANI_200_snapshot_sequence(
benchmark,
):
from ..analysis import setup_alchemical_system_and_energy_function
from ..ani import AlchemicalANI2x
torch.set_num_threads(4)
lambda_value = 0.1
name = "molDWRow_298"
(energy_function, tautomer, flipped,) = setup_alchemical_system_and_energy_function(
name=name,
env="droplet",
ANImodel=AlchemicalANI2x,
base_path=f"data/test_data/droplet/{name}",
diameter=10,
)
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)
coordinates = [x.xyz[0] for x in traj[:200]] * unit.nanometer
def wrap1():
for c in coordinates:
c = c.value_in_unit(unit.nanometer)
c = [c] * unit.nanometer
energy_function.calculate_energy(c, lambda_value)
benchmark(wrap1)
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_timing_for_single_energy_calculation_with_AlchemicalANI_20_snapshot_batch(
benchmark,
):
from ..analysis import setup_alchemical_system_and_energy_function
from ..ani import AlchemicalANI2x
lambda_value = 0.1
name = "molDWRow_298"
(energy_function, _, _,) = setup_alchemical_system_and_energy_function(
name=name,
env="droplet",
ANImodel=AlchemicalANI2x,
base_path=f"data/test_data/droplet/{name}",
diameter=10,
)
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, top = _get_traj(traj_path, top_path)
coordinates = [x.xyz[0] for x in traj[:20]] * unit.nanometer
def wrap1():
energy_function.calculate_energy(coordinates, lambda_value)
benchmark(wrap1)
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_timing_for_single_energy_calculation_with_AlchemicalANI_10_snapshots_batch(
benchmark,
):
from ..analysis import setup_alchemical_system_and_energy_function
from ..ani import AlchemicalANI2x
# NOTE: Sometimes this test fails? something wrong with molDWRow_68?
from ..constants import exclude_set_ANI, mols_with_charge, multiple_stereobonds
from ..constants import _get_names
from simtk import unit
torch.set_num_threads(1)
lambda_value = 0.1
name = "molDWRow_298"
(energy_function, _, _,) = setup_alchemical_system_and_energy_function(
name=name,
env="droplet",
ANImodel=AlchemicalANI2x,
base_path=f"data/test_data/droplet/{name}",
diameter=10,
)
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)
coordinates = [x.xyz[0] for x in traj[:10]] * unit.nanometer
def wrap1():
e = energy_function.calculate_energy(coordinates, lambda_value)
return e
e = benchmark(wrap1)
for e_pre, e_cal in zip(
[
-3515574.05857072,
-3515478.60995353,
-3515367.0878032,
-3515332.90224507,
-3515360.70976201,
-3515465.75272167,
-3515465.71963145,
-3515456.76306932,
-3515458.36516877,
-3515457.31727224,
],
e.energy.value_in_unit(unit.kilojoule_per_mole),
):
np.isclose(e_pre, e_cal, rtol=1e-5)
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)
def test_restraint_with_AlchemicalANI1ccx_for_batches_in_droplet():
import numpy as np
from ..constants import kT
from ..ani import AlchemicalANI2x
from ..analysis import setup_alchemical_system_and_energy_function
# generate smiles
name = "molDWRow_298"
energy_function, _, _ = setup_alchemical_system_and_energy_function(
name=name,
env="droplet",
ANImodel=AlchemicalANI2x,
diameter=10,
base_path=f"data/test_data/droplet/{name}",
)
# 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 = energy_function.calculate_energy(x0, lambda_value=0.0)
e_ = list(energy.energy.value_in_unit(unit.kilojoule_per_mole))
comp_ = [
-3515628.7141218423,
-3515543.79281359,
-3515512.8337717773,
-3515512.0131278695,
-3515511.2858314235,
-3515511.58966059,
-3515504.9864817774,
-3515505.5672234907,
-3515498.497307367,
-3515488.45506095,
]
for e, comp_e in zip(e_, comp_):
np.isclose(float(e), float(comp_e), rtol=1e-9)
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_droplet_bottom_restraint_with_AlchemicalANI1ccx_for_batches():
import numpy as np
from ..constants import kJ_mol_to_kT
from ..ani import AlchemicalANI2x
from ..analysis import setup_alchemical_system_and_energy_function
from ..restraints import CenterFlatBottomRestraint
name = "molDWRow_298"
energy_function, _, _ = setup_alchemical_system_and_energy_function(
name=name,
env="droplet",
ANImodel=AlchemicalANI2x,
diameter=10,
base_path=f"data/test_data/droplet/{name}",
)
# 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
coordinates = torch.tensor(
x0.value_in_unit(unit.nanometers),
requires_grad=True,
device=device,
dtype=torch.float32,
)
r = energy_function._compute_restraint_bias(coordinates, lambda_value=0.0)
r = r.detach().tolist()
full_restraints = [
0.4100764785669023,
0.03667892693933966,
0.0464554783605855,
1.6675558008972122,
1.6304614437658922,
2.7633326202414183,
3.7763598532698115,
9.16871744958779,
5.977379445002653,
3.429661035102869,
]
np.isclose(r, full_restraints)
coordinates = torch.tensor(
x0.value_in_unit(unit.angstrom),
requires_grad=True,
device=device,
dtype=torch.float32,
)
nr_of_mols = len(coordinates)
droplet_restraint = torch.tensor([0.0] * nr_of_mols,
device=device,
dtype=torch.float64)
for r in energy_function.list_of_lambda_restraints:
if isinstance(r, CenterFlatBottomRestraint):
c = r.restraint(coordinates)
droplet_restraint += c * kJ_mol_to_kT
print(droplet_restraint.tolist())
only_droplet_restrain_energies = [
0.19625303281218226,
0.03667892693933966,
0.0464554783605855,
0.4731166597911941,
1.6304614437658922,
2.7633326202414183,
3.7763598532698115,
4.700956016051697,
4.988977610418124,
3.429661035102869,
]
np.isclose(droplet_restraint.tolist(),
only_droplet_restrain_energies,
rtol=1e-6)
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,
)
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,
)
import torch.autograd.profiler as profiler
env = "droplet"
# env = "vacuum"
# system name
name = "molDWRow_298"
# read in pregenerated traj
traj_path = (
f"../data/test_data/{env}/molDWRow_298/molDWRow_298_lambda_0.0000_in_{env}.dcd"
)
top_path = f"../data/test_data/{env}/molDWRow_298/molDWRow_298_in_droplet.pdb"
# test energies with neutromeratio AlchemicalANI objects
# with ANI2x
traj, top = _get_traj(traj_path, top_path, None)
coordinates = [x.xyz[0] for x in traj[:800]] * unit.nanometer
energy_function, tautomer, flipped = setup_alchemical_system_and_energy_function(
name=name,
ANImodel=AlchemicalANI2x,
env=env,
diameter=10,
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"):
AlchemicalANI2x_energy_lambda_1 = energy_function.calculate_energy(
coordinates, lambda_value=1.0, include_restraint_energy_contribution=False
)
def test_COM_restraint_with_AlchemicalANI1ccx_for_batches_in_droplet():
import numpy as np
from ..constants import kT
from ..ani import AlchemicalANI2x
from ..analysis import setup_alchemical_system_and_energy_function
# generate smiles
name = "molDWRow_298"
energy_function, tautomer, flipped = setup_alchemical_system_and_energy_function(
name=name,
env="droplet",
ANImodel=AlchemicalANI2x,
diameter=10,
base_path=f"data/test_data/droplet/{name}",
)
# 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
coordinates = torch.tensor(
x0.value_in_unit(unit.nanometers),
requires_grad=True,
device=device,
dtype=torch.float32,
)
r = energy_function._compute_restraint_bias(coordinates, lambda_value=0.0)
r = r.detach().tolist()
full_restraints = [
0.4100764785669023,
0.03667892693933966,
0.0464554783605855,
1.6675558008972122,
1.6304614437658922,
2.7633326202414183,
3.7763598532698115,
9.16871744958779,
5.977379445002653,
3.429661035102869,
]
np.isclose(r, full_restraints)
print(energy_function.list_of_lambda_restraints[-1])
coordinates = torch.tensor(
x0.value_in_unit(unit.angstrom),
requires_grad=True,
device=device,
dtype=torch.float32,
)
com_restraint = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
com_r = (energy_function.list_of_lambda_restraints[-1].restraint(
coordinates).tolist())
np.isclose(com_restraint, com_r)
