Coverage for ase / utils / cube.py: 100.00%

1import numpy as np 

2from scipy.interpolate import interpn 

3 

4 

5def grid_2d_slice( 

6 spacings, 

7 array, 

8 u, 

9 v, 

10 o=(0, 0, 0), 

11 step=0.02, 

12 size_u=(-10, 10), 

13 size_v=(-10, 10), 

14): 

15 """Extract a 2D slice from a cube file using interpolation. 

16 

17 Works for non-orthogonal cells. 

18 

19 Parameters 

20 ---------- 

21 

22 cube: dict 

23 The cube dict as returned by ase.io.cube.read_cube 

24 

25 u: array_like 

26 The first vector defining the plane 

27 

28 v: array_like 

29 The second vector defining the plane 

30 

31 o: array_like 

32 The origin of the plane 

33 

34 step: float 

35 The step size of the interpolation grid in both directions 

36 

37 size_u: tuple 

38 The size of the interpolation grid in the u direction from the origin 

39 

40 size_v: tuple 

41 The size of the interpolation grid in the v direction from the origin 

42 

43 Returns 

44 ------- 

45 

46 X: np.ndarray 

47 The x coordinates of the interpolation grid 

48 

49 Y: np.ndarray 

50 The y coordinates of the interpolation grid 

51 

52 D: np.ndarray 

53 The interpolated data on the grid 

54 

55 Examples 

56 -------- 

57 

58 From a cube file, we can extract a 2D slice of the density along the 

59 the direction of the first three atoms in the file: 

60 

61 >>> from ase.io.cube import read_cube 

62 >>> from ase.utils.cube import grid_2d_slice 

63 >>> with open(..., 'r') as f: 

64 >>> cube = read_cube(f) 

65 >>> atoms = cube['atoms'] 

66 >>> spacings = cube['spacing'] 

67 >>> array = cube['data'] 

68 >>> u = atoms[1].position - atoms[0].position 

69 >>> v = atoms[2].position - atoms[0].position 

70 >>> o = atoms[0].position 

71 >>> X, Y, D = grid_2d_slice(spacings, array, u, v, o, size_u=(-1, 10), 

72 >>> size_v=(-1, 10)) 

73 >>> # We can now plot the data directly 

74 >>> import matplotlib.pyplot as plt 

75 >>> plt.pcolormesh(X, Y, D) 

76 """ 

77 

78 real_step = np.linalg.norm(spacings, axis=1) 

79 

80 u = np.array(u, dtype=np.float64) 

81 v = np.array(v, dtype=np.float64) 

82 o = np.array(o, dtype=np.float64) 

83 

84 size = array.shape 

85 

86 spacings = np.array(spacings) 

87 array = np.array(array) 

88 

89 cell = spacings * size 

90 

91 lengths = np.linalg.norm(cell, axis=1) 

92 

93 A = cell / lengths[:, None] 

94 

95 ox = np.arange(0, size[0]) * real_step[0] 

96 oy = np.arange(0, size[1]) * real_step[1] 

97 oz = np.arange(0, size[2]) * real_step[2] 

98 

99 u, v = u / np.linalg.norm(u), v / np.linalg.norm(v) 

100 

101 n = np.cross(u, v) 

102 n /= np.linalg.norm(n) 

103 

104 u_perp = np.cross(n, u) 

105 u_perp /= np.linalg.norm(u_perp) 

106 

107 # The basis of the plane 

108 B = np.array([u, u_perp, n]) 

109 Bo = np.dot(B, o) 

110 

111 det = u[0] * v[1] - v[0] * u[1] 

112 

113 if det == 0: 

114 zoff = 0 

115 else: 

116 zoff = ( 

117 (0 - o[1]) * (u[0] * v[2] - v[0] * u[2]) 

118 - (0 - o[0]) * (u[1] * v[2] - v[1] * u[2]) 

119 ) / det + o[2] 

120 

121 zoff = np.dot(B, [0, 0, zoff])[-1] 

122 

123 x, y = np.arange(*size_u, step), np.arange(*size_v, step) 

124 x += Bo[0] 

125 y += Bo[1] 

126 

127 X, Y = np.meshgrid(x, y) 

128 

129 Bvectors = np.stack((X, Y)).reshape(2, -1).T 

130 Bvectors = np.hstack((Bvectors, np.ones((Bvectors.shape[0], 1)) * zoff)) 

131 

132 vectors = np.dot(Bvectors, np.linalg.inv(B).T) 

133 # If the cell is not orthogonal, we need to rotate the vectors 

134 vectors = np.dot(vectors, np.linalg.inv(A)) 

135 

136 # We avoid nan values at boundary 

137 vectors = np.round(vectors, 12) 

138 

139 D = interpn( 

140 (ox, oy, oz), array, vectors, bounds_error=False, method='linear' 

141 ).reshape(X.shape) 

142 

143 return X - Bo[0], Y - Bo[1], D