IPKM/mdevaluate
6
1
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Added type hints

Browse Source
This commit is contained in:
Sebastian Kloth 2024-01-10 15:11:42 +01:00
parent 9e9c865cf9
commit 25cb3d38b3
1 changed files with 23 additions and 23 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

34
src/mdevaluate/extra/water.py
View File

@ -273,7 +273,9 @@ def LSI(
return pd.DataFrame({"I": I, "P": P})
def HDL_LDL_positions(frame, trajectory, selector=None):
def HDL_LDL_positions(
frame: CoordinateFrame, selector: Optional[Callable] = None
) -> Tuple[NDArray, NDArray]:
atoms_PBC = pbc_points(frame, frame.box, thickness=0.7)
atoms_tree = KDTree(atoms_PBC)
if selector:
@ -290,8 +292,12 @@ def HDL_LDL_positions(frame, trajectory, selector=None):
return pos_HDL, pos_LDL
def HDL_LDL_gr(trajectory, segments=10000, skip=0):
def gr_frame(frame, trajectory, bins):
def HDL_LDL_gr(
trajectory: Coordinates, segments: int = 10000, skip: float = 0.1
) -> pd.DataFrame:
def gr_frame(
frame: CoordinateFrame, trajectory: Coordinates, bins: ArrayLike
) -> NDArray:
atoms_ALL = frame
atoms_HDL, atoms_LDL = HDL_LDL_positions(frame, trajectory)
@ -334,20 +340,13 @@ def HDL_LDL_gr(trajectory, segments=10000, skip=0):
dist_HDL_LDL, bins=bins, range=(0, bins[-1]), density=False
)[0]
n = [len(atoms_ALL), len(atoms_HDL), len(atoms_LDL)] / np.prod(
np.diag(frame.box)
)
return (
np.array(
return np.array(
[
hist_ALL_ALL / len(atoms_ALL),
hist_HDL_HDL / len(atoms_HDL),
hist_LDL_LDL / len(atoms_LDL),
hist_HDL_LDL / len(atoms_HDL),
]
),
n,
)
start_frame = trajectory[int(len(trajectory) * skip)]
@ -358,24 +357,25 @@ def HDL_LDL_gr(trajectory, segments=10000, skip=0):
)
gr = []
ns = []
for frame_index in frame_indices:
hists, n = gr_frame(trajectory[frame_index], trajectory, bins)
hists = gr_frame(trajectory[frame_index], trajectory, bins)
gr.append(hists)
ns.append(n)
gr = np.mean(gr, axis=0)
gr = gr / (4 / 3 * np.pi * bins[1:] ** 3 - 4 / 3 * np.pi * bins[:-1] ** 3)
r = bins[1:] - (bins[1] - bins[0]) / 2
n = np.mean(ns, axis=0)
return pd.DataFrame(
{"r": r, "gr_ALL": [0], "gr_HDL": gr[1], "gr_LDL": gr[2], "gr_MIX": gr[3]}
)
def HDL_LDL_concentration(trajectory, segments=10000, skip=0):
def HDL_LDL_concentration_frame(frame, trajectrory, bins):
def HDL_LDL_concentration(
trajectory: Coordinates, segments: int = 10000, skip: float = 0.1
) -> pd.DataFrame:
def HDL_LDL_concentration_frame(
frame: CoordinateFrame, bins: ArrayLike
) -> Tuple[NDArray, NDArray]:
atoms_HDL, atoms_LDL = HDL_LDL_positions(frame, trajectory)
atoms_PBC_HDL = pbc_points(atoms_HDL, frame.box, thickness=0.61)
atoms_PBC_LDL = pbc_points(atoms_LDL, frame.box, thickness=0.61)