adga.py

import h5py
from shutil import copyfile
import numpy as np
import subprocess
import os
import time


# TODO: Exceptions
class Adga(object):
    """A class for abinitio DGA calculations.

    This class uses the ADGA project to calculate the non-local self-energy in
    the dynamical vertex approximation. It uses DMFT results from w2dynamics as
    input. Multiple worm samples in the input file are supported.

    The `jackknife` branch of the ADGA repository must be used. The
    location of the abinitiodga executable and the symmetrize script are
    assumed to be at `adga_root_path`/bin/abinitiodga and
    `adga_root_path`/scripts/symmetrize.py respectively.
    """

    def __init__(self, adga_root_path, adga_config_file_name,
                 two_particle_greens_function_file_name, n_processes=1):
        self.n_processes = n_processes
        self.executable = adga_root_path + "/bin/abinitiodga"
        self.symmetrize_script = adga_root_path + "/scripts/symmetrize.py"
        self.config_file_name = adga_config_file_name
        # Generate and a unique id for the names of the temporary files.
        # This allows to run multiple Jackknives at once.
        self.unique_id = int(time.time_ns() * 1e3)
        self.output_file_name = "adga{}.hdf5".format(self.unique_id)

        self.g2_file = h5py.File(two_particle_greens_function_file_name, "r")
        g4iw_file_name = "g4iw{}.hdf5".format(self.unique_id)
        copyfile(two_particle_greens_function_file_name, g4iw_file_name)
        self.g4iw_file = h5py.File(g4iw_file_name, "r+")

        self.worm_groups = [s + "/ineq-001/g4iw-worm/"
                            for s in self.g2_file.keys()
                            if (("worm-" in s) and (s != "worm-last"))]
        config = self.g2_file[".config"]
        self.n_meas = config.attrs["qmc.nmeas"]
        self.n_worm_samples = config.attrs["general.wormsteps"]
        self.compound_indexes = list(self.g2_file[self.worm_groups[0]].keys())

    def get_worm_sample_generator(self):
        """Return a generator yielding green's functions from different
        worm samples."""
        def worm_sample_generator():
            for worm_group in self.worm_groups:
                value_groups = [worm_group + i + "/value"
                                for i in self.compound_indexes]
                yield np.array([self.g2_file[value_group][...]
                                for value_group in value_groups])
        return worm_sample_generator

    # TODO: "Rename" this into calculate_self_energy()
    # TODO: Add calculate_susceptibility()
    def __call__(self, two_particle_greens_function):
        """Return the self-energy of the 2-particle Green's function.

        Write the 2-particle Green's function to the file g4iw.hdf5,
        which is used as an input file for the symmetrize.py script.
        Then do the abinitio DGA calculation and return the non-local
        self-energy as a 3D numpy array depending on k_x, k_y and
        omega."""

        g2 = two_particle_greens_function
        for i in range(len(self.compound_indexes)):
            data_set = self.g4iw_file["worm-001/ineq-001/g4iw-worm/" +
                                      self.compound_indexes[i] + "/value"]
            data_set[...] = g2[i]
        self.g4iw_file.flush()

        # TODO: Make a subroutine for symmetrizing g2
        # g4iw_sym.hdf5 must not exist when calling symmetrize.py
        symmetry_file = "g4iw{}_sym.hdf5".format(self.unique_id)
        if os.path.exists(symmetry_file):
            os.remove(symmetry_file)

        # Symmetrize the 2-particle Green's function
        print(self.symmetrize_script + " worm-001 " + self.g4iw_file.filename)
        subprocess.call([self.symmetrize_script, "worm-001",
                         self.g4iw_file.filename])

        # TODO: Make a subroutine for the whole DGA calculation
        # TODO: Values should be added if not existent
        # Make sure, abinitiodga uses the correct 2-particle and output
        # file
        self._replace_value_in_config_file("2PFile", symmetry_file)
        self._replace_value_in_config_file("Outfile", self.output_file_name)

        # Do the DGA calculation
        print("mpirun" + " -np " + str(self.n_processes) + self.executable
              + " " + self.config_file_name)
        subprocess.call(["mpirun", "-np", self.n_processes, self.executable,
                         self.config_file_name])

        # TODO: Make a subroutine for getting the self-energy
        output_file = h5py.File(self.output_file_name, "r")
        self_energy = np.array(output_file["selfenergy/nonloc/dga"][...])
        output_file.close()
        return self_energy[0, 0, :, :, 0, :]

    def _replace_value_in_config_file(self, key, value):
        # Replace the value of the key, if it exists in the config file.
        file_name = self.config_file_name
        lines = open(file_name).readlines()
        for line in lines:
            if line.startswith(key):
                old_line = line

        text = open(file_name).read()
        new_line = key + " = " + value + "\n"
        text = text.replace(old_line, new_line)
        f = open(file_name, "w")
        f.write(text)
        f.close()