Coverage for /builds/ase/ase/ase/calculators/test.py: 89.43%
123 statements
« prev ^ index » next coverage.py v7.5.3, created at 2025-08-02 00:12 +0000
« prev ^ index » next coverage.py v7.5.3, created at 2025-08-02 00:12 +0000
1# fmt: off
3from math import pi
5import numpy as np
7from ase.atoms import Atoms
8from ase.calculators.calculator import Calculator, kpts2ndarray
9from ase.calculators.fd import calculate_numerical_forces
10from ase.units import Bohr, Ha
13def make_test_dft_calculation():
14 a = b = 2.0
15 c = 6.0
16 atoms = Atoms(positions=[(0, 0, c / 2)],
17 symbols='H',
18 pbc=(1, 1, 0),
19 cell=(a, b, c),
20 calculator=TestCalculator())
21 return atoms
24class TestCalculator:
25 def __init__(self, nk=8):
26 assert nk % 2 == 0
27 bzk = []
28 weights = []
29 ibzk = []
30 w = 1.0 / nk**2
31 for i in range(-nk + 1, nk, 2):
32 for j in range(-nk + 1, nk, 2):
33 k = (0.5 * i / nk, 0.5 * j / nk, 0)
34 bzk.append(k)
35 if i >= j > 0:
36 ibzk.append(k)
37 if i == j:
38 weights.append(4 * w)
39 else:
40 weights.append(8 * w)
41 assert abs(sum(weights) - 1.0) < 1e-12
42 self.bzk = np.array(bzk)
43 self.ibzk = np.array(ibzk)
44 self.weights = np.array(weights)
46 # Calculate eigenvalues and wave functions:
47 self.init()
49 def init(self):
50 nibzk = len(self.weights)
51 nbands = 1
53 V = -1.0
54 self.eps = 2 * V * (np.cos(2 * pi * self.ibzk[:, 0]) +
55 np.cos(2 * pi * self.ibzk[:, 1]))
56 self.eps.shape = (nibzk, nbands)
58 self.psi = np.zeros((nibzk, 20, 20, 60), complex)
59 phi = np.empty((2, 2, 20, 20, 60))
60 z = np.linspace(-1.5, 1.5, 60, endpoint=False)
61 for i in range(2):
62 x = np.linspace(0, 1, 20, endpoint=False) - i
63 for j in range(2):
64 y = np.linspace(0, 1, 20, endpoint=False) - j
65 r = (((x[:, None]**2 +
66 y**2)[:, :, None] +
67 z**2)**0.5).clip(0, 1)
68 phi = 1.0 - r**2 * (3.0 - 2.0 * r)
69 phase = np.exp(pi * 2j * np.dot(self.ibzk, (i, j, 0)))
70 self.psi += phase[:, None, None, None] * phi
72 def get_pseudo_wave_function(self, band=0, kpt=0, spin=0):
73 assert spin == 0 and band == 0
74 return self.psi[kpt]
76 def get_eigenvalues(self, kpt=0, spin=0):
77 assert spin == 0
78 return self.eps[kpt]
80 def get_number_of_bands(self):
81 return 1
83 def get_k_point_weights(self):
84 return self.weights
86 def get_number_of_spins(self):
87 return 1
89 def get_fermi_level(self):
90 return 0.0
92 def get_pseudo_density(self):
93 n = 0.0
94 for w, eps, psi in zip(self.weights, self.eps[:, 0], self.psi):
95 if eps >= 0.0:
96 continue
97 n += w * (psi * psi.conj()).real
99 n[1:] += n[:0:-1].copy()
100 n[:, 1:] += n[:, :0:-1].copy()
101 n += n.transpose((1, 0, 2)).copy()
102 n /= 8
103 return n
106class TestPotential(Calculator):
107 implemented_properties = ['energy', 'forces']
109 def calculate(self, atoms, properties, system_changes):
110 Calculator.calculate(self, atoms, properties, system_changes)
111 E = 0.0
112 R = atoms.positions
113 F = np.zeros_like(R)
114 for a, r in enumerate(R):
115 D = R - r
116 d = (D**2).sum(1)**0.5
117 x = d - 1.0
118 E += np.vdot(x, x)
119 d[a] = 1
120 F -= (x / d)[:, None] * D
121 energy = 0.25 * E
122 self.results = {'energy': energy, 'forces': F}
125class FreeElectrons(Calculator):
126 """Free-electron band calculator.
128 Parameters:
130 nvalence: int
131 Number of electrons
132 kpts: dict
133 K-point specification.
135 Example:
136 >>> from ase.calculators.test import FreeElectrons
137 >>> calc = FreeElectrons(nvalence=1, kpts={'path': 'GXL'})
138 """
140 implemented_properties = ['energy']
141 default_parameters = {'kpts': np.zeros((1, 3)),
142 'nvalence': 0.0,
143 'nbands': 20,
144 'gridsize': 7}
146 def calculate(self, atoms, properties, system_changes):
147 Calculator.calculate(self, atoms)
148 self.kpts = kpts2ndarray(self.parameters.kpts, atoms)
149 icell = atoms.cell.reciprocal() * 2 * np.pi * Bohr
150 n = self.parameters.gridsize
151 offsets = np.indices((n, n, n)).T.reshape((n**3, 1, 3)) - n // 2
152 eps = 0.5 * (np.dot(self.kpts + offsets, icell)**2).sum(2).T
153 eps.sort()
154 self.eigenvalues = eps[:, :self.parameters.nbands] * Ha
155 self.results = {'energy': 0.0}
157 def get_eigenvalues(self, kpt, spin=0):
158 assert spin == 0
159 return self.eigenvalues[kpt].copy()
161 def get_fermi_level(self):
162 v = self.atoms.get_volume() / Bohr**3
163 kF = (self.parameters.nvalence / v * 3 * np.pi**2)**(1 / 3)
164 return 0.5 * kF**2 * Ha
166 def get_ibz_k_points(self):
167 return self.kpts.copy()
169 def get_number_of_spins(self):
170 return 1
173def gradient_test(atoms, indices=None):
174 """
175 Use numeric_force to compare analytical and numerical forces on atoms
177 If indices is None, test is done on all atoms.
178 """
179 if indices is None:
180 indices = range(len(atoms))
181 f = atoms.get_forces()[indices]
182 print('{:>16} {:>20}'.format('eps', 'max(abs(df))'))
183 for eps in np.logspace(-1, -8, 8):
184 fn = calculate_numerical_forces(atoms, eps, indices)
185 print(f'{eps:16.12f} {abs(fn - f).max():20.12f}')
186 return f, fn