def setUp(self):
self.max_workers = get_gpu_count()
self.cli = client.CUDAPoolClient(self.max_workers)
def verify(self):
"""
See abstract class for documentation.
"""
gpus = get_gpu_count()
assert self.max_workers <= gpus, f"More workers '{self.max_workers}' requested than GPUs '{gpus}'"
multiprocessing.set_start_method("spawn")
parser = argparse.ArgumentParser(
description="Absolute Hydration Free Energy Testing",
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)
parser.add_argument("--hosts",
nargs="*",
default=None,
help="Hosts running GRPC worker to use for compute")
parser.add_argument("--num_gpus",
type=int,
help="number of gpus",
default=get_gpu_count())
parser.add_argument("--num_windows",
type=int,
help="number of solvent lambda windows",
required=True)
parser.add_argument(
"--num_equil_steps",
type=int,
help="number of equilibration steps for each solvent lambda window",
required=True,
)
parser.add_argument(
"--num_prod_steps",
AbsoluteFreeEnergy,
RelativeFreeEnergy,
get_romol_conf,
setup_relative_restraints_by_distance,
)
from timemachine.fe.model_rabfe import AbsoluteConversionModel, AbsoluteStandardHydrationModel, RelativeBindingModel
from timemachine.ff import Forcefield
from timemachine.lib.potentials import Nonbonded, NonbondedInterpolated
from timemachine.md import builders, minimizer
from timemachine.parallel.client import CUDAPoolClient
from timemachine.parallel.utils import get_gpu_count
from timemachine.potentials import rmsd
from timemachine.testsystems.relative import hif2a_ligand_pair
DATA_DIR = os.path.join(os.path.dirname(__file__), "data")
NUM_GPUS = get_gpu_count()
class TestRABFEModels(TestCase):
def test_endpoint_parameters_match_decoupling_and_conversion_complex(self):
"""Verifies that the parameters at the endpoint of conversion match with the starting parameters of
the decoupling. Done on a complex model, as the hydration models differ
Conv: P_start -> P_independent
Decouple: P_independent -> P_arbitrary
"""
host_system, host_coords, host_box, host_topology = builders.build_water_system(
4.0)
num_host_atoms = host_coords.shape[0]
return np.mean(bonded_du_dl + nonbonded_du_dl)
dG_host = np.trapz([_mean_du_dlambda(x) for x in results],
lambda_schedule)
stage_dGs.append(dG_host)
pred = stage_dGs[0] - stage_dGs[1]
return pred
if __name__ == "__main__":
# command line argument parser generated the following variables:
# num_gpus, num_complex_windows, num_solvent_windows, num_equil_steps,
# num_prod_steps ...
num_gpus = get_gpu_count()
# TODO: can I just get all of the du/dlambda time-series information
# so that we can guess the variance of dG_estimate without having to run
# multiple times from scratch?
num_repeats = 5
client = CUDAPoolClient(max_workers=num_gpus)
# TODO: move this test system constructor into a test fixture sort of thing
suppl = Chem.SDMolSupplier(path_to_ligand, removeHs=False)
all_mols = [x for x in suppl]
mol_a = all_mols[1]
mol_b = all_mols[4]