grogu/src/grogupy/io.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.

from argparse import ArgumentParser
from pickle import dump, load

import numpy as np

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,
    automatic_ebot=False,
    eset=42,
    esetp=1000,
    calculate_charge=False,
    charges=[],
    parallel_solver_for_Gk=False,
    parallel_size=None,
    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):
    """_summary_

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


def load_pickle(infile, data):
    """_summary_

    Args:
        infile (_type_): _description_
        data (_type_): _description_

    Returns:
        _type_: _description_
    """
    with open(infile, "wb") as input_file:
        data = load(data, input_file)

    return data


def print_parameters(simulation_parameters):
    """_summary_

    Args:
        simulation_parameters (_type_): _description_
    """
    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(
        "================================================================================================================================================================"
    )
    if simulation_parameters["calculate_charge"]:
        print("The calculated charge of the Hamiltonian in the quantization axes: ")
        print(simulation_parameters["charges"])
        print(
            "================================================================================================================================================================"
        )


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

    Args:
        magnetic_entities (_type_): _description_
        pairs (_type_): _description_
    """
    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
    for mag_ent in 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(
        "----------------------------------------------------------------------------------------------------------------------------------------------------------------"
    )
    # iterate over magnetic entities
    for mag_ent in 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("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
    for pair in 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):
    """_summary_

    Args:
        times (_type_): _description_
    """
    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):
    """_summary_

    Args:
        simulation_parameters (_type_): _description_
    """

    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"],
    )
    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("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"])
    if simulation_parameters["automatic_ebot"]:
        print(
            "Ebot: ",
            simulation_parameters["ebot"],
            "        WARNING: This was automatically determined!",
        )
    else:
        print("Ebot: ", simulation_parameters["ebot"])
    print("Eset: ", simulation_parameters["eset"])
    print("Esetp: ", simulation_parameters["esetp"])
    print(
        "================================================================================================================================================================"
    )
    if simulation_parameters["calculate_charge"]:
        print(
            "Charge calculation required: True        WARNING: This causes a slowdown, because we have to use the complete Greens function!"
        )
    else:
        print("Charge calculation required: False")
    print(
        "================================================================================================================================================================"
    )