hrstm_from_wfn.py

#!/usr/bin/env python

# @author  Hillebrand, Fabian
# @date    2019
# @version 2.1.0

import argparse  # Terminal arguments
import os
import sys  # System access
import time  # Timing

# Python3 at least
assert sys.version_info >= (3, 0)

import numpy as np

ang2bohr = 1.88972612546
ev2hartree = 0.03674930814

# Own includes
# import hrstm_tools.ppstm_grid_orbitals as pgo
from mpi4py import MPI

import cp2k_spm_tools.cp2k_grid_orbitals as cgo
import hrstm_tools.cp2k_grid_matrix as cgm
import hrstm_tools.hrstm as hs
import hrstm_tools.hrstm_utils as hu
import hrstm_tools.tip_coeffs as tc
from cp2k_spm_tools.cube import Cube

mpi_comm = MPI.COMM_WORLD
mpi_size = mpi_comm.Get_size()
mpi_rank = mpi_comm.Get_rank()


parser = argparse.ArgumentParser(
    description="HR-STM for CP2K based on Chen's" + " Derivative Rule and the Probe Particle Model."
)

### ----------------------------------------------------------------------------
### Output file
parser.add_argument(
    "--output",
    type=str,
    metavar="OUTPUT",
    default="hrstm",
    required=False,
    help="Name for output file. File extension added automatically.",
)

### ----------------------------------------------------------------------------
### Input files for sample
parser.add_argument(
    "--cp2k_input_file", metavar="FILE", required=True, help="CP2K input file used for the sample calculation."
)

parser.add_argument(
    "--basis_set_file",
    metavar="FILE",
    required=True,
    help="File containing all basis sets used for sample calculation.",
)

parser.add_argument(
    "--xyz_file", metavar="FILE", required=True, help="File containing the atom positions for the sample."
)

parser.add_argument(
    "--wfn_file", metavar="FILE", required=True, help="Restart file containing the basis coefficients for the sample."
)

parser.add_argument("--hartree_file", metavar="FILE", required=True, help="Cube file with Hartree potential of sample.")

parser.add_argument(
    "--fwhm_sam",
    type=float,
    metavar="eV",
    default=0.05,
    required=False,
    help="Full width at half maximum for Gaussian broadening of DOS for sample.",
)

parser.add_argument(
    "--wn",
    type=float,
    metavar="eV",
    default=5.0,
    required=False,
    help="Work function for the sample. Used for the decay constant.",
)

### ----------------------------------------------------------------------------
### Energy range for sample
parser.add_argument(
    "--emin",
    type=float,
    metavar="eV",
    default=-2.0,
    required=False,
    help="Lower energy used for cut-off with respect to Fermi energy for sample.",
)

parser.add_argument(
    "--emax",
    type=float,
    metavar="eV",
    default=2.0,
    required=False,
    help="Upper energy used for cut-off with respect to Fermi energy for sample.",
)

### ----------------------------------------------------------------------------
### Parameters for putting sample orbitals on grid
parser.add_argument(
    "--dx_wfn",
    type=float,
    metavar="Ang",
    default=0.2,
    required=False,
    help="Spacing for wavefunction grid used by interpolation.",
)

parser.add_argument(
    "--rcut",
    type=float,
    metavar="Ang",
    default=15.0,
    required=False,
    help="Cut-off radius used when computing sample wavefunction.",
)

parser.add_argument(
    "--extrap_dist",
    type=float,
    metavar="Ang",
    default=0.5,
    required=False,
    help="Starting distance from highest atom where extrapolation is used from.",
)

### ----------------------------------------------------------------------------
### Parameters for tip
parser.add_argument(
    "--tip_pos_files",
    metavar="FILE",
    nargs="+",
    required=False,
    help="File paths to positions of probe particles (without '_<x,y,z>.npy').",
)

parser.add_argument(
    "--tip_shift",
    type=float,
    metavar="Ang",
    required=False,
    help="z-distance shift for metal tip with respect to apex probe particle." + " Only used with relaxed positions.",
)

parser.add_argument(
    "--eval_region",
    type=float,
    metavar="Ang",
    nargs="+",
    required=False,
    help="Tip positions in case of non-relaxed positions.",
)

parser.add_argument(
    "--dx_tip", type=float, metavar="Ang", required=False, help="Spacing for a uniform grid for non-relaxed positions."
)

parser.add_argument(
    "--pdos_list",
    metavar="FILE",
    nargs="+",
    required=True,
    help="List of PDOS files for the different tip apexes used as tip"
    + " coefficients. Or, alternatively, four numbers corresponding to"
    + " [s py pz px] for uniform PDOS values (constant DOS).",
)

parser.add_argument(
    "--orbs_tip",
    type=int,
    metavar="l",
    default=1,
    required=False,
    help="Integer indicating which orbitals of the tip are used following the" + " angular momentum quantum number.",
)

parser.add_argument(
    "--fwhm_tip",
    type=float,
    metavar="eV",
    default=0.0,
    required=False,
    help="Full width at half maximum for Gaussian broadening of DOS for tip.",
)

parser.add_argument(
    "--rotate", action="store_true", default=False, required=False, help="If set, tip coefficients will be rotated."
)

### ----------------------------------------------------------------------------
### Parameters for HR-STM images
parser.add_argument("--voltages", type=float, metavar="eV", nargs="+", required=True, help="Voltages used for STM.")

