def ref_allpairs(conf):
return prepare_reference_nonbonded(
params=params,
exclusion_idxs=np.array([], dtype=np.int32),
scales=np.zeros((0, 2), dtype=np.float64),
lambda_plane_idxs=lambda_plane_idxs,
lambda_offset_idxs=lambda_offset_idxs,
beta=beta,
cutoff=cutoff,
)(conf, params, example_box, lamb)
def make_reference_nonbonded(lambda_plane_idxs):
return prepare_reference_nonbonded(
params=params,
exclusion_idxs=np.array([], dtype=np.int32),
scales=np.zeros((0, 2), dtype=np.float64),
lambda_plane_idxs=lambda_plane_idxs,
lambda_offset_idxs=lambda_offset_idxs,
beta=beta,
cutoff=cutoff,
)
def test_correctness(self):
"""
Test against the reference jax platform for correctness.
"""
# we can't go bigger than this due to memory limitations in the the reference platform.
for N in [33, 65, 231, 1050, 4080]:
np.random.seed(2022)
test_conf = self.host_conf[:N]
# strip out parts of the system
test_exclusions = []
test_scales = []
for (i, j), (sa, sb) in zip(self.nonbonded_fn.get_exclusion_idxs(), self.nonbonded_fn.get_scale_factors()):
if i < N and j < N:
test_exclusions.append((i, j))
test_scales.append((sa, sb))
test_exclusions = np.array(test_exclusions, dtype=np.int32)
test_scales = np.array(test_scales, dtype=np.float64)
test_params = self.nonbonded_fn.params[:N, :]
test_lambda_plane_idxs = np.random.randint(low=-2, high=2, size=N, dtype=np.int32)
test_lambda_offset_idxs = np.random.randint(low=-2, high=2, size=N, dtype=np.int32)
test_nonbonded_fn = potentials.Nonbonded(
test_exclusions, test_scales, test_lambda_plane_idxs, test_lambda_offset_idxs, self.beta, self.cutoff
)
ref_nonbonded_fn = prepare_reference_nonbonded(
test_params,
test_exclusions,
test_scales,
test_lambda_plane_idxs,
test_lambda_offset_idxs,
self.beta,
self.cutoff,
)
for precision, rtol, atol in [(np.float64, 1e-8, 1e-8), (np.float32, 1e-4, 5e-4)]:
self.compare_forces(
test_conf,
test_params,
self.box,
self.lamb,
ref_nonbonded_fn,
test_nonbonded_fn,
rtol=rtol,
atol=atol,
precision=precision,
)
def test_nblist_box_resize(self):
# test that running the coordinates under two different boxes produces correct results
# since we should be rebuilding the nblist when the box sizes change.
host_system, host_coords, box, _ = builders.build_water_system(3.0)
host_fns, host_masses = openmm_deserializer.deserialize_system(host_system, cutoff=1.0)
for f in host_fns:
if isinstance(f, potentials.Nonbonded):
test_nonbonded_fn = f
host_conf = []
for x, y, z in host_coords:
host_conf.append([to_md_units(x), to_md_units(y), to_md_units(z)])
host_conf = np.array(host_conf)
lamb = 0.1
ref_nonbonded_fn = prepare_reference_nonbonded(
test_nonbonded_fn.params,
test_nonbonded_fn.get_exclusion_idxs(),
test_nonbonded_fn.get_scale_factors(),
test_nonbonded_fn.get_lambda_plane_idxs(),
test_nonbonded_fn.get_lambda_offset_idxs(),
test_nonbonded_fn.get_beta(),
test_nonbonded_fn.get_cutoff(),
)
big_box = box + np.eye(3) * 1000
# print(big_box, small_box)
# (ytz): note the ordering should be from large box to small box. though in the current code
# the rebuild is triggered as long as the box *changes*.
for test_box in [big_box, box]:
for precision, rtol, atol in [(np.float64, 1e-8, 1e-10), (np.float32, 1e-4, 3e-5)]:
self.compare_forces(
host_conf,
test_nonbonded_fn.params,
test_box,
lamb,
ref_nonbonded_fn,
test_nonbonded_fn,
rtol=rtol,
atol=atol,
precision=precision,
)
def test_dhfr(self):
pdb_path = 'tests/data/5dfr_solv_equil.pdb'
host_pdb = app.PDBFile(pdb_path)
protein_ff = app.ForceField('amber99sbildn.xml', 'tip3p.xml')
host_system = protein_ff.createSystem(host_pdb.topology,
nonbondedMethod=app.NoCutoff,
constraints=None,
rigidWater=False)
host_coords = host_pdb.positions
box = host_pdb.topology.getPeriodicBoxVectors()
box = np.asarray(box / box.unit)
host_fns, host_masses = openmm_deserializer.deserialize_system(
host_system, cutoff=1.0)
for f in host_fns:
if isinstance(f, potentials.Nonbonded):
nonbonded_fn = f
host_conf = []
for x, y, z in host_coords:
host_conf.append([to_md_units(x), to_md_units(y), to_md_units(z)])
host_conf = np.array(host_conf)
beta = 2.0
cutoff = 1.1
lamb = 0.1
max_N = host_conf.shape[0]
for N in [33, 65, 231, 1050, 4080]:
print("N", N)
test_conf = host_conf[:N]
# process exclusions
test_exclusions = []
test_scales = []
for (i, j), (sa, sb) in zip(nonbonded_fn.get_exclusion_idxs(),
nonbonded_fn.get_scale_factors()):
if i < N and j < N:
test_exclusions.append((i, j))
test_scales.append((sa, sb))
test_exclusions = np.array(test_exclusions, dtype=np.int32)
test_scales = np.array(test_scales, dtype=np.float64)
test_params = nonbonded_fn.params[:N, :]
test_lambda_plane_idxs = np.random.randint(low=-2,
high=2,
size=N,
dtype=np.int32)
test_lambda_offset_idxs = np.random.randint(low=-2,
high=2,
size=N,
dtype=np.int32)
test_nonbonded_fn = potentials.Nonbonded(test_exclusions,
test_scales,
test_lambda_plane_idxs,
test_lambda_offset_idxs,
beta, cutoff)
ref_nonbonded_fn = prepare_reference_nonbonded(
test_params, test_exclusions, test_scales,
test_lambda_plane_idxs, test_lambda_offset_idxs, beta, cutoff)
for precision, rtol in [(np.float64, 1e-8), (np.float32, 1e-4)]:
self.compare_forces(test_conf,
test_params,
box,
lamb,
ref_nonbonded_fn,
test_nonbonded_fn,
rtol,
precision=precision)