Source code for directupsampling.effqh.utils

from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np

from directupsampling.utils import set_volume

if TYPE_CHECKING:
    from ase.atoms import Atoms




def get_ref_atoms(atoms: Atoms | list[Atoms], volume: float) -> Atoms:
    """Return atoms scaled to *volume* per atom, or the list entry closest to it.

    Returns
    -------
    Atoms

    Raises
    ------
    TypeError
        If no entry in the list matches *volume*.

    """
    if hasattr(atoms, "get_masses"):
        ref_atoms = atoms.copy()
        set_volume(ref_atoms, volume * len(atoms))
    else:
        volumes = [_.get_volume() / len(_) for _ in atoms]
        inds = np.isclose(volumes, volume).nonzero()[0]
        if len(inds) == 0:
            msg = (
                f"No reference atoms found for volume {volume}. "
                f"(The reference atoms have the following volumes: {volumes})"
            )
            raise TypeError(msg)
        ref_atoms = atoms[inds[0]].copy()
    return ref_atoms





def reshape_sphinx_to_phonopy(fc_matrix: np.ndarray) -> np.ndarray:
    n0 = int(fc_matrix.shape[0] / 3)
    n1 = int(fc_matrix.shape[1] / 3)
    return fc_matrix.reshape(n0, 3, n1, 3).transpose([0, 2, 1, 3])





def reshape_phonopy_to_sphinx(fc_matrix: np.ndarray) -> np.ndarray:
    n0, n1 = fc_matrix.shape[0:2]
    return fc_matrix.transpose([0, 2, 1, 3]).reshape(n0 * 3, n1 * 3)