diff --git a/src/mdevaluate/distribution.py b/src/mdevaluate/distribution.py index dd069a7..501b71a 100644 --- a/src/mdevaluate/distribution.py +++ b/src/mdevaluate/distribution.py @@ -120,7 +120,7 @@ def gr( def distance_distribution( - atoms: CoordinateFrame, bins: Optional[int, ArrayLike] + atoms: CoordinateFrame, bins: Union[int, ArrayLike] ) -> NDArray: connection_vectors = atoms[:-1, :] - atoms[1:, :] connection_lengths = (connection_vectors**2).sum(axis=1) ** 0.5 diff --git a/src/mdevaluate/pbc.py b/src/mdevaluate/pbc.py index 16678d8..8755b43 100644 --- a/src/mdevaluate/pbc.py +++ b/src/mdevaluate/pbc.py @@ -1,5 +1,6 @@ +from __future__ import annotations from collections import OrderedDict -from typing import Optional, Union +from typing import Optional, Union, TYPE_CHECKING import numpy as np from numpy.typing import ArrayLike, NDArray @@ -7,7 +8,8 @@ from numpy.typing import ArrayLike, NDArray from itertools import product from .logging import logger -from .coordinates import CoordinateFrame +if TYPE_CHECKING: + from mdevaluate.coordinates import CoordinateFrame def pbc_diff( @@ -70,7 +72,7 @@ def pbc_diff_tric(coords_a: NDArray, coords_b: NDArray, box: NDArray) -> NDArray def pbc_dist( - atoms_a: CoordinateFrame, atoms_b: CoordinateFrame, box: Optional[NDArray] = None + atoms_a: NDArray, atoms_b: NDArray, box: Optional[NDArray] = None ) -> ArrayLike: return ((pbc_diff(atoms_a, atoms_b, box) ** 2).sum(axis=1)) ** 0.5