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

1import numpy as np 

2from scipy.interpolate import interpn 

3 

4 

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

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

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

8 

9 Works for non-orthogonal cells. 

10 

11 Parameters: 

12 

13 cube: dict 

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

15 

16 u: array_like 

17 The first vector defining the plane 

18 

19 v: array_like 

20 The second vector defining the plane 

21 

22 o: array_like 

23 The origin of the plane 

24 

25 step: float 

26 The step size of the interpolation grid in both directions 

27 

28 size_u: tuple 

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

30 

31 size_v: tuple 

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

33 

34 Returns: 

35 

36 X: np.ndarray 

37 The x coordinates of the interpolation grid 

38 

39 Y: np.ndarray 

40 The y coordinates of the interpolation grid 

41 

42 D: np.ndarray 

43 The interpolated data on the grid 

44 

45 Examples: 

46 

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

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

49 

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

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

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

53 >>> cube = read_cube(f) 

54 >>> atoms = cube['atoms'] 

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

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

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

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

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

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

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

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

63 >>> import matplotlib.pyplot as plt 

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

65 """ 

66 

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

68 

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

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

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

72 

73 size = array.shape 

74 

75 spacings = np.array(spacings) 

76 array = np.array(array) 

77 

78 cell = spacings * size 

79 

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

81 

82 A = cell / lengths[:, None] 

83 

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

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

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

87 

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

89 

90 n = np.cross(u, v) 

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

92 

93 u_perp = np.cross(n, u) 

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

95 

96 # The basis of the plane 

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

98 Bo = np.dot(B, o) 

99 

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

101 

102 if det == 0: 

103 zoff = 0 

104 else: 

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

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

107 / det + o[2] 

108 

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

110 

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

112 x += Bo[0] 

113 y += Bo[1] 

114 

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

116 

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

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

119 

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

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

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

123 

124 # We avoid nan values at boundary 

125 vectors = np.round(vectors, 12) 

126 

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

128 array, 

129 vectors, 

130 bounds_error=False, 

131 method='linear' 

132 ).reshape(X.shape) 

133 

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