# Source code for MDAnalysis.lib.mdamath

    # -*- Mode: python; tab-width: 4; indent-tabs-mode:nil; coding:utf-8 -*-
    # vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4
    #
    # MDAnalysis --- https://www.mdanalysis.org
    # Copyright (c) 2006-2017 The MDAnalysis Development Team and contributors
    # (see the file AUTHORS for the full list of names)
    #
    # Released under the Lesser GNU Public Licence, v2.1 or any higher version
    #
    # Please cite your use of MDAnalysis in published work:
    #
    # R. J. Gowers, M. Linke, J. Barnoud, T. J. E. Reddy, M. N. Melo, S. L. Seyler,
    # D. L. Dotson, J. Domanski, S. Buchoux, I. M. Kenney, and O. Beckstein.
    # MDAnalysis: A Python package for the rapid analysis of molecular dynamics
    # simulations. In S. Benthall and S. Rostrup editors, Proceedings of the 15th
    # Python in Science Conference, pages 102-109, Austin, TX, 2016. SciPy.
    # doi: 10.25080/majora-629e541a-00e
    #
    # N. Michaud-Agrawal, E. J. Denning, T. B. Woolf, and O. Beckstein.
    # MDAnalysis: A Toolkit for the Analysis of Molecular Dynamics Simulations.
    # J. Comput. Chem. 32 (2011), 2319--2327, doi:10.1002/jcc.21787
    #

    """
    Mathematical helper functions --- :mod:`MDAnalysis.lib.mdamath`
    ===============================================================

    Helper functions for common mathematical operations. Some of these functions
    are written in C/cython for higher performance.

    Linear algebra
    --------------

    .. autofunction:: normal
    .. autofunction:: norm
    .. autofunction:: pdot
    .. autofunction:: pnorm
    .. autofunction:: angle
    .. autofunction:: dihedral
    .. autofunction:: stp
    .. autofunction:: sarrus_det
    .. autofunction:: triclinic_box
    .. autofunction:: triclinic_vectors
    .. autofunction:: box_volume

    Connectivity
    ------------

    .. autofunction:: make_whole
    .. autofunction:: find_fragments

    .. versionadded:: 0.11.0
    .. versionchanged: 1.0.0
       Unused function :func:`_angle()` has now been removed.

    """
    import numpy as np

    from ..exceptions import NoDataError
    from . import util
    from ._cutil import (
        make_whole,
        find_fragments,
        _sarrus_det_single,
        _sarrus_det_multiple,
    )
    import numpy.typing as npt
    from typing import Union

    # geometric functions

    def norm(v: npt.ArrayLike) -> float:
        r"""Calculate the norm of a vector v.

        .. math:: v = \sqrt{\mathbf{v}\cdot\mathbf{v}}

        This version is faster then numpy.linalg.norm because it only works for a
        single vector and therefore can skip a lot of the additional fuss
        linalg.norm does.

        Parameters
        ----------
        v : array_like
            1D array of shape (N) for a vector of length N

        Returns
        -------
        float
            norm of the vector

        """
        return np.sqrt(np.dot(v, v))

    def normal(vec1: npt.ArrayLike, vec2: npt.ArrayLike) -> npt.NDArray:
        r"""Returns the unit vector normal to two vectors.

        .. math::

           \hat{\mathbf{n}} = \frac{\mathbf{v}_1 \times \mathbf{v}_2}{|\mathbf{v}_1 \times \mathbf{v}_2|}

        If the two vectors are collinear, the vector :math:`\mathbf{0}` is returned.

        .. versionchanged:: 0.11.0
           Moved into lib.mdamath
        """
        # TODO: enable typing when https://github.com/python/mypy/issues/11347 done
        normal: npt.NDArray = np.cross(vec1, vec2)  # type: ignore
        n = norm(normal)
        if n == 0.0:
            return normal  # returns [0,0,0] instead of [nan,nan,nan]
        # ... could also use numpy.nan_to_num(normal/norm(normal))
        return normal / n

    def pdot(a: npt.NDArray, b: npt.NDArray) -> npt.NDArray:
        """Pairwise dot product.

        ``a`` must be the same shape as ``b``.

        Parameters
        ----------
        a: :class:`numpy.ndarray` of shape (N, M)
        b: :class:`numpy.ndarray` of shape (N, M)

        Returns
        -------
        :class:`numpy.ndarray` of shape (N,)
        """
        return np.einsum("ij,ij->i", a, b)

    def pnorm(a: npt.NDArray) -> npt.NDArray:
        """Euclidean norm of each vector in a matrix

        Parameters
        ----------
        a: :class:`numpy.ndarray` of shape (N, M)

        Returns
        -------
        :class:`numpy.ndarray` of shape (N,)
        """
        return pdot(a, a) ** 0.5

    def angle(a: npt.ArrayLike, b: npt.ArrayLike) -> float:
        """Returns the angle between two vectors in radians

        .. versionchanged:: 0.11.0
           Moved into lib.mdamath
        .. versionchanged:: 1.0.0
           Changed rounding-off code to use `np.clip()`. Values lower than
           -1.0 now return `np.pi` instead of `-np.pi`
        """
        x = np.dot(a, b) / (norm(a) * norm(b))
        # catch roundoffs that lead to nan otherwise
        x = np.clip(x, -1.0, 1.0)
        return np.arccos(x)

    def stp(
        vec1: npt.ArrayLike, vec2: npt.ArrayLike, vec3: npt.ArrayLike
    ) -> float:
        r"""Takes the scalar triple product of three vectors.

        Returns the volume *V* of the parallel epiped spanned by the three
        vectors

        .. math::

            V = \mathbf{v}_3 \cdot (\mathbf{v}_1 \times \mathbf{v}_2)

        .. versionchanged:: 0.11.0
           Moved into lib.mdamath
        """
        # TODO: enable typing when https://github.com/python/mypy/issues/11347 done
        return np.dot(vec3, np.cross(vec1, vec2))  # type: ignore

    def dihedral(ab: npt.ArrayLike, bc: npt.ArrayLike, cd: npt.ArrayLike) -> float:
        r"""Returns the dihedral angle in radians between vectors connecting A,B,C,D.

        The dihedral measures the rotation around bc::

             ab
           A---->B
                  \ bc
                  _\'
                    C---->D
                      cd

        The dihedral angle is restricted to the range -π <= x <= π.

        .. versionadded:: 0.8
        .. versionchanged:: 0.11.0
           Moved into lib.mdamath
        """
        x = angle(normal(ab, bc), normal(bc, cd))
        return x if stp(ab, bc, cd) <= 0.0 else -x

    def sarrus_det(matrix: npt.NDArray) -> Union[float, npt.NDArray]:
        """Computes the determinant of a 3x3 matrix according to the
        `rule of Sarrus`_.

        If an array of 3x3 matrices is supplied, determinants are computed per
        matrix and returned as an appropriately shaped numpy array.

        .. _rule of Sarrus:
           https://en.wikipedia.org/wiki/Rule_of_Sarrus

        Parameters
        ----------
        matrix : numpy.ndarray
            An array of shape ``(..., 3, 3)`` with the 3x3 matrices residing in the
            last two dimensions.

        Returns
        -------
        det : float or numpy.ndarray
            The determinant(s) of `matrix`.
            If ``matrix.shape == (3, 3)``, the determinant will be returned as a
            scalar. If ``matrix.shape == (..., 3, 3)``, the determinants will be
            returned as a :class:`numpy.ndarray` of shape ``(...,)`` and dtype
            ``numpy.float64``.

        Raises
        ------
        ValueError:
            If `matrix` has less than two dimensions or its last two dimensions
            are not of shape ``(3, 3)``.

        .. versionadded:: 0.20.0
        """
        m = matrix.astype(np.float64)
        shape = m.shape
        ndim = m.ndim
        if ndim < 2 or shape[-2:] != (3, 3):
            raise ValueError(
                "Invalid matrix shape: must be (3, 3) or (..., 3, 3), "
                "got {}.".format(shape)
            )
        if ndim == 2:
            return _sarrus_det_single(m)
        return _sarrus_det_multiple(m.reshape((-1, 3, 3))).reshape(shape[:-2])

    def triclinic_box(
        x: npt.ArrayLike, y: npt.ArrayLike, z: npt.ArrayLike
    ) -> npt.NDArray:
        """Convert the three triclinic box vectors to
        ``[lx, ly, lz, alpha, beta, gamma]``.

        If the resulting box is invalid, i.e., any box length is zero or negative,
        or any angle is outside the open interval ``(0, 180)``, a zero vector will
        be returned.

        All angles are in degrees and defined as follows:

        * ``alpha = angle(y,z)``
        * ``beta  = angle(x,z)``
        * ``gamma = angle(x,y)``

        Parameters
        ----------
        x : array_like
            Array of shape ``(3,)`` representing the first box vector
        y : array_like
            Array of shape ``(3,)`` representing the second box vector
        z : array_like
            Array of shape ``(3,)`` representing the third box vector

        Returns
        -------
        numpy.ndarray
            A numpy array of shape ``(6,)`` and dtype ``np.float32`` providing the
            unitcell dimensions in the same format as returned by
            :attr:`MDAnalysis.coordinates.timestep.Timestep.dimensions`:\n
            ``[lx, ly, lz, alpha, beta, gamma]``.\n
            Invalid boxes are returned as a zero vector.

        Note
        ----
        Definition of angles: http://en.wikipedia.org/wiki/Lattice_constant

        See Also
        --------
        :func:`~MDAnalysis.lib.mdamath.triclinic_vectors`

        .. versionchanged:: 0.20.0
           Calculations are performed in double precision and invalid box vectors
           result in an all-zero box.
        """
        x = np.asarray(x, dtype=np.float64)
        y = np.asarray(y, dtype=np.float64)
        z = np.asarray(z, dtype=np.float64)
        lx = norm(x)
        ly = norm(y)
        lz = norm(z)
        with np.errstate(invalid="ignore"):
            alpha = np.rad2deg(np.arccos(np.dot(y, z) / (ly * lz)))
            beta = np.rad2deg(np.arccos(np.dot(x, z) / (lx * lz)))
            gamma = np.rad2deg(np.arccos(np.dot(x, y) / (lx * ly)))
        box = np.array([lx, ly, lz, alpha, beta, gamma], dtype=np.float32)
        # Only positive edge lengths and angles in (0, 180) are allowed:
        if np.all(box > 0.0) and alpha < 180.0 and beta < 180.0 and gamma < 180.0:
            return box
        # invalid box, return zero vector:
        return np.zeros(6, dtype=np.float32)

    def triclinic_vectors(
        dimensions: npt.ArrayLike, dtype: npt.DTypeLike = np.float32
    ) -> npt.NDArray:
        """Convert ``[lx, ly, lz, alpha, beta, gamma]`` to a triclinic matrix
        representation.

        Original `code by Tsjerk Wassenaar`_ posted on the Gromacs mailinglist.

        If `dimensions` indicates a non-periodic system (i.e., all lengths are
        zero), zero vectors are returned. The same applies for invalid `dimensions`,
        i.e., any box length is zero or negative, or any angle is outside the open
        interval ``(0, 180)``.

        .. _code by Tsjerk Wassenaar:
           http://www.mail-archive.com/gmx-users@gromacs.org/msg28032.html

        Parameters
        ----------
        dimensions : array_like
            Unitcell dimensions provided in the same format as returned by
            :attr:`MDAnalysis.coordinates.timestep.Timestep.dimensions`:\n
            ``[lx, ly, lz, alpha, beta, gamma]``.
        dtype: numpy.dtype
            The data type of the returned box matrix.

        Returns
        -------
        box_matrix : numpy.ndarray
            A numpy array of shape ``(3, 3)`` and dtype `dtype`,
            with ``box_matrix[0]`` containing the first, ``box_matrix[1]`` the
            second, and ``box_matrix[2]`` the third box vector.

        Notes
        -----
        * The first vector is guaranteed to point along the x-axis, i.e., it has the
          form ``(lx, 0, 0)``.
        * The second vector is guaranteed to lie in the x/y-plane, i.e., its
          z-component is guaranteed to be zero.
        * If any box length is negative or zero, or if any box angle is zero, the
          box is treated as invalid and an all-zero-matrix is returned.

        .. versionchanged:: 0.7.6
           Null-vectors are returned for non-periodic (or missing) unit cell.
        .. versionchanged:: 0.20.0
           * Calculations are performed in double precision and zero vectors are
             also returned for invalid boxes.
           * Added optional output dtype parameter.
        """
        dim = np.asarray(dimensions, dtype=np.float64)
        lx, ly, lz, alpha, beta, gamma = dim
        # Only positive edge lengths and angles in (0, 180) are allowed:
        if not (
            np.all(dim > 0.0) and alpha < 180.0 and beta < 180.0 and gamma < 180.0
        ):
            # invalid box, return zero vectors:
            box_matrix = np.zeros((3, 3), dtype=dtype)
        # detect orthogonal boxes:
        elif alpha == beta == gamma == 90.0:
            # box is orthogonal, return a diagonal matrix:
            box_matrix = np.diag(dim[:3].astype(dtype, copy=False))
        # we have a triclinic box:
        else:
            box_matrix = np.zeros((3, 3), dtype=np.float64)
            box_matrix[0, 0] = lx
            # Use exact trigonometric values for right angles:
            if alpha == 90.0:
                cos_alpha = 0.0
            else:
                cos_alpha = np.cos(np.deg2rad(alpha))
            if beta == 90.0:
                cos_beta = 0.0
            else:
                cos_beta = np.cos(np.deg2rad(beta))
            if gamma == 90.0:
                cos_gamma = 0.0
                sin_gamma = 1.0
            else:
                gamma = np.deg2rad(gamma)
                cos_gamma = np.cos(gamma)
                sin_gamma = np.sin(gamma)
            box_matrix[1, 0] = ly * cos_gamma
            box_matrix[1, 1] = ly * sin_gamma
            box_matrix[2, 0] = lz * cos_beta
            box_matrix[2, 1] = lz * (cos_alpha - cos_beta * cos_gamma) / sin_gamma
            box_matrix[2, 2] = np.sqrt(
                lz * lz - box_matrix[2, 0] ** 2 - box_matrix[2, 1] ** 2
            )
            # The discriminant of the above square root is only negative or zero for
            # triplets of box angles that lead to an invalid box (i.e., the sum of
            # any two angles is less than or equal to the third).
            # We don't need to explicitly test for np.nan here since checking for a
            # positive value already covers that.
            if box_matrix[2, 2] > 0.0:
                # all good, convert to correct dtype:
                box_matrix = box_matrix.astype(dtype, copy=False)
            else:
                # invalid box, return zero vectors:
                box_matrix = np.zeros((3, 3), dtype=dtype)
        return box_matrix

    def box_volume(dimensions: npt.ArrayLike) -> float:
        """Return the volume of the unitcell described by `dimensions`.

        The volume is computed as the product of the box matrix trace, with the
        matrix obtained from :func:`triclinic_vectors`.
        If the box is invalid, i.e., any box length is zero or negative, or any
        angle is outside the open interval ``(0, 180)``, the resulting volume will
        be zero.

        Parameters
        ----------
        dimensions : array_like
            Unitcell dimensions provided in the same format as returned by
            :attr:`MDAnalysis.coordinates.timestep.Timestep.dimensions`:\n
            ``[lx, ly, lz, alpha, beta, gamma]``.

        Returns
        -------
        volume : float
            The volume of the unitcell. Will be zero for invalid boxes.

        .. versionchanged:: 0.20.0
            Calculations are performed in double precision and zero is returned
            for invalid dimensions.
        """
        dim = np.asarray(dimensions, dtype=np.float64)
        lx, ly, lz, alpha, beta, gamma = dim
        if alpha == beta == gamma == 90.0 and lx > 0 and ly > 0 and lz > 0:
            # valid orthogonal box, volume is the product of edge lengths:
            volume = lx * ly * lz
        else:
            # triclinic or invalid box, volume is trace product of box matrix
            # (invalid boxes are set to all zeros by triclinic_vectors):
            tri_vecs = triclinic_vectors(dim, dtype=np.float64)
            volume = tri_vecs[0, 0] * tri_vecs[1, 1] * tri_vecs[2, 2]
        return volume