shnitsel.geo.alignment
======================

.. py:module:: shnitsel.geo.alignment


Functions
---------

.. autoapisummary::

   shnitsel.geo.alignment.get_centered_geometry
   shnitsel.geo.alignment.rotational_procrustes_np
   shnitsel.geo.alignment.rotational_procrustes
   shnitsel.geo.alignment.kabsch


Module Contents
---------------

.. py:function:: get_centered_geometry(atXYZ, by_mass = False)

   Helper function to set the center of the geometry (i.e. mean along the `atom` axis) to zero.

   :param atXYZ: Array of positional data
   :type atXYZ: AtXYZ
   :param by_mass: Flag whether the centering/average should be center of mass or just plain average of positions. Defaults to False.
   :type by_mass: Literal[False], optional

   :raises NotImplementedError: Centering the COM instead of the mean is currently not implemented.

   :returns: Resulting positions after centering.
   :rtype: AtXYZ


.. py:function:: rotational_procrustes_np(A, B, weight = None)

   Rotationally align the geometrie(s) in A to the single geometry in B.

   This helper function is specifically tailored to work directly on numpy arrays
   in contrast to its namesake `rotational_procrustes()`, which accepts
   `xarray.DataArray` parameters.

   :param A: The geometries to process with shape
             ``(n_geometries, n_points, n_coordinates)``
   :type A: np.ndarray
   :param B: The reference geometry with shape
             ``(n_points, n_coordinates)``
   :type B: np.ndarray
   :param weight: How much importance should be given to the alignment of
                  each point, by default equal importance
   :type weight: Sequence[float], optional

   :returns: An array with the same shape as A
   :rtype: np.ndarray

   .. rubric:: Notes

   We are solving the following minimization problem for each geometry `A(i)`:
   .. math::
       \min_R \Vert A(i)R - B\Vert^2_F

   For this, we calculte the weighted cross-covariance matrix:
   We are solving the minimization problem
   .. math::
       C = A^T B

   And then we find the SVD
   .. math::
       C = U \Sigma V^T

   Which provides us with the optimum rotation:
   .. math::
       R = V U^T

   If the resulting R has negative determinant, the rotation is instead
   a mirroring operation and we invert the sign of the last column of vt
   to restore rotational properties.


.. py:function:: rotational_procrustes(A, B, dim0 = 'atom', dim1 = 'direction', weight = None)

   Rotationally align the geometry or geometries in A to the single geometry in B.

   :param A: The (optionally multiple) geometries to process
   :type A: xr.DataArray
   :param B: The reference geometry
   :type B: xr.DataArray
   :param dim0: The name of the dimension over points to be rotated;
                must be present in ``A`` and ``B`; by default 'atom'
   :type dim0: str, optional
   :param dim1: The name of the dimension over the coordinates of the aforementioned
                points; must be present in ``A`` and ``B`; by default 'direction'
   :type dim1: str, optional
   :param weight: How much importance should be given to the alignment of
                  each point (atom), by default equal importance
   :type weight: Sequence[float], optional

   :returns: An xr.DataArray with the same shape as ``A`` but with entries aligned to the overall geometry of ``B``
   :rtype: xr.DataArray


.. py:function:: kabsch(atXYZ, reference_or_indexers = None, **indexers_kwargs)

   Rotationally align the molecular geometries in ``atXYZ`` to a single molecular geometry.

   If no `reference_or_indexers` argument (or the `indexers_kwargs` option) is passed,
   this function will try to use the first frame or first timestep in `atXYZ` as a reference.

   :param atXYZ: The geometries to process (with dims 'atom', 'direction')
   :type atXYZ: xr.DataArray
   :param reference_or_indexers: Either a reference geometry (with dims 'atom', 'direction')
                                 or an indexer dictionary which will be passed to ``atXYZ.sel()``
                                 to indetify a single geometry in the `atXYZ` parameter to
                                 use as a reference point.
   :type reference_or_indexers: xr.DataArray | dict, optional
   :param \*\*indexer_kwargs: The keyword-argument form of the indexer to be passed to ``atXYZ.sel()``

   :returns: The aligned geometries
   :rtype: xr.DataArray

   :raises ValueError: If nothing is done to indicate a reference geometry, i.e.
       neither reference_or_indexers nor indexer_kwargs are passed