1"""This module defines an ASE interface to deMon.
2
3http://www.demon-software.com
4
5"""
6import os
7import os.path as op
8import subprocess
9import shutil
10
11import numpy as np
12
13from ase.units import Bohr, Hartree
14import ase.data
15from ase.calculators.calculator import FileIOCalculator, ReadError
16from ase.calculators.calculator import Parameters, all_changes
17from ase.calculators.calculator import equal
18import ase.io
19from .demon_io import parse_xray
20
21m_e_to_amu = 1822.88839
22
23
24class Parameters_deMon(Parameters):
25    """Parameters class for the calculator.
26    Documented in Base_deMon.__init__
27
28    The options here are the most important ones that the user needs to be
29    aware of. Further options accepted by deMon can be set in the dictionary
30    input_arguments.
31
32    """
33    def __init__(
34            self,
35            label='rundir',
36            atoms=None,
37            command=None,
38            restart=None,
39            basis_path=None,
40            ignore_bad_restart_file=FileIOCalculator._deprecated,
41            deMon_restart_path='.',
42            title='deMon input file',
43            scftype='RKS',
44            forces=False,
45            dipole=False,
46            xc='VWN',
47            guess='TB',
48            print_out='MOE',
49            basis={},
50            ecps={},
51            mcps={},
52            auxis={},
53            augment={},
54            input_arguments=None):
55        kwargs = locals()
56        kwargs.pop('self')
57        Parameters.__init__(self, **kwargs)
58
59
60class Demon(FileIOCalculator):
61    """Calculator interface to the deMon code. """
62
63    implemented_properties = [
64        'energy',
65        'forces',
66        'dipole',
67        'eigenvalues']
68
69    def __init__(self, **kwargs):
70        """ASE interface to the deMon code.
71
72        The deMon2k code can be obtained from http://www.demon-software.com
73
74        The DEMON_COMMAND environment variable must be set to run the executable, in bash it would be set along the lines of
75        export DEMON_COMMAND="deMon.4.3.6.std > deMon_ase.out 2>&1"
76
77        Parameters:
78
79        label : str
80            relative path to the run directory
81        atoms : Atoms object
82            the atoms object
83        command  : str
84            Command to run deMon. If not present the environment varable DEMON_COMMAND will be used
85        restart  : str
86            Relative path to ASE restart directory for parameters and atoms object and results
87        basis_path  : str
88            Relative path to the directory containing BASIS, AUXIS, ECPS, MCPS and AUGMENT
89        ignore_bad_restart_file : bool
90            Ignore broken or missing ASE restart files
91            By default, it is an error if the restart
92            file is missing or broken.
93        deMon_restart_path  : str
94            Relative path to the deMon restart dir
95        title : str
96            Title in the deMon input file.
97        scftype : str
98            Type of scf
99        forces  : bool
100            If True a force calculation will be enforced.
101        dipole  : bool
102            If True a dipole calculation will be enforced
103        xc : str
104            xc-functional
105        guess : str
106            guess for initial density and wave functions
107        print_out : str | list
108            Options for the printing in deMon
109        basis : dict
110            Definition of basis sets.
111        ecps  : dict
112            Definition of ECPs
113        mcps  : dict
114            Definition of MCPs
115        auxis  : dict
116            Definition of AUXIS
117        augment : dict
118            Definition of AUGMENT
119        input_arguments : dict
120            Explicitly given input arguments. The key is the input keyword
121            and the value is either a str, a list of str (will be written on the same line as the keyword),
122            or a list of lists of str (first list is written on the first line, the others on following lines.)
123
124        For example usage, see the tests h2o.py and h2o_xas_xes.py in the directory ase/test/demon
125
126        """
127
128        parameters = Parameters_deMon(**kwargs)
129
130        # Setup the run command
131        command = parameters['command']
132        if command is None:
133            command = os.environ.get('DEMON_COMMAND')
134
135        if command is None:
136            mess = 'The "DEMON_COMMAND" environment is not defined.'
137            raise ValueError(mess)
138        else:
139            parameters['command'] = command
140
141        # Call the base class.
142        FileIOCalculator.__init__(
143            self,
144            **parameters)
145
146    def __getitem__(self, key):
147        """Convenience method to retrieve a parameter as
148        calculator[key] rather than calculator.parameters[key]
149
150            Parameters:
151                key       : str, the name of the parameters to get.
152        """
153        return self.parameters[key]
154
155    def set(self, **kwargs):
156        """Set all parameters.
157
158        Parameters:
159            kwargs  : Dictionary containing the keywords for deMon
160        """
161        # Put in the default arguments.
162        kwargs = self.default_parameters.__class__(**kwargs)
163
164        if 'parameters' in kwargs:
165            filename = kwargs.pop('parameters')
166            parameters = Parameters.read(filename)
167            parameters.update(kwargs)
168            kwargs = parameters
169
170        changed_parameters = {}
171
172        for key, value in kwargs.items():
173            oldvalue = self.parameters.get(key)
174            if key not in self.parameters or not equal(value, oldvalue):
175                changed_parameters[key] = value
176                self.parameters[key] = value
177
178        return changed_parameters
179
180    def link_file(self, fromdir, todir, filename):
181
182        if op.exists(todir + '/' + filename):
183            os.remove(todir + '/' + filename)
184
185        if op.exists(fromdir + '/' + filename):
186            os.symlink(fromdir + '/' + filename,
187                       todir + '/' + filename)
188        else:
189            raise RuntimeError(
190                "{0} doesn't exist".format(fromdir + '/' + filename))
191
192    def calculate(self,
193                  atoms=None,
194                  properties=['energy'],
195                  system_changes=all_changes):
196        """Capture the RuntimeError from FileIOCalculator.calculate
197        and add a little debug information from the deMon output.
198
199        See base FileIocalculator for documentation.
200        """
201
202        if atoms is not None:
203            self.atoms = atoms.copy()
204
205        self.write_input(self.atoms, properties, system_changes)
206        if self.command is None:
207            raise RuntimeError('Please set $%s environment variable ' %
208                               ('DEMON_COMMAND') +
209                               'or supply the command keyword')
210        command = self.command  # .replace('PREFIX', self.prefix)
211
212        # basis path
213        basis_path = self.parameters['basis_path']
214        if basis_path is None:
215            basis_path = os.environ.get('DEMON_BASIS_PATH')
216
217        if basis_path is None:
218            raise RuntimeError('Please set basis_path keyword,' +
219                               ' or the DEMON_BASIS_PATH' +
220                               ' environment variable')
221
222        # link restart file
223        value = self.parameters['guess']
224        if value.upper() == 'RESTART':
225            value2 = self.parameters['deMon_restart_path']
226
227            if op.exists(self.directory + '/deMon.rst')\
228                    or op.islink(self.directory + '/deMon.rst'):
229                os.remove(self.directory + '/deMon.rst')
230            abspath = op.abspath(value2)
231
232            if op.exists(abspath + '/deMon.mem') \
233                    or op.islink(abspath + '/deMon.mem'):
234
235                shutil.copy(abspath + '/deMon.mem',
236                            self.directory + '/deMon.rst')
237            else:
238                raise RuntimeError(
239                    "{0} doesn't exist".format(abspath + '/deMon.rst'))
240
241        abspath = op.abspath(basis_path)
242
243        for name in ['BASIS', 'AUXIS', 'ECPS', 'MCPS', 'FFDS']:
244            self.link_file(abspath, self.directory, name)
245
246        subprocess.check_call(command, shell=True, cwd=self.directory)
247
248        try:
249            self.read_results()
250        except Exception:  # XXX Which kind of exception?
251            with open(self.directory + '/deMon.out', 'r') as fd:
252                lines = fd.readlines()
253            debug_lines = 10
254            print('##### %d last lines of the deMon.out' % debug_lines)
255            for line in lines[-20:]:
256                print(line.strip())
257            print('##### end of deMon.out')
258            raise RuntimeError
259
260    def set_label(self, label):
261        """Set label directory """
262
263        self.label = label
264
265        # in our case self.directory = self.label
266        self.directory = self.label
267        if self.directory == '':
268            self.directory = os.curdir
269
270    def write_input(self, atoms, properties=None, system_changes=None):
271        """Write input (in)-file.
272        See calculator.py for further details.
273
274        Parameters:
275             atoms        : The Atoms object to write.
276             properties   : The properties which should be calculated.
277             system_changes : List of properties changed since last run.
278
279        """
280        # Call base calculator.
281        FileIOCalculator.write_input(
282            self,
283            atoms=atoms,
284            properties=properties,
285            system_changes=system_changes)
286
287        if system_changes is None and properties is None:
288            return
289
290        filename = self.label + '/deMon.inp'
291
292        add_print = ''
293
294        # Start writing the file.
295        with open(filename, 'w') as fd:
296
297            # write keyword argument keywords
298            value = self.parameters['title']
299            self._write_argument('TITLE', value, fd)
300
301            fd.write('#\n')
302
303            value = self.parameters['scftype']
304            self._write_argument('SCFTYPE', value, fd)
305
306            value = self.parameters['xc']
307            self._write_argument('VXCTYPE', value, fd)
308
309            value = self.parameters['guess']
310            self._write_argument('GUESS', value, fd)
311
312            # obtain forces through a single BOMD step
313            # only if forces is in properties, or if keyword forces is True
314            value = self.parameters['forces']
315            if 'forces' in properties or value:
316
317                self._write_argument('DYNAMICS',
318                                     ['INT=1', 'MAX=0', 'STEP=0'], fd)
319                self._write_argument('TRAJECTORY', 'FORCES', fd)
320                self._write_argument('VELOCITIES', 'ZERO', fd)
321                add_print = add_print + ' ' + 'MD OPT'
322
323            # if dipole is True, enforce dipole calculation.
324            # Otherwise only if asked for
325            value = self.parameters['dipole']
326            if 'dipole' in properties or value:
327                self._write_argument('DIPOLE', '', fd)
328
329            # print argument, here other options could change this
330            value = self.parameters['print_out']
331            assert(type(value) is str)
332            value = value + add_print
333
334            if not len(value) == 0:
335                self._write_argument('PRINT', value, fd)
336                fd.write('#\n')
337
338            # write general input arguments
339            self._write_input_arguments(fd)
340
341            fd.write('#\n')
342
343            # write basis set, ecps, mcps, auxis, augment
344            basis = self.parameters['basis']
345            if 'all' not in basis:
346                basis['all'] = 'DZVP'
347            self._write_basis(fd, atoms, basis, string='BASIS')
348
349            ecps = self.parameters['ecps']
350            if not len(ecps) == 0:
351                self._write_basis(fd, atoms, ecps, string='ECPS')
352
353            mcps = self.parameters['mcps']
354            if not len(mcps) == 0:
355                self._write_basis(fd, atoms, mcps, string='MCPS')
356
357            auxis = self.parameters['auxis']
358            if not len(auxis) == 0:
359                self._write_basis(fd, atoms, auxis, string='AUXIS')
360
361            augment = self.parameters['augment']
362            if not len(augment) == 0:
363                self._write_basis(fd, atoms, augment, string='AUGMENT')
364
365            # write geometry
366            self._write_atomic_coordinates(fd, atoms)
367
368            # write xyz file for good measure.
369            ase.io.write(self.label + '/deMon_atoms.xyz', self.atoms)
370
371    def read(self, restart_path):
372        """Read parameters from directory restart_path."""
373
374        self.set_label(restart_path)
375
376        if not op.exists(restart_path + '/deMon.inp'):
377            raise ReadError('The restart_path file {0} does not exist'
378                            .format(restart_path))
379
380        self.atoms = self.deMon_inp_to_atoms(restart_path + '/deMon.inp')
381
382        self.read_results()
383
384    def _write_input_arguments(self, fd):
385        """Write directly given input-arguments."""
386        input_arguments = self.parameters['input_arguments']
387
388        # Early return
389        if input_arguments is None:
390            return
391
392        for key, value in input_arguments.items():
393            self._write_argument(key, value, fd)
394
395    def _write_argument(self, key, value, fd):
396        """Write an argument to file.
397         key :  a string coresponding to the input keyword
398         value : the arguments, can be a string, a number or a list
399         f :  and open file
400        """
401
402        # for only one argument, write on same line
403        if not isinstance(value, (tuple, list)):
404            line = key.upper()
405            line += '    ' + str(value).upper()
406            fd.write(line)
407            fd.write('\n')
408
409        # for a list, write first argument on the first line,
410        # then the rest on new lines
411        else:
412            line = key
413            if not isinstance(value[0], (tuple, list)):
414                for i in range(len(value)):
415                    line += '  ' + str(value[i].upper())
416                fd.write(line)
417                fd.write('\n')
418            else:
419                for i in range(len(value)):
420                    for j in range(len(value[i])):
421                        line += '  ' + str(value[i][j]).upper()
422                    fd.write(line)
423                    fd.write('\n')
424                    line = ''
425
426    def _write_atomic_coordinates(self, fd, atoms):
427        """Write atomic coordinates.
428
429        Parameters:
430        - f:     An open file object.
431        - atoms: An atoms object.
432        """
433
434        fd.write('#\n')
435        fd.write('# Atomic coordinates\n')
436        fd.write('#\n')
437        fd.write('GEOMETRY CARTESIAN ANGSTROM\n')
438
439        for i in range(len(atoms)):
440            xyz = atoms.get_positions()[i]
441            chem_symbol = atoms.get_chemical_symbols()[i]
442            chem_symbol += str(i + 1)
443
444            # if tag is set to 1 then we have a ghost atom,
445            # set nuclear charge to 0
446            if(atoms.get_tags()[i] == 1):
447                nuc_charge = str(0)
448            else:
449                nuc_charge = str(atoms.get_atomic_numbers()[i])
450
451            mass = atoms.get_masses()[i]
452
453            line = '{0:6s}'.format(chem_symbol).rjust(10) + ' '
454            line += '{0:.5f}'.format(xyz[0]).rjust(10) + ' '
455            line += '{0:.5f}'.format(xyz[1]).rjust(10) + ' '
456            line += '{0:.5f}'.format(xyz[2]).rjust(10) + ' '
457            line += '{0:5s}'.format(nuc_charge).rjust(10) + ' '
458            line += '{0:.5f}'.format(mass).rjust(10) + ' '
459
460            fd.write(line)
461            fd.write('\n')
462
463    # routine to write basis set inormation, including ecps and auxis
464    def _write_basis(self, fd, atoms, basis={}, string='BASIS'):
465        """Write basis set, ECPs, AUXIS, or AUGMENT basis
466
467        Parameters:
468        - f:     An open file object.
469        - atoms: An atoms object.
470        - basis: A dictionary specifying the basis set
471        - string: 'BASIS', 'ECP','AUXIS' or 'AUGMENT'
472        """
473
474        # basis for all atoms
475        line = '{0}'.format(string).ljust(10)
476
477        if 'all' in basis:
478            default_basis = basis['all']
479            line += '({0})'.format(default_basis).rjust(16)
480
481        fd.write(line)
482        fd.write('\n')
483
484        # basis for all atomic species
485        chemical_symbols = atoms.get_chemical_symbols()
486        chemical_symbols_set = set(chemical_symbols)
487
488        for i in range(chemical_symbols_set.__len__()):
489            symbol = chemical_symbols_set.pop()
490
491            if symbol in basis:
492                line = '{0}'.format(symbol).ljust(10)
493                line += '({0})'.format(basis[symbol]).rjust(16)
494                fd.write(line)
495                fd.write('\n')
496
497        # basis for individual atoms
498        for i in range(len(atoms)):
499
500            if i in basis:
501                symbol = str(chemical_symbols[i])
502                symbol += str(i + 1)
503
504                line = '{0}'.format(symbol).ljust(10)
505                line += '({0})'.format(basis[i]).rjust(16)
506                fd.write(line)
507                fd.write('\n')
508
509    # Analysis routines
510    def read_results(self):
511        """Read the results from output files."""
512        self.read_energy()
513        self.read_forces(self.atoms)
514        self.read_eigenvalues()
515        self.read_dipole()
516        self.read_xray()
517
518    def read_energy(self):
519        """Read energy from deMon's text-output file."""
520        with open(self.label + '/deMon.out', 'r') as fd:
521            text = fd.read().upper()
522
523        lines = iter(text.split('\n'))
524
525        for line in lines:
526            if line.startswith(' TOTAL ENERGY                ='):
527                self.results['energy'] = float(line.split()[-1]) * Hartree
528                break
529        else:
530            raise RuntimeError
531
532    def read_forces(self, atoms):
533        """Read the forces from the deMon.out file."""
534
535        natoms = len(atoms)
536        filename = self.label + '/deMon.out'
537
538        if op.isfile(filename):
539            with open(filename, 'r') as fd:
540                lines = fd.readlines()
541
542                # find line where the orbitals start
543                flag_found = False
544                for i in range(len(lines)):
545                    if lines[i].rfind('GRADIENTS OF TIME STEP 0 IN A.U.') > -1:
546                        start = i + 4
547                        flag_found = True
548                        break
549
550                if flag_found:
551                    self.results['forces'] = np.zeros((natoms, 3), float)
552                    for i in range(natoms):
553                        line = [s for s in lines[i + start].strip().split(' ')
554                                if len(s) > 0]
555                        f = -np.array([float(x) for x in line[2:5]])
556                        self.results['forces'][i, :] = f * (Hartree / Bohr)
557
558    def read_eigenvalues(self):
559        """Read eigenvalues from the 'deMon.out' file."""
560        assert os.access(self.label + '/deMon.out', os.F_OK)
561
562        # Read eigenvalues
563        with open(self.label + '/deMon.out', 'r') as fd:
564            lines = fd.readlines()
565
566        # try  PRINT MOE
567        eig_alpha, occ_alpha = self.read_eigenvalues_one_spin(
568            lines, 'ALPHA MO ENERGIES', 6)
569        eig_beta, occ_beta = self.read_eigenvalues_one_spin(
570            lines, 'BETA MO ENERGIES', 6)
571
572        # otherwise try PRINT MOS
573        if len(eig_alpha) == 0 and len(eig_beta) == 0:
574            eig_alpha, occ_alpha = self.read_eigenvalues_one_spin(
575                lines, 'ALPHA MO COEFFICIENTS', 5)
576            eig_beta, occ_beta = self.read_eigenvalues_one_spin(
577                lines, 'BETA MO COEFFICIENTS', 5)
578
579        self.results['eigenvalues'] = np.array([eig_alpha, eig_beta]) * Hartree
580        self.results['occupations'] = np.array([occ_alpha, occ_beta])
581
582    def read_eigenvalues_one_spin(self, lines, string, neigs_per_line):
583        """Utility method for retreiving eigenvalues after the string "string"
584        with neigs_per_line eigenvlaues written per line
585        """
586        eig = []
587        occ = []
588
589        skip_line = False
590        more_eigs = False
591
592        # find line where the orbitals start
593        for i in range(len(lines)):
594            if lines[i].rfind(string) > -1:
595                ii = i
596                more_eigs = True
597                break
598
599        while more_eigs:
600            # search for two empty lines in a row preceding a line with
601            # numbers
602            for i in range(ii + 1, len(lines)):
603                if len(lines[i].split()) == 0 and \
604                        len(lines[i + 1].split()) == 0 and \
605                        len(lines[i + 2].split()) > 0:
606                    ii = i + 2
607                    break
608
609            # read eigenvalues, occupations
610            line = lines[ii].split()
611            if len(line) < neigs_per_line:
612                # last row
613                more_eigs = False
614            if line[0] != str(len(eig) + 1):
615                more_eigs = False
616                skip_line = True
617
618            if not skip_line:
619                line = lines[ii + 1].split()
620                for l in line:
621                    eig.append(float(l))
622                line = lines[ii + 3].split()
623                for l in line:
624                    occ.append(float(l))
625                ii = ii + 3
626
627        return eig, occ
628
629    def read_dipole(self):
630        """Read dipole moment."""
631        dipole = np.zeros(3)
632        with open(self.label + '/deMon.out', 'r') as fd:
633            lines = fd.readlines()
634
635            for i in range(len(lines)):
636                if lines[i].rfind('DIPOLE') > -1 and lines[i].rfind('XAS') == -1:
637                    dipole[0] = float(lines[i + 1].split()[3])
638                    dipole[1] = float(lines[i + 2].split()[3])
639                    dipole[2] = float(lines[i + 3].split()[3])
640
641                    # debye to e*Ang
642                    self.results['dipole'] = dipole * 0.2081943482534
643
644                    break
645
646    def read_xray(self):
647        """Read deMon.xry if present."""
648
649        # try to read core IP from, .out file
650        filename = self.label + '/deMon.out'
651        core_IP = None
652        if op.isfile(filename):
653            with open(filename, 'r') as fd:
654                lines = fd.readlines()
655
656            for i in range(len(lines)):
657                if lines[i].rfind('IONIZATION POTENTIAL') > -1:
658                    core_IP = float(lines[i].split()[3])
659
660        try:
661            mode, ntrans, E_trans, osc_strength, trans_dip = parse_xray(self.label + '/deMon.xry')
662        except ReadError:
663            pass
664        else:
665            xray_results = {'xray_mode': mode,
666                            'ntrans': ntrans,
667                            'E_trans': E_trans,
668                            'osc_strength': osc_strength,  # units?
669                            'trans_dip': trans_dip,  # units?
670                            'core_IP': core_IP}
671
672            self.results['xray'] = xray_results
673
674    def deMon_inp_to_atoms(self, filename):
675        """Routine to read deMon.inp and convert it to an atoms object."""
676
677        with open(filename, 'r') as fd:
678            lines = fd.readlines()
679
680        # find line where geometry starts
681        for i in range(len(lines)):
682            if lines[i].rfind('GEOMETRY') > -1:
683                if lines[i].rfind('ANGSTROM'):
684                    coord_units = 'Ang'
685                elif lines.rfind('Bohr'):
686                    coord_units = 'Bohr'
687                ii = i
688                break
689
690        chemical_symbols = []
691        xyz = []
692        atomic_numbers = []
693        masses = []
694
695        for i in range(ii + 1, len(lines)):
696            try:
697                line = lines[i].split()
698
699                if(len(line) > 0):
700                    for symbol in ase.data.chemical_symbols:
701                        found = None
702                        if line[0].upper().rfind(symbol.upper()) > -1:
703                            found = symbol
704                            break
705
706                        if found is not None:
707                            chemical_symbols.append(found)
708                        else:
709                            break
710
711                        xyz.append([float(line[1]), float(line[2]), float(line[3])])
712
713                if len(line) > 4:
714                    atomic_numbers.append(int(line[4]))
715
716                if len(line) > 5:
717                    masses.append(float(line[5]))
718
719            except Exception:  # XXX Which kind of exception?
720                raise RuntimeError
721
722        if coord_units == 'Bohr':
723            xyz = xyz * Bohr
724
725        natoms = len(chemical_symbols)
726
727        # set atoms object
728        atoms = ase.Atoms(symbols=chemical_symbols, positions=xyz)
729
730        # if atomic numbers were read in, set them
731        if(len(atomic_numbers) == natoms):
732            atoms.set_atomic_numbers(atomic_numbers)
733
734        # if masses were read in, set them
735        if(len(masses) == natoms):
736            atoms.set_masses(masses)
737
738        return atoms
739