adga.py

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


# TODO: Exceptions
# TODO: Documentation
# TODO: Document member variables
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.

    .. note::
        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.
    
    .. _ADGA project: https://github.com/AbinitioDGA/ADGA
    .. _w2dynamics: https://github.com/w2dynamics/w2dynamics
    """

    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"

        # Generate a unique id for the names of the temporary files.
        # Together with a unique copy of the config file, this allows to
        # run multiple jackknives at once.
        self._unique_id = int(time.time() * 1e6)
        self._output_file_name = "adga{}.hdf5".format(self._unique_id)
        g4iw_file_name = "g4iw{}.hdf5".format(self._unique_id)
        self._symmetry_file_name = "g4iw{}_sym.hdf5".format(self._unique_id)
        self._config_file_name = adga_config_file_name.split(".")[0] \
            + str(self._unique_id) + ".conf"
        copyfile(adga_config_file_name, self._config_file_name)

        self._g2_file = h5py.File(two_particle_greens_function_file_name, "r")
        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_qmc_measurements = config.attrs["qmc.nmeas"]
        self._n_worm_samples = config.attrs["general.wormsteps"]
        self._compound_indexes = list(
            self._g2_file[self._worm_groups[0]].keys())

    @property
    def n_qmc_measurements(self):
        return self._n_qmc_measurements

    @property
    def n_worm_samples(self):
        return self._n_worm_samples

    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

    def calculate_self_energy(self, two_particle_greens_function):
        """Return the self-energy on the 2-particle level.

        Symmetrize the given 2-particle Green's function. Then do the
        abinitio DGA calculation and return the non-local self-energy as
        a 3D numpy array depending on :math:`k_x,\\ k_y` and
        :math:`\\nu`."""

        self._symmetrize_g2(two_particle_greens_function)
        self._do_dga_calculation()
        return self._get_self_energy()

    def calculate_susceptibilities(self, two_particle_greens_function):
        """Return the susceptibilities on the 2-particle level.

        Symmetrize the given 2-particle Green's function. Then do the
        abinitio DGA calculation and return the non-local
        susceptibilities in 3 channels as a 4D numpy array depending on
        :math:`k_x, k_y` and :math:`\\nu`. The first index selects the
        channel: bubble_nl, density or magnetic."""

        self._symmetrize_g2(two_particle_greens_function)
        self._do_dga_calculation()
        return self._get_susceptibilities()

    def _symmetrize_g2(self, g2):
        # Copy the two-particle Green's function to the g4iw file and
        # call symmetrize.py worm-001 g4iw_file.filename
        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()

        # g4iw{}_sym.hdf5 must not exist when calling symmetrize.py
        if os.path.exists(self._symmetry_file_name):
            os.remove(self._symmetry_file_name)

        print(self._symmetrize_script + " worm-001 " + self._g4iw_file.filename)
        subprocess.call([self._symmetrize_script, "worm-001",
                         self._g4iw_file.filename])
        return

    def _do_dga_calculation(self):
        # Copy the config file, add the unique ID and set the correct
        # 2-particle and output file. This allows to run multiple
        # jackknives at the same time without interference.
        self._set_value_in_config_file("Two-Particle", "2PFile",
                                       self._symmetry_file_name)
        self._set_value_in_config_file("General", "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", str(self.n_processes),
                         self._executable, self._config_file_name])
        return

    def _set_value_in_config_file(self, section, key, value):
        # Replace or add the key-value pair in the given section of the
        # config file
        file_name = self._config_file_name
        new_line = key + " = " + value + "\n"

        lines = open(file_name).readlines()
        line_is_in_section = False
        for line in lines:
            section_is_found = line.startswith("[" + section + "]")
            key_is_found = line_is_in_section and line.startswith(key)
            end_of_section_is_found = line and line_is_in_section \
                and line[0] == "[" and line[1].isalpha()
            if section_is_found:
                line_is_in_section = True
                continue
            elif key_is_found:
                # Replace the key-value pair with the new one
                lines[lines.index(line)] = new_line
                break
            elif end_of_section_is_found:
                # Add the key-value pair at the end of the section
                lines.insert(lines.index(line) - 1, new_line)
                break

        with open(file_name, "w") as config_file:
            for line in lines:
                config_file.write(line)
        return

    def _get_self_energy(self):
        # Return the non-local DGA self-energy from the output file as a
        # 3D numpy array
        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 _get_susceptibilities(self):
        # Return the non-local susceptibility from the output file as a
        # 4D numpy array. The first index selects the channel:
        # bubble_nl, density or magnetic.
        output_file = h5py.File(self._output_file_name, "r")
        susceptibilities = []
        susceptibilities.append(
            output_file["susceptibility/nonloc/bubble_nl"][...])
        susceptibilities.append(output_file["susceptibility/nonloc/dens"][...])
        susceptibilities.append(output_file["susceptibility/nonloc/magn"][...])
        output_file.close()
        susceptibilities = np.array(susceptibilities)
        return susceptibilities[:, 0, 0, :, :, 0, :]