### ----------------------------------------------------------------------------
### Parse args one one rank and broadcast it
### ----------------------------------------------------------------------------

args = None
args_success = False
try:
    if mpi_rank == 0:
        args = parser.parse_args()
        args_success = True
finally:
    args_success = mpi_comm.bcast(args_success, root=0)
if not args_success:
    raise ImportError("Failed to read input arguments.")

args = mpi_comm.bcast(args, root=0)


### ----------------------------------------------------------------------------
### Read tip positions on ranks
### ----------------------------------------------------------------------------

start = time.time()
if args.tip_pos_files is not None and args.tip_shift is not None:
    pos_local, dim_pos, eval_region_wfn, lVec = hu.read_tip_positions(
        args.tip_pos_files, args.tip_shift, args.dx_wfn, mpi_rank, mpi_size, mpi_comm
    )
elif args.eval_region is not None and args.dx_tip is not None and not args.rotate:
    pos_local, dim_pos, eval_region_wfn, lVec = hu.create_tip_positions(
        args.eval_region, args.dx_tip, mpi_rank, mpi_size, mpi_comm
    )
else:
    raise ImportError("Missing tip positions.")
end = time.time()
print("Reading tip positions in {} seconds for rank {}.".format(end - start, mpi_rank))

### ----------------------------------------------------------------------------
### Setup tip coefficients
### ----------------------------------------------------------------------------

start = time.time()
tip_coeffs = tc.TipCoefficients(mpi_rank, mpi_size, mpi_comm)
tip_coeffs.read_coefficients(
    args.orbs_tip,
    args.pdos_list,
    min(-max(args.voltages), 0) - 4.0 * args.fwhm_tip,
    max(-min(args.voltages), 0) + 4.0 * args.fwhm_tip,
)
tip_coeffs.initialize(pos_local, args.rotate)
end = time.time()
print("Reading tip coefficients in {} seconds for rank {}.".format(end - start, mpi_rank))

### ----------------------------------------------------------------------------
###  Evaluate sample orbitals
### ----------------------------------------------------------------------------

start = time.time()
wfn_grid_orb = cgo.Cp2kGridOrbitals(mpi_rank, mpi_size, mpi_comm, single_precision=False)
# wfn_grid_orb = pgo.PPSTMGridOrbitals(mpi_rank, mpi_size, mpi_comm,
#    single_precision=False)
wfn_grid_orb.read_cp2k_input(args.cp2k_input_file)
wfn_grid_orb.read_xyz(args.xyz_file)
wfn_grid_orb.read_basis_functions(args.basis_set_file)
wfn_grid_orb.load_restart_wfn_file(
    args.wfn_file, emin=args.emin - 4.0 * args.fwhm_sam, emax=args.emax + 4.0 * args.fwhm_sam
)
wfn_grid_orb.calc_morbs_in_region(
    args.dx_wfn,
    z_eval_region=eval_region_wfn[2],
    # Hack for PPSTM: actually used for workfunction
    #    reserve_extrap = args.wn,
    reserve_extrap=args.extrap_dist,
    eval_cutoff=args.rcut,
)
end = time.time()
print("Building CP2K wave function matrix in {} seconds for rank {}.".format(end - start, mpi_rank))

### ----------------------------------------------------------------------------
### Extrapolate orbitals
### ----------------------------------------------------------------------------

start = time.time()
hart_cube = Cube()
hart_cube.read_cube_file(args.hartree_file)
extrap_plane_z = eval_region_wfn[2][1] - np.max(wfn_grid_orb.ase_atoms.positions[:, 2])
hart_plane = hart_cube.get_plane_above_topmost_atom(extrap_plane_z) - wfn_grid_orb.ref_energy * ev2hartree
del hart_cube, extrap_plane_z
wfn_grid_orb.extrapolate_morbs(hart_plane=hart_plane)
end = time.time()
print("Extrapolating CP2K wave function matrix in {} seconds for rank {}.".format(end - start, mpi_rank))

### ----------------------------------------------------------------------------
### Divide grid orbitals along space and put in wrapper
### ----------------------------------------------------------------------------

start = time.time()
wfn_grid_matrix = cgm.Cp2kGridMatrix(
    wfn_grid_orb, eval_region_wfn, pos_local[1:], args.orbs_tip, args.wn, mpi_rank, mpi_size, mpi_comm
)
del wfn_grid_orb, pos_local
wfn_grid_matrix.divide()
end = time.time()
print("Setting up wave function object in {} seconds for rank {}.".format(end - start, mpi_rank))

### ----------------------------------------------------------------------------
### Evaluate HRSTM and write output
### ----------------------------------------------------------------------------

# Meta information
if mpi_rank == 0:
    meta = {"dimGrid": dim_pos, "lVec": lVec, "voltages": args.voltages, "input": args}
    if os.path.dirname(args.output) != "":
        os.makedirs(os.path.dirname(args.output), exist_ok=True)
    np.save(args.output + "_meta.npy", meta)
start = time.time()
hrstm = hs.Hrstm(tip_coeffs, dim_pos, wfn_grid_matrix, args.fwhm_sam, args.fwhm_tip, mpi_rank, mpi_size, mpi_comm)
hrstm.run(args.voltages)
hrstm.write_compressed(args.output)
end = time.time()
print("Evaluating HRSTM-run method in {} seconds for rank {}.".format(end - start, mpi_rank))