Coverage for gpaw/new/external_potential.py: 28%

1from __future__ import annotations 

2 

3import numpy as np 

4from ase.units import Ha, Bohr 

5 

6from gpaw.typing import Array1D, Array2D, ArrayLike1D 

7from gpaw.core import UGArray 

8from gpaw.new.density import Density 

9 

10 

11def create_external_potential(params: dict) -> ExternalPotential: 

12 if not params: 

13 return ExternalPotential() 

14 params = params.copy() 

15 name = params.pop('name') 

16 if name == 'BField': 

17 return BField(**params) 

18 if name == 'ConstantElectricField': 

19 return ConstantElectricField(**params) 

20 raise ValueError 

21 

22 

23class ExternalPotential: 

24 def update_potential(self, 

25 vt_sR: UGArray, 

26 density) -> float: 

27 return 0.0 

28 

29 def add_paw_correction(self, Delta_p: Array1D, dH_sp: Array2D) -> float: 

30 return 0.0 

31 

32 

33class ConstantElectricField(ExternalPotential): 

34 def __init__(self, strength, direction=[0, 0, 1], tolerance=1e-7): 

35 """External constant electric field. 

36 

37 strength: float 

38 Field strength in V/Ang. 

39 direction: vector 

40 Polarization direction. 

41 """ 

42 self.strength = strength * Bohr / Ha 

43 self.direction_v = np.array(direction, dtype=float) 

44 self.direction_v /= np.linalg.norm(self.direction_v) 

45 self.field_v = self.strength * self.direction_v 

46 self.tolerance = tolerance 

47 self.name = 'ConstantElectricField' 

48 

49 def __str__(self): 

50 return ('Constant electric field: ' 

51 '({:.3f}, {:.3f}, {:.3f}) V/Ang' 

52 .format(*(self.field_v * Ha / Bohr))) 

53 

54 def update_potential(self, 

55 vt_sR: UGArray, 

56 density) -> float: 

57 grid = vt_sR.desc 

58 L_c = grid.cell_cv @ self.direction_v 

59 (axis,) = np.where(abs(L_c) > self.tolerance)[0] 

60 # assert not grid.pbc_c[axis] 

61 relpos_r = np.linspace(grid.start_c[axis], 

62 grid.end_c[axis], 

63 grid.mysize_c[axis], 

64 endpoint=False) / grid.size_c[axis] 

65 v_r = L_c[axis] * (relpos_r - 0.5) * self.strength 

66 if grid.start_c[axis] == 0: 

67 v_r[0] = 0.0 

68 vt_sR.data += v_r.reshape([1] + 

69 [-1 if c == axis else 1 for c in range(3)]) 

70 return 0.0 

71 

72 

73class BField(ExternalPotential): 

74 def __init__(self, field: ArrayLike1D): 

75 """Constant magnetic field. 

76 

77 field: 

78 B-field vector in units of Ha/bohr-magnoton. 

79 """ 

80 self.field_v = np.array(field) / Ha 

81 assert self.field_v.shape == (3,) 

82 

83 def update_potential(self, 

84 vt_sR: UGArray, 

85 density: Density) -> float: 

86 magmom_v, _ = density.calculate_magnetic_moments() 

87 eext = -self.field_v @ magmom_v 

88 ncomponents = len(vt_sR) 

89 if ncomponents == 2: 

90 assert (self.field_v[:2] == 0.0).all() 

91 vt_sR.data[0] -= self.field_v[2] 

92 vt_sR.data[1] += self.field_v[2] 

93 elif ncomponents == 4: 

94 vt_sR.data[1:] = -self.field_v.reshape((3, 1, 1, 1)) 

95 else: 

96 1 / 0 

97 return eext 

98 

99 def add_paw_correction(self, Delta_p: Array1D, dH_sp: Array2D) -> float: 

100 if len(dH_sp) == 2: 

101 c = (4 * np.pi)**0.5 * self.field_v[2] 

102 dH_sp[0] -= c * Delta_p 

103 dH_sp[1] += c * Delta_p 

104 else: 

105 c_vp = (4 * np.pi)**0.5 * self.field_v[:, np.newaxis] 

106 dH_sp[1:] -= c_vp * Delta_p 

107 return 0.0