grogu/src/grogupy/core.py

# Copyright (c) [2024] []
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import numpy as np
from numpy.linalg import inv

from grogupy.utils import blow_up_orbindx, commutator, parse_magnetic_entity


def parallel_Gk(HK, SK, eran, eset):
    return inv(SK * eran.reshape(eset, 1, 1) - HK)


def sequential_GK(HK, SK, eran, eset):
    Gk = np.zeros(shape=(eset, HK.shape[0], HK.shape[1]), dtype="complex128")
    for j in range(eset):
        Gk[j] = inv(SK * eran[j] - HK)

    return Gk


def calc_Vu(H, Tu):
    """_summary_

    Args:
        H (_type_): _description_
        Tu (_type_): _description_

    Returns:
        _type_: _description_
    """

    Vu1 = 1j / 2 * commutator(H, Tu)  # equation 100
    Vu2 = 1 / 8 * commutator(commutator(Tu, H), Tu)  # equation 100

    return Vu1, Vu2


def remove_clutter_for_save(pairs, magnetic_entities):
    """_summary_

    Args:
        pairs (_type_): _description_
        magnetic_entities (_type_): _description_

    Returns:
        _type_: _description_
    """
    # remove clutter from magnetic entities and pair information
    for pair in pairs:
        del pair["Gij"]
        del pair["Gij_tmp"]
        del pair["Gji"]
        del pair["Gji_tmp"]
    for mag_ent in magnetic_entities:
        del mag_ent["Gii"]
        del mag_ent["Gii_tmp"]
        del mag_ent["Vu1"]
        del mag_ent["Vu2"]

    return pairs, magnetic_entities


def build_hh_ss(dh):
    """_summary_

    Args:
        dh (_type_): _description_

    Returns:
        _type_: _description_
    """
    NO = dh.no  # shorthand for number of orbitals in the unit cell

    # preprocessing Hamiltonian and overlap matrix elements
    h11 = dh.tocsr(dh.M11r)
    h11 += dh.tocsr(dh.M11i) * 1.0j
    h11 = h11.toarray().reshape(NO, dh.n_s, NO).transpose(0, 2, 1).astype("complex128")

    h22 = dh.tocsr(dh.M22r)
    h22 += dh.tocsr(dh.M22i) * 1.0j
    h22 = h22.toarray().reshape(NO, dh.n_s, NO).transpose(0, 2, 1).astype("complex128")

    h12 = dh.tocsr(dh.M12r)
    h12 += dh.tocsr(dh.M12i) * 1.0j
    h12 = h12.toarray().reshape(NO, dh.n_s, NO).transpose(0, 2, 1).astype("complex128")

    h21 = dh.tocsr(dh.M21r)
    h21 += dh.tocsr(dh.M21i) * 1.0j
    h21 = h21.toarray().reshape(NO, dh.n_s, NO).transpose(0, 2, 1).astype("complex128")

    sov = (
        dh.tocsr(dh.S_idx)
        .toarray()
        .reshape(NO, dh.n_s, NO)
        .transpose(0, 2, 1)
        .astype("complex128")
    )

    # Reorganization of Hamiltonian and overlap matrix elements to SPIN BOX representation
    U = np.vstack(
        [np.kron(np.eye(NO, dtype=int), [1, 0]), np.kron(np.eye(NO, dtype=int), [0, 1])]
    )
    # This is the permutation that transforms ud1ud2 to u12d12
    # That is this transforms FROM SPIN BOX to ORBITAL BOX => U
    # the inverse transformation is U.T u12d12 to ud1ud2
    # That is FROM ORBITAL BOX to SPIN BOX => U.T

    # From now on everything is in SPIN BOX!!
    hh, ss = np.array(
        [
            U.T
            @ np.block([[h11[:, :, i], h12[:, :, i]], [h21[:, :, i], h22[:, :, i]]])
            @ U
            for i in range(dh.lattice.nsc.prod())
        ]
    ), np.array(
        [
            U.T
            @ np.block(
                [[sov[:, :, i], sov[:, :, i] * 0], [sov[:, :, i] * 0, sov[:, :, i]]]
            )
            @ U
            for i in range(dh.lattice.nsc.prod())
        ]
    )

    return hh, ss, NO


