1#! /usr/bin/env python3
2
3import argparse
4import h5py
5import json
6from mpi4py import MPI
7import numpy
8import os
9import scipy.sparse
10import sys
11import time
12from afqmctools.inputs.from_pyscf import write_qmcpack
13from afqmctools.utils.misc import get_git_hash
14
15def parse_args(args, comm):
16    """Parse command-line arguments.
17
18    Parameters
19    ----------
20    args : list of strings
21        command-line arguments.
22
23    Returns
24    -------
25    options : :class:`argparse.ArgumentParser`
26        Command line arguments.
27    """
28
29    if comm.rank == 0:
30        parser = argparse.ArgumentParser(description = __doc__)
31        parser.add_argument('-i', '--input', dest='chk_file', type=str,
32                            default=None, help='Input pyscf .chk file.')
33        parser.add_argument('-o', '--output', dest='hamil_file',
34                            type=str, default='hamiltonian.h5',
35                            help='Output file name for QMCPACK hamiltonian.')
36        parser.add_argument('-w', '--wavefunction', dest='wfn_file',
37                            type=str, default=None,
38                            help='Output file name for QMCPACK wavefunction. '
39                                 'By default will write to hamil_file.')
40        parser.add_argument('-q', '--qmcpack-input', dest='qmc_input',
41                            type=str, default=None,
42                            help='Generate skeleton QMCPACK input xml file.')
43        parser.add_argument('-t', '--cholesky-threshold', dest='thresh',
44                            type=float, default=1e-5,
45                            help='Cholesky convergence threshold.')
46        parser.add_argument('-k', '--kpoint', dest='kpoint_sym',
47                            action='store_true', default=False,
48                            help='Generate explicit kpoint dependent integrals.')
49        parser.add_argument('--density-fit', dest='df',
50                            action='store_true', default=False,
51                            help='Use density fitting integrals stored in '
52                            'input pyscf chkpoint file.')
53        parser.add_argument('-a', '--ao', '--ortho-ao', dest='ortho_ao',
54                            action='store_true', default=False,
55                            help='Transform to ortho AO basis. Default assumes '
56                            'we work in MO basis')
57        parser.add_argument('-c', '--cas',
58                            help='Specify a CAS in the form of N,M.',
59                            type=lambda s: [int(item) for item in s.split(',')],
60                            default=None)
61        parser.add_argument('-d', '--disable-ham', dest='disable_ham',
62                            action='store_true', default=False,
63                            help='Disable hamiltonian generation.')
64        parser.add_argument('-n', '--num-dets', dest='ndet_max',
65                            type=int, default=1,
66                            help='Set upper limit on number of determinants to '
67                            'generate.')
68        parser.add_argument('-r', '--real-ham', dest='real_chol',
69                            action='store_true', default=False,
70                            help='Write integrals as real numbers.')
71        parser.add_argument('-p', '--phdf', dest='phdf',
72                            action='store_true', default=False,
73                            help='Use parallel hdf5.')
74        parser.add_argument('--low', dest='low_thresh',
75                            type=float, default=0.1,
76                            help='Lower threshold for non-integer occupancies'
77                            'to include in multi-determinant exansion.')
78        parser.add_argument('--high', dest='high_thresh',
79                            type=float, default=0.95,
80                            help='Upper threshold for non-integer occupancies'
81                            'to include in multi-determinant exansion.')
82        parser.add_argument('--dense', dest='dense',
83                            action='store_true', default=False,
84                            help='Write dense Hamiltonian.')
85        parser.add_argument('-v', '--verbose', action='count', default=0,
86                            help='Verbose output.')
87
88        options = parser.parse_args(args)
89    else:
90        options = None
91    options = comm.bcast(options, root=0)
92
93    if not options.chk_file:
94        if comm.rank == 0:
95            parser.print_help()
96        sys.exit()
97
98    return options
99
100def write_metadata(options, sha1, cwd, date_time):
101    op_dict = vars(options)
102    op_dict['git_hash'] = sha1
103    op_dict['working_directory'] = os.getcwd()
104    op_dict['date_time'] = date_time
105    if options.wfn_file != options.hamil_file:
106        with h5py.File(options.wfn_file, 'a') as fh5:
107            try:
108                fh5['metadata'] = json.dumps(op_dict)
109            except RuntimeError:
110                del fh5['metadata']
111                fh5['metadata'] = json.dumps(op_dict)
112    if not options.disable_ham:
113        with h5py.File(options.hamil_file, 'a') as fh5:
114            fh5['metadata'] = json.dumps(op_dict)
115
116def main(args):
117    """Generate QMCPACK input from pyscf checkpoint file.
118
119    Parameters
120    ----------
121    args : list of strings
122        command-line arguments.
123    """
124    comm = MPI.COMM_WORLD
125    options = parse_args(args, comm)
126    if comm.rank == 0:
127        cwd = os.getcwd()
128        sha1 = get_git_hash()
129        date_time = time.asctime()
130        print(" # Generating QMCPACK input from PYSCF checkpoint file.")
131        print(" # git sha1: {}".format(sha1))
132        print(" # Date/Time: {}".format(date_time))
133        print(" # Working directory: {}".format(cwd))
134
135    if options.wfn_file is None:
136        if options.disable_ham:
137            options.wfn_file = 'wfn.dat'
138        else:
139            options.wfn_file = options.hamil_file
140    write_qmcpack(options.chk_file, options.hamil_file, options.thresh,
141                  comm=comm,
142                  ortho_ao=options.ortho_ao,
143                  kpoint=options.kpoint_sym, df=options.df,
144                  verbose=options.verbose, cas=options.cas,
145                  qmc_input=options.qmc_input,
146                  wfn_file=options.wfn_file,
147                  write_hamil=(not options.disable_ham),
148                  ndet_max=options.ndet_max,
149                  real_chol=options.real_chol,
150                  phdf=options.phdf,
151                  low=options.low_thresh,
152                  high=options.high_thresh,
153                  dense=options.dense)
154    if comm.rank == 0:
155        write_metadata(options, sha1, cwd, date_time)
156
157    if comm.rank == 0:
158        print("\n # Finished.")
159
160if __name__ == '__main__':
161
162    main(sys.argv[1:])
163