1#!/usr/local/bin/python3.8
2
3# Script to create a kgrid.inp file automatically from a Quantum Espresso
4# data-file.xml file.
5#
6# Felipe H. da Jornada (May, 2015)
7
8
9import xml.etree.cElementTree as ET
10
11
12def create_kgrid_inp(datafile, f_kgrid, kgrid, kshift, qshift, use_trs,
13    out_cart, out_oct):
14
15    f_kgrid.write((3*'{} '+'\n').format(*kgrid))
16    f_kgrid.write((3*'{:.12f} '+'\n').format(*kshift))
17    f_kgrid.write((3*'{:.12f} '+'\n').format(*qshift))
18
19    tree = ET.parse(datafile)
20    a1 = tree.find('CELL/DIRECT_LATTICE_VECTORS/a1').text.strip()
21    a2 = tree.find('CELL/DIRECT_LATTICE_VECTORS/a2').text.strip()
22    a3 = tree.find('CELL/DIRECT_LATTICE_VECTORS/a3').text.strip()
23    f_kgrid.write(' {}\n'.format(a1))
24    f_kgrid.write(' {}\n'.format(a2))
25    f_kgrid.write(' {}\n'.format(a3))
26
27    nat = int(tree.find('IONS/NUMBER_OF_ATOMS').text.strip())
28    f_kgrid.write('{}\n'.format(nat))
29
30    for iat in range(nat):
31        node = tree.find('IONS/ATOM.{}'.format(iat+1))
32        ityp = node.get('INDEX')
33        tau = node.get('tau')
34        f_kgrid.write((' {} {}\n').format(ityp, tau))
35
36    node = tree.find('PLANE_WAVES/FFT_GRID')
37    fftgrid  = [node.get(nr) for nr in ('nr1', 'nr2', 'nr3')]
38    f_kgrid.write('{} {} {}\n'.format(*fftgrid))
39
40    def write_bool(b):
41        if b:
42            f_kgrid.write('.true.\n')
43        else:
44            f_kgrid.write('.false.\n')
45
46    write_bool(use_trs)
47    write_bool(out_cart)
48    write_bool(out_oct)
49    write_bool(False)
50
51
52if __name__=="__main__":
53    from argparse import ArgumentParser
54
55    desc = ('Creates a kgrid.inp file given a data-file.xml from a Quantum'
56    ' Espresso save directory.')
57    parser = ArgumentParser(description=desc)
58
59    group = parser.add_argument_group('required input and output files')
60    group.add_argument('datafile', help='input data-file.xml from a QE save directory')
61    group.add_argument('kgrid_inp', help='output kgrid.inp file to generate')
62
63    group = parser.add_argument_group('kgrid specification')
64    group.add_argument('--kgrid', type=int, default=[1,1,1], nargs=3,
65    metavar=('nk1', 'nk2', 'nk3'),
66        help='k-point density in each reciprocal-lattice direction. Defaults to 1.')
67    group.add_argument('--kshift', type=float, default=[0.,0.,0.], nargs=3,
68        metavar=('ks1', 'ks2', 'ks3'),
69        help='k-shift in each reciprocal-lattice direction. Defaults to 0.0.')
70    group.add_argument('--qshift', type=float, default=[0.,0.,0.], nargs=3,
71        metavar=('qs1', 'qs2', 'qs3'),
72        help='q0-shift in each reciprocal-lattice direction. Defaults to 0.0.')
73    group.add_argument('--use_trs', default=False, action='store_true',
74        help='Whether to use time-reversal symmetries. Defaults to false.')
75
76    group = parser.add_argument_group('output control')
77    group.add_argument('--output_cartesian', dest='out_cart', default=False,
78        action='store_true', help=('Write output in Cartesian coordinates.'
79        ' Default is false, which writes in crystal coordinates.'))
80    group.add_argument('--output_octopus', dest='out_oct', default=False,
81        action='store_true',  help=('Write output in Octopus format.'
82        ' Default is false, which writes in Quantum Espresso format.'))
83
84    args = parser.parse_args()
85
86    with open(args.kgrid_inp, 'w') as f_kgrid:
87        create_kgrid_inp(args.datafile, f_kgrid, args.kgrid, args.kshift,
88            args.qshift, args.use_trs, args.out_cart, args.out_oct)
89