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

1# fmt: off 

2 

3import numpy as np 

4from scipy.interpolate import interpn 

5 

6 

7def grid_2d_slice(spacings, array, u, v, o=(0, 0, 0), step=0.02, 

8 size_u=(-10, 10), size_v=(-10, 10)): 

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

10 

11 Works for non-orthogonal cells. 

12 

13 Parameters: 

14 

15 cube: dict 

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

17 

18 u: array_like 

19 The first vector defining the plane 

20 

21 v: array_like 

22 The second vector defining the plane 

23 

24 o: array_like 

25 The origin of the plane 

26 

27 step: float 

28 The step size of the interpolation grid in both directions 

29 

30 size_u: tuple 

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

32 

33 size_v: tuple 

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

35 

36 Returns: 

37 

38 X: np.ndarray 

39 The x coordinates of the interpolation grid 

40 

41 Y: np.ndarray 

42 The y coordinates of the interpolation grid 

43 

44 D: np.ndarray 

45 The interpolated data on the grid 

46 

47 Examples: 

48 

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

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

51 

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

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

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

55 >>> cube = read_cube(f) 

56 >>> atoms = cube['atoms'] 

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

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

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

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

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

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

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

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

65 >>> import matplotlib.pyplot as plt 

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

67 """ 

68 

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

70 

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

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

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

74 

75 size = array.shape 

76 

77 spacings = np.array(spacings) 

78 array = np.array(array) 

79 

80 cell = spacings * size 

81 

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

83 

84 A = cell / lengths[:, None] 

85 

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

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

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

89 

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

91 

92 n = np.cross(u, v) 

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

94 

95 u_perp = np.cross(n, u) 

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

97 

98 # The basis of the plane 

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

100 Bo = np.dot(B, o) 

101 

102 det = (u[0] * v[1] - v[0] * u[1]) 

103 

104 if det == 0: 

105 zoff = 0 

106 else: 

107 zoff = ((0 - o[1]) * (u[0] * v[2] - v[0] * u[2]) - 

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

109 / det + o[2] 

110 

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

112 

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

114 x += Bo[0] 

115 y += Bo[1] 

116 

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

118 

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

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

121 

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

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

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

125 

126 # We avoid nan values at boundary 

127 vectors = np.round(vectors, 12) 

128 

129 D = interpn((ox, oy, oz), 

130 array, 

131 vectors, 

132 bounds_error=False, 

133 method='linear' 

134 ).reshape(X.shape) 

135 

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