def test_vhf_from_pvhf():
trj = md.load(get_fn("spce.xtc"), top=get_fn("spce.gro"))
unique_atoms = get_unique_atoms(trj)
tuples_combination = combinations_with_replacement(unique_atoms, 2)
# obtaining g_r_t from total func
chunk_length = 20
r, t, g_r_t = compute_van_hove(trj, chunk_length=chunk_length)
partial_dict = compute_van_hove(trj, chunk_length=chunk_length, partial=True)
# obtating dict of np.array of pvhf
partial_dict = {}
for pairs in tuples_combination:
pair1 = pairs[0]
pair2 = pairs[1]
# Set in alphabetical order
if pairs[0].name > pairs[1].name:
pair2 = pairs[0]
pair1 = pairs[1]
x = compute_partial_van_hove(
trj,
chunk_length=chunk_length,
selection1=f"name {pair1.name}",
selection2=f"name {pair2.name}",
)
partial_dict[pairs] = x[1]
# obtaining total_grt from partial
total_g_r_t = vhf_from_pvhf(trj, partial_dict)
assert np.allclose(g_r_t, total_g_r_t)
def test_pvhf_error_is_tuple():
trj = md.load(get_fn("spce.xtc"), top=get_fn("spce.gro"))
unique_atoms = get_unique_atoms(trj)
key = frozenset({unique_atoms[0], unique_atoms[1]})
partial_dict = {}
x = compute_partial_van_hove(
trj,
chunk_length=20,
selection1="name O",
selection2="name O",
)
partial_dict[key] = x[1]
with pytest.raises(ValueError, match="Dictionary key not valid. Must be a tuple"):
vhf_from_pvhf(trj, partial_dict)
def test_pvhf_error_is_atom_type():
trj = md.load(get_fn("spce.xtc"), top=get_fn("spce.gro"))
unique_atoms = get_unique_atoms(trj)
partial_dict = {}
x = compute_partial_van_hove(
trj,
chunk_length=20,
selection1="name O",
selection2="name O",
)
partial_dict[("H", "O")] = x[1]
with pytest.raises(
ValueError, match="Dictionary key not valid. Must be type"
):
vhf_from_pvhf(trj, partial_dict)
def test_pvhf_error_2_atoms_per_pair():
trj = md.load(get_fn("spce.xtc"), top=get_fn("spce.gro"))
unique_atoms = get_unique_atoms(trj)
partial_dict = {}
x = compute_partial_van_hove(
trj,
chunk_length=20,
selection1="name O",
selection2="name O",
)
partial_dict[(unique_atoms[0], unique_atoms[1], unique_atoms[1])] = x[1]
with pytest.raises(
ValueError, match="Dictionary key not valid. Must only have 2 atoms per pair"
):
vhf_from_pvhf(trj, partial_dict)
def test_pvhf_error_atoms_in_trj():
trj = md.load(get_fn("spce.xtc"), top=get_fn("spce.gro"))
unique_atoms = get_unique_atoms(trj)
atom = md.core.topology.Atom(
name="Na", element=md.core.element.sodium, index=0, residue=1
)
partial_dict = {}
x = compute_partial_van_hove(
trj,
chunk_length=20,
selection1="name O",
selection2="name O",
)
partial_dict[(atom, unique_atoms[0])] = x[1]
with pytest.raises(
ValueError, match="Dictionary key not valid, `Atom`"
):
vhf_from_pvhf(trj, partial_dict)
def vhf_from_pvhf(trj, partial_dict, water=False):
"""
Compute the total Van Hove function from partial Van Hove functions
Parameters
----------
trj : mdtrj.Trajectory
trajectory on which partial vhf were calculated form
partial_dict : dict
dictionary containing partial vhf as a np.array.
Key is a tuple of len 2 with 2 atom types
Return
-------
total_grt : numpy.ndarray
Total Van Hove Function generated from addition of partial Van Hove Functions
"""
unique_atoms = get_unique_atoms(trj)
all_atoms = [atom for atom in trj.topology.atoms]
norm_coeff = 0
dict_shape = list(partial_dict.values())[0][0].shape
total_grt = np.zeros(dict_shape)
for atom_pair in partial_dict.keys():
# checks if key is a tuple
if isinstance(atom_pair, tuple) == False:
raise ValueError("Dictionary key not valid. Must be a tuple.")
for atom in atom_pair:
# checks if the atoms in tuple pair are atom types
if type(atom) != type(unique_atoms[0]):
raise ValueError(
"Dictionary key not valid. Must be type `MDTraj.Atom`.")
# checks if atoms are in the trajectory
if atom not in all_atoms:
raise ValueError(
f"Dictionary key not valid, `Atom` {atom} not in MDTraj trajectory."
)
# checks if key has two atoms
if len(atom_pair) != 2:
raise ValueError(
"Dictionary key not valid. Must only have 2 atoms per pair.")
atom1 = atom_pair[0]
atom2 = atom_pair[1]
coeff = (get_form_factor(element_name=f"{atom1.element.symbol}",
water=False) *
get_form_factor(element_name=f"{atom2.element.symbol}",
water=False) *
len(trj.topology.select(f"name {atom1.name}")) /
(trj.n_atoms) *
len(trj.topology.select(f"name {atom2.name}")) /
(trj.n_atoms))
normalized_pvhf = coeff * partial_dict[atom_pair]
norm_coeff += coeff
total_grt = np.add(total_grt, normalized_pvhf)
total_grt /= norm_coeff
return total_grt