1import os
2import re
3from subprocess import call, TimeoutExpired
4from copy import deepcopy
5
6import numpy as np
7
8from ase import Atoms
9from ase.utils import workdir
10from ase.units import Hartree, Bohr, Debye
11from ase.calculators.singlepoint import SinglePointCalculator
12
13
14def _format_value(val):
15    if isinstance(val, bool):
16        return '.t.' if val else '.f.'
17    return str(val).upper()
18
19
20def _write_block(name, args):
21    out = [' ${}'.format(name.upper())]
22    for key, val in args.items():
23        out.append('  {}={}'.format(key.upper(), _format_value(val)))
24    out.append(' $END')
25    return '\n'.join(out)
26
27
28def _write_geom(atoms, basis_spec):
29    out = [' $DATA', atoms.get_chemical_formula(), 'C1']
30    for i, atom in enumerate(atoms):
31        out.append('{:<3} {:>3} {:20.13e} {:20.13e} {:20.13e}'
32                   .format(atom.symbol, atom.number, *atom.position))
33        if basis_spec is not None:
34            basis = basis_spec.get(i)
35            if basis is None:
36                basis = basis_spec.get(atom.symbol)
37            if basis is None:
38                raise ValueError('Could not find an appropriate basis set '
39                                 'for atom number {}!'.format(i))
40            out += [basis, '']
41    out.append(' $END')
42    return '\n'.join(out)
43
44
45def _write_ecp(atoms, ecp):
46    out = [' $ECP']
47    for i, symbol in enumerate(atoms.symbols):
48        if i in ecp:
49            out.append(ecp[i])
50        elif symbol in ecp:
51            out.append(ecp[symbol])
52        else:
53            raise ValueError('Could not find an appropriate ECP for '
54                             'atom number {}!'.format(i))
55    out.append(' $END')
56    return '\n'.join(out)
57
58
59_xc = dict(LDA='SVWN')
60
61
62def write_gamess_us_in(fd, atoms, properties=None, **params):
63    params = deepcopy(params)
64
65    if properties is None:
66        properties = ['energy']
67
68    # set RUNTYP from properties iff value not provided by the user
69    contrl = params.pop('contrl', dict())
70    if 'runtyp' not in contrl:
71        if 'forces' in properties:
72            contrl['runtyp'] = 'gradient'
73        else:
74            contrl['runtyp'] = 'energy'
75
76    # Support xc keyword for functional specification
77    xc = params.pop('xc', None)
78    if xc is not None and 'dfttyp' not in contrl:
79        contrl['dfttyp'] = _xc.get(xc.upper(), xc.upper())
80
81    # Automatically determine multiplicity from magnetic moment
82    magmom_tot = int(round(atoms.get_initial_magnetic_moments().sum()))
83    if 'mult' not in contrl:
84        contrl['mult'] = abs(magmom_tot) + 1
85
86    # Since we're automatically determining multiplicity, we also
87    # need to automatically switch to UHF when the multiplicity
88    # is not 1
89    if 'scftyp' not in contrl:
90        contrl['scftyp'] = 'rhf' if contrl['mult'] == 1 else 'uhf'
91
92    # effective core potentials
93    ecp = params.pop('ecp', None)
94    if ecp is not None and 'pp' not in contrl:
95        contrl['pp'] = 'READ'
96
97    # If no basis set is provided, use 3-21G by default.
98    basis_spec = None
99    if 'basis' not in params:
100        params['basis'] = dict(gbasis='N21', ngauss=3)
101    else:
102        keys = set(params['basis'])
103        # Check if the user is specifying a literal per-atom basis.
104        # We assume they are passing a per-atom basis if the keys of the
105        # basis dict are atom symbols, or if they are atom indices, or
106        # a mixture of both.
107        if (keys.intersection(set(atoms.symbols))
108                or any(map(lambda x: isinstance(x, int), keys))):
109            basis_spec = params.pop('basis')
110
111    out = [_write_block('contrl', contrl)]
112    out += [_write_block(*item) for item in params.items()]
113    out.append(_write_geom(atoms, basis_spec))
114    if ecp is not None:
115        out.append(_write_ecp(atoms, ecp))
116    fd.write('\n\n'.join(out))
117
118
119_geom_re = re.compile(r'^\s*ATOM\s+ATOMIC\s+COORDINATES')
120_atom_re = re.compile(r'^\s*(\S+)\s+(\S+)\s+(\S+)\s+(\S+)\s+(\S+)\s*\n')
121_energy_re = re.compile(r'^\s*FINAL [\S\s]+ ENERGY IS\s+(\S+) AFTER')
122_grad_re = re.compile(r'^\s*GRADIENT OF THE ENERGY\s*')
123_dipole_re = re.compile(r'^\s+DX\s+DY\s+DZ\s+\/D\/\s+\(DEBYE\)')
124
125
126def read_gamess_us_out(fd):
127    atoms = None
128    energy = None
129    forces = None
130    dipole = None
131    for line in fd:
132        # Geometry
133        if _geom_re.match(line):
134            fd.readline()
135            symbols = []
136            pos = []
137            while True:
138                atom = _atom_re.match(fd.readline())
139                if atom is None:
140                    break
141                symbol, _, x, y, z = atom.groups()
142                symbols.append(symbol.capitalize())
143                pos.append(list(map(float, [x, y, z])))
144            atoms = Atoms(symbols, np.array(pos) * Bohr)
145            continue
146
147        # Energy
148        ematch = _energy_re.match(line)
149        if ematch is not None:
150            energy = float(ematch.group(1)) * Hartree
151
152        # MPn energy. Supplants energy parsed above.
153        elif line.strip().startswith('TOTAL ENERGY'):
154            energy = float(line.strip().split()[-1]) * Hartree
155
156        # Higher-level energy (e.g. coupled cluster)
157        # Supplants energies parsed above.
158        elif line.strip().startswith('THE FOLLOWING METHOD AND ENERGY'):
159            energy = float(fd.readline().strip().split()[-1]) * Hartree
160
161        # Gradients
162        elif _grad_re.match(line):
163            for _ in range(3):
164                fd.readline()
165            grad = []
166            while True:
167                atom = _atom_re.match(fd.readline())
168                if atom is None:
169                    break
170                grad.append(list(map(float, atom.groups()[2:])))
171            forces = -np.array(grad) * Hartree / Bohr
172        elif _dipole_re.match(line):
173            dipole = np.array(list(map(float, fd.readline().split()[:3])))
174            dipole *= Debye
175
176    atoms.calc = SinglePointCalculator(atoms, energy=energy,
177                                       forces=forces, dipole=dipole)
178    return atoms
179
180
181def read_gamess_us_punch(fd):
182    atoms = None
183    energy = None
184    forces = None
185    dipole = None
186    for line in fd:
187        if line.strip() == '$DATA':
188            symbols = []
189            pos = []
190            while line.strip() != '$END':
191                line = fd.readline()
192                atom = _atom_re.match(line)
193                if atom is None:
194                    # The basis set specification is interlaced with the
195                    # molecular geometry. We don't care about the basis
196                    # set, so ignore lines that don't match the pattern.
197                    continue
198                symbols.append(atom.group(1).capitalize())
199                pos.append(list(map(float, atom.group(3, 4, 5))))
200            atoms = Atoms(symbols, np.array(pos))
201        elif line.startswith('E('):
202            energy = float(line.split()[1][:-1]) * Hartree
203        elif line.strip().startswith('DIPOLE'):
204            dipole = np.array(list(map(float, line.split()[1:]))) * Debye
205        elif line.strip() == '$GRAD':
206            # The gradient block also contains the energy, which we prefer
207            # over the energy obtained above because it is more likely to
208            # be consistent with the gradients. It probably doesn't actually
209            # make a difference though.
210            energy = float(fd.readline().split()[1]) * Hartree
211            grad = []
212            while line.strip() != '$END':
213                line = fd.readline()
214                atom = _atom_re.match(line)
215                if atom is None:
216                    continue
217                grad.append(list(map(float, atom.group(3, 4, 5))))
218            forces = -np.array(grad) * Hartree / Bohr
219
220    atoms.calc = SinglePointCalculator(atoms, energy=energy, forces=forces,
221                                       dipole=dipole)
222
223    return atoms
224
225
226def clean_userscr(userscr, prefix):
227    for fname in os.listdir(userscr):
228        tokens = fname.split('.')
229        if tokens[0] == prefix and tokens[-1] != 'bak':
230            fold = os.path.join(userscr, fname)
231            os.rename(fold, fold + '.bak')
232
233
234def get_userscr(prefix, command):
235    prefix_test = prefix + '_test'
236    command = command.replace('PREFIX', prefix_test)
237    with workdir(prefix_test, mkdir=True):
238        try:
239            call(command, shell=True, timeout=2)
240        except TimeoutExpired:
241            pass
242
243        try:
244            with open(prefix_test + '.log') as fd:
245                for line in fd:
246                    if line.startswith('GAMESS supplementary output files'):
247                        return ' '.join(line.split(' ')[8:]).strip()
248        except FileNotFoundError:
249            return None
250
251    return None
252