Coverage for gpaw/fd_operators.py: 79%

1"""Finite difference operators.

3This file defines a series of finite difference operators used in grid mode.

4"""

5from __future__ import annotations

7from math import factorial as fact

8from math import pi

9from typing import TYPE_CHECKING

11import numpy as np

12from ase.geometry.minkowski_reduction import reduction_full

13from numpy.fft import fftn, ifftn

14from scipy.spatial import Voronoi

16import gpaw.cgpaw as cgpaw

17from gpaw import debug

18from gpaw.gpu import cupy_is_fake

19from gpaw.grid_descriptor import GridDescriptor

20from gpaw.typing import Array2D, ArrayLike2D

22if TYPE_CHECKING:

23 from gpaw.core import UGDesc

25# Expansion coefficients for finite difference Laplacian. The numbers are

26# from J. R. Chelikowsky et al., Phys. Rev. B 50, 11355 (1994):

27laplace = [[0],

28 [-2, 1],

29 [-5 / 2, 4 / 3, -1 / 12],

30 [-49 / 18, 3 / 2, -3 / 20, 1 / 90],

31 [-205 / 72, 8 / 5, -1 / 5, 8 / 315, -1 / 560],

32 [-5269 / 1800, 5 / 3, -5 / 21, 5 / 126, -5 / 1008, 1 / 3150],

33 [-5369 / 1800, 12 / 7, -15 / 56, 10 / 189, -1 / 112, 2 / 1925,

34 -1 / 16632]]

36derivatives = [[1 / 2],

37 [2 / 3, -1 / 12],

38 [3 / 4, -3 / 20, 1 / 60],

39 [4 / 5, -1 / 5, 4 / 105, -1 / 280]]

42class FDOperator:

43 def __init__(self, coef_p, offset_pc, gd, dtype=float,

44 description=None, xp=np):

45 """FDOperator(coefs, offsets, gd, dtype) -> FDOperator object.

46 """

48 # Is this a central finite-difference type of stencil?

49 cfd = True

50 for offset_c in offset_pc:

51 if sum([offset != 0 for offset in offset_c]) >= 2:

52 cfd = False

54 mp = np.abs(offset_pc).max() # padding

56 # If stencil is strictly larger than any domain, some CPU will

57 # need data from the second-nearest-neighbour domain.

58 # We don't support that. (Typically happens with 1D distributions)

59 assert (mp <= gd.n_c).all(), 'Stencil longer than domain'

60 n_c = gd.n_c

61 M_c = n_c + 2 * mp

62 stride_c = np.array([M_c[1] * M_c[2], M_c[2], 1])

63 offset_p = np.dot(offset_pc, stride_c)

64 coef_p = np.ascontiguousarray(coef_p, float)

65 neighbor_cd = gd.neighbor_cd

67 assert coef_p.ndim == 1

68 assert coef_p.shape == offset_p.shape

69 assert dtype in [float, complex]

71 self.dtype = dtype

72 self.shape = tuple(n_c)

74 if gd.comm.size > 1:

75 comm = gd.comm.get_c_object()

76 else:

77 comm = None

79 assert neighbor_cd.flags.c_contiguous and offset_p.flags.c_contiguous

80 self.mp = mp

81 self.gd = gd

82 self.npoints = len(coef_p)

83 self.coef_p = coef_p

84 self.offset_p = offset_p

85 self.offset_pc = offset_pc

86 self.comm = comm

87 self.cfd = cfd

89 self.xp = xp

90 gpu = xp is not np and not cupy_is_fake

91 self.operator = cgpaw.Operator(coef_p, offset_p, n_c, mp,

92 neighbor_cd, dtype == float,

93 comm, cfd, gpu)

95 if description is None:

96 description = '%d point finite-difference stencil' % self.npoints

97 self.description = description

99 def __str__(self):

100 return '<' + self.description + '>'

101

102 def __call__(self, in_xR, out_xR=None):

