Coverage for /builds/ase/ase/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 cube: dict 

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

23 

24 u: array_like 

25 The first vector defining the plane 

26 

27 v: array_like 

28 The second vector defining the plane 

29 

30 o: array_like 

31 The origin of the plane 

32 

33 step: float 

34 The step size of the interpolation grid in both directions 

35 

36 size_u: tuple 

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

38 

39 size_v: tuple 

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

41 

42 Returns: 

43 

44 X: np.ndarray 

45 The x coordinates of the interpolation grid 

46 

47 Y: np.ndarray 

48 The y coordinates of the interpolation grid 

49 

50 D: np.ndarray 

51 The interpolated data on the grid 

52 

53 Examples: 

54 

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

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

57 

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

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

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

61 >>> cube = read_cube(f) 

62 >>> atoms = cube['atoms'] 

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

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

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

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

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

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

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

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

71 >>> import matplotlib.pyplot as plt 

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

73 """ 

74 

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

76 

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

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

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

80 

81 size = array.shape 

82 

83 spacings = np.array(spacings) 

84 array = np.array(array) 

85 

86 cell = spacings * size 

87 

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

89 

90 A = cell / lengths[:, None] 

91 

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

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

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

95 

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

97 

98 n = np.cross(u, v) 

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

100 

101 u_perp = np.cross(n, u) 

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

103 

104 # The basis of the plane 

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

106 Bo = np.dot(B, o) 

107 

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

109 

110 if det == 0: 

111 zoff = 0 

112 else: 

113 zoff = ( 

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

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

116 ) / det + o[2] 

117 

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

119 

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

121 x += Bo[0] 

122 y += Bo[1] 

123 

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

125 

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

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

128 

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

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

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

132 

133 # We avoid nan values at boundary 

134 vectors = np.round(vectors, 12) 

135 

136 D = interpn( 

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

138 ).reshape(X.shape) 

139 

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