def test_save_and_load_single_restraint(clean_files):
""" Test we can save a simple restraint """
rest = DAT_restraint()
rest.amber_index = True
rest.continuous_apr = False
rest.auto_apr = False
rest.topology = os.path.join(os.path.dirname(__file__),
"../data/cb6-but/cb6-but-notcentered.pdb")
rest.mask1 = ":CB6@O,O2,O4,O6,O8,O10"
rest.mask2 = ":BUT@C3"
rest.attach["target"] = 3.0
rest.attach["num_windows"] = 4
rest.attach["fc_initial"] = 0.0
rest.attach["fc_final"] = 3.0
rest.pull["fc"] = rest.attach["fc_final"]
rest.pull["num_windows"] = 4
rest.pull["target_initial"] = rest.attach["target"]
rest.pull["target_final"] = 6.0
rest.release["target"] = rest.pull["target_final"]
rest.release["num_windows"] = rest.attach["num_windows"]
rest.release["fc_initial"] = rest.attach["fc_initial"]
rest.release["fc_final"] = rest.attach["fc_final"]
rest.initialize()
save_restraints([rest], os.path.join("tmp", "rest.json"))
restraints = load_restraints(os.path.join("tmp", "rest.json"))
assert rest == restraints[0]
def test_save_and_load_list_restraint(clean_files):
""" Test we can save and load a list of restraints """
rest1 = DAT_restraint()
rest1.amber_index = True
rest1.continuous_apr = False
rest1.auto_apr = False
rest1.topology = os.path.join(os.path.dirname(__file__),
"../data/cb6-but/cb6-but-notcentered.pdb")
rest1.mask1 = ":CB6@O,O2,O4,O6,O8,O10"
rest1.mask2 = ":BUT@C3"
rest1.attach["target"] = 3.0
rest1.attach["num_windows"] = 4
rest1.attach["fc_initial"] = 0.0
rest1.attach["fc_final"] = 3.0
rest1.pull["fc"] = rest1.attach["fc_final"]
rest1.pull["num_windows"] = 4
rest1.pull["target_initial"] = rest1.attach["target"]
rest1.pull["target_final"] = 6.0
rest1.release["target"] = rest1.pull["target_final"]
rest1.release["num_windows"] = rest1.attach["num_windows"]
rest1.release["fc_initial"] = rest1.attach["fc_initial"]
rest1.release["fc_final"] = rest1.attach["fc_final"]
rest1.initialize()
rest2 = DAT_restraint()
rest2.amber_index = True
rest2.continuous_apr = False
rest2.auto_apr = False
rest2.topology = os.path.join(os.path.dirname(__file__),
"../data/cb6-but/cb6-but-notcentered.pdb")
rest2.mask1 = ":CB6@O,O2,O4,O6,O8,O10"
rest2.mask2 = ":BUT@C3"
rest2.mask3 = ":BUT@C"
rest2.attach["target"] = 180.0
rest2.attach["num_windows"] = 4
rest2.attach["fc_final"] = 75.0
rest2.pull["fc"] = rest2.attach["fc_final"]
rest2.pull["num_windows"] = 4
rest2.pull["target_final"] = 180.0
rest2.release["target"] = rest2.pull["target_final"]
rest2.release["num_windows"] = rest2.attach["num_windows"]
rest2.release["fc_final"] = rest2.attach["fc_final"]
rest2.initialize()
save_restraints([rest1, rest2], os.path.join("tmp", "restraints.json"))
restraints = load_restraints(os.path.join("tmp", "restraints.json"))
assert rest1 == restraints[0]
assert rest2 == restraints[1]
def _create_dummy_restraint(self, initial_structure):
if self.guest != "release":
windows = [
self.host_yaml["calculation"]["windows"]["attach"],
self.host_yaml["calculation"]["windows"]["pull"],
None,
]
else:
windows = [
None, None, self.host_yaml["calculation"]["windows"]["release"]
]
guest_restraint = DAT_restraint()
guest_restraint.auto_apr = True
guest_restraint.continuous_apr = True
guest_restraint.amber_index = False if self.backend == "openmm" else True
guest_restraint.topology = str(initial_structure)
guest_restraint.mask1 = "@1"
guest_restraint.mask2 = "@2"
if self.guest != "release":
restraint = self.guest_yaml["restraints"]["guest"][0]
guest_restraint.attach["target"] = restraint["restraint"][
"attach"]["target"]
guest_restraint.attach["fc_final"] = restraint["restraint"][
"attach"]["force_constant"]
guest_restraint.attach["fraction_list"] = self.host_yaml[
"calculation"]["lambda"]["attach"]
guest_restraint.pull["target_final"] = self.host_yaml[
"calculation"]["target"]["pull"]
guest_restraint.pull["num_windows"] = windows[1]
else:
# Remember, the purpose of this *fake* restraint is *only* to figure out how many windows to make,
# so we can use the OpenFF Evaluator to solvate the structures for us. To figure out how many winodws
# we need, just setting the lambda values should be sufficient.
guest_restraint.auto_apr = False
guest_restraint.continuous_apr = False
guest_restraint.release["target"] = 1.0
guest_restraint.release["fc_final"] = 1.0
guest_restraint.release["fraction_list"] = self.host_yaml[
"calculation"]["lambda"]["release"]
guest_restraint.initialize()
return guest_restraint
def build_guest_restraints(
cls,
coordinate_path: str,
attach_lambdas: List[float],
n_pull_windows: Optional[int],
restraint_schemas: List[Dict[str, Any]],
use_amber_indices: bool = False,
) -> List[DAT_restraint]:
"""A method to convert a set of guest restraints defined by their 'schemas'
into corresponding ``DAT_restraint``objects.
Each 'schema' should be a dictionary with:
* an ``atoms`` entry with a value of the atom selection make which specifies
which atoms the restraint will apply to
and additionally a nested ``attach`` and ``pull`` dictionary with
* a ``force_constant`` entry which specifies the force constant of the
restraint.
* a ``target`` entry which specifies the target value of the restraint.
These 'schemas` map directly to the 'restraints -> guest -> restraint'
dictionaries specified in the `taproom` guest YAML files.
Parameters
----------
coordinate_path
The path to the coordinate file which the restraints will be applied to.
This should contain either the host or the complex, the dummy atoms and
and solvent.
attach_lambdas
The values 'lambda' being used during the attach phase of the APR
calculation.
n_pull_windows
The total number of pull windows being used in the APR calculation.
restraint_schemas
The list of dictionaries which provide the settings to use for each
wall restraint to add.
use_amber_indices
Whether to use amber based (i.e. starting from 1) restraint indices or
OpenMM based (i.e. starting from 0) indices.
Returns
-------
The constructed wall restraint objects.
"""
restraints = []
for restraint_schema in restraint_schemas:
mask = restraint_schema["atoms"].split()
guest_restraint = DAT_restraint()
guest_restraint.auto_apr = True
guest_restraint.continuous_apr = False
guest_restraint.amber_index = use_amber_indices
guest_restraint.topology = coordinate_path
guest_restraint.mask1 = mask[0]
guest_restraint.mask2 = mask[1]
guest_restraint.mask3 = mask[2] if len(mask) > 2 else None
guest_restraint.mask4 = mask[3] if len(mask) > 3 else None
guest_restraint.attach["target"] = restraint_schema["attach"]["target"]
guest_restraint.attach["fc_final"] = restraint_schema["attach"][
"force_constant"
]
guest_restraint.attach["fraction_list"] = attach_lambdas
if n_pull_windows:
guest_restraint.pull["target_final"] = restraint_schema["pull"][
"target"
]
guest_restraint.pull["num_windows"] = n_pull_windows
guest_restraint.initialize()
restraints.append(guest_restraint)
return restraints
def build_wall_restraints(
cls,
coordinate_path: str,
n_attach_windows: int,
restraint_schemas: List[Dict[str, Any]],
use_amber_indices: bool = False,
) -> List[DAT_restraint]:
"""A method to convert a set of wall restraints defined by their 'schemas'
into corresponding ``DAT_restraint``objects.
Each 'schema' should be a dictionary with:
* an ``atoms`` entry with a value of the atom selection make which specifies
which atoms the restraint will apply to
* a ``force_constant`` entry which specifies the force constant of the
restraint.
* a ``target`` entry which specifies the target value of the restraint.
These 'schemas` map directly to the 'restraints -> wall_restraints -> restraint'
dictionaries specified in the `taproom` guest YAML files.
Parameters
----------
coordinate_path
The path to the coordinate file which the restraints will be applied to.
This should contain either the host or the complex, the dummy atoms and
and solvent.
n_attach_windows
The total number of attach windows being used in the APR calculation.
restraint_schemas
The list of dictionaries which provide the settings to use for each
wall restraint to add.
use_amber_indices
Whether to use amber based (i.e. starting from 1) restraint indices or
OpenMM based (i.e. starting from 0) indices.
Returns
-------
The constructed wall restraint objects.
"""
restraints = []
for restraint_schema in restraint_schemas:
restraint = DAT_restraint()
restraint.auto_apr = True
restraint.continuous_apr = False
restraint.amber_index = use_amber_indices
restraint.topology = coordinate_path
restraint.mask1 = restraint_schema["atoms"].split()[0]
restraint.mask2 = restraint_schema["atoms"].split()[1]
restraint.attach["fc_final"] = restraint_schema["force_constant"]
restraint.attach["fraction_list"] = [1.0] * n_attach_windows
restraint.attach["target"] = restraint_schema["target"]
# Minimum distance is 0 Angstrom
restraint.custom_restraint_values["r1"] = 0
restraint.custom_restraint_values["r2"] = 0
# Harmonic force constant beyond target distance.
restraint.custom_restraint_values["rk2"] = restraint_schema[
"force_constant"
]
restraint.custom_restraint_values["rk3"] = restraint_schema[
"force_constant"
]
restraint.initialize()
restraints.append(restraint)
return restraints
def test_evaluator_analyze(clean_files):
input_pdb = os.path.join(os.path.dirname(__file__),
"../data/cb6-but/vac.pdb")
structure = pmd.load_file(input_pdb, structure=True)
guest_atom_indices = []
for atom in structure.topology.atoms():
if atom.residue.name == "BUT":
guest_atom_indices.append(atom.index)
host_guest_structure = Setup.prepare_complex_structure(
input_pdb,
guest_atom_indices,
":BUT@C :BUT@C3",
24.0,
0,
46,
)
Setup.add_dummy_atoms_to_structure(
host_guest_structure,
[
np.array([0, 0, 0]),
np.array([0, 0, -3.0]),
np.array([0, 2.2, -5.2]),
],
np.zeros(3),
)
# Distance restraint
rest1 = DAT_restraint()
rest1.continuous_apr = True
rest1.amber_index = True
rest1.topology = host_guest_structure
rest1.mask1 = ":DM1"
rest1.mask2 = ":BUT@C"
rest1.attach["target"] = 6.0
rest1.attach["fraction_list"] = [0.00, 0.04, 0.181, 0.496, 1.000]
rest1.attach["fc_final"] = 5.0
rest1.pull["fc"] = rest1.attach["fc_final"]
rest1.pull["target_initial"] = rest1.attach["target"]
rest1.pull["target_final"] = 24.0
rest1.pull["num_windows"] = 19
rest1.initialize()
# Angle 1 restraint
rest2 = DAT_restraint()
rest2.continuous_apr = True
rest2.amber_index = True
rest2.topology = input_pdb
rest2.mask1 = ":DM2"
rest2.mask2 = ":DM1"
rest2.mask3 = ":BUT@C"
rest2.attach["target"] = 180.0
rest2.attach["fraction_list"] = [0.00, 0.04, 0.181, 0.496, 1.000]
rest2.attach["fc_final"] = 100.0
rest2.pull["fc"] = rest2.attach["fc_final"]
rest2.pull["target_initial"] = rest2.attach["target"]
rest2.pull["target_final"] = rest2.attach["target"]
rest2.pull["num_windows"] = 19
rest2.initialize()
# Angle 2
rest3 = DAT_restraint()
rest3.continuous_apr = True
rest3.amber_index = True
rest3.topology = input_pdb
rest3.mask1 = ":DM1"
rest3.mask2 = ":BUT@C"
rest3.mask3 = ":BUT@C3"
rest3.attach["target"] = 180.0
rest3.attach["fraction_list"] = [0.00, 0.04, 0.181, 0.496, 1.000]
rest3.attach["fc_final"] = 100.0
rest3.pull["fc"] = rest2.attach["fc_final"]
rest3.pull["target_initial"] = rest2.attach["target"]
rest3.pull["target_final"] = rest2.attach["target"]
rest3.pull["num_windows"] = 19
rest3.initialize()
temperature = 298.15
guest_restraints = [rest1, rest2, rest3]
ref_state_work = Analyze.compute_ref_state_work(temperature,
guest_restraints)
assert pytest.approx(ref_state_work, abs=1e-3) == -7.14151
fe_sym = Analyze.symmetry_correction(n_microstates=1, temperature=298.15)
assert fe_sym == 0.0
fe_sym = Analyze.symmetry_correction(n_microstates=2, temperature=298.15)
assert pytest.approx(fe_sym, abs=1e-3) == -0.410679
def initialize_restraints(self, structure="output.pdb"):
if self.guest != "release":
windows = [
self.host_yaml["calculation"]["windows"]["attach"],
self.host_yaml["calculation"]["windows"]["pull"],
None,
]
else:
windows = [
None, None, self.host_yaml["calculation"]["windows"]["release"]
]
static_restraints = []
for restraint in self.host_yaml["restraints"]["static"]:
static = static_DAT_restraint(
restraint_mask_list=restraint["restraint"]["atoms"].split(),
num_window_list=windows,
ref_structure=str(structure),
force_constant=restraint["restraint"]["force_constant"],
amber_index=False if self.backend == "openmm" else True,
)
static_restraints.append(static)
conformational_restraints = []
if self.host_yaml["restraints"]["conformational"]:
for conformational in self.host_yaml["restraints"][
"conformational"]:
mask = conformational["restraint"]["atoms"].split()
conformational_restraint = DAT_restraint()
conformational_restraint.auto_apr = True
conformational_restraint.continuous_apr = True
conformational_restraint.amber_index = False if self.backend == "openmm" else True
conformational_restraint.topology = str(structure)
conformational_restraint.mask1 = mask[0]
conformational_restraint.mask2 = mask[1]
conformational_restraint.mask3 = mask[2] if len(
mask) > 2 else None
conformational_restraint.mask4 = mask[3] if len(
mask) > 3 else None
if self.guest != "release":
conformational_restraint.attach["target"] = conformational[
"restraint"]["target"]
conformational_restraint.attach[
"fc_final"] = conformational["restraint"][
"force_constant"]
conformational_restraint.attach[
"fraction_list"] = self.host_yaml["calculation"][
"lambda"]["attach"]
conformational_restraint.pull[
"target_final"] = conformational["restraint"]["target"]
conformational_restraint.pull["num_windows"] = windows[1]
else:
conformational_restraint.auto_apr = False
conformational_restraint.continuous_apr = False
conformational_restraint.release[
"target"] = conformational["restraint"]["target"]
conformational_restraint.release[
"fc_final"] = conformational["restraint"][
"force_constant"]
conformational_restraint.release[
"fraction_list"] = self.host_yaml["calculation"][
"lambda"]["release"]
conformational_restraint.initialize()
conformational_restraints.append(conformational_restraint)
else:
logger.debug("Skipping conformational restraints...")
symmetry_restraints = []
if self.guest != "release" and "symmetry_correction" in self.guest_yaml:
for symmetry in self.guest_yaml["symmetry_correction"][
"restraints"]:
symmetry_restraint = DAT_restraint()
symmetry_restraint.auto_apr = True
symmetry_restraint.continuous_apr = True
symmetry_restraint.amber_index = False if self.backend == "openmm" else True
symmetry_restraint.topology = str(structure)
symmetry_restraint.mask1 = symmetry["atoms"].split()[0]
symmetry_restraint.mask2 = symmetry["atoms"].split()[1]
symmetry_restraint.mask3 = symmetry["atoms"].split()[2]
symmetry_restraint.attach["fc_final"] = symmetry[
"force_constant"]
symmetry_restraint.attach["fraction_list"] = [1.0] * len(
self.host_yaml["calculation"]["lambda"]["attach"])
# This target should be overridden by the custom values.
symmetry_restraint.attach["target"] = 999.99
symmetry_restraint.custom_restraint_values["r2"] = 91
symmetry_restraint.custom_restraint_values["r3"] = 91
# 0 force constant between 91 degrees and 180 degrees.
symmetry_restraint.custom_restraint_values["rk3"] = 0.0
symmetry_restraint.initialize()
symmetry_restraints.append(symmetry_restraint)
else:
logger.debug("Skipping symmetry restraints...")
wall_restraints = []
if self.guest != "release" and "wall_restraints" in self.guest_yaml[
'restraints']:
for wall in self.guest_yaml["restraints"]["wall_restraints"]:
wall_restraint = DAT_restraint()
wall_restraint.auto_apr = True
wall_restraint.continuous_apr = True
wall_restraint.amber_index = False if self.backend == "openmm" else True
wall_restraint.topology = str(structure)
wall_restraint.mask1 = wall["restraint"]["atoms"].split()[0]
wall_restraint.mask2 = wall["restraint"]["atoms"].split()[1]
wall_restraint.attach["fc_final"] = wall["restraint"][
"force_constant"]
wall_restraint.attach["fraction_list"] = [1.0] * len(
self.host_yaml["calculation"]["lambda"]["attach"])
wall_restraint.attach["target"] = wall["restraint"]["target"]
# Minimum distance is 0 Angstrom
wall_restraint.custom_restraint_values["r1"] = 0
wall_restraint.custom_restraint_values["r2"] = 0
# Harmonic force constant beyond target distance.
wall_restraint.custom_restraint_values["rk2"] = wall[
"restraint"]["force_constant"]
wall_restraint.custom_restraint_values["rk3"] = wall[
"restraint"]["force_constant"]
wall_restraint.initialize()
wall_restraints.append(wall_restraint)
else:
logger.debug("Skipping wall restraints...")
guest_restraints = []
for restraint in [] if not hasattr(
self,
'guest_yaml') else self.guest_yaml["restraints"]["guest"]:
mask = restraint["restraint"]["atoms"].split()
guest_restraint = DAT_restraint()
guest_restraint.auto_apr = True
guest_restraint.continuous_apr = True
guest_restraint.amber_index = False if self.backend == "openmm" else True
guest_restraint.topology = str(structure)
guest_restraint.mask1 = mask[0]
guest_restraint.mask2 = mask[1]
guest_restraint.mask3 = mask[2] if len(mask) > 2 else None
guest_restraint.mask4 = mask[3] if len(mask) > 3 else None
guest_restraint.attach["target"] = restraint["restraint"][
"attach"]["target"]
guest_restraint.attach["fc_final"] = restraint["restraint"][
"attach"]["force_constant"]
guest_restraint.attach["fraction_list"] = self.host_yaml[
"calculation"]["lambda"]["attach"]
guest_restraint.pull["target_final"] = restraint["restraint"][
"pull"]["target"]
guest_restraint.pull["num_windows"] = windows[1]
guest_restraint.initialize()
guest_restraints.append(guest_restraint)
return (
static_restraints,
conformational_restraints,
symmetry_restraints,
wall_restraints,
guest_restraints,
)