1"""This module defines an ASE interface to NWchem
2
3https://nwchemgit.github.io
4"""
5import os
6import numpy as np
7
8from ase import io
9from ase.units import Hartree
10from ase.calculators.calculator import FileIOCalculator
11from ase.spectrum.band_structure import BandStructure
12
13
14class NWChem(FileIOCalculator):
15    implemented_properties = ['energy', 'forces', 'stress', 'dipole']
16    command = 'nwchem PREFIX.nwi > PREFIX.nwo'
17    accepts_bandpath_keyword = True
18    discard_results_on_any_change = True
19
20    def __init__(self, restart=None,
21                 ignore_bad_restart_file=FileIOCalculator._deprecated,
22                 label='nwchem', atoms=None, command=None, **kwargs):
23        """
24        NWChem keywords are specified using (potentially nested)
25        dictionaries. Consider the following input file block:
26
27        >>> dft
28        >>>   odft
29        >>>   mult 2
30        >>>   convergence energy 1e-9 density 1e-7 gradient 5e-6
31        >>> end
32
33        This can be generated by the NWChem calculator by using the
34        following settings:
35
36        >>> calc = NWChem(dft={'odft': None,
37        >>>                    'mult': 2,
38        >>>                    'convergence': {'energy': 1e-9,
39        >>>                                    'density': 1e-7,
40        >>>                                    'gradient': 5e-6,
41        >>>                                    },
42        >>>                    },
43        >>>               )
44
45        In addition, the calculator supports several special keywords:
46
47        theory: str
48            Which NWChem module should be used to calculate the
49            energies and forces. Supported values are ``'dft'``,
50            ``'scf'``, ``'mp2'``, ``'ccsd'``, ``'tce'``, ``'tddft'``,
51            ``'pspw'``, ``'band'``, and ``'paw'``. If not provided, the
52            calculator will attempt to guess which theory to use based
53            on the keywords provided by the user.
54        xc: str
55            The exchange-correlation functional to use. Only relevant
56            for DFT calculations.
57        task: str
58            What type of calculation is to be performed, e.g.
59            ``'energy'``, ``'gradient'``, ``'optimize'``, etc. When
60            using ``'SocketIOCalculator'``, ``task`` should be set
61            to ``'optimize'``. In most other circumstances, ``task``
62            should not be set manually.
63        basis: str or dict
64            Which basis set to use for gaussian-type orbital
65            calculations. Set to a string to use the same basis for all
66            elements. To use a different basis for different elements,
67            provide a dict of the form:
68
69            >>> calc = NWChem(...,
70            >>>               basis={'O': '3-21G',
71            >>>                      'Si': '6-31g'})
72
73        basispar: str
74            Additional keywords to put in the NWChem ``basis`` block,
75            e.g. ``'rel'`` for relativistic bases.
76        symmetry: int or str
77            The point group (for gaussian-type orbital calculations) or
78            space group (for plane-wave calculations) of the system.
79            Supports both group names (e.g. ``'c2v'``, ``'Fm3m'``) and
80            numbers (e.g. ``225``).
81        autosym: bool
82            Whether NWChem should automatically determine the symmetry
83            of the structure (defaults to ``False``).
84        center: bool
85            Whether NWChem should automatically center the structure
86            (defaults to ``False``). Enable at your own risk.
87        autoz: bool
88            Whether NWChem should automatically construct a Z-matrix
89            for your molecular system (defaults to ``False``).
90        geompar: str
91            Additional keywords to put in the NWChem `geometry` block,
92            e.g. ``'nucleus finite'`` for gaussian-shaped nuclear
93            charges. Do not set ``'autosym'``, ``'center'``, or
94            ``'autoz'`` in this way; instead, use the appropriate
95            keyword described above for these settings.
96        set: dict
97            Used to manually create or modify entries in the NWChem
98            rtdb. For example, the following settings enable
99            pseudopotential filtering for plane-wave calculations:
100
101            >>> set nwpw:kbpp_ray .true.
102            >>> set nwpw:kbpp_filter .true.
103
104            These settings are generated by the NWChem calculator by
105            passing the arguments:
106
107            >>> calc = NWChem(...,
108            >>>               set={'nwpw:kbpp_ray': True,
109            >>>                    'nwpw:kbpp_filter': True})
110
111        kpts: (int, int, int), or dict
112            Indicates which k-point mesh to use. Supported syntax is
113            similar to that of GPAW. Implies ``theory='band'``.
114        bandpath: BandPath object
115            The band path to use for a band structure calculation.
116            Implies ``theory='band'``.
117        """
118        FileIOCalculator.__init__(self, restart, ignore_bad_restart_file,
119                                  label, atoms, command, **kwargs)
120        self.calc = None
121
122    def write_input(self, atoms, properties=None, system_changes=None):
123        FileIOCalculator.write_input(self, atoms, properties, system_changes)
124
125        # Prepare perm and scratch directories
126        perm = os.path.abspath(self.parameters.get('perm', self.label))
127        scratch = os.path.abspath(self.parameters.get('scratch', self.label))
128        os.makedirs(perm, exist_ok=True)
129        os.makedirs(scratch, exist_ok=True)
130
131        io.write(self.label + '.nwi', atoms, properties=properties,
132                 label=self.label, **self.parameters)
133
134    def read_results(self):
135        output = io.read(self.label + '.nwo')
136        self.calc = output.calc
137        self.results = output.calc.results
138
139    def band_structure(self):
140        self.calculate()
141        perm = self.parameters.get('perm', self.label)
142        if self.calc.get_spin_polarized():
143            alpha = np.loadtxt(os.path.join(perm, self.label + '.alpha_band'))
144            beta = np.loadtxt(os.path.join(perm, self.label + '.beta_band'))
145            energies = np.array([alpha[:, 1:], beta[:, 1:]]) * Hartree
146        else:
147            data = np.loadtxt(os.path.join(perm,
148                                           self.label + '.restricted_band'))
149            energies = data[np.newaxis, :, 1:] * Hartree
150        eref = self.calc.get_fermi_level()
151        if eref is None:
152            eref = 0.
153        return BandStructure(self.parameters.bandpath, energies, eref)
154