103 """New UGArray interface."""

104 if out_xR is None:

105 out_xR = in_xR.new()

106 if self.xp is np:

107 self.operator.apply(in_xR.data, out_xR.data,

108 in_xR.desc.phase_factor_cd)

109 elif cupy_is_fake:

110 self.operator.apply(in_xR.data._data, out_xR.data._data,

111 in_xR.desc.phase_factor_cd)

112 else:

113 assert isinstance(in_xR.data, self.xp.ndarray)

114 assert isinstance(out_xR.data, self.xp.ndarray)

115 self.operator.apply_gpu(in_xR.data.data.ptr, out_xR.data.data.ptr,

116 in_xR.data.shape, in_xR.data.dtype,

117 in_xR.desc.phase_factor_cd)

118 return out_xR

119

120 def apply(self, in_xg, out_xg, phase_cd=None):

121 """Old NumPy interface."""

122 if self.xp is np:

123 self.operator.apply(in_xg, out_xg, phase_cd)

124 elif cupy_is_fake:

125 self.operator.apply(in_xg._data, out_xg._data,

126 phase_cd)

127 else:

128 self.operator.apply_gpu(in_xg.data.ptr,

129 out_xg.data.ptr,

130 in_xg.shape, in_xg.dtype, phase_cd)

131

132 def relax(self, relax_method, f_g, s_g, n, w=None):

133 if self.xp is np:

134 self.operator.relax(relax_method, f_g, s_g, n, w)

135 elif cupy_is_fake:

136 self.operator.relax(relax_method, f_g._data, s_g._data, n, w)

137 else:

138 self.operator.relax_gpu(relax_method,

139 f_g.data.ptr,

140 s_g.data.ptr,

141 n, w)

142

143 def get_diagonal_element(self):

144 return self.operator.get_diagonal_element()

145

146 def get_async_sizes(self):

147 return self.operator.get_async_sizes()

148

149

150if debug:

151 _FDOperator = FDOperator

152

153 class FDOperator(_FDOperator): # type: ignore

154 def apply(self, in_xg, out_xg, phase_cd=None):

155 assert in_xg.shape == out_xg.shape

156 assert in_xg.shape[-3:] == self.shape

157 assert in_xg.flags.c_contiguous

158 assert in_xg.dtype == self.dtype

159 assert out_xg.flags.c_contiguous

160 assert out_xg.dtype == self.dtype

161 assert (self.dtype == float or

162 (phase_cd.dtype == complex and

163 phase_cd.shape == (3, 2)))

164 _FDOperator.apply(self, in_xg, out_xg, phase_cd)

165

166 def relax(self, relax_method, f_g, s_g, n, w=None):

167 assert f_g.shape == self.shape

168 assert s_g.shape == self.shape

169 assert f_g.flags.c_contiguous

170 assert f_g.dtype == float

171 assert s_g.flags.c_contiguous

172 assert s_g.dtype == float

173 assert self.dtype == float

174 _FDOperator.relax(self, relax_method, f_g, s_g, n, w)

175

176

177def Laplace(gd, scale=1.0, n=1, dtype=float, xp=np):

178 if n == 9:

179 return FTLaplace(gd, scale, dtype)

180 else:

181 return GUCLaplace(gd, scale, n, dtype, xp=xp)

182

183

184class GUCLaplace(FDOperator):

185 def __init__(self, gd, scale=1.0, n=1, dtype=float, xp=np):

186 """Laplacian for general non orthorhombic grid.

187

188 gd: GridDescriptor

189 Descriptor for grid.

190 scale: float

191 Scaling factor. Use scale=-0.5 for a kinetic energy operator.

192 n: int

193 Range of stencil. Stencil has O(h^(2n)) error.

194 dtype: float or complex

195 Data-type to work on.

196 """

197

198 # Order the 26 neighbor grid points after length

199 # (reduced to 13 inequivalent):

