def test_raise(self):
# if func is not support pmap
self.assertRaises(ValueError, lambda: pt.pmap(pt.bfactors, self.traj))
# run time: openmp vs pmap
if 'OPENMP' in pt.compiled_info():
self.assertRaises(RuntimeError,
lambda: pt.pmap(pt.watershell, self.traj))
# if traj is not TrajectoryIterator
self.assertRaises(ValueError, lambda: pt.pmap(pt.radgyr, self.traj[:]))
# raise if a given method does not support pmap
def need_to_raise(traj=self.traj):
pt.pmap(2, pt.bfactors, traj)
self.assertRaises(ValueError, lambda: need_to_raise())
# raise if a traj is not TrajectoryIterator
def need_to_raise_2(traj=self.traj):
pt.pmap(pt.bfactors, traj[:], n_cores=2)
# raise if turn off pmap by setting _is_parallelizable to False
pt.radgyr._is_parallelizable = False
self.assertRaises(ValueError, lambda: pt.pmap(pt.radgyr, self.traj))
pt.radgyr._is_parallelizable = True
def test_progress(self):
self.traj = pt.iterload(fn('Tc5b.x'), fn('Tc5b.top'))
pt.pmap(pt.radgyr,
self.traj,
progress=True,
progress_params=dict(every=2))
pt.pmap(pt.radgyr, self.traj, progress=True)
def test_general(self):
traj = pt.iterload("./data/Tc5b.x", "./data/Tc5b.top")
# with mask
saved_data = pt.radgyr(traj, '@CA')
data = pt.pmap(pt.radgyr, traj, '@CA')
data = pt.tools.dict_to_ndarray(data)
aa_eq(saved_data, data)
# with a series of functions
func_list = [pt.radgyr, pt.molsurf, pt.rmsd]
ref = traj[-3]
for n_cores in [2, 3]:
for func in func_list:
if func in [pt.rmsd, ]:
pout = pt.tools.dict_to_ndarray(pt.pmap(func=func,
traj=traj,
ref=ref,
n_cores=n_cores))
serial_out = flatten(func(traj, ref=ref))
else:
pout = pt.tools.dict_to_ndarray(pt.pmap(n_cores=n_cores,
func=func,
traj=traj))
serial_out = flatten(func(traj))
aa_eq(pout[0], serial_out)
# test worker
# need to test this since coverages seems not recognize partial func
from pytraj.parallel.multiprocessing_ import worker_byfunc
data = worker_byfunc(rank=2, n_cores=8, func=pt.radgyr, traj=traj, args=(), kwd={'mask': '@CA'}, iter_options={})
assert data[0] == 2, 'rank must be 2'
assert data[2] == 1, 'n_frames for rank=2 should be 1 (only 10 frames in total)'
def test_raise(self):
# if func is not support pmap
self.assertRaises(ValueError, lambda: pt.pmap(pt.bfactors, self.traj))
# run time: openmp vs pmap
if 'OPENMP' in pt.compiled_info():
self.assertRaises(RuntimeError,
lambda: pt.pmap(pt.watershell, self.traj))
# if traj is not TrajectoryIterator
self.assertRaises(ValueError, lambda: pt.pmap(pt.radgyr, self.traj[:]))
# raise if a given method does not support pmap
def need_to_raise(traj=self.traj):
pt.pmap(2, pt.bfactors, traj)
self.assertRaises(ValueError, lambda: need_to_raise())
# raise if a traj is not TrajectoryIterator
def need_to_raise_2(traj=self.traj):
pt.pmap(pt.bfactors, traj[:], n_cores=2)
# raise if turn off pmap by setting _is_parallelizable to False
pt.radgyr._is_parallelizable = False
self.assertRaises(ValueError, lambda: pt.pmap(pt.radgyr, self.traj))
pt.radgyr._is_parallelizable = True
def test_volmap(self):
traj = pt.iterload("data/tz2.ortho.nc", "data/tz2.ortho.parm7")
# raise if does not provide size
self.assertRaises(
AssertionError, lambda: pt.pmap(pt.volmap,
traj,
mask=':WAT@O',
grid_spacing=(0.5, 0.5, 0.5),
n_cores=2))
mask = ':WAT@O'
grid_spacing = (0.5, 0.5, 0.5)
for n_cores in [1, 2, 3]:
for size in [(20, 20, 20), (20, 40, 60)]:
serial_out = pt.volmap(traj,
mask=mask,
grid_spacing=grid_spacing,
size=size)
parallel_out = pt.pmap(pt.volmap,
traj,
mask=mask,
grid_spacing=grid_spacing,
size=size,
n_cores=n_cores)
self.assertEqual(serial_out.shape, tuple(2 * x for x in size))
aa_eq(serial_out, parallel_out)
def test_nh_paramters(self):
parmfile = cpptraj_test_dir + '/Test_IRED/1IEE_A_prot.prmtop'
trajfile = cpptraj_test_dir + '/Test_IRED/1IEE_A_test.mdcrd'
traj = pt.iterload(trajfile, parmfile)
h_indices = pt.select_atoms('@H', traj.top)
n_indices = h_indices - 1
nh_indices = list(zip(n_indices, h_indices))
# single core
orders = pt.NH_order_parameters(traj, nh_indices, tcorr=8000.)
saved_S2 = np.loadtxt(cpptraj_test_dir +
'/Test_IRED/orderparam.save').T[-1]
aa_eq(orders, saved_S2)
# multiple core
# default 2
orders = pt.pmap(pt.NH_order_parameters, traj, nh_indices, tcorr=8000.)
aa_eq(orders, saved_S2)
for n_cores in [1, 2, 3, 4, -1]:
orders = pt.NH_order_parameters(traj,
nh_indices,
tcorr=8000.,
n_cores=n_cores)
aa_eq(orders, saved_S2)
orders = pt.pmap(pt.NH_order_parameters,
traj,
nh_indices,
tcorr=8000.,
n_cores=n_cores)
aa_eq(orders, saved_S2)
def test_nh_paramters(self):
parmfile = cpptraj_test_dir + '/Test_IRED/1IEE_A_prot.prmtop'
trajfile = cpptraj_test_dir + '/Test_IRED/1IEE_A_test.mdcrd'
traj = pt.iterload(trajfile, parmfile)
h_indices = pt.select_atoms('@H', traj.top)
n_indices = h_indices - 1
nh_indices = list(zip(n_indices, h_indices))
# single core
orders = pt.NH_order_parameters(traj, nh_indices, tcorr=8000.)
saved_S2 = np.loadtxt(cpptraj_test_dir +
'/Test_IRED/orderparam.save').T[-1]
aa_eq(orders, saved_S2)
# multiple core
# default 2
orders = pt.pmap(pt.NH_order_parameters, traj, nh_indices, tcorr=8000.)
aa_eq(orders, saved_S2)
for n_cores in [1, 2, 3, 4, -1]:
orders = pt.NH_order_parameters(traj,
nh_indices,
tcorr=8000.,
n_cores=n_cores)
aa_eq(orders, saved_S2)
orders = pt.pmap(pt.NH_order_parameters,
traj,
nh_indices,
tcorr=8000.,
n_cores=n_cores)
aa_eq(orders, saved_S2)
def test_pmap_hbond(self):
traj = pt.iterload(fn('tz2.nc'), fn('tz2.parm7'))
hbond_data_serial = pt.hbond(traj, dtype='dict')
hbond_data_pmap = pt.pmap(pt.hbond, traj, n_cores=3)
assert sorted(hbond_data_serial.keys()) == sorted(
hbond_data_pmap.keys())
for key, value in hbond_data_serial.items():
aa_eq(hbond_data_serial[key], hbond_data_pmap[key])
assert value.dtype == np.int32
with pytest.raises(ValueError):
# cpptraj style
# not support yet.
pt.pmap(['radgyr', 'hbond'], traj, n_cores=3)
def test_matrix_module(self):
traj = pt.iterload("data/tz2.nc", "data/tz2.parm7")
for n_cores in [2, 3]:
for func in [matrix.dist, matrix.idea]:
x = pt.pmap(func, traj, '@CA', n_cores=n_cores)
aa_eq(x, func(traj, '@CA'))
def test_matrix_module(self):
traj = pt.iterload("data/tz2.nc", "data/tz2.parm7")
for n_cores in [2, 3]:
for func in [matrix.dist, matrix.idea]:
x = pt.pmap(func, traj, '@CA', n_cores=n_cores)
aa_eq(x, func(traj, '@CA'))
def test_pmap_average_structure(self):
traj = pt.iterload("data/tz2.nc", "data/tz2.parm7")
saved_frame = pt.mean_structure(traj, '@CA')
saved_xyz = saved_frame.xyz
for n_cores in [2, 3, 4]:
frame = pt.pmap(pt.mean_structure, traj, '@CA', n_cores=n_cores)
aa_eq(frame.xyz, saved_xyz)
def test_frame_indices(self):
traj = pt.iterload(fn('tz2.nc'), fn('tz2.parm7'))
# frame_indices could be a list, range
frame_indices_list = [[0, 8, 9, 3, 2, 5], range(6)]
for frame_indices in frame_indices_list:
for n_cores in [2, 3]:
serial_out = pt.radgyr(
traj, '@CA', frame_indices=frame_indices)
parallel_out = pt.pmap(
pt.radgyr, traj, '@CA', frame_indices=frame_indices)
parallel_out_c_style = pt.pmap(
['radgyr @CA nomax'], traj, frame_indices=frame_indices)
aa_eq([serial_out], pt.tools.dict_to_ndarray(parallel_out))
aa_eq([serial_out],
pt.tools.dict_to_ndarray(parallel_out_c_style))
def pmap_(traj=traj, n_cores=4):
print('n_cores = ', n_cores)
print(traj)
x = pt.pmap(n_cores, pt.matrix.dist, traj, '@P', dtype='ndarray')
# make sure to reproduce serial verion
#data = np.sum((v[1] * v[2] for v in x), axis=0)
#return data / traj.n_frames
return x
def test_rotation_matrix_in_rmsd_calculation(self):
traj = pt.iterload("data/tz2.nc", "data/tz2.parm7")
saved_mat = pt.rotation_matrix(traj, ref=traj[3], mask='@CA')
saved_rmsd = pt.rmsd(traj, ref=traj[3], mask='@CA')
for n_cores in [2, 3]:
out = pt.pmap(pt.rotation_matrix, traj, ref=traj[3], mask='@CA')
out_with_rmsd = pt.pmap(pt.rotation_matrix,
traj,
ref=traj[3],
mask='@CA',
with_rmsd=True)
mat = out[list(out.keys())[0]]
mat2, rmsd_ = out_with_rmsd[list(out_with_rmsd.keys())[0]]
aa_eq(saved_mat, mat)
aa_eq(saved_mat, mat2)
aa_eq(saved_rmsd, rmsd_)
def test_pmap_average_structure(self):
traj = pt.iterload("data/tz2.nc", "data/tz2.parm7")
saved_frame = pt.mean_structure(traj, '@CA')
saved_xyz = saved_frame.xyz
for n_cores in [2, 3, 4]:
frame = pt.pmap(pt.mean_structure, traj, '@CA', n_cores=n_cores)
aa_eq(frame.xyz, saved_xyz)
def test_rotation_matrix_in_rmsd_calculation(self):
traj = pt.iterload("data/tz2.nc", "data/tz2.parm7")
saved_mat = pt.rotation_matrix(traj, ref=traj[3], mask='@CA')
saved_rmsd = pt.rmsd(traj, ref=traj[3], mask='@CA')
for n_cores in [2, 3]:
out = pt.pmap(pt.rotation_matrix, traj, ref=traj[3], mask='@CA')
out_with_rmsd = pt.pmap(pt.rotation_matrix,
traj,
ref=traj[3],
mask='@CA',
with_rmsd=True)
mat = out[list(out.keys())[0]]
mat2, rmsd_ = out_with_rmsd[list(out_with_rmsd.keys())[0]]
aa_eq(saved_mat, mat)
aa_eq(saved_mat, mat2)
aa_eq(saved_rmsd, rmsd_)
def test_iter_options(self):
traj = pt.iterload("data/tz2.ortho.nc", "data/tz2.ortho.parm7")
t0 = traj[:].autoimage().rmsfit(ref=0)
saved_avg = pt.mean_structure(t0)
saved_radgyr = pt.radgyr(traj, '@CA')
# perform autoimage, then rms fit to 1st frame, then compute mean structure
iter_options = {'autoimage': True, 'rmsfit': 0}
for n_cores in [2, 3]:
avg = pt.pmap(pt.mean_structure,
traj,
iter_options=iter_options,
n_cores=n_cores)
aa_eq(saved_avg.xyz, avg.xyz)
radgyr_ = pt.tools.dict_to_ndarray(
pt.pmap(pt.radgyr, traj, iter_options={'mask': '@CA'}))
aa_eq(radgyr_[0], saved_radgyr)
def test_volmap(self):
traj = pt.iterload("data/tz2.ortho.nc", "data/tz2.ortho.parm7")
# raise if does not provide size
self.assertRaises(AssertionError, lambda: pt.pmap(pt.volmap, traj, mask=':WAT@O',
grid_spacing=(0.5, 0.5, 0.5),
n_cores=2))
mask = ':WAT@O'
grid_spacing = (0.5, 0.5, 0.5)
for n_cores in [1, 2, 3]:
for size in [(20, 20, 20), (20, 40, 60)]:
serial_out = pt.volmap(traj, mask=mask, grid_spacing=grid_spacing, size=size)
parallel_out = pt.pmap(pt.volmap, traj, mask=mask, grid_spacing=grid_spacing,
size=size, n_cores=n_cores)
self.assertEqual(serial_out.shape, tuple(2 * x for x in size))
aa_eq(serial_out, parallel_out)
def test_sander_pmap_simple(self):
traj = pt.iterload(fn('Tc5b.x'), fn('Tc5b.top'))
fname = traj.top.filename
serial = pt.energy_decomposition(traj, prmtop=fname)['dihedral']
parallel = pt.pmap(n_cores=4,
func=pt.energy_decomposition,
traj=traj,
prmtop=fname)['dihedral']
aa_eq(serial, parallel)
def test_sander_pmap_simple(self):
traj = pt.iterload('./data/Tc5b.x', './data/Tc5b.top')
fname = traj.top.filename
serial = pt.energy_decomposition(traj, prmtop=fname)['dihedral']
parallel = pt.pmap(n_cores=4,
func=pt.energy_decomposition,
traj=traj,
prmtop=fname)['dihedral']
aa_eq(serial, parallel)
def test_iter_options(self):
traj = pt.iterload("data/tz2.ortho.nc", "data/tz2.ortho.parm7")
t0 = traj[:].autoimage().rmsfit(ref=0)
saved_avg = pt.mean_structure(t0)
saved_radgyr = pt.radgyr(traj, '@CA')
# perform autoimage, then rms fit to 1st frame, then compute mean structure
iter_options = {'autoimage': True, 'rmsfit': 0}
for n_cores in [2, 3]:
avg = pt.pmap(pt.mean_structure,
traj,
iter_options=iter_options,
n_cores=n_cores)
aa_eq(saved_avg.xyz, avg.xyz)
radgyr_ = pt.tools.dict_to_ndarray(pt.pmap(pt.radgyr,
traj,
iter_options={'mask':
'@CA'}))
aa_eq(radgyr_[0], saved_radgyr)
def test_c_command_style(self):
traj = pt.iterload(fn('tz2.nc'), fn('tz2.parm7'))
angle_ = pt.angle(traj, ':3 :4 :5')
distance_ = pt.distance(traj, '@10 @20')
data = pt.pmap(['angle :3 :4 :5', 'distance @10 @20'], traj, n_cores=2)
assert isinstance(data, OrderedDict), 'must be OrderDict'
arr = pt.tools.dict_to_ndarray(data)
aa_eq(angle_, arr[0])
aa_eq(distance_, arr[1])
# as whole text, case 1
data = pt.pmap(
'''angle :3 :4 :5
distance @10 @20''', traj, n_cores=2)
assert isinstance(data, OrderedDict), 'must be OrderDict'
arr = pt.tools.dict_to_ndarray(data)
aa_eq(angle_, arr[0])
aa_eq(distance_, arr[1])
def test_insert_new_function(self):
traj = pt.iterload("data/tz2.nc", "data/tz2.parm7")
# create mutable Trajectory
t0 = traj[:]
for frame in t0:
frame.xyz[:10] += 1.
data_parallel = pt.pmap(pt.radgyr, traj, n_cores=2, apply=change_10_atoms)
data_serial = pt.radgyr(t0)
aa_eq(data_parallel['RoG_00000'], data_serial)
def test_general(self):
traj = pt.iterload("./data/Tc5b.x", "./data/Tc5b.top")
# with mask
saved_data = pt.radgyr(traj, '@CA')
data = pt.pmap(pt.radgyr, traj, '@CA')
data = pt.tools.dict_to_ndarray(data)
aa_eq(saved_data, data)
# with a series of functions
func_list = [pt.radgyr, pt.molsurf, pt.rmsd]
ref = traj[-3]
for n_cores in [2, 3]:
for func in func_list:
if func in [
pt.rmsd,
]:
pout = pt.tools.dict_to_ndarray(
pt.pmap(func=func, traj=traj, ref=ref,
n_cores=n_cores))
serial_out = flatten(func(traj, ref=ref))
else:
pout = pt.tools.dict_to_ndarray(
pt.pmap(n_cores=n_cores, func=func, traj=traj))
serial_out = flatten(func(traj))
aa_eq(pout[0], serial_out)
# test worker
# need to test this since coverages seems not recognize partial func
from pytraj.parallel.multiprocessing_ import worker_byfunc
data = worker_byfunc(rank=2,
n_cores=8,
func=pt.radgyr,
traj=traj,
args=(),
kwd={'mask': '@CA'},
iter_options={})
assert data[0] == 2, 'rank must be 2'
assert data[
2] == 1, 'n_frames for rank=2 should be 1 (only 10 frames in total)'
def test_trajectoryiterator_with_transformation(self):
traj_on_disk = pt.iterload(fn('tz2.nc'), fn('tz2.parm7'))
traj_on_mem = pt.load(fn('tz2.nc'), fn('tz2.parm7'))
traj_on_disk.superpose(mask='@CA', ref=3)
traj_on_mem.superpose(mask='@CA', ref=3)
rmsd0_dict = pt.pmap(
pt.rmsd_nofit, traj_on_disk, mask='@CB', n_cores=2, ref=0)
rmsd1 = pt.rmsd_nofit(traj_on_mem, mask='@CB', ref=0)
aa_eq(pt.tools.dict_to_ndarray(rmsd0_dict), [rmsd1])
def test_ired_vector_and_matrix_pmap(self):
traj = pt.iterload(fn('tz2.nc'), fn('tz2.parm7'))
h = traj.top.select('@H')
n = h - 1
nh = list(zip(n, h))
exptected_vecs, exptected_mat = pt.ired_vector_and_matrix(traj, nh)
for n_cores in [2, 3]:
vecs, mat = pt.pmap(
pt.ired_vector_and_matrix, traj, nh, n_cores=n_cores)
aa_eq(exptected_vecs, vecs, decimal=7)
aa_eq(exptected_mat, mat, decimal=7)
def test_c_command_style(self):
traj = pt.iterload("data/tz2.nc", "data/tz2.parm7")
angle_ = pt.angle(traj, ':3 :4 :5')
distance_ = pt.distance(traj, '@10 @20')
data = pt.pmap(['angle :3 :4 :5', 'distance @10 @20'], traj, n_cores=2)
assert isinstance(data, OrderedDict), 'must be OrderDict'
arr = pt.tools.dict_to_ndarray(data)
aa_eq(angle_, arr[0])
aa_eq(distance_, arr[1])
# as whole text, case 1
data = pt.pmap('''angle :3 :4 :5
distance @10 @20''',
traj,
n_cores=2)
assert isinstance(data, OrderedDict), 'must be OrderDict'
arr = pt.tools.dict_to_ndarray(data)
aa_eq(angle_, arr[0])
aa_eq(distance_, arr[1])
def test_insert_new_function(self):
traj = pt.iterload(fn('tz2.nc'), fn('tz2.parm7'))
# create mutable Trajectory
t0 = traj[:]
for frame in t0:
frame.xyz[:10] += 1.
data_parallel = pt.pmap(
pt.radgyr, traj, n_cores=2, apply=change_10_atoms)
data_serial = pt.radgyr(t0)
aa_eq(data_parallel['RoG_00000'], data_serial)
def test_different_references(self):
traj = self.traj
func = pt.rmsd
for i in range(0, 8, 2):
ref = self.traj[i]
for n_cores in [2, 3, ]:
pout = pt.tools.dict_to_ndarray(pt.pmap(n_cores=n_cores,
func=func,
traj=traj,
ref=ref))
serial_out = flatten(func(traj, ref=ref))
aa_eq(pout[0], serial_out)
def test_check_valid_command(self):
from pytraj.parallel.base import check_valid_command
assert check_valid_command(['rms', ]) == (['rms refindex 0 '], True)
assert check_valid_command(['distrmsd', ]) == (['distrmsd refindex 0 '], True)
assert check_valid_command(['nativecontacts', ]) == (['nativecontacts refindex 0 '], True)
assert check_valid_command(['nastruct', ]) == (['nastruct refindex 0 '], True)
assert check_valid_command(['symmetricrmsd', ]) == (['symmetricrmsd refindex 0 '], True)
traj = pt.iterload("data/tz2.nc", "data/tz2.parm7")
aa_eq(pt.tools.dict_to_ndarray(
pt.pmap(['rmsd'], traj, ref=traj[3], n_cores=3)),
pt.rmsd(traj, ref=traj[3]))
# provide refindex
aa_eq(pt.tools.dict_to_ndarray(
pt.pmap(['rmsd refindex 0'], traj, ref=traj[3], n_cores=3)),
pt.rmsd(traj, ref=traj[3]))
aa_eq(pt.tools.dict_to_ndarray(
pt.pmap(['rmsd refindex 0'], traj, ref=[traj[3], traj[0]], n_cores=3)),
pt.rmsd(traj, ref=traj[3]))
# if user does not provide reference, need to give it to them
aa_eq(pt.tools.dict_to_ndarray(
pt.pmap(['rmsd'], traj, n_cores=3)),
pt.rmsd(traj, ref=traj[0]))
# does not support matrix
self.assertRaises(ValueError, lambda: pt.pmap(['matrix'], traj, n_cores=2))
# do not accept any c analysis command
for word in c_commands.analysis_commands:
self.assertRaises(ValueError, lambda: pt.pmap(word, traj, n_cores=2))
def test_different_references(self):
traj = self.traj
func = pt.rmsd
for i in range(0, 8, 2):
ref = self.traj[i]
for n_cores in [
2,
3,
]:
pout = pt.tools.dict_to_ndarray(
pt.pmap(n_cores=n_cores, func=func, traj=traj, ref=ref))
serial_out = flatten(func(traj, ref=ref))
aa_eq(pout[0], serial_out)
def test_frame_indices(self):
traj = pt.iterload("data/tz2.nc", "data/tz2.parm7")
# frame_indices could be a list, range
frame_indices_list = [[0, 8, 9, 3, 2, 5], range(6)]
for frame_indices in frame_indices_list:
for n_cores in [2, 3]:
serial_out = pt.radgyr(traj,
'@CA',
frame_indices=frame_indices)
parallel_out = pt.pmap(pt.radgyr,
traj,
'@CA',
frame_indices=frame_indices)
parallel_out_c_style = pt.pmap(
['radgyr @CA nomax'],
traj,
frame_indices=frame_indices)
aa_eq(serial_out, pt.tools.dict_to_ndarray(parallel_out))
aa_eq(serial_out,
pt.tools.dict_to_ndarray(parallel_out_c_style))
def test_ired_vector_and_matrix_pmap(self):
traj = pt.iterload("data/tz2.nc", "data/tz2.parm7")
h = traj.top.select('@H')
n = h - 1
nh = list(zip(n, h))
exptected_vecs, exptected_mat = pt.ired_vector_and_matrix(traj, nh)
for n_cores in [2, 3]:
vecs, mat = pt.pmap(pt.ired_vector_and_matrix,
traj,
nh,
n_cores=n_cores)
aa_eq(exptected_vecs, vecs, decimal=7)
aa_eq(exptected_mat, mat, decimal=7)
def test_check_valid_command(self):
from pytraj.parallel.base import check_valid_command
assert check_valid_command([
'rms',
]) == (['rms refindex 0 '], True)
assert check_valid_command([
'distrmsd',
]) == (['distrmsd refindex 0 '], True)
assert check_valid_command([
'nativecontacts',
]) == (['nativecontacts refindex 0 '], True)
assert check_valid_command([
'nastruct',
]) == (['nastruct refindex 0 '], True)
assert check_valid_command([
'symmetricrmsd',
]) == (['symmetricrmsd refindex 0 '], True)
traj = pt.iterload(fn('tz2.nc'), fn('tz2.parm7'))
aa_eq(
pt.tools.dict_to_ndarray(
pt.pmap(['rmsd'], traj, ref=traj[3], n_cores=3)),
[pt.rmsd(traj, ref=traj[3])])
# provide refindex
aa_eq(
pt.tools.dict_to_ndarray(
pt.pmap(['rmsd refindex 0'], traj, ref=traj[3], n_cores=3)),
[pt.rmsd(traj, ref=traj[3])])
aa_eq(
pt.tools.dict_to_ndarray(
pt.pmap(
['rmsd refindex 0'],
traj,
ref=[traj[3], traj[0]],
n_cores=3)), [pt.rmsd(traj, ref=traj[3])])
# if user does not provide reference, need to give it to them
aa_eq(
pt.tools.dict_to_ndarray(pt.pmap(['rmsd'], traj, n_cores=3)),
[pt.rmsd(traj, ref=traj[0])])
# does not support matrix
self.assertRaises(ValueError,
lambda: pt.pmap(['matrix'], traj, n_cores=2))
# do not accept any c analysis command
for word in c_commands.ANALYSIS_COMMANDS:
with pytest.raises(ValueError):
pt.pmap(word, traj, n_cores=2)
def test_reference(self):
traj = pt.iterload("./data/tz2.nc", "./data/tz2.parm7")
for n_cores in [2, 3]:
# use 4-th Frame for reference
data = pt.pmap(['rms @CA refindex 0'],
traj,
ref=traj[3],
n_cores=n_cores)
# another way to get reference
data2 = pt.pmap(['rms @CA reference'],
traj,
ref=traj[3],
n_cores=n_cores)
# use int for ref
data3 = pt.pmap(pt.rmsd,
traj,
ref=3,
mask='@CA',
n_cores=n_cores)
# use int for ref: use cpptraj's commmand style
data4 = pt.pmap(['rms @CA reference'],
traj,
ref=3,
n_cores=n_cores)
arr = pt.tools.dict_to_ndarray(data)[0]
arr2 = pt.tools.dict_to_ndarray(data2)[0]
arr3 = pt.tools.dict_to_ndarray(data3)[0]
arr4 = pt.tools.dict_to_ndarray(data4)[0]
aa_eq(arr, pt.rmsd(traj, ref=3, mask='@CA'))
aa_eq(arr2, pt.rmsd(traj, ref=3, mask='@CA'))
aa_eq(arr3, pt.rmsd(traj, ref=3, mask='@CA'))
aa_eq(arr4, pt.rmsd(traj, ref=3, mask='@CA'))
# use 4-th and 5-th Frame for reference
data = pt.pmap(
['rms @CA refindex 0', 'rms @CB refindex 1'],
traj,
ref=[traj[3], traj[4]],
n_cores=n_cores)
arr = pt.tools.dict_to_ndarray(data)[0]
aa_eq(arr, pt.rmsd(traj, '@CA', 3))
arr1 = pt.tools.dict_to_ndarray(data)[1]
aa_eq(arr1, pt.rmsd(traj, ref=4, mask='@CB'))
# no ref
data = pt.pmap(['radgyr', ], traj, n_cores=n_cores)
arr = pt.tools.dict_to_ndarray(data)[0]
aa_eq(arr, pt.radgyr(traj))
def test_pmap_hbond_with_solvent_bridge(self):
for trajin_fn, parm_fn in [('tz2.ortho.nc', 'tz2.ortho.parm7'),
('tz2.truncoct.nc', 'tz2.truncoct.parm7')]:
traj = pt.iterload(fn(trajin_fn), fn(parm_fn))
kwargs = dict(solvent_donor=':WAT@O', solvent_acceptor=':WAT')
hbond_data_serial = pt.hbond(traj, dtype='dict', **kwargs)
hbond_data_pmap = pt.pmap(
pt.hbond, traj, dtype='dict', n_cores=3, **kwargs)
assert sorted(hbond_data_serial.keys()) == sorted(
hbond_data_pmap.keys())
for key, value in hbond_data_serial.items():
if key.endswith('HB[ID]'):
assert hbond_data_pmap[key].tolist() == hbond_data_serial[
key].tolist()
print(hbond_data_pmap[key].tolist(),
hbond_data_serial[key].tolist())
else:
aa_eq(hbond_data_serial[key], hbond_data_pmap[key])
assert value.dtype == np.int32
def test_reference(self):
traj = pt.iterload("./data/tz2.nc", "./data/tz2.parm7")
for n_cores in [2, 3]:
# use 4-th Frame for reference
data = pt.pmap(['rms @CA refindex 0'],
traj,
ref=traj[3],
n_cores=n_cores)
# another way to get reference
data2 = pt.pmap(['rms @CA reference'],
traj,
ref=traj[3],
n_cores=n_cores)
# use int for ref
data3 = pt.pmap(pt.rmsd, traj, ref=3, mask='@CA', n_cores=n_cores)
# use int for ref: use cpptraj's commmand style
data4 = pt.pmap(['rms @CA reference'],
traj,
ref=3,
n_cores=n_cores)
arr = pt.tools.dict_to_ndarray(data)[0]
arr2 = pt.tools.dict_to_ndarray(data2)[0]
arr3 = pt.tools.dict_to_ndarray(data3)[0]
arr4 = pt.tools.dict_to_ndarray(data4)[0]
aa_eq(arr, pt.rmsd(traj, ref=3, mask='@CA'))
aa_eq(arr2, pt.rmsd(traj, ref=3, mask='@CA'))
aa_eq(arr3, pt.rmsd(traj, ref=3, mask='@CA'))
aa_eq(arr4, pt.rmsd(traj, ref=3, mask='@CA'))
# use 4-th and 5-th Frame for reference
data = pt.pmap(['rms @CA refindex 0', 'rms @CB refindex 1'],
traj,
ref=[traj[3], traj[4]],
n_cores=n_cores)
arr = pt.tools.dict_to_ndarray(data)[0]
aa_eq(arr, pt.rmsd(traj, '@CA', 3))
arr1 = pt.tools.dict_to_ndarray(data)[1]
aa_eq(arr1, pt.rmsd(traj, ref=4, mask='@CB'))
# no ref
data = pt.pmap([
'radgyr',
], traj, n_cores=n_cores)
arr = pt.tools.dict_to_ndarray(data)[0]
aa_eq(arr, pt.radgyr(traj))
def test_sander_pmap_with_options(self):
'''need to write mm_options as text
'''
code_local = '''
mm_options = sander.gas_input(8)
'''
traj = pt.iterload(fn('Tc5b.x'), fn('Tc5b.top'))
for code in [code_global, code_local]:
data_parallel = pt.pmap(pt.energy_decomposition,
traj,
mm_options=code,
n_cores=3,
dtype='dict')
mm_options = sander.gas_input(8)
data_serial = pt.energy_decomposition(traj,
mm_options=mm_options,
dtype='dict')
aa_eq(pt.tools.dict_to_ndarray(data_parallel),
pt.tools.dict_to_ndarray(data_serial))
def test_lie():
topology_fn = cpptraj_test_dir + '/FtuFabI.NAD.TCL.parm7'
trajin_fn = cpptraj_test_dir + '/FtuFabI.NAD.TCL.nc'
traj = pt.iterload(trajin_fn, topology_fn)
cpptraj_in = """
parm {}
trajin {}
lie LIE :TCS cutvdw 12 cutelec 12
""".format(topology_fn, trajin_fn)
state = pt.load_cpptraj_state(cpptraj_in)
state.run()
mask = ':TCS'
options = 'cutvdw 12 cutelec 12'
data = pt.lie(traj, mask, options=options)
for key in ['LIE[EELEC]', 'LIE[EVDW]']:
aa_eq(data[key], state.data[key])
if sys.platform.startswith('linux'):
data_parallel = pt.pmap(pt.lie, traj, mask, options, n_cores=3)
for key in ['LIE[EELEC]', 'LIE[EVDW]']:
aa_eq(data_parallel[key], state.data[key])
def test_sander_pmap_with_options(self):
'''need to write mm_options as text
'''
code_local = '''
mm_options = sander.gas_input(8)
'''
traj = pt.iterload('./data/Tc5b.x', './data/Tc5b.top')
for code in [code_global, code_local]:
data_parallel = pt.pmap(pt.energy_decomposition,
traj,
mm_options=code,
n_cores=3,
dtype='dict')
mm_options = sander.gas_input(8)
data_serial = pt.energy_decomposition(traj,
mm_options=mm_options,
dtype='dict')
aa_eq(
pt.tools.dict_to_ndarray(data_parallel),
pt.tools.dict_to_ndarray(data_serial))
def pmap_(n_cores):
print('n_cores = ', n_cores)
traj = pt.iterload(fname, tname)
print('traj size = %s (GB)' % traj._estimated_GB)
x = pt.pmap(n_cores, pt.rmsd, traj, mask='!:WAT', ref=traj[0])
return x
def need_to_raise_2(traj=self.traj):
pt.pmap(pt.bfactors, traj[:], n_cores=2)
def test_user_defined_method(self):
traj = pt.iterload("./data/tz2.nc", "./data/tz2.parm7")
for n_cores in [3, 5, 7]:
data = pt.pmap(method, traj, n_cores=n_cores)
joint_data = pt.tools.flatten(x[1] for x in data)
aa_eq(joint_data, pt.tools.flatten(method(traj)))
def pmap_(traj=traj, nh=nh, n_cores=4):
print('n_cores = ', n_cores)
print(traj)
return pt.pmap(pt.ired_vector_and_matrix, traj, nh, n_cores=n_cores)
def need_to_raise(traj=self.traj):
pt.pmap(2, pt.bfactors, traj)
def need_to_raise(traj=self.traj):
pt.pmap(2, pt.bfactors, traj)
def need_to_raise_2(traj=self.traj):
pt.pmap(pt.bfactors, traj[:], n_cores=2)
import pytraj as pt
# use `iterload` to save memory
traj = pt.iterload("../tests/data/tz2.ortho.nc",
"../tests/data/tz2.ortho.parm7")
# use 4 available cores
# calculate radgyr for all atoms
result = pt.pmap(n_cores=4, func=pt.radgyr, traj=traj)
print(result)
# serial version
result = pt.radgyr(traj)
print(result)
def pmap_(traj=traj, n_cores=4):
print('n_cores = ', n_cores)
print(traj)
x = pt.pmap(n_cores, pt.rmsd, traj, mask='!:WAT', ref=traj[0])
return x
import pytraj as pt
from pytraj.testing import Timer
from time import time
traj = pt.iterload('md*.nc', 'prmtop', frame_slice=(0, 1000))
print(traj)
mask = '!:WAT,Na+'
t0 = time()
pt.rmsd(traj, mask=mask)
t_serial = time() - t0
for n_cores in [1, 2, 3, 4, 5, 6, 7, 8]:
t0 = time()
pt.pmap(pt.rmsd, traj, mask=mask, ref=traj[0], n_cores=n_cores)
t_par = time() - t0
efficiency = (t_serial / t_par) / n_cores
print('n_cores = {}, time = {}, speed up = {}, efficiency = {}'.format(
n_cores, t_par, t_serial / t_par, efficiency))