def setup_pairs_and_magnetic_entities(
    magnetic_entities, pairs, dh, simulation_parameters
):
    """_summary_

    Args:
        magnetic_entities (_type_): _description_
        pairs (_type_): _description_
        dh (_type_): _description_
        simulation_parameters (_type_): _description_

    Returns:
        _type_: _description_
    """
    # for every site we have to store 3 Greens function (and the associated _tmp-s) in the 3 reference directions
    for mag_ent in magnetic_entities:
        parsed = parse_magnetic_entity(dh, **mag_ent)  # parse orbital indexes
        mag_ent["orbital_indices"] = parsed
        mag_ent["spin_box_indices"] = blow_up_orbindx(
            parsed
        )  # calculate spin box indexes
        # if orbital is not set use all
        if "l" not in mag_ent.keys():
            mag_ent["l"] = "all"
        if isinstance(mag_ent["atom"], int):
            mag_ent["tags"] = [
                f"[{mag_ent['atom']}]{dh.atoms[mag_ent['atom']].tag}({mag_ent['l']})"
            ]
            mag_ent["xyz"] = [dh.xyz[mag_ent["atom"]]]
        if isinstance(mag_ent["atom"], list):
            mag_ent["tags"] = []
            mag_ent["xyz"] = []
            # iterate over atoms
            for atom_idx in mag_ent["atom"]:
                mag_ent["tags"].append(
                    f"[{atom_idx}]{dh.atoms[atom_idx].tag}({mag_ent['l']})"
                )
                mag_ent["xyz"].append(dh.xyz[atom_idx])

        # calculate size for Greens function generation
        spin_box_shape = len(mag_ent["spin_box_indices"])

        mag_ent["energies"] = (
            []
        )  # we will store the second order energy derivations here

        # These will be the perturbed potentials from eq. 100
        mag_ent["Vu1"] = []  # so they are independent in memory
        mag_ent["Vu2"] = []

        mag_ent["Gii"] = []  # Greens function
        mag_ent["Gii_tmp"] = []  # Greens function for parallelization
        for i in simulation_parameters["ref_xcf_orientations"]:
            # Rotations for every quantization axis
            mag_ent["Vu1"].append([])
            mag_ent["Vu2"].append([])
            # Greens functions for every quantization axis
            mag_ent["Gii"].append(
                np.zeros(
                    (simulation_parameters["eset"], spin_box_shape, spin_box_shape),
                    dtype="complex128",
                )
            )
            mag_ent["Gii_tmp"].append(
                np.zeros(
                    (simulation_parameters["eset"], spin_box_shape, spin_box_shape),
                    dtype="complex128",
                )
            )

    # for every site we have to store 2x3 Greens function (and the associated _tmp-s)
    # in the 3 reference directions, because G_ij and G_ji are both needed
    for pair in pairs:
        # calculate distance
        xyz_ai = magnetic_entities[pair["ai"]]["xyz"]
        xyz_aj = magnetic_entities[pair["aj"]]["xyz"]
        xyz_aj = xyz_aj + pair["Ruc"] @ simulation_parameters["cell"]
        pair["dist"] = np.linalg.norm(xyz_ai - xyz_aj)

        # calculate size for Greens function generation
        spin_box_shape_i = len(magnetic_entities[pair["ai"]]["spin_box_indices"])
        spin_box_shape_j = len(magnetic_entities[pair["aj"]]["spin_box_indices"])
        pair["tags"] = []
        for mag_ent in [magnetic_entities[pair["ai"]], magnetic_entities[pair["aj"]]]:
            tag = ""
            # get atoms of magnetic entity
            atoms_idx = mag_ent["atom"]
            orbitals = mag_ent["l"]

            # if magnetic entity contains one atoms
            if isinstance(atoms_idx, int):
                tag += f"[{atoms_idx}]{dh.atoms[atoms_idx].tag}({orbitals})"

            # if magnetic entity contains more than one atoms
            if isinstance(atoms_idx, list):
                # iterate over atoms
                atom_group = "{"
                for atom_idx in atoms_idx:
                    atom_group += f"[{atom_idx}]{dh.atoms[atom_idx].tag}({orbitals})--"
                # end {} of the atoms in the magnetic entity
                tag += atom_group[:-2] + "}"
            pair["tags"].append(tag)
        pair["energies"] = []  # we will store the second order energy derivations here

        pair["Gij"] = []  # Greens function
        pair["Gji"] = []
        pair["Gij_tmp"] = []  # Greens function for parallelization
        pair["Gji_tmp"] = []
        for i in simulation_parameters["ref_xcf_orientations"]:
            # Greens functions for every quantization axis
            pair["Gij"].append(
                np.zeros(
                    (simulation_parameters["eset"], spin_box_shape_i, spin_box_shape_j),
                    dtype="complex128",
                )
            )
            pair["Gij_tmp"].append(
                np.zeros(
                    (simulation_parameters["eset"], spin_box_shape_i, spin_box_shape_j),
                    dtype="complex128",
                )
            )
            pair["Gji"].append(
                np.zeros(
                    (simulation_parameters["eset"], spin_box_shape_j, spin_box_shape_i),
                    dtype="complex128",
                )
            )
            pair["Gji_tmp"].append(
                np.zeros(
                    (simulation_parameters["eset"], spin_box_shape_j, spin_box_shape_i),
                    dtype="complex128",
                )
            )

    return pairs, magnetic_entities


def onsite_projection(matrix, idx1, idx2):
    """_summary_

    Args:
        matrix (_type_): _description_
        idx (_type_): _description_

    Returns:
        _type_: _description_
    """

    return matrix[..., idx1, :][..., idx2]