Source code for qmmm_pme.common.utils

#! /usr/bin/env python3
"""A module containing helper functions accessed by multiple classes.
"""
from __future__ import annotations

from typing import Any

import numpy as np
from numpy.typing import NDArray


def align_dict(
        dictionary: dict[str, Any],
) -> dict[str, float]:
    """Create a 'flat' version of energy dictionary generated by
    :class:`Simulation` calculations.

    :param energy: The energy calculated by the :class:`Simulation`.
    :return: A flattend dictionary of energies.
    """
    flat = {}
    for key, val in dictionary.items():
        flat.update(
            align_dict(val) if isinstance(val, dict) else {key: val},
        )
    return flat


def compute_least_mirror(
        i_vector: NDArray[np.float64],
        j_vector: NDArray[np.float64],
        box: NDArray[np.float64],
) -> NDArray[np.float64]:
    """Calculates a least mirror vector.

    Returns the least mirror coordinates of i_vector with respect to
    j_vector given a set of box vectors from a periodic triclinic
    system.

    :param i_vector: Position vector, in Angstroms.
    :param j_vector: Reference vector, in Angstroms.
    :param box: |box|
    :return: Least mirror vector of the position vector with respect to
        the reference vector.
    """
    r_vector = i_vector - j_vector
    r_vector -= box[2] * np.floor(r_vector[2]/box[2][2] + 0.5)
    r_vector -= box[1] * np.floor(r_vector[1]/box[1][1] + 0.5)
    r_vector -= box[0] * np.floor(r_vector[0]/box[0][0] + 0.5)
    return r_vector


def compute_lattice_constants(
        box: NDArray[np.float64],
) -> tuple[float, float, float, float, float, float]:
    """Calculates length and angle lattice constants.

    Returns the lattice constants a, b, c, alpha, beta, and gamma using
    a set of box vectors for a periodic triclinic system.

    :param box: |box|
    :return: The characteristic lengths of a triclinic box, in
        Angstroms, and the characteristic angles of a triclinic box,
        in degrees.
    """
    vec_a = box[:, 0]
    vec_b = box[:, 1]
    vec_c = box[:, 2]
    len_a = np.linalg.norm(vec_a)
    len_b = np.linalg.norm(vec_b)
    len_c = np.linalg.norm(vec_c)
    alpha = 180*np.arccos(np.dot(vec_b, vec_c)/len_b/len_c)/np.pi
    beta = 180*np.arccos(np.dot(vec_a, vec_c)/len_a/len_c)/np.pi
    gamma = 180*np.arccos(np.dot(vec_a, vec_b)/len_a/len_b)/np.pi
    return len_a, len_b, len_c, alpha, beta, gamma