ldos_example.py

import argparse
import os
import signal
from mpi4py import MPI
import time
from functools import partial
import json
from loguru import logger

from qcat.pp import LDOS
from qcat.io_kernel import qe_io, qbox_io
from qcat.utils import utils
from qcat.utils import setLogger

setLogger()

comm = MPI.COMM_WORLD

if __name__ == "__main__":

    signal.signal(signal.SIGINT, partial(utils.handler, comm))
    # ----------------------------------Prepare------------------------------------
    rank = comm.Get_rank()
    if rank == 0:
        utils.time_now()

    parser = argparse.ArgumentParser()
    parser.add_argument("-a", "--abinitio", type=str,
            help="abinitio software: qe/qbox. Default: qe")
    parser.add_argument("-s", "--saveFileFolder", type=str,
            help="*.save folder generated by QE or the folder that store qbox output. Default: ./scf.save")
    parser.add_argument("-d", "--delta", type=float,
            help="delta that control local VB/CB. Default: 0.001")
    parser.add_argument("-o", "--orientation", type=str,
            help="LDOS orientation. Default: z")
    args = parser.parse_args()

    # default values
    if not args.abinitio:
        args.abinitio = "qe"
    if not args.saveFileFolder:
        args.saveFileFolder = "./scf.save"
    assert(os.path.exists(args.saveFileFolder))
    if not args.delta:
        args.delta = 0.001 
    if not args.orientation:
        args.orientation = "z"

    conf_tab = {"software": args.abinitio,
                "saveFileFolder": args.saveFileFolder,
                "delta": args.delta,
                "orientation": args.orientation,
                "MPI size": comm.Get_size()}
    utils.print_conf(conf_tab)

    # ---------------------------------Compute LDOS------------------------------------

    rank = comm.Get_rank()
    st = time.time()

    # initialize reader instance
    abinitioRead = None
    if args.abinitio.lower() == "qbox":
        abinitioRead = qbox_io.QBOXRead(args.saveFileFolder, comm)
    elif args.abinitio.lower() == "qe":
        abinitioRead = qe_io.QERead(args.saveFileFolder, comm)

    # delta     : parameter to control the location of LCBM/LVBM (checkout eq.3 of https://doi.org/10.1063/1.4811481)
    # comm      : MPI COMM_WORLD
    localDensityOfState = LDOS(read_obj=abinitioRead,
                               delta=args.delta,
                               comm=comm)

    # compute LDOS along z orientation (can be 'x' or 'y')
    localDensityOfState.computeLDOS(axis=args.orientation)

    # obtain lcbm and lvbm
    # lcbm: local CBM
    # lvbm: local VBM
    # lcbm and lvbm are 1d numpy.ndarray with size of FFT grid along chosen axis (z in this case).
    lcbm, lvbm = localDensityOfState.localBandEdge()

    if rank == 0:
        # write lcbm and lvbm to json file
        data = {'lcbm': lcbm.tolist(), 'lvbm': lvbm.tolist()}
        fname = "ldos.json"
        with open(fname, 'w') as json_file:
            json.dump(data, json_file)
        logger.info(f"Data has been written to {fname}")

    # get the execution time
    # get the end time
    et = time.time()
    elapsed_time = et - st
    comm.Barrier()
    if rank == 0:
        print('Execution time:', elapsed_time, 'seconds')