1# distutils: libraries = STD_LIBS 2# distutils: sources = yt/utilities/lib/cykdtree/c_utils.cpp 3# distutils: depends = yt/utilities/lib/cykdtree/c_utils.hpp 4# distutils: language = c++ 5# distutils: extra_compile_args = CPP03_FLAG 6import numpy as np 7 8cimport cython 9cimport numpy as np 10from libc.stdint cimport int32_t, int64_t, uint32_t, uint64_t 11from libcpp cimport bool as cbool 12from libcpp.pair cimport pair 13from libcpp.vector cimport vector 14 15import copy 16 17 18def py_max_pts(np.ndarray[np.float64_t, ndim=2] pos): 19 r"""Get the maximum of points along each coordinate. 20 21 Args: 22 pos (np.ndarray of float64): (n,m) array of n m-D coordinates. 23 24 Returns: 25 np.ndarray of float64: Maximum of pos along each coordinate. 26 27 """ 28 cdef uint64_t n = <uint64_t>pos.shape[0] 29 cdef uint32_t m = <uint32_t>pos.shape[1] 30 cdef np.float64_t* cout = max_pts(&pos[0,0], n, m) 31 cdef uint32_t i = 0 32 cdef np.ndarray[np.float64_t] out = np.zeros(m, 'float64') 33 for i in range(m): 34 out[i] = cout[i] 35 #free(cout) 36 return out 37 38def py_min_pts(np.ndarray[np.float64_t, ndim=2] pos): 39 r"""Get the minimum of points along each coordinate. 40 41 Args: 42 pos (np.ndarray of float64): (n,m) array of n m-D coordinates. 43 44 Returns: 45 np.ndarray of float64: Minimum of pos along each coordinate. 46 47 """ 48 cdef uint64_t n = <uint64_t>pos.shape[0] 49 cdef uint32_t m = <uint32_t>pos.shape[1] 50 cdef np.float64_t* cout = min_pts(&pos[0,0], n, m) 51 cdef uint32_t i = 0 52 cdef np.ndarray[np.float64_t] out = np.zeros(m, 'float64') 53 for i in range(m): 54 out[i] = cout[i] 55 #free(cout) 56 return out 57 58def py_argmax_pts_dim(np.ndarray[np.float64_t, ndim=2] pos, 59 uint64_t[:] idx, 60 np.uint32_t d, int Lidx0, int Ridx0): 61 r"""Get the maximum of points along one dimension for a subset of the 62 point indices. This is essentially max(pos[idx[Lidx:(Ridx+1)], d]). 63 64 Args: 65 pos (np.ndarray of float64): (n,m) array of n m-D coordinates. 66 idx (np.ndarray of uint64_t): (n,) array of indices for positions. 67 d (uint32_t): Dimension to compute maximum along. 68 Lidx (int): Index in idx that search should begin at. 69 Ridx (int): Index in idx that search should end at. 70 71 Returns: 72 uint64_t: Index in idx that provides maximum position in the subset 73 indices along dimension d. 74 75 """ 76 cdef np.intp_t n = pos.shape[0] 77 cdef uint32_t m = <uint32_t>pos.shape[1] 78 cdef uint64_t Lidx = 0 79 cdef uint64_t Ridx = <uint64_t>(n-1) 80 if (Lidx0 < 0): 81 Lidx = <uint64_t>(n + Lidx0) 82 elif Lidx0 >= n: 83 raise Exception("Left index (%d) exceeds size of positions array (%d)." 84 % (Lidx0, n)) 85 else: 86 Lidx = <uint64_t>Lidx0 87 if (Ridx0 < 0): 88 Ridx = <uint64_t>(n + Ridx0) 89 elif Ridx0 >= n: 90 raise Exception("Right index (%d) exceeds size of positions array (%d)." 91 % (Ridx0, n)) 92 else: 93 Ridx = <uint64_t>Ridx0 94 cdef np.uint64_t cout = Lidx 95 if (Ridx > Lidx): 96 cout = argmax_pts_dim(&pos[0,0], &idx[0], m, d, Lidx, Ridx) 97 return cout 98 99def py_argmin_pts_dim(np.ndarray[np.float64_t, ndim=2] pos, 100 uint64_t[:] idx, 101 np.uint32_t d, int Lidx0, int Ridx0): 102 r"""Get the minimum of points along one dimension for a subset of the 103 point indices. This is essentially min(pos[idx[Lidx:(Ridx+1)], d]). 104 105 Args: 106 pos (np.ndarray of float64): (n,m) array of n m-D coordinates. 107 idx (np.ndarray of uint64_t): (n,) array of indices for positions. 108 d (uint32_t): Dimension to compute minimum along. 109 Lidx (int): Index in idx that search should begin at. 110 Ridx (int): Index in idx that search should end at. 111 112 Returns: 113 uint64_t: Index in idx that provides minimum position in the subset 114 indices along dimension d. 115 116 """ 117 cdef uint64_t n = <uint64_t>pos.shape[0] 118 cdef uint32_t m = <uint32_t>pos.shape[1] 119 cdef uint64_t Lidx = 0 120 cdef uint64_t Ridx = n 121 if (Lidx0 < 0): 122 Lidx = <uint64_t>(<int>n + Lidx0) 123 else: 124 Lidx = <uint64_t>Lidx0 125 if (Ridx0 < 0): 126 Ridx = <uint64_t>(<int>n + Ridx0) 127 else: 128 Ridx = <uint64_t>Ridx0 129 cdef np.uint64_t cout = Lidx 130 if (Ridx > Lidx): 131 cout = argmin_pts_dim(&pos[0,0], &idx[0], m, d, Lidx, Ridx) 132 return cout 133 134def py_quickSort(np.ndarray[np.float64_t, ndim=2] pos, np.uint32_t d): 135 r"""Get the indices required to sort coordinates along one dimension. 136 137 Args: 138 pos (np.ndarray of float64): (n,m) array of n m-D coordinates. 139 d (np.uint32_t): Dimension that pos should be sorted along. 140 141 Returns: 142 np.ndarray of uint64: Indices that sort pos along dimension d. 143 144 """ 145 cdef np.intp_t ndim = pos.shape[1] 146 cdef int64_t l = 0 147 cdef int64_t r = pos.shape[0]-1 148 cdef uint64_t[:] idx 149 idx = np.arange(pos.shape[0]).astype('uint64') 150 cdef double *ptr_pos = NULL 151 cdef uint64_t *ptr_idx = NULL 152 if pos.shape[0] != 0: 153 ptr_pos = &pos[0,0] 154 ptr_idx = &idx[0] 155 quickSort(ptr_pos, ptr_idx, ndim, d, l, r) 156 return idx 157 158def py_insertSort(np.ndarray[np.float64_t, ndim=2] pos, np.uint32_t d): 159 r"""Get the indices required to sort coordinates along one dimension. 160 161 Args: 162 pos (np.ndarray of float64): (n,m) array of n m-D coordinates. 163 d (np.uint32_t): Dimension that pos should be sorted along. 164 165 Returns: 166 np.ndarray of uint64: Indices that sort pos along dimension d. 167 168 """ 169 cdef np.intp_t ndim = pos.shape[1] 170 cdef int64_t l = 0 171 cdef int64_t r = pos.shape[0]-1 172 cdef uint64_t[:] idx 173 idx = np.arange(pos.shape[0]).astype('uint64') 174 cdef double *ptr_pos = NULL 175 cdef uint64_t *ptr_idx = NULL 176 if pos.shape[0] != 0: 177 ptr_pos = &pos[0,0] 178 ptr_idx = &idx[0] 179 insertSort(ptr_pos, ptr_idx, ndim, d, l, r) 180 return idx 181 182def py_pivot(np.ndarray[np.float64_t, ndim=2] pos, np.uint32_t d): 183 r"""Get the index of the median of medians along one dimension and indices 184 that partition pos according to the median of medians. 185 186 Args: 187 pos (np.ndarray of float64): (n,m) array of n m-D coordinates. 188 d (np.uint32_t): Dimension that pos should be partitioned along. 189 190 Returns: 191 tuple of int64 and np.ndarray of uint64: Index q of idx that is the 192 pivot. All elements of idx before the pivot will be less than 193 the pivot. If there is an odd number of points, the pivot will 194 be the median. 195 196 """ 197 cdef np.intp_t ndim = pos.shape[1] 198 cdef int64_t l = 0 199 cdef int64_t r = pos.shape[0]-1 200 cdef uint64_t[:] idx 201 idx = np.arange(pos.shape[0]).astype('uint64') 202 cdef double *ptr_pos = NULL 203 cdef uint64_t *ptr_idx = NULL 204 if pos.shape[0] != 0: 205 ptr_pos = &pos[0,0] 206 ptr_idx = &idx[0] 207 cdef int64_t q = pivot(ptr_pos, ptr_idx, ndim, d, l, r) 208 return q, idx 209 210def py_partition(np.ndarray[np.float64_t, ndim=2] pos, np.uint32_t d, 211 np.int64_t p): 212 r"""Get the indices required to partition coordinates along one dimension. 213 214 Args: 215 pos (np.ndarray of float64): (n,m) array of n m-D coordinates. 216 d (np.uint32_t): Dimension that pos should be partitioned along. 217 p (np.int64_t): Element of pos[:,d] that should be used as the pivot 218 to partition pos. 219 220 Returns: 221 tuple of int64 and np.ndarray of uint64: Location of the pivot in the 222 partitioned array and the indices required to partition the array 223 such that elements before the pivot are smaller and elements after 224 the pivot are larger. 225 226 """ 227 cdef np.intp_t ndim = pos.shape[1] 228 cdef int64_t l = 0 229 cdef int64_t r = pos.shape[0]-1 230 cdef uint64_t[:] idx 231 idx = np.arange(pos.shape[0]).astype('uint64') 232 cdef double *ptr_pos = NULL 233 cdef uint64_t *ptr_idx = NULL 234 if pos.shape[0] != 0: 235 ptr_pos = &pos[0,0] 236 ptr_idx = &idx[0] 237 cdef int64_t q = partition(ptr_pos, ptr_idx, ndim, d, l, r, p) 238 return q, idx 239 240def py_partition_given_pivot(np.ndarray[np.float64_t, ndim=2] pos, 241 np.uint32_t d, np.float64_t pval): 242 r"""Get the indices required to partition coordinates along one dimension. 243 244 Args: 245 pos (np.ndarray of float64): (n,m) array of n m-D coordinates. 246 d (np.uint32_t): Dimension that pos should be partitioned along. 247 pval (np.float64_t): Value that should be used to partition pos. 248 249 Returns: 250 tuple of int64 and np.ndarray of uint64: Location of the largest value 251 that is smaller than pval in partitioned array and the indices 252 required to partition the array such that elements before the pivot 253 are smaller and elements after the pivot are larger. 254 255 """ 256 cdef np.intp_t ndim = pos.shape[1] 257 cdef int64_t l = 0 258 cdef int64_t r = pos.shape[0]-1 259 cdef uint64_t[:] idx 260 idx = np.arange(pos.shape[0]).astype('uint64') 261 cdef double *ptr_pos = NULL 262 cdef uint64_t *ptr_idx = NULL 263 if pos.shape[0] != 0: 264 ptr_pos = &pos[0,0] 265 ptr_idx = &idx[0] 266 cdef int64_t q = partition_given_pivot(ptr_pos, ptr_idx, ndim, d, l, r, 267 pval) 268 return q, idx 269 270def py_select(np.ndarray[np.float64_t, ndim=2] pos, np.uint32_t d, 271 np.int64_t t): 272 r"""Get the indices required to partition coordiantes such that the first 273 t elements in pos[:,d] are the smallest t elements in pos[:,d]. 274 275 Args: 276 pos (np.ndarray of float64): (n,m) array of n m-D coordinates. 277 d (np.uint32_t): Dimension that pos should be partitioned along. 278 t (np.int64_t): Number of smallest elements in pos[:,d] that should be 279 partitioned. 280 281 Returns: 282 tuple of int64 and np.ndarray of uint64: Location of element t in the 283 partitioned array and the indices required to partition the array 284 such that elements before element t are smaller and elements after 285 the pivot are larger. 286 287 """ 288 cdef np.intp_t ndim = pos.shape[1] 289 cdef int64_t l = 0 290 cdef int64_t r = pos.shape[0]-1 291 cdef uint64_t[:] idx 292 idx = np.arange(pos.shape[0]).astype('uint64') 293 cdef double *ptr_pos = NULL 294 cdef uint64_t *ptr_idx = NULL 295 if pos.shape[0] != 0: 296 ptr_pos = &pos[0,0] 297 ptr_idx = &idx[0] 298 cdef int64_t q = select(ptr_pos, ptr_idx, ndim, d, l, r, t) 299 return q, idx 300 301 302def py_split(np.ndarray[np.float64_t, ndim=2] pos, 303 np.ndarray[np.float64_t, ndim=1] mins = None, 304 np.ndarray[np.float64_t, ndim=1] maxs = None, 305 bool use_sliding_midpoint = False): 306 r"""Get the indices required to split the positions equally along the 307 largest dimension. 308 309 Args: 310 pos (np.ndarray of float64): (n,m) array of n m-D coordinates. 311 mins (np.ndarray of float64, optional): (m,) array of mins. Defaults 312 to None and is set to mins of pos along each dimension. 313 maxs (np.ndarray of float64, optional): (m,) array of maxs. Defaults 314 to None and is set to maxs of pos along each dimension. 315 use_sliding_midpoint (bool, optional): If True, the sliding midpoint 316 rule is used to split the positions. Defaults to False. 317 318 Returns: 319 tuple(int64, uint32, np.ndarray of uint64): The index of the split in 320 the partitioned array, the dimension of the split, and the indices 321 required to partition the array. 322 323 """ 324 cdef np.intp_t npts = pos.shape[0] 325 cdef np.intp_t ndim = pos.shape[1] 326 cdef uint64_t Lidx = 0 327 cdef uint64_t[:] idx 328 idx = np.arange(pos.shape[0]).astype('uint64') 329 cdef double *ptr_pos = NULL 330 cdef uint64_t *ptr_idx = NULL 331 cdef double *ptr_mins = NULL 332 cdef double *ptr_maxs = NULL 333 if (npts != 0) and (ndim != 0): 334 if mins is None: 335 mins = np.min(pos, axis=0) 336 if maxs is None: 337 maxs = np.max(pos, axis=0) 338 ptr_pos = &pos[0,0] 339 ptr_idx = &idx[0] 340 ptr_mins = &mins[0] 341 ptr_maxs = &maxs[0] 342 cdef int64_t q = 0 343 cdef double split_val = 0.0 344 cdef cbool c_midpoint_flag = <cbool>use_sliding_midpoint 345 cdef uint32_t dsplit = split(ptr_pos, ptr_idx, Lidx, npts, ndim, 346 ptr_mins, ptr_maxs, q, split_val, 347 c_midpoint_flag) 348 return q, dsplit, idx 349