import numpy as np
from cython.parallel import prange
cimport cython
cimport numpy as cnp
cnp.import_array()
ctypedef cnp.npy_uint32 uint32_t
@cython.boundscheck(False)
@cython.wraparound(False)
@cython.nonecheck(False)
cdef void filter_inner(
uint32_t[::1] I,
double[::1] Z,
double[::1] M,
double[::1] T,
double Z1,
double Z2,
double T1,
double T2,
double M1,
double M2,
double[::1] T_out,
uint32_t[::1] I_out,
double[::1] M_out,
object[:, :] output_nested,
Py_ssize_t n,
Py_ssize_t p
) nogil:
cdef Py_ssize_t i, out_idx
cdef cython.bint in_range = 0
out_idx = 0
for i in range(n):
in_range = (
(Z1 < Z[i]) & (Z[i] < Z2) &
(T1 < T[i]) & (T[i] < T2) &
(M1 < M[i]) & (M[i] < M2)
)
if in_range:
T_out[out_idx] = T[i]
I_out[out_idx] = I[i]
M_out[out_idx] = M[i]
out_idx += 1
# I don't know if this hack works, but it allows use of python objects
# and compiles
with gil:
output_nested[p, 0] = T_out[:out_idx]
output_nested[p, 1] = I_out[:out_idx]
output_nested[p, 2] = M_out[:out_idx]
@cython.boundscheck(False)
@cython.wraparound(False)
@cython.nonecheck(False)
cpdef filter_outer(
uint32_t[::1] I,
double[::1] Z,
double[::1] M,
double[::1] T,
double Z1,
double Z2,
double T1,
double T2,
double M1,
double M2
):
cdef Py_ssize_t p, r, n
r = Z1.shape[0]
n = Z.shape[0]
cdef double[::1] T_out = np.empty_like(T)
cdef uint32_t[::1] I_out = np.empty_like(I)
cdef double[::1] M_out = np.empty_like(M)
cdef object[:, :] output_nested = np.empty((r, 3), dtype=object)
for p in prange(r, nogil=True, schedule="static"):
filter_inner(
I,
Z,
M,
T,
Z1[p],
Z2[p],
T1[p],
T2[p],
M1[p],
M2[p],
T_out,
I_out,
M_out,
output_nested,
n,
p
)
return np.asarray(output_nested)