Coverage for gpaw/jellium.py: 81%

1"""Helper classes for doing jellium calculations.""" 

2from math import pi 

3 

4import numpy as np 

5from ase.units import Bohr 

6from gpaw.new.environment import Jellium as JelliumEnv 

7 

8 

9def create_background_charge(**kwargs): 

10 if 'z1' in kwargs: 

11 return JelliumSlab(**kwargs) 

12 return Jellium(**kwargs) 

13 

14 

15class Jellium(): 

16 """ The Jellium object """ 

17 def __init__(self, charge): 

18 """ Initialize the Jellium object 

19 

20 Input: charge, the total Jellium background charge. 

21 """ 

22 self.charge = charge 

23 self.rs = None # the Wigner-Seitz radius 

24 self.volume = None 

25 self.mask_g = None 

26 self.gd = None 

27 

28 def todict(self): 

29 return {'charge': self.charge} 

30 

31 def set_grid_descriptor(self, gd): 

32 """ Set the grid descriptor for the Jellium background charge""" 

33 self.gd = gd 

34 self.mask_g = self.get_mask().astype(float) 

35 self.volume = self.gd.comm.sum_scalar(self.mask_g.sum()) * self.gd.dv 

36 if self.charge != 0.: 

37 self.rs = (3 / pi / 4 * self.volume / abs(self.charge))**(1 / 3) 

38 else: 

39 self.rs = np.inf # Avoid having to see numpy's warning. 

40 

41 def get_mask(self): 

42 """Choose which grid points are inside the jellium. 

43 

44 gd: grid descriptor 

45 

46 Return ndarray of ones and zeros indicating where the jellium 

47 is. This implementation will put the positive background in the 

48 whole cell. Overwrite this method in subclasses.""" 

49 

50 return self.gd.zeros() + 1.0 

51 

52 def add_charge_to(self, rhot_g): 

53 """ Add Jellium background charge to pseudo charge density rhot_g""" 

54 rhot_g -= self.mask_g * (self.charge / self.volume) 

55 

56 def add_fourier_space_charge_to(self, pd, rhot_q): 

57 rhot_g = pd.gd.zeros() 

58 self.add_charge_to(rhot_g) 

59 rhot_q += pd.fft(rhot_g) 

60 

61 def build(self, 

62 setups, 

63 grid, 

64 relpos_ac, 

65 log, 

66 comm, 

67 **kwargs) -> JelliumEnv: 

68 self.set_grid_descriptor(grid._gd) 

69 return JelliumEnv(self, len(relpos_ac), grid) 

70 

71 

72class JelliumSlab(Jellium): 

73 """ The Jellium slab object """ 

74 def __init__(self, charge, z1, z2): 

75 """Put the positive background charge where z1 < z < z2. 

76 

77 z1: float 

78 Position of lower surface in Angstrom units. 

79 z2: float 

80 Position of upper surface in Angstrom units.""" 

81 Jellium.__init__(self, charge) 

82 self.z1 = (z1 - 0.0001) / Bohr 

83 self.z2 = (z2 - 0.0001) / Bohr 

84 

85 def todict(self): 

86 dct = Jellium.todict(self) 

87 dct.update(z1=self.z1 * Bohr + 0.0001, 

88 z2=self.z2 * Bohr + 0.0001) 

89 return dct 

90 

91 def get_mask(self): 

92 r_gv = self.gd.get_grid_point_coordinates().transpose((1, 2, 3, 0)) 

93 # r_gv: 4-dimensional ndarray 

94 # positions of the grid points in Bohr units. 

95 return np.logical_and(r_gv[:, :, :, 2] > self.z1, 

96 r_gv[:, :, :, 2] < self.z2)