Coverage for tcvx21/grillix_post/lineouts/csr_linear_interpolation_m.py: 82%

1import numpy as np 

2from scipy.sparse import csr_matrix 

3 

4 

5def _ij_to_l_index(x_unstructured: np.ndarray, y_unstructured: np.ndarray, x_index: int, y_index: int) -> int: 

6 """Converts (i, j) x/y indices to (l) grid indices""" 

7 

8 grid_index = np.where(np.logical_and(x_unstructured[x_index] == x_unstructured, y_unstructured[y_index] == y_unstructured))[0] 

9 

10 if len(grid_index) == 1: 

11 return grid_index[0] 

12 elif len(grid_index) == 0: 

13 return -1 

14 else: 

15 raise ValueError("Multiple grid_indices cannot match a single (x_index, y_index)") 

16 

17def _find_nearest_neighbor(x_unstructured: np.ndarray, y_unstructured: np.ndarray, x_query: float, y_query: float) -> int: 

18 """Finds the nearest point to a given query point""" 

19 

20 cartesian_distance = np.sqrt( (x_unstructured - x_query)**2 + (y_unstructured - y_query)**2 ) 

21 

22 return np.argmin(cartesian_distance) 

23 

24def _weightings_interpolate(x_unstructured: np.ndarray, y_unstructured: np.ndarray, x_query: float, y_query: float) -> tuple: 

25 """ 

26 Returns the indices and weightings of neighbours for linear interpolation 

27 on the unstructured (x,y,z) tricolumn data 

28 """ 

29 assert(x_unstructured.ndim == 1) 

30 assert(y_unstructured.ndim == 1) 

31 assert(x_unstructured.size == y_unstructured.size) 

32 

33 if np.any([ 

34 x_query < x_unstructured.min(), 

35 x_query > x_unstructured.max(), 

36 y_query < y_unstructured.min(), 

37 y_query > y_unstructured.max() 

38 ]): 

39 return 0, 0, 0, 0, 0.0, 0.0, 0.0, 0.0 

40 

41 diff_x = x_query - x_unstructured 

42 diff_y = y_query - y_unstructured 

43 

44 ix0 = np.nanargmin(np.where(diff_x >= 0, diff_x, np.nan)) 

45 ix1 = np.nanargmax(np.where(diff_x <= 0, diff_x, np.nan)) 

46 

47 iy0 = np.nanargmin(np.where(diff_y >= 0, diff_y, np.nan)) 

48 iy1 = np.nanargmax(np.where(diff_y <= 0, diff_y, np.nan)) 

49 

50 i00 = _ij_to_l_index(x_unstructured, y_unstructured, ix0, iy0) 

51 i01 = _ij_to_l_index(x_unstructured, y_unstructured, ix0, iy1) 

52 i10 = _ij_to_l_index(x_unstructured, y_unstructured, ix1, iy0) 

53 i11 = _ij_to_l_index(x_unstructured, y_unstructured, ix1, iy1) 

54 

55 # If any points not found 

56 if np.any(np.array([i00, i01, i10, i11]) < 0): 

57 nearest_neighbour = _find_nearest_neighbor(x_unstructured, y_unstructured, x_query, y_query) 

58 # Set the interpolation weighting to take a single point value 

59 return nearest_neighbour, 0, 0, 0, 1.0, 0.0, 0.0, 0.0 

60 

61 # Find the distance from the cell edges to the query point 

62 dx0 = x_query - x_unstructured[ix0] 

63 dy0 = y_query - y_unstructured[iy0] 

64 dx1 = x_unstructured[ix1] - x_query 

65 dy1 = y_unstructured[iy1] - y_query 

66 

67 # If edge-of-grid in x-direction, collapse to nearest neighbour in x 

68 if ix0 == ix1: 

69 wx0 = 1.0 

70 wx1 = 0.0 

71 else: 

72 wx0 = dx1/(dx0+dx1) 

73 wx1 = dx0/(dx0+dx1) 

74 

75 # If edge-of-grid in y-direction, collapse to nearest neighbour in y 

76 if iy0 == iy1: 

77 wy0 = 1.0 

78 wy1 = 0.0 

79 else: 

80 wy0 = dy1/(dy0+dy1) 

81 wy1 = dy0/(dy0+dy1) 

82 

83 return i00, i10, i01, i11, wx0*wy0, wx1*wy0, wx0*wy1, wx1*wy1 

84 

85def make_matrix_interp(x_unstructured: np.ndarray, y_unstructured: np.ndarray, x_queries: np.ndarray, y_queries: np.ndarray) -> \ 

86 csr_matrix: 

87 """ 

88 Makes a csr matrix which extracts the values at points defined 

89 by x_queries, y_queries. 

90 csr_matrix * unstructured_data = values at queries 

91 """ 

92 assert(x_unstructured.ndim == 1) 

93 assert(y_unstructured.ndim == 1) 

94 assert(x_unstructured.size == y_unstructured.size) 

95 assert(x_queries.ndim == 1) 

96 assert(y_queries.ndim == 1) 

97 assert(x_queries.size == y_queries.size) 

98 

99 nz = 0 

100 indi = [] 

101 indj = [] 

102 val = [] 

103 

104 # For each query, find the bilinear interpolation stencil 

105 for l in range(x_queries.size): 

106 indi.append(nz) 

107 

108 i00, i10, i01, i11, w00, w10, w01, w11 = \ 

109 _weightings_interpolate(x_unstructured, y_unstructured, x_queries[l], y_queries[l]) 

110 

111 for index, weight in zip([i00, i10, i01, i11], [w00, w10, w01, w11]): 

112 if weight > 0.0: 

113 nz += 1 

114 indj.append(index) 

115 val.append(weight) 

116 

117 indi.append(nz) 

118 

119 return csr_matrix((val, indj, indi), shape=(x_queries.size, x_unstructured.size))