cube_from_wfn.py

#!/usr/bin/env python
import os
import numpy as np
import time
import copy
import sys

import argparse

ang_2_bohr = 1.0/0.52917721067
hart_2_ev = 27.21138602

import cp2k_spm_tools.cp2k_grid_orbitals as cgo
from cp2k_spm_tools import common, cube

from mpi4py import MPI

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

parser = argparse.ArgumentParser(
    description='Creates Gaussian cube files from cp2k .wfn file.')

parser.add_argument(
    '--cp2k_input_file',
    metavar='FILENAME',
    required=True,
    help='CP2K input of the SCF calculation.')
parser.add_argument(
    '--basis_set_file',
    metavar='FILENAME',
    required=True,
    help='File containing the used basis sets.')
parser.add_argument(
    '--xyz_file',
    metavar='FILENAME',
    required=True,
    help='.xyz file containing the geometry.')
parser.add_argument(
    '--wfn_file',
    metavar='FILENAME',
    required=True,
    help='cp2k restart file containing the wavefunction.')

parser.add_argument(
    '--output_dir',
    metavar='DIR',
    required=True,
    help='directory where to output the cubes.')
### -----------------------------------------------------------
parser.add_argument(
    '--dx',
    type=float,
    metavar='DX',
    default=0.2,
    help='Spatial step for the grid (angstroms).')
parser.add_argument(
    '--eval_cutoff',
    type=float,
    metavar='D',
    default=14.0,
    help=("Size of the region around the atom where each"
          " orbital is evaluated (only used for 'G' region).")
)
parser.add_argument(
    '--eval_region',
    type=str,
    nargs=6,
    metavar='X',
    required=False,
    default = ['G', 'G', 'G', 'G', 'G', 'G'],
    help=common.eval_region_description
)
parser.add_argument(
    '--pbc',
    type=int,
    nargs=3,
    metavar='X',
    required=False,
    default = [1, 1, 1],
    help="periodic boundary conditions in directions [x,y,z]. (1=on, 0=off)"
)
### -----------------------------------------------------------
parser.add_argument(
    '--n_homo',
    type=int,
    metavar='N',
    default=0,
    help="Number of HOMO orbitals to export.")
parser.add_argument(
    '--n_lumo',
    type=int,
    metavar='N',
    default=0,
    help="Number of LUMO orbitals to export.")
parser.add_argument(
    '--orb_square',
    action='store_true',
    help=("Additionally generate the square (RHO) for each MO.")
)
### -----------------------------------------------------------
parser.add_argument(
    '--charge_dens',
    action='store_true',
    help=("Calculate charge density (all occupied orbitals are evaluated).")
)
parser.add_argument(
    '--charge_dens_artif_core',
    action='store_true',
    help=("Calculate charge density with 'fake' artificial core (all occ orbitals are evaluated).")
)
parser.add_argument(
    '--spin_dens',
    action='store_true',
    help=("Calculate spin density (all occupied orbitals are evaluated).")
)
### -----------------------------------------------------------
parser.add_argument(
    '--do_not_center_atoms',
    action='store_true',
    help=("Center atoms to cell.")
)
### -----------------------------------------------------------

time0 = time.time()

### ------------------------------------------------------
### Parse args for only one rank to suppress duplicate stdio
### ------------------------------------------------------

args = None
args_success = False
try:
    if mpi_rank == 0:
        args = parser.parse_args()
        args_success = True
finally:
    args_success = comm.bcast(args_success, root=0)

if not args_success:
    print(mpi_rank, "exiting")
    exit(0)

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

output_dir = args.output_dir if args.output_dir[-1] == '/' else args.output_dir+"/"

### ------------------------------------------------------
### Evaluate orbitals on the real-space grid
### ------------------------------------------------------

n_homo = args.n_homo
n_lumo = args.n_lumo

n_homo_range = n_homo
if args.charge_dens or args.spin_dens:
    n_homo_range = None

mol_grid_orb = cgo.Cp2kGridOrbitals(mpi_rank, mpi_size, comm, single_precision=False)
mol_grid_orb.read_cp2k_input(args.cp2k_input_file)
mol_grid_orb.read_xyz(args.xyz_file)
if not args.do_not_center_atoms:
    mol_grid_orb.center_atoms_to_cell()
mol_grid_orb.read_basis_functions(args.basis_set_file)
mol_grid_orb.load_restart_wfn_file(args.wfn_file, n_occ=n_homo_range, n_virt=n_lumo)

eval_reg = common.parse_eval_region_input(args.eval_region, mol_grid_orb.ase_atoms, mol_grid_orb.cell)

mol_grid_orb.calc_morbs_in_region(args.dx,
                                x_eval_region = eval_reg[0],
                                y_eval_region = eval_reg[1],
                                z_eval_region = eval_reg[2],
                                pbc=np.array(args.pbc, dtype=bool),
                                reserve_extrap = 0.0,
                                eval_cutoff = args.eval_cutoff)


### ------------------------------------------------------
### Export the data
### ------------------------------------------------------

ase_atoms = mol_grid_orb.ase_atoms
origin = mol_grid_orb.origin
cell = mol_grid_orb.eval_cell*np.eye(3)
vol_elem = np.prod(mol_grid_orb.dv)

for imo in np.arange(n_homo+n_lumo):
    i_rel_homo = imo - n_homo + 1
    for ispin in range(mol_grid_orb.nspin):

        if imo >= len(mol_grid_orb.cwf.global_morb_indexes[ispin]):
            continue

        global_index = mol_grid_orb.cwf.global_morb_indexes[ispin][imo]

        if i_rel_homo < 0:
            hl_label = "HOMO%+d" % i_rel_homo
        elif i_rel_homo == 0:
            hl_label = "HOMO"
        elif i_rel_homo == 1:
            hl_label = "LUMO"
        else:
            hl_label = "LUMO%+d" % (i_rel_homo-1)

        name = "S%d_%d_%s" % (ispin, global_index, hl_label)
        mol_grid_orb.write_cube(output_dir + name + ".cube", i_rel_homo, spin=ispin)

        if args.orb_square:
            mol_grid_orb.write_cube(output_dir + name + "_sq.cube", i_rel_homo, spin=ispin, square=True)

if args.charge_dens:
    mol_grid_orb.calculate_and_save_charge_density(output_dir + "charge_density.cube")
if args.charge_dens_artif_core:
    mol_grid_orb.calculate_and_save_charge_density(output_dir + "charge_density_artif.cube", artif_core=True)

if args.spin_dens:
    mol_grid_orb.calculate_and_save_spin_density(output_dir + "spin_density.cube")

print("R%d/%d: finished, total time: %.2fs"%(mpi_rank, mpi_size, (time.time() - time0)))