# Source code for MDAnalysis.transformations.fit

    # -*- 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.
    #
    # 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
    #

    """\
    Fitting transformations --- :mod:`MDAnalysis.transformations.fit`
    =================================================================

    Translate and/or rotates the coordinates of a given trajectory to align
    a given AtomGroup to a reference structure.

    .. autoclass:: fit_translation

    .. autoclass:: fit_rot_trans

    """
    import numpy as np

    from ..analysis import align
    from ..lib.transformations import euler_from_matrix, euler_matrix

    from .base import TransformationBase

    class fit_translation(TransformationBase):
        """Translates a given AtomGroup so that its center of geometry/mass matches
        the respective center of the given reference. A plane can be given by the
        user using the option `plane`, and will result in the removal of
        the translation motions of the AtomGroup over that particular plane.

        Example
        -------
        Removing the translations of a given AtomGroup `ag` on the XY plane by
        fitting its center of mass to the center of mass of a reference `ref`:

        .. code-block:: python

            ag = u.select_atoms("protein")
            ref = mda.Universe("reference.pdb")
            transform = mda.transformations.fit_translation(ag, ref, plane="xy",
                                                            weights="mass")
            u.trajectory.add_transformations(transform)

        Parameters
        ----------
        ag : Universe or AtomGroup
           structure to translate, a
           :class:`~MDAnalysis.core.groups.AtomGroup` or a whole
           :class:`~MDAnalysis.core.universe.Universe`
        reference : Universe or AtomGroup
           reference structure, a :class:`~MDAnalysis.core.groups.AtomGroup` or a
           whole :class:`~MDAnalysis.core.universe.Universe`
        plane: str, optional
            used to define the plane on which the translations will be removed.
            Defined as a string of the plane.
            Suported planes are yz, xz and xy planes.
        weights : {"mass", ``None``} or array_like, optional
           choose weights. With ``"mass"`` uses masses as weights; with ``None``
           weigh each atom equally. If a float array of the same length as
           `ag` is provided, use each element of the `array_like` as a
           weight for the corresponding atom in `ag`.

        Returns
        -------
        MDAnalysis.coordinates.timestep.Timestep

        .. versionchanged:: 2.0.0
           The transformation was changed from a function/closure to a class
           with ``__call__``.
        .. versionchanged:: 2.0.0
           The transformation was changed to inherit from the base class for
           limiting threads and checking if it can be used in parallel analysis.
        """

        def __init__(
            self,
            ag,
            reference,
            plane=None,
            weights=None,
            max_threads=None,
            parallelizable=True,
        ):
            super().__init__(
                max_threads=max_threads, parallelizable=parallelizable
            )

            self.ag = ag
            self.reference = reference
            self.plane = plane
            self.weights = weights

            if self.plane is not None:
                axes = {"yz": 0, "xz": 1, "xy": 2}
                try:
                    self.plane = axes[self.plane]
                except (TypeError, KeyError):
                    raise ValueError(
                        f"{self.plane} is not a valid plane"
                    ) from None
            try:
                if self.ag.atoms.n_residues != self.reference.atoms.n_residues:
                    errmsg = (
                        f"{self.ag} and {self.reference} have mismatched"
                        f"number of residues"
                    )

                    raise ValueError(errmsg)
            except AttributeError:
                errmsg = (
                    f"{self.ag} or {self.reference} is not valid"
                    f"Universe/AtomGroup"
                )
                raise AttributeError(errmsg) from None
            self.ref, self.mobile = align.get_matching_atoms(
                self.reference.atoms, self.ag.atoms
            )
            self.weights = align.get_weights(self.ref.atoms, weights=self.weights)
            self.ref_com = self.ref.center(self.weights)

        def _transform(self, ts):
            mobile_com = np.asarray(
                self.mobile.atoms.center(self.weights), np.float32
            )
            vector = self.ref_com - mobile_com
            if self.plane is not None:
                vector[self.plane] = 0
            ts.positions += vector

            return ts

    class fit_rot_trans(TransformationBase):
        """Perform a spatial superposition by minimizing the RMSD.

        Spatially align the group of atoms `ag` to `reference` by doing a RMSD
        fit.

        This fit works as a way to remove translations and rotations of a given
        AtomGroup in a trajectory. A plane can be given using the flag `plane`
        so that only translations and rotations in that particular plane are
        removed. This is useful for protein-membrane systems to where the membrane
        must remain in the same orientation.

        Note
        ----
        ``max_threads`` is set to 1 for this transformation
        with which it performs better.

        Example
        -------
        Removing the translations and rotations of a given AtomGroup `ag` on the XY plane
        by fitting it to a reference `ref`, using the masses as weights for the RMSD fit:

        .. code-block:: python

            ag = u.select_atoms("protein")
            ref = mda.Universe("reference.pdb")
            transform = mda.transformations.fit_rot_trans(ag, ref, plane="xy",
                                                          weights="mass")
            u.trajectory.add_transformations(transform)

        Parameters
        ----------
        ag : Universe or AtomGroup
           structure to translate and rotate, a
           :class:`~MDAnalysis.core.groups.AtomGroup` or a whole
           :class:`~MDAnalysis.core.universe.Universe`
        reference : Universe or AtomGroup
           reference structure, a :class:`~MDAnalysis.core.groups.AtomGroup` or a whole
           :class:`~MDAnalysis.core.universe.Universe`
        plane: str, optional
            used to define the plane on which the rotations and translations will be removed.
            Defined as a string of the plane. Supported planes are "yz", "xz" and "xy" planes.
        weights : {"mass", ``None``} or array_like, optional
           choose weights. With ``"mass"`` uses masses as weights; with ``None``
           weigh each atom equally. If a float array of the same length as
           `ag` is provided, use each element of the `array_like` as a
           weight for the corresponding atom in `ag`.

        Returns
        -------
        MDAnalysis.coordinates.timestep.Timestep

        .. versionchanged:: 2.0.0
           The transformation was changed from a function/closure to a class
           with ``__call__``.
        .. versionchanged:: 2.0.0
           The transformation was changed to inherit from the base class for
           limiting threads and checking if it can be used in parallel analysis.
        """

        def __init__(
            self,
            ag,
            reference,
            plane=None,
            weights=None,
            max_threads=1,
            parallelizable=True,
        ):
            super().__init__(
                max_threads=max_threads, parallelizable=parallelizable
            )

            self.ag = ag
            self.reference = reference
            self.plane = plane
            self.weights = weights

            if self.plane is not None:
                axes = {"yz": 0, "xz": 1, "xy": 2}
                try:
                    self.plane = axes[self.plane]
                except (TypeError, KeyError):
                    raise ValueError(
                        f"{self.plane} is not a valid plane"
                    ) from None
            try:
                if self.ag.atoms.n_residues != self.reference.atoms.n_residues:
                    errmsg = (
                        f"{self.ag} and {self.reference} have mismatched "
                        f"number of residues"
                    )
                    raise ValueError(errmsg)
            except AttributeError:
                errmsg = (
                    f"{self.ag} or {self.reference} is not valid "
                    f"Universe/AtomGroup"
                )
                raise AttributeError(errmsg) from None
            self.ref, self.mobile = align.get_matching_atoms(
                self.reference.atoms, self.ag.atoms
            )
            self.weights = align.get_weights(self.ref.atoms, weights=self.weights)
            self.ref_com = self.ref.center(self.weights)
            self.ref_coordinates = self.ref.atoms.positions - self.ref_com

        def _transform(self, ts):
            mobile_com = self.mobile.atoms.center(self.weights)
            mobile_coordinates = self.mobile.atoms.positions - mobile_com
            rotation, dump = align.rotation_matrix(
                mobile_coordinates, self.ref_coordinates, weights=self.weights
            )
            vector = self.ref_com
            if self.plane is not None:
                matrix = np.r_[rotation, np.zeros(3).reshape(1, 3)]
                matrix = np.c_[matrix, np.zeros(4)]
                euler_angs = np.asarray(
                    euler_from_matrix(matrix, axes="sxyz"), np.float32
                )
                for i in range(0, euler_angs.size):
                    euler_angs[i] = (
                        euler_angs[self.plane] if i == self.plane else 0
                    )
                rotation = euler_matrix(
                    euler_angs[0], euler_angs[1], euler_angs[2], axes="sxyz"
                )[:3, :3]
                vector[self.plane] = mobile_com[self.plane]
            ts.positions = ts.positions - mobile_com
            ts.positions = np.dot(ts.positions, rotation.T)
            ts.positions = ts.positions + vector
            return ts