Coverage for gpaw/utilities/adjust_cell.py: 97%

1import numpy as np 

2 

3from ase import Atoms 

4 

5from gpaw.utilities import h2gpts 

6from gpaw.grid_descriptor import GridDescriptor 

7 

8 

9def adjust_cell(atoms: Atoms, 

10 border: float, 

11 h: float = 0.2, 

12 idiv: int = 4) -> None: 

13 """Adjust the cell such that 

14 1. The vacuum around all atoms is at least border 

15 in non-periodic directions 

16 2. The grid spacing chosen by GPAW will be as similar 

17 as possible to h in all directions 

18 """ 

19 n_pbc = atoms.pbc.sum() 

20 if n_pbc == 3: 

21 return 

22 

23 pos_ac = atoms.get_positions() 

24 lowest_c = np.minimum.reduce(pos_ac) 

25 largest_c = np.maximum.reduce(pos_ac) 

26 

27 for i, v_c, in enumerate(atoms.cell): 

28 if (v_c == 0).all(): 

29 assert not atoms.pbc[i] # pbc with zero cell size make no sense 

30 atoms.cell[i, i] = 1 

31 

32 if n_pbc: 

33 N_c = h2gpts(h, atoms.cell, idiv) 

34 gd = GridDescriptor(N_c, atoms.cell, atoms.pbc) 

35 h_c = gd.get_grid_spacings() 

36 h = 0 

37 for pbc, h1 in zip(atoms.pbc, h_c): 

38 if pbc: 

39 h += h1 / n_pbc 

40 

41 # the optimal h to be set to non-periodic directions 

42 h_c = np.array([h, h, h]) 

43 

44 shift_c = np.zeros(3) 

45 

46 # adjust each cell direction 

47 for i in range(3): 

48 if atoms.pbc[i]: 

49 continue 

50 

51 # cell direction 

52 u_c = atoms.cell[i] / np.linalg.norm(atoms.cell[i]) 

53 

54 extension = (largest_c - lowest_c) * u_c 

55 min_size = extension + 2 * border 

56 

57 h = h_c[i] 

58 N = min_size / h 

59 N = -(N // -idiv) * idiv # ceil div 

60 size = N * h 

61 

62 atoms.cell[i] = size * u_c 

63 

64 # shift structure to the center 

65 shift_c += (size - extension) / 2 * u_c 

66 shift_c -= lowest_c * u_c 

67 

68 atoms.translate(shift_c)