main.py

#!/usr/bin/env python
# TODO: Use Argparse for command line arguments and --help
# """
# Usage:
#     python main.py [config file] [number of processes]

# Everything after the second command line argument will be ignored. If
# the number of processes is not specified, the ADGA calculation will be
# done with a single process. If no config file is specified, the default
# value "jackknife.ini" will be used. For an example of such a config file
# see `example_config.ini`.
# """

import jackknife as jk
import adga
import numpy as np
import sys
import h5py
import configparser
from distutils.util import strtobool


def main():
    """Do jackknife resampling for ADGA calculations.
    
    Determine the config file and number of processe to run ADGA with
    from the command line arguments. If all or some arguments are
    missing, use the default values. They are "jackknife.ini" for
    the config file and "1" for the number of processes. Read the config
    file and check if obligatory keys are missing. If not, perform a
    jackknife estimation of the quantity chosen in the config file. The
    results as well as some useful meta data are stored in an HDF5 file.
    
    .. note::
        Right now it is only possible to calculate either the
        self-energy *or* the susceptibilities. If both
        ``calculate_self_energy`` and ``calculate_susceptibilities`` are
        set to "yes" or "true" in the config file, only the self-energy
        is calculated.
    
    :raise ValueError: if an obligatory key is missing in the config
                       file
    """
    
    config_file_name = "jackknife.ini"
    n_processes = 1
    if len(sys.argv) == 1:
        print("No configuration file specified. Default is " +
              config_file_name)
    if len(sys.argv) >= 2:
        config_file_name = sys.argv[1]
    if len(sys.argv) >= 3:
        n_processes = int(sys.argv[2])
    if len(sys.argv) > 3:
        print("""Too many command line arguments. Please specify only the
              jackknife configuration file and the number of processes.
              Everything after that will be ignored.""")

    # Define the dictionary for the [General] section of the INI file.
    # Obligatory keys are set to `None`.
    # TODO: Enable parameter lists (especially for files)
    general = {"adga_root_directory": None,
               "adga_config_file": None,
               "two_particle_file": None,
               "output_file_prefix": "jackknife",
               "store_output_samples": "no",
               "calculate_self_energy": "yes",
               "calculate_susceptibilities": "no"}

    # Read config file and check if obligatory keys are missing
    config = configparser.ConfigParser()
    config.read(config_file_name)
    for key in config["General"]:
        general[key] = config["General"][key]
    for key in general:
        if general[key] is None:
            raise ValueError("Error while parsing conifg file.\n"
                             + "Obligatory key \"{}\" not found.".format(key))

    abinitio_dga = adga.Adga(general["adga_root_directory"],
                             general["adga_config_file"],
                             general["two_particle_file"], n_processes)
    x_generator = abinitio_dga.get_worm_sample_generator()
    n = abinitio_dga.n_worm_samples
    if strtobool(general["calculate_self_energy"]):
        f = abinitio_dga.calculate_self_energy
    elif strtobool(general["calculate_susceptibilities"]):
        f = abinitio_dga.calculate_susceptibilities
    else:
        print("Since you don't want to calculate anything we're done.")
        sys.exit()

    # Do the Jackknife estimation
    jackknife = jk.Jackknife(x_generator, n, f, -1,
                             general["output_file_prefix"],
                             strtobool(general["store_output_samples"]))
    jackknife.do_estimation()
    jackknife.write_results_to_file()

    # Add some useful metadata to the output file
    try:
        with h5py.File(jackknife.output_file_name, "r+") as output_file:
            config = output_file[".config"].attrs
            config.create("qmc.nmeas", abinitio_dga.n_qmc_measurements)
            config.create("adga_root_directory",
                          np.string_(general["adga_root_directory"]))
            config.create("adga_config_file",
                          np.string_(general["adga_config_file"]))
            config.create("two_particle_file",
                          np.string_(general["two_particle_file"]))
            config.create("n_processes", n_processes)
            config.create("store_output_samples",
                          bool(strtobool(general["store_output_samples"])))
            config.create("calculate_self_energy",
                          bool(strtobool(general["calculate_self_energy"])))
            config.create("calculate_susceptibilities",
                          bool(strtobool(
                              general["calculate_susceptibilities"])))
    except Exception:
        print("""-- Something went wrong when adding metadata to the output
              file.""")
        raise
    return

if __name__ == "__main__":
    main()