1import numpy as np
2
3from ase.calculators.calculator import Calculator, all_properties
4from ase.calculators.calculator import PropertyNotImplementedError
5
6
7class SinglePointCalculator(Calculator):
8    """Special calculator for a single configuration.
9
10    Used to remember the energy, force and stress for a given
11    configuration.  If the positions, atomic numbers, unit cell, or
12    boundary conditions are changed, then asking for
13    energy/forces/stress will raise an exception."""
14
15    name = 'unknown'
16
17    def __init__(self, atoms, **results):
18        """Save energy, forces, stress, ... for the current configuration."""
19        Calculator.__init__(self)
20        self.results = {}
21        for property, value in results.items():
22            assert property in all_properties
23            if value is None:
24                continue
25            if property in ['energy', 'magmom', 'free_energy']:
26                self.results[property] = value
27            else:
28                self.results[property] = np.array(value, float)
29        self.atoms = atoms.copy()
30
31    def __str__(self):
32        tokens = []
33        for key, val in sorted(self.results.items()):
34            if np.isscalar(val):
35                txt = '{}={}'.format(key, val)
36            else:
37                txt = '{}=...'.format(key)
38            tokens.append(txt)
39        return '{}({})'.format(self.__class__.__name__, ', '.join(tokens))
40
41    def get_property(self, name, atoms=None, allow_calculation=True):
42        if atoms is None:
43            atoms = self.atoms
44        if name not in self.results or self.check_state(atoms):
45            if allow_calculation:
46                raise PropertyNotImplementedError(
47                    'The property "{0}" is not available.'.format(name))
48            return None
49
50        result = self.results[name]
51        if isinstance(result, np.ndarray):
52            result = result.copy()
53        return result
54
55
56class SinglePointKPoint:
57    def __init__(self, weight, s, k, eps_n=[], f_n=[]):
58        self.weight = weight
59        self.s = s  # spin index
60        self.k = k  # k-point index
61        self.eps_n = eps_n
62        self.f_n = f_n
63
64
65def arrays_to_kpoints(eigenvalues, occupations, weights):
66    """Helper function for building SinglePointKPoints.
67
68    Convert eigenvalue, occupation, and weight arrays to list of
69    SinglePointKPoint objects."""
70    nspins, nkpts, nbands = eigenvalues.shape
71    assert eigenvalues.shape == occupations.shape
72    assert len(weights) == nkpts
73    kpts = []
74    for s in range(nspins):
75        for k in range(nkpts):
76            kpt = SinglePointKPoint(
77                weight=weights[k], s=s, k=k,
78                eps_n=eigenvalues[s, k], f_n=occupations[s, k])
79            kpts.append(kpt)
80    return kpts
81
82
83class SinglePointDFTCalculator(SinglePointCalculator):
84    def __init__(self, atoms,
85                 efermi=None, bzkpts=None, ibzkpts=None, bz2ibz=None,
86                 **results):
87        self.bz_kpts = bzkpts
88        self.ibz_kpts = ibzkpts
89        self.bz2ibz = bz2ibz
90        self.eFermi = efermi
91
92        SinglePointCalculator.__init__(self, atoms, **results)
93        self.kpts = None
94
95    def get_fermi_level(self):
96        """Return the Fermi-level(s)."""
97        return self.eFermi
98
99    def get_bz_to_ibz_map(self):
100        return self.bz2ibz
101
102    def get_bz_k_points(self):
103        """Return the k-points."""
104        return self.bz_kpts
105
106    def get_number_of_spins(self):
107        """Return the number of spins in the calculation.
108
109        Spin-paired calculations: 1, spin-polarized calculation: 2."""
110        if self.kpts is not None:
111            nspin = set()
112            for kpt in self.kpts:
113                nspin.add(kpt.s)
114            return len(nspin)
115        return None
116
117    def get_spin_polarized(self):
118        """Is it a spin-polarized calculation?"""
119        nos = self.get_number_of_spins()
120        if nos is not None:
121            return nos == 2
122        return None
123
124    def get_ibz_k_points(self):
125        """Return k-points in the irreducible part of the Brillouin zone."""
126        return self.ibz_kpts
127
128    def get_kpt(self, kpt=0, spin=0):
129        if self.kpts is not None:
130            counter = 0
131            for kpoint in self.kpts:
132                if kpoint.s == spin:
133                    if kpt == counter:
134                        return kpoint
135                    counter += 1
136        return None
137
138    def get_k_point_weights(self):
139        """ Retunrs the weights of the k points """
140        if self.kpts is not None:
141            weights = []
142            for kpoint in self.kpts:
143                if kpoint.s == 0:
144                    weights.append(kpoint.weight)
145            return np.array(weights)
146        return None
147
148    def get_occupation_numbers(self, kpt=0, spin=0):
149        """Return occupation number array."""
150        kpoint = self.get_kpt(kpt, spin)
151        if kpoint is not None:
152            return kpoint.f_n
153        return None
154
155    def get_eigenvalues(self, kpt=0, spin=0):
156        """Return eigenvalue array."""
157        kpoint = self.get_kpt(kpt, spin)
158        if kpoint is not None:
159            return kpoint.eps_n
160        return None
161
162    def get_homo_lumo(self):
163        """Return HOMO and LUMO energies."""
164        if self.kpts is None:
165            raise RuntimeError('No kpts')
166        eHs = []
167        eLs = []
168        for kpt in self.kpts:
169            eH, eL = self.get_homo_lumo_by_spin(kpt.s)
170            eHs.append(eH)
171            eLs.append(eL)
172        return np.array(eHs).max(), np.array(eLs).min()
173
174    def get_homo_lumo_by_spin(self, spin=0):
175        """Return HOMO and LUMO energies for a given spin."""
176        if self.kpts is None:
177            raise RuntimeError('No kpts')
178        for kpt in self.kpts:
179            if kpt.s == spin:
180                break
181        else:
182            raise RuntimeError('No k-point with spin {0}'.format(spin))
183        if self.eFermi is None:
184            raise RuntimeError('Fermi level is not available')
185        eH = -1.e32
186        eL = 1.e32
187        for kpt in self.kpts:
188            if kpt.s == spin:
189                for e in kpt.eps_n:
190                    if e <= self.eFermi:
191                        eH = max(eH, e)
192                    else:
193                        eL = min(eL, e)
194        return eH, eL
195