200 M_ic = np.indices((3, 3, 3)).reshape((3, -3)).T[-13:] - 1

201 u_cv = gd.h_cv / (gd.h_cv**2).sum(1)[:, np.newaxis]**0.5

202 u2_i = (np.dot(M_ic, u_cv)**2).sum(1)

203 i_d = u2_i.argsort()

204

205 # x^2, y^2, z^2, yz, xz, xy:

206 m_mv = np.array([(2, 0, 0), (0, 2, 0), (0, 0, 2),

207 (0, 1, 1), (1, 0, 1), (1, 1, 0)])

208 # Try 3, 4, 5 and 6 of the shortest directions:

209 for D in range(3, 7):

210 h_dv = np.dot(M_ic[i_d[:D]], gd.h_cv)

211 A_md = (h_dv**m_mv[:, np.newaxis, :]).prod(2)

212 # Find best stencil coefficients:

213 a_d, residual, rank, s = np.linalg.lstsq(A_md, [1, 1, 1, 0, 0, 0],

214 rcond=-1)

215 if residual.sum() < 1e-14:

216 if rank != D:

217 raise ValueError(

218 'You have a weird unit cell! '

219 'Try to use the maximally reduced Niggli cell. '

220 'See the ase.build.niggli_reduce() function.')

221 # D directions was OK

222 break

223

224 a_d *= scale

225 offsets = [(0, 0, 0)]

226 coefs = [laplace[n][0] * a_d.sum()]

227 for d in range(D):

228 M_c = M_ic[i_d[d]]

229 offsets.extend(np.arange(1, n + 1)[:, np.newaxis] * M_c)

230 coefs.extend(a_d[d] * np.array(laplace[n][1:]))

231 offsets.extend(np.arange(-1, -n - 1, -1)[:, np.newaxis] * M_c)

232 coefs.extend(a_d[d] * np.array(laplace[n][1:]))

233

234 FDOperator.__init__(self, coefs, offsets, gd, dtype, xp=xp)

235

