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