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)