236 self.description = (

237 '%d*%d+1=%d point O(h^%d) finite-difference Laplacian' %

238 ((self.npoints - 1) // n, n, self.npoints, 2 * n))

239

240

241def find_neighbors(h_cv: ArrayLike2D) -> Array2D:

242 """Find nearest neighbors sorted after distance.

243

244 With ``M_dc = find_neighbors(h_cv)``, the neighbors will be at

245 ``M_dc @ h_cv``, where ``M_dc`` is a (#neighbors, 3)-shaped

246 ndarray of int. If h is a vector pointing at a

247 neighbor grid-point then we don't also include -h in the list.

248

249 Examples:

250

251 >>> h_cv = [[0.1, 0, 0], [0, 0.11, 0], [0, 0, 0.12]]

252 >>> find_neighbors(h_cv)

253 array([[1, 0, 0],

254 [0, 1, 0],

255 [0, 0, 1]])

256 >>> h_cv = [[0.1, 0, 0], [0.2, 0.11, 0], [0, 0, 0.12]]

257 >>> find_neighbors(h_cv)

258 array([[ 1, 0, 0],

259 [-2, 1, 0],

260 [ 0, 0, 1]])

261 """

262 # Do Minkowski reduction:

263 h_bv, U_bc = reduction_full(h_cv) # h_bv = U_bc @ h_cv

264

265 # 27 points: (-1,0,1)x(-1,0,1)x(-1,0,1)

266 M_ib = np.indices((3, 3, 3)).reshape((3, -1)).T - 1

267 h_iv = M_ib.dot(h_bv)

268 voro = Voronoi(h_iv)

269 i_d = [] # List[int]

270 for i1, i2 in voro.ridge_points:

271 if i1 == 13 and i2 > 13: # 13 is the point in the middle

272 i_d.append(i2)

273 elif i2 == 13 and i1 > 13:

274 i_d.append(i1)

275 h2_d = (h_iv[i_d]**2).sum(1)

276 M_db = M_ib[[i_d[d] for d in h2_d.argsort()]]

277

278 # Transform back to origial unreduced space:

279 # h_dv = M_db @ h_bv = M_db @ U_bc @ h_cv = M_dc @ h_cv

280 M_dc = M_db @ U_bc

281 return M_dc

282

283

284class Gradient(FDOperator):

285 def __init__(self,

286 grid: UGDesc | GridDescriptor,

287 v: int,

288 scale=1.0,

289 n=1,

290 dtype=float,

291 xp=np):

292 """Symmetric gradient for general non orthorhombic grid.

293

294 gd: GridDescriptor

295 Descriptor for grid.

296 v: int

297 Direction of gradient: 0, 1, or 2 for x, y and z.

298 scale: float

299 Scaling factor.

300 n: int

301 Range of stencil. Stencil has O(h^(2n)) error.

302 dtype: float or complex

303 Data-type to work on.

304 """

305 if isinstance(grid, GridDescriptor):

306 gd = grid

307 else:

308 gd = grid._gd

309

310 M_dc = find_neighbors(gd.h_cv)

311 h_dv = M_dc @ gd.h_cv # vectors pointing at neighbor grid-points

312 D = len(h_dv) # number of neighbors (3, 4, 5, 6 or 7)

313

314 # Find gradient along 3 directions (n_cv):

315 invh_vc = np.linalg.inv(gd.h_cv)

316 n_cv = (invh_vc / (invh_vc**2).sum(axis=0)**0.5).T

317

318 coef_cd = np.zeros((3, D)) # unknown weights

319 for c, n_v in enumerate(n_cv):

320 # Point on the plane perpendicular to the directions get a

321 # weight of zero

322 ok_d = abs(h_dv.dot(n_v)) > 1e-10

323 h_mv = h_dv[ok_d]

324

325 # The theee equations: A_jm.dot(coef_m) = [1, 0, 0]

326 A_jm = np.array([h_mv.dot(n_v),

327 h_mv.dot(gd.h_cv[c - 1]),

328 h_mv.dot(gd.h_cv[c - 2])])

329

330 U_jm, S_m, V_mm = np.linalg.svd(A_jm, full_matrices=False)

331 coef_m = V_mm.T.dot(np.diag(S_m**-1).dot(U_jm.T.dot([1, 0, 0])))

332 coef_cd[c, ok_d] = coef_m

333

334 # Now get the v-gradient:

335 coef_d = np.linalg.inv(n_cv)[v].dot(coef_cd) * scale

336

337 # Create stencil:

338 offsets = []

339 coefs: list[float] = []

340 stencil = np.array(derivatives[n - 1])

341 for d, c in enumerate(coef_d):

342 if abs(c) < 1e-10:

343 continue

344 M_c = M_dc[d]

345 offsets.extend(np.arange(1, n + 1)[:, np.newaxis] * M_c)

346 coefs.extend(c * stencil)

347 offsets.extend(np.arange(-1, -n - 1, -1)[:, np.newaxis] * M_c)

348 coefs.extend(-c * stencil)

349

350 FDOperator.__init__(self, coefs, offsets, gd, dtype, xp=xp)

351

352 self.description = (

353 'Finite-difference {}-derivative with O(h^{}) error ({} points)'

354 .format('xyz'[v], 2 * n, self.npoints))

355

356

357class LaplaceA(FDOperator):

358 def __init__(self, gd, scale, dtype=float, xp=np):

359 assert gd.orthogonal

360 c = np.divide(-1 / 12, gd.h_cv.diagonal()**2) * scale # Why divide?

361 c0 = c[1] + c[2]

362 c1 = c[0] + c[2]

363 c2 = c[1] + c[0]

364 a = -16.0 * np.sum(c)

365 b = 10.0 * c + 0.125 * a

366 FDOperator.__init__(self,

367 [a,

368 b[0], b[0],

369 b[1], b[1],

370 b[2], b[2],

371 c0, c0, c0, c0,

372 c1, c1, c1, c1,

373 c2, c2, c2, c2],

374 [(0, 0, 0),

375 (-1, 0, 0), (1, 0, 0),

376 (0, -1, 0), (0, 1, 0),

377 (0, 0, -1), (0, 0, 1),

378 (0, -1, -1), (0, -1, 1), (0, 1, -1), (0, 1, 1),

379 (-1, 0, -1), (-1, 0, 1), (1, 0, -1), (1, 0, 1),

380 (-1, -1, 0), (-1, 1, 0), (1, -1, 0), (1, 1, 0)],

381 gd, dtype,

382 'O(h^4) Mehrstellen Laplacian (A)',

383 xp=xp)

384

385

386class LaplaceB(FDOperator):

387 def __init__(self, gd, dtype=float, xp=np):

388 a = 0.5

389 b = 1.0 / 12.0

390 FDOperator.__init__(self,

391 [a,

392 b, b, b, b, b, b],

393 [(0, 0, 0),

394 (-1, 0, 0), (1, 0, 0),

395 (0, -1, 0), (0, 1, 0),

396 (0, 0, -1), (0, 0, 1)],

397 gd, dtype,

398 'O(h^4) Mehrstellen Laplacian (B)',

399 xp=xp)

400

401

402class FTLaplace:

403 def __init__(self, gd, scale, dtype):

404 assert gd.comm.size == 1 and gd.pbc_c.all()

405

406 N_c1 = gd.N_c[:, np.newaxis]

407 i_cq = np.indices(gd.N_c).reshape((3, -1))

408 i_cq += N_c1 // 2

409 i_cq %= N_c1

410 i_cq -= N_c1 // 2

411 B_vc = 2.0 * pi * gd.icell_cv.T

412 k_vq = np.dot(B_vc, i_cq)

413 k_vq *= k_vq

414 self.k2_Q = k_vq.sum(axis=0).reshape(gd.N_c)

415 self.k2_Q *= -scale

416 self.d = 6.0 / gd.h_cv[0, 0]**2

417 self.npoints = 1000

418

419 def apply(self, in_xg, out_xg, phase_cd=None):

420 if in_xg.ndim > 3:

421 for in_g, out_g in zip(in_xg, out_xg):

422 out_g[:] = ifftn(fftn(in_g) * self.k2_Q).real

423 else:

424 out_xg[:] = ifftn(fftn(in_xg) * self.k2_Q).real

425

426 def get_diagonal_element(self):

427 return self.d

428

429

430class OldGradient(FDOperator):

431 def __init__(self, gd, v, scale=1.0, n=1, dtype=float, xp=np):

432 h = (gd.h_cv**2).sum(1)**0.5

433 d = gd.xxxiucell_cv[:, v]

434 A = np.zeros((2 * n + 1, 2 * n + 1))

435 for i, io in enumerate(range(-n, n + 1)):

436 for j in range(2 * n + 1):

437 A[i, j] = io**j / float(fact(j))

438 A[n, 0] = 1.0

439 coefs = np.linalg.inv(A)[1]

440 coefs = np.delete(coefs, len(coefs) // 2)

441 offs = np.delete(np.arange(-n, n + 1), n)

442 coef_p = []

443 offset_pc = []

444 for i in range(3):

445 if abs(d[i]) > 1e-11:

446 coef_p.extend(list(coefs * d[i] / h[i] * scale))

447 offset = np.zeros((2 * n, 3), int)

448 offset[:, i] = offs

449 offset_pc.extend(offset)

450

451 FDOperator.__init__(self, coef_p, offset_pc, gd, dtype,

452 'O(h^%d) %s-gradient stencil' % (2 * n, 'xyz'[v]),

453 xp=xp)