Relativistic magnetic interactions from nonorthogonal basis sets

# 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.magnetism import blow_up_orbindx, parse_magnetic_entity
from grogupy.utilities import commutator

def parallel_Gk(HK, SK, eran, eset):
    """Calculates the Greens function by inversion.

    It calculates the Greens function on all the energy levels at the same time.

    Args:
        HK : (NO, NO), np.array_like
            Hamiltonian at a given k point
        SK : (NO, NO), np.array_like
            Overlap Matrix at a given k point
        eran : (eset) np.array_like
            Energy sample along the contour
        eset : int
            Number of energy samples along the contour

    Returns:
        Gk : (eset, NO, NO), np.array_like
            Green's function at a given k point
    """

    # Calculates the Greens function on all the energy levels
    return inv(SK * eran.reshape(eset, 1, 1) - HK)

def sequential_GK(HK, SK, eran, eset):

    It calculates sequentially over the energy levels.

    Args:
        eset : int

    Returns:
    """

    # creates an empty holder
    Gk = np.zeros(shape=(eset, HK.shape[0], HK.shape[1]), dtype="complex128")
    # fills the holder sequentially by the Greens function on a given energy
    for j in range(eset):
        Gk[j] = inv(SK * eran[j] - HK)

    return Gk

def calc_Vu(H, Tu):
    """Calculates the local perturbation in case of a spin rotation.

    Args:
        H : (NO, NO) np.array_like
            Hamiltonian
        Tu : (NO, NO) array_like
            Rotation around u

    Returns:
        Vu1 : (NO, NO) np.array_like
            First order perturbed matrix
        Vu2 : (NO, NO) np.array_like
            Second order perturbed matrix
    """

    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):
    """Removes unimportant data from the dictionaries.

    It is used before saving to throw away data that
    is not needed for post processing.

    Args:
        pairs : dict
            Contains all the pair information
        magnetic_entities : dict
            Contains all the magnetic entity information

    Returns:
            Contains all the reduced pair information
            Contains all the reduced magnetic entity information
    """

    # 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):
    """It builds the Hamiltonian and Overlap matrix from the sisl.dh class.

    It restructures the data in the SPIN BOX representation, where NS is
    the number of supercells and NO is the number of orbitals.

    Args:
        dh : sisl.physics.Hamiltonian
            Hamiltonian read in by sisl

    Returns:
        hh : (NS, NO, NO) np.array_like
            Hamiltonian in SPIN BOX representation
        ss : (NS, NO, NO) np.array_like
            Overlap matrix in SPIN BOX representation
    """

    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
        ]
    )

    return hh, ss

def setup_pairs_and_magnetic_entities(
    magnetic_entities, pairs, dh, simulation_parameters
):
    """It creates the complete structure of the dictionaries and fills some basic data.

    It creates orbital indexes, spin box indexes, coordinates and tags for magnetic entities.
    Furthermore it creates the structures for the energies, the perturbed potentials and
    the Greens function calculation. It dose the same for the pairs.

    Args:
            Contains the initial pair information
            Contains the initial magnetic entity information
        simulation_parameters : dict
            A set of parameters from the simulation

    Returns:
            Contains the initial information and the complete structure
    """

    # for every site we have to store 3 Greens function (and the associated _tmp-s) in the 3 reference directions
        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"

        # tag creation for one atom
        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"]]]
        # tag creation for more atoms
        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"])

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

        # 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 _ 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(
                )
            )

    # 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
        # 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)

        spin_box_shape_i = len(magnetic_entities[pair["ai"]]["spin_box_indices"])
        spin_box_shape_j = len(magnetic_entities[pair["aj"]]["spin_box_indices"])
        # tag generation
        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"] = []
            pair["Gij"].append(
                np.zeros(
                    (simulation_parameters["eset"], spin_box_shape_i, spin_box_shape_j),
                )
            )
            pair["Gij_tmp"].append(
                np.zeros(
                )
            )
            pair["Gji"].append(
                np.zeros(
                    (simulation_parameters["eset"], spin_box_shape_j, spin_box_shape_i),
                )
            )
            pair["Gji_tmp"].append(
                np.zeros(
                )
            )

def onsite_projection(matrix, idx1, idx2):
    """It produces the slices of a matrix for the on site projection.

    The slicing is along the last two axes as these contains the orbital indexing.

    Args:
        matrix : (..., :, :) np.array_like
            Some matrix
        idx : np.array_like
            The indexes of the orbitals

    Returns:
        np.array_like
            Reduced matrix based on the projection
    """

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

#
#
#

import warnings
from sys import getsizeof
from timeit import default_timer as timer

# use numpy number of threads one
try:
    from threadpoolctl import threadpool_info, threadpool_limits

    user_api = threadpool_info()["user_api"]
    threadpool_limits(limits=1, user_api=user_api)
except:
    print("Warning: threadpoolctl could not make numpy use single thread!")

import sisl
from mpi4py import MPI

try:
    from tqdm import tqdm

    tqdm_imported = True
except:
    print("Please install tqdm for nice progress bar.")
    tqdm_imported = False

from grogupy import *

def main():
    # runtime information
    times = dict()
    times["start_time"] = timer()

    # input output stuff
    ######################################################################

    infile = "/Users/danielpozsar/Downloads/nojij/Fe3GeTe2/monolayer/soc/lat3_791/Fe3GeTe2.fdf"
    outfile = "./Fe3GeTe2_notebook"

    magnetic_entities = [
        dict(atom=3, l=2),
        dict(atom=4, l=2),
        dict(atom=5, l=2),
    ]
    pairs = [
        dict(ai=0, aj=1, Ruc=np.array([0, 0, 0])),
        dict(ai=0, aj=2, Ruc=np.array([0, 0, 0])),
        dict(ai=1, aj=2, Ruc=np.array([0, 0, 0])),
        dict(ai=0, aj=2, Ruc=np.array([-1, -1, 0])),
        dict(ai=1, aj=2, Ruc=np.array([-1, -1, 0])),
        dict(ai=0, aj=2, Ruc=np.array([-1, 0, 0])),
        dict(ai=1, aj=2, Ruc=np.array([-1, 0, 0])),
        dict(ai=1, aj=2, Ruc=np.array([-2, 0, 0])),
        dict(ai=1, aj=2, Ruc=np.array([-3, 0, 0])),
    ]
    simulation_parameters = default_args
    simulation_parameters["infile"] = infile
    simulation_parameters["outfile"] = outfile
    simulation_parameters["kset"] = 20
    simulation_parameters["kdirs"] = "xy"
    simulation_parameters["eset"] = 600
    simulation_parameters["esetp"] = 10000

    # MPI parameters
    comm = MPI.COMM_WORLD
    size = comm.Get_size()
    rank = comm.Get_rank()
    root_node = 0

    if rank == root_node:
        # include parallel size in simulation parameters
        simulation_parameters["parallel_size"] = size

        # check versions for debugging
        try:
            print("sisl version: ", sisl.__version__)
        except:
            print("sisl version unknown.")
        try:
            print("numpy version: ", np.__version__)
        except:
            print("numpy version unknown.")

        # rename outfile
        if not simulation_parameters["outfile"].endswith(".pickle"):
            simulation_parameters["outfile"] += ".pickle"

        # if ebot is not given put it 0.1 eV under the smallest energy
        if simulation_parameters["ebot"] is None:
            try:
                eigfile = simulation_parameters["infile"][:-3] + "EIG"
                simulation_parameters["ebot"] = read_siesta_emin(eigfile) - 0.1
                simulation_parameters["automatic_ebot"] = True
            except:
                print("Could not determine ebot.")
                print("Parameter was not given and .EIG file was not found.")
        else:
            simulation_parameters["automatic_ebot"] = False

    # read sile
    fdf = sisl.get_sile(simulation_parameters["infile"])

    # read in hamiltonian
    dh = fdf.read_hamiltonian()

    # read unit cell vectors
    simulation_parameters["cell"] = fdf.read_geometry().cell

    # unit cell index
    uc_in_sc_idx = dh.lattice.sc_index([0, 0, 0])

        print("\n\n\n\n\n")
        print(
            "#################################################################### JOB INFORMATION ###########################################################################"
        )
        print_job_description(simulation_parameters)
        print(
            "################################################################################################################################################################"
        )

        times["setup_time"] = timer()
        print(f"Setup done. Elapsed time: {times['setup_time']} s")
        print(
            "================================================================================================================================================================"
        )

    # reformat Hamltonian and Overlap matrix for manipulations
    hh, ss = build_hh_ss(dh)

    # symmetrizing Hamiltonian and Overlap matrix to make them hermitian
    for i in range(dh.lattice.sc_off.shape[0]):
        j = dh.lattice.sc_index(-dh.lattice.sc_off[i])
        h1, h1d = hh[i], hh[j]
        hh[i], hh[j] = (h1 + h1d.T.conj()) / 2, (h1d + h1.T.conj()) / 2
        s1, s1d = ss[i], ss[j]
        ss[i], ss[j] = (s1 + s1d.T.conj()) / 2, (s1d + s1.T.conj()) / 2

    # identifying TRS and TRB parts of the Hamiltonian
    TAUY = np.kron(np.eye(NO), tau_y)
    hTR = np.array([TAUY @ hh[i].conj() @ TAUY for i in range(dh.lattice.nsc.prod())])
    hTRS = (hh + hTR) / 2
    hTRB = (hh - hTR) / 2

    # extracting the exchange field
    traced = [spin_tracer(hTRB[i]) for i in range(dh.lattice.nsc.prod())]  # equation 77
    XCF = np.array(
        [
            np.array([f["x"] / 2 for f in traced]),
            np.array([f["y"] / 2 for f in traced]),
            np.array([f["z"] / 2 for f in traced]),
        ]
    )

    # check if exchange field has scalar part
    max_xcfs = abs(np.array(np.array([f["c"] / 2 for f in traced]))).max()
    if max_xcfs > 1e-12:
        warnings.warn(
            f"Exchange field has non negligible scalar part. Largest value is {max_xcfs}"
        )

        times["H_and_XCF_time"] = timer()
        print(
            f"Hamiltonian and exchange field rotated. Elapsed time: {times['H_and_XCF_time']} s"
        )
        print(
        )

    # initialize pairs and magnetic entities based on input information
    pairs, magnetic_entities = setup_pairs_and_magnetic_entities(
        magnetic_entities, pairs, dh, simulation_parameters
    )

        times["site_and_pair_dictionaries_time"] = timer()
        print(
            f"Site and pair dictionaries created. Elapsed time: {times['site_and_pair_dictionaries_time']} s"
        )
        print(
        )

    # generate k space sampling
    kset = make_kset(
        dirs=simulation_parameters["kdirs"], NUMK=simulation_parameters["kset"]
    )

    # generate weights for k points
    wkset = np.ones(len(kset)) / len(kset)

    # split the k points based on MPI size
    kpcs = np.array_split(kset, size)

    # use progress bar if available
    if rank == root_node and tqdm_imported:
        kpcs[root_node] = tqdm(kpcs[root_node], desc="k loop")

        times["k_set_time"] = timer()
        print(f"k set created. Elapsed time: {times['k_set_time']} s")
        print(
        )

    # this will contain the three Hamiltonian in the
    # reference directions needed to calculate the energy
    # variations upon rotation
    hamiltonians = []

    # iterate over the reference directions (quantization axes)
    for i, orient in enumerate(simulation_parameters["ref_xcf_orientations"]):
        # obtain rotated exchange field and Hamiltonian
        R = RotMa2b(simulation_parameters["scf_xcf_orientation"], orient["o"])
        rot_XCF = np.einsum("ij,jklm->iklm", R, XCF)
        rot_H_XCF = sum(
            [np.kron(rot_XCF[i], tau) for i, tau in enumerate([tau_x, tau_y, tau_z])]
        )
        rot_H_XCF_uc = rot_H_XCF[uc_in_sc_idx]

        # obtain total Hamiltonian with the rotated exchange field
        rot_H = hTRS + rot_H_XCF  # equation 76

        # store the relevant information of the Hamiltonian
        hamiltonians.append(dict(orient=orient["o"], H=rot_H))

        # these are the rotations (for now) perpendicular to the quantization axis
        for u in orient["vw"]:
            # section 2.H
            Tu = np.kron(np.eye(NO, dtype=int), tau_u(u))
            Vu1, Vu2 = calc_Vu(rot_H_XCF_uc, Tu)

            for mag_ent in magnetic_entities:
                idx = mag_ent["spin_box_indices"]
                # fill up the perturbed potentials (for now) based on the on-site projections
                mag_ent["Vu1"][i].append(onsite_projection(Vu1, idx, idx))
                mag_ent["Vu2"][i].append(onsite_projection(Vu2, idx, idx))

        times["reference_rotations_time"] = timer()
        print(
            f"Rotations done perpendicular to quantization axis. Elapsed time: {times['reference_rotations_time']} s"
        )
        print(
        )

    # provide helpful information to estimate the runtime and memory
    # requirements of the Greens function calculations
        print("Starting matrix inversions.")
        if simulation_parameters["padawan_mode"]:
            print("Padawan mode: ")
            print(f"Total number of k points: {kset.shape[0]}")
            print(
                f"Number of energy samples per k point: {simulation_parameters['eset']}"
            )
            print(f"Total number of directions: {len(hamiltonians)}")
            print(
                f"Total number of matrix inversions: {kset.shape[0] * len(hamiltonians) * simulation_parameters['eset']}"
            )
            print(
                f"The shape of the Hamiltonian and the Greens function is {NO}x{NO}={NO*NO}"
            )
            # https://stackoverflow.com/questions/70746660/how-to-predict-memory-requirement-for-np-linalg-inv
            # memory is O(64 n**2) for complex matrices
            memory_size = getsizeof(hamiltonians[0]["H"].base) / 1024
            print(
                f"Memory taken by a single Hamiltonian is: {getsizeof(hamiltonians[0]['H'].base) / 1024} KB"
            )
            print(f"Expected memory usage per matrix inversion: {memory_size * 32} KB")
            print(
                f"Expected memory usage per k point for parallel inversion: {memory_size * len(hamiltonians) * simulation_parameters['eset'] * 32} KB"
            )
            print(
                f"Expected memory usage on root node: {len(np.array_split(kset, size)[0]) * memory_size * len(hamiltonians) * simulation_parameters['eset'] * 32 / 1024} MB"
            )
        print(
        )

    # MPI barrier
    comm.Barrier()

    # make energy contour
    cont = make_contour(
        emin=simulation_parameters["ebot"],
        enum=simulation_parameters["eset"],
        p=simulation_parameters["esetp"],
    )
    eran = cont.ze

    # sampling the integrand on the contour and the BZ
    for k in kpcs[rank]:
        # weight of k point in BZ integral
        wk = wkset[rank]

        # iterate over reference directions
        for i, hamiltonian_orientation in enumerate(hamiltonians):
            # calculate Hamiltonian and Overlap matrix in a given k point
            H = hamiltonian_orientation["H"]
            HK, SK = hsk(H, ss, dh.sc_off, k)

            if simulation_parameters["parallel_solver_for_Gk"]:
                Gk = parallel_Gk(HK, SK, eran, simulation_parameters["eset"])
            else:
                # solve Greens function sequentially for the energies, because of memory bound
                Gk = sequential_GK(HK, SK, eran, simulation_parameters["eset"])

            # store the Greens function slice of the magnetic entities
                mag_ent["Gii_tmp"][i] += onsite_projection(Gk, idx, idx) * wk

            for pair in pairs:
                # add phase shift based on the cell difference
                phase = np.exp(1j * 2 * np.pi * k @ pair["Ruc"].T)

                # get the pair orbital sizes from the magnetic entities
                ai = magnetic_entities[pair["ai"]]["spin_box_indices"]
                aj = magnetic_entities[pair["aj"]]["spin_box_indices"]

                # store the Greens function slice of the magnetic entities
                pair["Gij_tmp"][i] += onsite_projection(Gk, ai, aj) * phase * wk
                pair["Gji_tmp"][i] += onsite_projection(Gk, aj, ai) / phase * wk

    # summ reduce partial results of mpi nodes
    for i in range(len(hamiltonians)):
        for mag_ent in magnetic_entities:
            comm.Reduce(mag_ent["Gii_tmp"][i], mag_ent["Gii"][i], root=root_node)

        for pair in pairs:
            comm.Reduce(pair["Gij_tmp"][i], pair["Gij"][i], root=root_node)
            comm.Reduce(pair["Gji_tmp"][i], pair["Gji"][i], root=root_node)

        times["green_function_inversion_time"] = timer()
        print(
            f"Calculated Greens functions. Elapsed time: {times['green_function_inversion_time']} s"
        )
        print(
        )

        # iterate over the magnetic entities
        for tracker, mag_ent in enumerate(magnetic_entities):
            # iterate over the quantization axes
            for i, Gii in enumerate(mag_ent["Gii"]):
                storage = []
                # iterate over the first and second order local perturbations
                for Vu1, Vu2 in zip(mag_ent["Vu1"][i], mag_ent["Vu2"][i]):
                    # The Szunyogh-Lichtenstein formula
                    traced = np.trace(
                        (Vu2 @ Gii + 0.5 * Gii @ Vu1 @ Gii), axis1=1, axis2=2
                    )  # this is the on site projection
                    # evaluation of the contour integral
                    storage.append(int_de_ke(traced, cont.we))
                # fill up the magnetic entities dictionary with the energies
                magnetic_entities[tracker]["energies"].append(storage)
            # convert to np array
            magnetic_entities[tracker]["energies"] = np.array(
                magnetic_entities[tracker]["energies"]
            )

        # iterate over the pairs
        for tracker, pair in enumerate(pairs):
            for i, (Gij, Gji) in enumerate(zip(pair["Gij"], pair["Gji"])):
                site_i = magnetic_entities[pair["ai"]]
                site_j = magnetic_entities[pair["aj"]]
                # iterate over the first order local perturbations in all possible orientations for the two sites
                for Vui in site_i["Vu1"][i]:
                    for Vuj in site_j["Vu1"][i]:
                        # The Szunyogh-Lichtenstein formula
                        traced = np.trace(
                            (Vui @ Gij @ Vuj @ Gji), axis1=1, axis2=2
                        )  # this is the on site projection
                        # evaluation of the contour integral
                        storage.append(int_de_ke(traced, cont.we))
                # fill up the pairs dictionary with the energies
                pairs[tracker]["energies"].append(storage)
            pairs[tracker]["energies"] = np.array(pairs[tracker]["energies"])

        # calculate magnetic parameters
            Kxx, Kyy, Kzz, consistency = calculate_anisotropy_tensor(mag_ent)
            mag_ent["K"] = np.array([Kxx, Kyy, Kzz]) * sisl.unit_convert("eV", "meV")
            mag_ent["K_consistency"] = consistency

            J_iso, J_S, D, J = calculate_exchange_tensor(pair)
            pair["J_iso"] = J_iso * sisl.unit_convert("eV", "meV")
            pair["J_S"] = J_S * sisl.unit_convert("eV", "meV")
            pair["D"] = D * sisl.unit_convert("eV", "meV")
            pair["J"] = J * sisl.unit_convert("eV", "meV")

        times["end_time"] = timer()
        print(
            "##################################################################### GROGU OUTPUT #############################################################################"
        )

        print_parameters(simulation_parameters)
        print_atoms_and_pairs(magnetic_entities, pairs)
        print(
        )
        print_runtime_information(times)
        print("")

        # remove unwanted stuff before saving
        pairs, magnetic_entities = remove_clutter_for_save(pairs, magnetic_entities)
        # create output dictionary with all the relevant data
        results = dict(
            parameters=simulation_parameters,
            magnetic_entities=magnetic_entities,
            pairs=pairs,
            runtime=times,
        )

        # save results
        save_pickle(simulation_parameters["outfile"], results)

if __name__ == "__main__":
    main()

#
#
#

from argparse import ArgumentParser
from pickle import dump, load

default_args = dict(
    infile=None,
    outfile=None,
    scf_xcf_orientation=np.array([0, 0, 1]),
    ref_xcf_orientations=[
        dict(o=np.array([1, 0, 0]), vw=[np.array([0, 1, 0]), np.array([0, 0, 1])]),
        dict(o=np.array([0, 1, 0]), vw=[np.array([1, 0, 0]), np.array([0, 0, 1])]),
        dict(o=np.array([0, 0, 1]), vw=[np.array([1, 0, 0]), np.array([0, 1, 0])]),
    ],
    kset=2,
    kdirs="xyz",
    ebot=None,
    eset=42,
    esetp=1000,
    parallel_solver_for_Gk=False,
    padawan_mode=True,
)

# parser = ArgumentParser()

# parser.add_argument('--input'   , dest = 'infile' , default=None                       , help = 'Input file name')
# parser.add_argument('--output'  , dest = 'outfile', default=None                       , help = 'Output file name')

# parser.add_argument('--kset'    , dest = 'kset'   , default  = 2           , type=int  , help = 'k-space resolution of Jij calculation')
# parser.add_argument('--kdirs'   , dest = 'kdirs'  , default  = 'xyz'                   , help = 'Definition of k-space dimensionality')
# parser.add_argument('--ebot'    , dest = 'ebot'   , default  = None        , type=float, help = 'Bottom energy of the contour')
# parser.add_argument('--eset'    , dest = 'eset'   , default  = 42          , type=int  , help = 'Number of energy points on the contour')
# parser.add_argument('--eset-p'  , dest = 'esetp'  , default  = 1000        , type=int  , help = 'Parameter tuning the distribution on the contour')

# cmd_line_args = parser.parse_args()

def save_pickle(outfile, data):
    """Saves the data in the outfile with pickle.

    Args:
        outfile : str
            Path to outfile
        data : dict
            Contains the data
    """

    # save dictionary
    with open(outfile, "wb") as output_file:
        dump(data, output_file)

def load_pickle(infile):
    """Loads the data from the infile with pickle.

    Args:
        infile : str
            Path to infile

    Returns:
        data : dict
            A dictionary of data
    """

    # open and read file
    with open(infile, "wb") as input_file:
        data = load(data, input_file)

    return data

def print_parameters(simulation_parameters):
    """It prints the simulation parameters for the grogu out.

    Args:
            It contains the simulations parameters
    """

    print(
        "================================================================================================================================================================"
    )
    print("Input file: ")
    print(simulation_parameters["infile"])
    print("Output file: ")
    print(simulation_parameters["outfile"])
    print(
        "Number of nodes in the parallel cluster: ",
        simulation_parameters["parallel_size"],
    )
    print(
    )
    print("Cell [Ang]: ")
    print(simulation_parameters["cell"])
    print(
    )
    print("DFT axis: ")
    print(simulation_parameters["scf_xcf_orientation"])
    print("Quantization axis and perpendicular rotation directions:")
    for ref in simulation_parameters["ref_xcf_orientations"]:
        print(ref["o"], " --» ", ref["vw"])
    print(
    )
    print("Parameters for the contour integral:")
    print("Number of k points: ", simulation_parameters["kset"])
    print("k point directions: ", simulation_parameters["kdirs"])
    print("Ebot: ", simulation_parameters["ebot"])
    print("Eset: ", simulation_parameters["eset"])
    print("Esetp: ", simulation_parameters["esetp"])
    print(
    )

def print_atoms_and_pairs(magnetic_entities, pairs):
    """It prints the pair and magnetic entity information for the grogu out.

    Args:
            It contains the data on the magnetic entities
            It contains the data on the pairs
    """

    print("Atomic information: ")
    print(
        "----------------------------------------------------------------------------------------------------------------------------------------------------------------"
    )
    print(
        "[atom index]Element(orbitals)        x [Ang]        y [Ang]        z [Ang]        Sx        Sy        Sz        Q        Lx        Ly        Lz        Jx        Jy        Jz"
    )
    print(
    )
    # iterate over magnetic entities
        # iterate over atoms
        for tag, xyz in zip(mag_ent["tags"], mag_ent["xyz"]):
            # coordinates and tag
            print(f"{tag}        {xyz[0]}        {xyz[1]}        {xyz[2]}")
        print("")

    print(
    )
    print("Anisotropy [meV]")
    print(
    )
    print("Magnetic entity        x [Ang]        y [Ang]        z [Ang]")
    print(
    )
            print("Consistency check: ", mag_ent["K_consistency"])
            print("Anisotropy diag: ", mag_ent["K"])
        print("")

    print(
    )
    print("Exchange [meV]")
    print(
    )
    print("Magnetic entity1        Magnetic entity2       [i  j  k]       d [Ang]")
    print(
    )
    # iterate over pairs
        # print pair parameters
        print(
            f"{pair['tags'][0]}        {pair['tags'][1]}        {pair['Ruc']}        d [Ang] {pair['dist']}"
        )
        # print magnetic parameters
        print("Isotropic: ", pair["J_iso"], "        # Tr[J] / 3")
        print("")
        print("DMI: ", pair["D"], "        # Dx, Dy, Dz")
        print("")
        print(
            "Symmetric-anisotropy: ",
            pair["J_S"],
            "        # J_S = J - J_iso * I ––> Jxx, Jyy, Jxy, Jxz, Jyz",
        )
        print("")
        print("J:        # Jxx, Jxy, Jxz, Jyx, Jyy, Jyz, Jzx, Jzy, Jzz")
        print(pair["J"])
        print(
            "----------------------------------------------------------------------------------------------------------------------------------------------------------------"
        )

    print(
    )

def print_runtime_information(times):
    """It prints the runtime information for the grogu out.

    Args:
        times : dict
            It contains the runtime data
    """

    print("Runtime information: ")
    print(f"Total runtime: {times['end_time'] - times['start_time']} s")
    print(
    )
    print(f"Initial setup: {times['setup_time'] - times['start_time']} s")
    print(
        f"Hamiltonian conversion and XC field extraction: {times['H_and_XCF_time'] - times['setup_time']:.3f} s"
    )
    print(
        f"Pair and site datastructure creatrions: {times['site_and_pair_dictionaries_time'] - times['H_and_XCF_time']:.3f} s"
    )
    print(
        f"k set cration and distribution: {times['k_set_time'] - times['site_and_pair_dictionaries_time']:.3f} s"
    )
    print(
        f"Rotating XC potential: {times['reference_rotations_time'] - times['k_set_time']:.3f} s"
    )
    print(
        f"Greens function inversion: {times['green_function_inversion_time'] - times['reference_rotations_time']:.3f} s"
    )
    print(
        f"Calculate energies and magnetic components: {times['end_time'] - times['green_function_inversion_time']:.3f} s"
    )

def print_job_description(simulation_parameters):
    """It prints the parameters and the description of the job.

    Args:
    """

    print(
    )
    print(
    )
    if simulation_parameters["parallel_solver_for_Gk"]:
        print("solver used for Greens function calculation: parallel")
    else:
        print("solver used for Greens function calculation: sequential")
    print(
    )
    print(
    )
    print(
    )
    if simulation_parameters["automatic_ebot"]:
        print(
            "Ebot: ",
            simulation_parameters["ebot"],
            "        WARNING: This was automatically determined!",
        )
    else:
        print("Ebot: ", simulation_parameters["ebot"])
    print(
    )

#
#
#

def blow_up_orbindx(orb_indices):
    """Function to blow up orbital indices to make SPIN BOX indices.

    Args:
        orb_indices : np.array_like
            These are the indices in ORBITAL BOX

    Returns:
            These are the indices in SPIN BOX
    """

    orb_indices = np.array([[2 * o, 2 * o + 1] for o in orb_indices]).flatten()

    return orb_indices

def spin_tracer(M):
    """Spin tracer utility.

    This takes an operator with the orbital-spin sequence:
    orbital 1 up,
    orbital 1 down,
    orbital 2 up,
    orbital 2 down,
    that is in the SPIN-BOX representation,
    and extracts orbital dependent Pauli traces.

    Args:
        M : np.array_like
            Traceable matrix

    Returns:
        dict
            It contains the traced matrix with "x", "y", "z" and "c"
    """
    M11 = M[0::2, 0::2]
    M12 = M[0::2, 1::2]
    M21 = M[1::2, 0::2]
    M22 = M[1::2, 1::2]

    M_o = dict()
    M_o["x"] = M12 + M21
    M_o["y"] = 1j * (M12 - M21)
    M_o["z"] = M11 - M22
    M_o["c"] = M11 + M22

    return M_o

def parse_magnetic_entity(dh, atom=None, l=None, **kwargs):
    """Function to define orbital indexes of a given magnetic entity.

    Args:
            Hamiltonian from sisl
        atom : integer or list of integers, optional
            Defining atom (or atoms) in the unit cell forming the magnetic entity. Defaults to None
        l : integer, optional
            Defining the angular momentum channel. Defaults to None

    Returns:
        list
            The orbital indexes of the given magnetic entity
    """
    # case where we deal with more than one atom defining the magnetic entity
    if type(atom) == list:
        dat = []
        for a in atom:
            a_orb_idx = dh.geometry.a2o(a, all=True)
            if (
                type(l) == int
            ):  # if specified we restrict to given l angular momentum channel inside each atom
                a_orb_idx = a_orb_idx[[o.l == l for o in dh.geometry.atoms[a].orbitals]]
            dat.append(a_orb_idx)
        orbital_indeces = np.hstack(dat)
    # case where we deal with a singel atom magnetic entity
    elif type(atom) == int:
        orbital_indeces = dh.geometry.a2o(atom, all=True)
        if (
            type(l) == int
        ):  # if specified we restrict to given l angular momentum channel
            orbital_indeces = orbital_indeces[
                [o.l == l for o in dh.geometry.atoms[atom].orbitals]
            ]

    return orbital_indeces  # numpy array containing integers labeling orbitals associated to a magnetic entity.

def calculate_anisotropy_tensor(mag_ent):
    """Calculates the renormalized anisotropy tensor from the energies.

    It uses the grogu convention for output.

    Args:
        mag_ent : dict
            An element from the magnetic entities

    Returns:
        K : np.array_like
            elements of the anisotropy tensor
    """

    # get the energies
    energies = mag_ent["energies"]

    # calculate the diagonal tensor elements
    Kxx = energies[1, 1] - energies[1, 0]
    Kyy = energies[0, 1] - energies[0, 0]
    Kzz = 0

    # perform consistency check
    calculated_diff = Kyy - Kxx
    expected_diff = energies[2, 0] - energies[2, 1]
    consistency_check = abs(calculated_diff - expected_diff)

    return Kxx, Kyy, Kzz, consistency_check

def calculate_exchange_tensor(pair):
    """Calculates the exchange tensor from the energies.

    It produces the isotropic exchange, the relevant elements
    from the Dzyaloshinskii-Morilla (Dm) tensor, the symmetric-anisotropy
    and the complete exchange tensor.

    Args:
        pair : dict
            An element from the pairs

    Returns:
        J_iso : float
            Isotropic exchange (Tr[J] / 3)
        J_S : np.array_like
            Symmetric-anisotropy (J_S = J - J_iso * I ––> Jxx, Jyy, Jxy, Jxz, Jyz)
        D : np.array_like
            DM elements (Dx, Dy, Dz)
        J : np.array_like
            Complete exchange tensor flattened (Jxx, Jxy, Jxz, Jyx, Jyy, Jyz, Jzx, Jzy, Jzz)
    """

    # energies from rotations
    energies = pair["energies"]
    # Initialize output arrays
    J = np.zeros((3, 3))
    D = np.zeros(3)

    # J matrix calculations
    # J(1,1) = mean([DEij(2,2,2), DEij(2,2,3)])
    J[0, 0] = np.mean([energies[1, 3], energies[2, 3]])

    # J(1,2) = -mean([DEij(1,2,3), DEij(2,1,3)])
    J[0, 1] = -np.mean([energies[2, 1], energies[2, 2]])
    J[1, 0] = J[0, 1]

    # J(1,3) = -mean([DEij(1,2,2), DEij(2,1,2)])
    J[0, 2] = -np.mean([energies[1, 1], energies[1, 2]])
    J[2, 0] = J[0, 2]

    # J(2,2) = mean([DEij(2,2,1), DEij(1,1,3)])
    J[1, 1] = np.mean([energies[0, 3], energies[2, 0]])

    # J(2,3) = -mean([DEij(1,2,1), DEij(2,1,1)])
    J[1, 2] = -np.mean([energies[0, 1], energies[0, 2]])
    J[2, 1] = J[1, 2]

    # J(3,3) = mean([DEij(1,1,1), DEij(1,1,2)])
    J[2, 2] = np.mean([energies[0, 0], energies[1, 0]])

    # D vector calculations
    # D(1) = mean([DEij(1,2,1), -DEij(2,1,1)])
    D[0] = np.mean([energies[0, 1], -energies[0, 2]])

    # D(2) = mean([DEij(2,1,2), -DEij(1,2,2)])
    D[1] = np.mean([energies[1, 2], -energies[1, 1]])

    # D(3) = mean([DEij(1,2,3), -DEij(2,1,3)])
    D[2] = np.mean([energies[2, 1], -energies[2, 2]])

    J_iso = np.trace(J) / 3
    # based on the grogu output pdf
    # traceless symmetric exchange matrix:
    # Jxx, Jyy, Jxy, Jxz, Jyz
    J_S = np.array([J[0, 0] - J_iso, J[1, 1] - J_iso, J[0, 1], J[0, 2], J[1, 2]])

    return J_iso, J_S, D, J

#
#
#

from scipy.special import roots_legendre
from sisl.io.siesta import eigSileSiesta

# Pauli matrices
tau_x = np.array([[0, 1], [1, 0]])
tau_y = np.array([[0, -1j], [1j, 0]])
tau_z = np.array([[1, 0], [0, -1]])
tau_0 = np.array([[1, 0], [0, 1]])

def commutator(a, b):
    """Shorthand for commutator.

    Commutator of two matrices in the mathematical sense.

    Args:
        a : np.array_like
            The first matrix
        b : np.array_like
            The second matrix

    Returns:
            The commutator of a and b
    """

    return a @ b - b @ a

# define some useful functions
def hsk(H, ss, sc_off, k=(0, 0, 0)):
    """Speed up Hk and Sk generation.

    Calculates the Hamiltonian and the Overlap matrix at a given k point. It is faster that the sisl version.

    Args:
        H : np.array_like
            Hamiltonian in spin box form
        ss : np.array_like
            Overlap matrix in spin box form
        sc_off : list
            supercell indexes of the Hamiltonian
        k : tuple, optional
            The k point where the matrices are set up. Defaults to (0, 0, 0)

    Returns:
            Hamiltonian at the given k point
            Overlap matrix at the given k point
    """

    # this two conversion lines are from the sisl source
    k = np.asarray(k, np.float64)
    k.shape = (-1,)

    # this generates the list of phases
    phases = np.exp(-1j * 2 * np.pi * k @ sc_off.T)

    # phases applied to the hamiltonian
    HK = np.einsum("abc,a->bc", H, phases)
    SK = np.einsum("abc,a->bc", ss, phases)

    return HK, SK

def make_kset(dirs="xyz", NUMK=20):
    """Simple k-grid generator to sample the Brillouin zone.

    Depending on the value of the dirs
    argument k sampling in 1,2 or 3 dimensions is generated.
    If dirs argument does not contain either of x, y or z
    a kset of a single k-pont at the origin is returned. The total number of k points is the NUMK**(dimensions)

    Args:
        dirs : str, optional
            Directions of the k points in the Brillouin zone. They are the three lattice vectors. Defaults to "xyz"
        NUMK : int, optional
            The number of k points in a direction. Defaults to 20

    Returns:
            An array of k points that uniformly sample the Brillouin zone in the given directions
    """

    # if there is no xyz in dirs return the Gamma point
    if not (sum([d in dirs for d in "xyz"])):
        return np.array([[0, 0, 0]])

    kran = len(dirs) * [np.linspace(0, 1, NUMK, endpoint=False)]
    mg = np.meshgrid(*kran)
    dirsdict = dict()

    for d in enumerate(dirs):
        dirsdict[d[1]] = mg[d[0]].flatten()
    for d in "xyz":
        if not (d in dirs):
            dirsdict[d] = 0 * dirsdict[dirs[0]]
    kset = np.array([dirsdict[d] for d in "xyz"]).T

    return kset

def make_contour(emin=-20, emax=0.0, enum=42, p=150):
    """A more sophisticated contour generator.

    Calculates the parameters for the complex contour integral. It uses the
    Legendre-Gauss quadrature method.   It returns a class that contains
    the information for the contour integral.

    Args:
        emin : int, optional
            Energy minimum of the contour. Defaults to -20
        emax : float, optional
            Energy maximum of the contour. Defaults to 0.0, so the Fermi level
        enum : int, optional
            Number of sample points along the contour. Defaults to 42
        p : int, optional
            Shape parameter that describes the distribution of the sample points. Defaults to 150

    Returns:
        ccont
            Contains all the information for the contour integral
    """
    x, wl = roots_legendre(enum)
    R = (emax - emin) / 2  # radius
    z0 = (emax + emin) / 2  # center point
    y1 = -np.log(1 + np.pi * p)  # lower bound
    y2 = 0  # upper bound

    y = (y2 - y1) / 2 * x + (y2 + y1) / 2
    phi = (np.exp(-y) - 1) / p  # angle parameter
    ze = z0 + R * np.exp(1j * phi)  # complex points for path
    we = -(y2 - y1) / 2 * np.exp(-y) / p * 1j * (ze - z0) * wl

    # just an empty container class
    class ccont:
        pass

    cont = ccont()
    cont.R = R
    cont.z0 = z0
    cont.ze = ze
    cont.we = we
    cont.enum = enum

    return cont

def tau_u(u):
    """Pauli matrix in direction u.

    Returns the vector u in the basis of the Pauli matrices.

    Args:
        u : list or np.array_like
            The direction

    Returns:
            Arbitrary direction in the base of the Pauli matrices
    """

    # u is force to be of unit length
    u = u / np.linalg.norm(u)

    return u[0] * tau_x + u[1] * tau_y + u[2] * tau_z

#
def crossM(u):
    """Definition for the cross-product matrix.

    It acts as a cross product with vector u.

    Args:
            The second vector in the cross product

    Returns:
            The matrix that represents teh cross product with a vector
    """
    return np.array([[0, -u[2], u[1]], [u[2], 0, -u[0]], [-u[1], u[0], 0]])

def RotM(theta, u, eps=1e-10):
    """Definition of rotation matrix with angle theta around direction u.

    Args:
        theta : float
            The angle of rotation
        u : np.array_like
            The rotation axis
        eps : float, optional
            Cutoff for small elements in the resulting matrix. Defaults to 1e-10

    Returns:
            The rotation matrix
    """

    M = (
        np.cos(theta) * np.eye(3)
        + np.sin(theta) * crossM(u)
        + (1 - np.cos(theta)) * np.outer(u, u)
    )

    # kill off small numbers
    M[abs(M) < eps] = 0.0
    return M

def RotMa2b(a, b, eps=1e-10):
    """Definition of rotation matrix rotating unit vector a to unit vector b.

    Function returns array R such that R @ a = b holds.

    Args:
            First vector
            Second vector

    Returns:
            The rotation matrix with the above property
    """

    v = np.cross(a, b)
    c = a @ b
    M = np.eye(3) + crossM(v) + crossM(v) @ crossM(v) / (1 + c)

    return M

def read_siesta_emin(eigfile):
    """It reads the lowest energy level from the siesta run.

    It uses the .EIG file from siesta that contains the eigenvalues.

    Args:
        eigfile : str
            The path to the .EIG file

    Returns:
        float
            The energy minimum
    """

    # read the file
    eigs = eigSileSiesta(eigfile).read_data()

    return eigs.min()

def int_de_ke(traced, we):
    """It numerically integrates the traced matrix.

    It is a wrapper from numpy.trapz and it contains the
    relevant constants to calculate the energy integral from
    equation 93 or 96.

    Args:
        traced : np.array_like
            The trace of a matrix or a matrix product
        we : float
            The weight of a point on the contour

    Returns:
        float
            The energy calculated from the integral formula
    """

    return np.trapz(-1 / np.pi * np.imag(traced * we))