mephisto/mephisto.cpp
/**
 * @brief Return the indices of the bins to which each value in input array belongs.
 *
 * @see https://numpy.org/doc/stable/reference/generated/numpy.digitize.html
 * @see https://github.com/numpy/numpy/blob/v1.18.1/numpy/lib/function_base.py#L4700-L4808
 *
 * @param self
 * @param args
 * @param kwds
 * @return PyObject*
 */
@brief Return the indices of the bins to which each value in input array belongs.


@brief Return the indices of the bins to which each value in input array belongs .
static PyObject*
digitize(PyObject *self, PyObject *args, PyObject *kwds) {
    PyObject *x = nullptr;
    PyObject *bins = nullptr;
    npy_intp right = false;
    static char *kwlist[] = {"x", "bins", "right", NULL};
    PyArrayObject *indices = nullptr;
    PyArrayObject *x_np = nullptr;
    PyArrayObject *bins_np = nullptr;
    double *bins_data = nullptr;
    npy_intp bins_length;
    bool increasing;
    npy_intp *x_shape = nullptr;
    double *x_data = nullptr;
    npy_intp x_length;
    int64_t *indices_data = nullptr;
    if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO|p", kwlist, &x, &bins, &right))
    {
        PyErr_SetString(PyExc_ValueError, "Wrong Arguments");
        goto fail;
    }
    x_np = reinterpret_cast< PyArrayObject* >(PyArray_FROM_OTF(x, NPY_DOUBLE, NPY_ARRAY_C_CONTIGUOUS | NPY_ARRAY_ALIGNED));
    if (!x_np)
    {
        PyErr_SetString(PyExc_ValueError, "PyArray_FROM_OTF failed for x_np");
        goto fail;
    }
    bins_np = reinterpret_cast< PyArrayObject* >(PyArray_FROM_OTF(bins, NPY_DOUBLE, NPY_ARRAY_C_CONTIGUOUS | NPY_ARRAY_ALIGNED));
    if (!bins_np)
    {
        PyErr_SetString(PyExc_ValueError, "PyArray_FROM_OTF failed for bins_np");
        goto fail;
    }
    if (PyArray_NDIM(bins_np) != 1)
    {
        PyErr_SetString(PyExc_ValueError, "`bins` must be 1d array");
        goto fail;
    }
    bins_data = (double*)PyArray_DATA(bins_np);
    bins_length = PyArray_SIZE(bins_np);
    if (bins_length < 2)
    {
        PyErr_SetString(PyExc_ValueError, "`bins` must contain at least 2 values");
        goto fail;
    }
    increasing = bins_data[0] < bins_data[1];
    for (npy_intp i = 1; i < bins_length - 1; ++i)
        if ((bins_data[i] >= bins_data[i+1]) == increasing)
        {
            PyErr_SetString(PyExc_ValueError, "`bins` must be monotonic");
            goto fail;
        }
    x_shape = PyArray_DIMS(x_np);
    if (!x_shape)
        goto fail;
    indices = reinterpret_cast< PyArrayObject* >(PyArray_SimpleNew(PyArray_NDIM(x_np), x_shape, NPY_INT64));
    if (!indices)
        goto fail;
    x_data = (double*)PyArray_DATA(x_np);
    if (!x_data)
        goto fail;
    x_length = PyArray_SIZE(x_np);
    indices_data = (int64_t*)PyArray_DATA(indices);
    if (!indices_data)
        goto fail;
    if (increasing && !right)
    {
        #pragma omp parallel for schedule(static, THREAD_CHUNK_SIZE)
        for (npy_intp i = 0; i < x_length; ++i)
        {
            auto const tmp = x_data[i];
            int64_t lower = 0;
            int64_t upper = bins_length;
            int64_t bin = 0;
            while (true)
            {
                bin = lower + (upper - lower) / 2;
                if (bin == 0 || bin == bins_length || (bins_data[bin-1] <= tmp && bins_data[bin] > tmp))
                    break;
                if (bins_data[bin] > tmp)
                    upper = bin;
                else
                    lower = bin + 1;
            }
            indices_data[i] = bin;
        }
    } else if (increasing && right)
    {
        #pragma omp parallel for schedule(static, THREAD_CHUNK_SIZE)
        for (npy_intp i = 0; i < x_length; ++i)
        {
            auto const tmp = x_data[i];
            int64_t lower = 0;
            int64_t upper = bins_length;
            int64_t bin = 0;
            while (true)
            {
                bin = lower + (upper - lower) / 2;
                if (bin == 0 || bin == bins_length || (bins_data[bin-1] < tmp && bins_data[bin] >= tmp))
                    break;
                if (bins_data[bin] >= tmp)
                    upper = bin;
                else
                    lower = bin;
            }
            indices_data[i] = bin;
        }
    } else if (!increasing && !right)
    {
        #pragma omp parallel for schedule(static, THREAD_CHUNK_SIZE)
        for (npy_intp i = 0; i < x_length; ++i)
        {
            auto const tmp = x_data[i];
            int64_t lower = 0;
            int64_t upper = bins_length;
            int64_t bin = 0;
            while (true)
            {
                bin = lower + (upper - lower) / 2;
                if (bin == 0 || bin == bins_length || (bins_data[bin-1] > tmp && bins_data[bin] <= tmp))
                    break;
                if (bins_data[bin] <= tmp)
                    upper = bin;
                else
                    lower = bin;
            }
            indices_data[i] = bin;
        }
    } else
    {
        #pragma omp parallel for schedule(static, THREAD_CHUNK_SIZE)
        for (npy_intp i = 0; i < x_length; ++i)
        {
            auto const tmp = x_data[i];
            int64_t lower = 0;
            int64_t upper = bins_length;
            int64_t bin = 0;
            while (true)
            {
                bin = lower + (upper - lower) / 2;
                if (bin == 0 || bin == bins_length || (bins_data[bin-1] >= tmp && bins_data[bin] < tmp))
                    break;
                if (bins_data[bin] < tmp)
                    upper = bin;
                else
                    lower = bin + 1;
            }
            indices_data[i] = bin;
        }
    }
    goto success;
success:
    Py_DECREF(x_shape);
    Py_DECREF(bins_np);
    Py_DECREF(x_np);
    return Py_BuildValue("N", indices);
fail:
    Py_XDECREF(x_shape);
    Py_XDECREF(bins_np);
    Py_XDECREF(x_np);
    Py_XDECREF(indices);
    return NULL;
}
static PyObject * digitize ( PyObject * self , PyObject * args , PyObject * kwds ) { PyObject * x = nullptr ; PyObject * bins = nullptr ; npy_intp right = false ; static char * kwlist [ ] = { " " , " " , " " , NULL } ; PyArrayObject * indices = nullptr ; PyArrayObject * x_np = nullptr ; PyArrayObject * bins_np = nullptr ; double * bins_data = nullptr ; npy_intp bins_length ; bool increasing ; npy_intp * x_shape = nullptr ; double * x_data = nullptr ; npy_intp x_length ; int64_t * indices_data = nullptr ; if ( ! PyArg_ParseTupleAndKeywords ( args , kwds , " " , kwlist , & x , & bins , & right ) ) { PyErr_SetString ( PyExc_ValueError , " " ) ; goto fail ; } x_np = reinterpret_cast < PyArrayObject * > ( PyArray_FROM_OTF ( x , NPY_DOUBLE , NPY_ARRAY_C_CONTIGUOUS | NPY_ARRAY_ALIGNED ) ) ; if ( ! x_np ) { PyErr_SetString ( PyExc_ValueError , " " ) ; goto fail ; } bins_np = reinterpret_cast < PyArrayObject * > ( PyArray_FROM_OTF ( bins , NPY_DOUBLE , NPY_ARRAY_C_CONTIGUOUS | NPY_ARRAY_ALIGNED ) ) ; if ( ! bins_np ) { PyErr_SetString ( PyExc_ValueError , " " ) ; goto fail ; } if ( PyArray_NDIM ( bins_np ) != 1 ) { PyErr_SetString ( PyExc_ValueError , " " ) ; goto fail ; } bins_data = ( double * ) PyArray_DATA ( bins_np ) ; bins_length = PyArray_SIZE ( bins_np ) ; if ( bins_length < 2 ) { PyErr_SetString ( PyExc_ValueError , " " ) ; goto fail ; } increasing = bins_data [ 0 ] < bins_data [ 1 ] ; for ( npy_intp i = 1 ; i < bins_length - 1 ; ++ i ) if ( ( bins_data [ i ] >= bins_data [ i + 1 ] ) == increasing ) { PyErr_SetString ( PyExc_ValueError , " " ) ; goto fail ; } x_shape = PyArray_DIMS ( x_np ) ; if ( ! x_shape ) goto fail ; indices = reinterpret_cast < PyArrayObject * > ( PyArray_SimpleNew ( PyArray_NDIM ( x_np ) , x_shape , NPY_INT64 ) ) ; if ( ! indices ) goto fail ; x_data = ( double * ) PyArray_DATA ( x_np ) ; if ( ! x_data ) goto fail ; x_length = PyArray_SIZE ( x_np ) ; indices_data = ( int64_t * ) PyArray_DATA ( indices ) ; if ( ! indices_data ) goto fail ; if ( increasing && ! right ) { #pragma  omp parallel for schedule(static, THREAD_CHUNK_SIZE) for ( npy_intp i = 0 ; i < x_length ; ++ i ) { auto const tmp = x_data [ i ] ; int64_t lower = 0 ; int64_t upper = bins_length ; int64_t bin = 0 ; while ( true ) { bin = lower + ( upper - lower ) / 2 ; if ( bin == 0 || bin == bins_length || ( bins_data [ bin - 1 ] <= tmp && bins_data [ bin ] > tmp ) ) break ; if ( bins_data [ bin ] > tmp ) upper = bin ; else lower = bin + 1 ; } indices_data [ i ] = bin ; } } else if ( increasing && right ) { #pragma  omp parallel for schedule(static, THREAD_CHUNK_SIZE) for ( npy_intp i = 0 ; i < x_length ; ++ i ) { auto const tmp = x_data [ i ] ; int64_t lower = 0 ; int64_t upper = bins_length ; int64_t bin = 0 ; while ( true ) { bin = lower + ( upper - lower ) / 2 ; if ( bin == 0 || bin == bins_length || ( bins_data [ bin - 1 ] < tmp && bins_data [ bin ] >= tmp ) ) break ; if ( bins_data [ bin ] >= tmp ) upper = bin ; else lower = bin ; } indices_data [ i ] = bin ; } } else if ( ! increasing && ! right ) { #pragma  omp parallel for schedule(static, THREAD_CHUNK_SIZE) for ( npy_intp i = 0 ; i < x_length ; ++ i ) { auto const tmp = x_data [ i ] ; int64_t lower = 0 ; int64_t upper = bins_length ; int64_t bin = 0 ; while ( true ) { bin = lower + ( upper - lower ) / 2 ; if ( bin == 0 || bin == bins_length || ( bins_data [ bin - 1 ] > tmp && bins_data [ bin ] <= tmp ) ) break ; if ( bins_data [ bin ] <= tmp ) upper = bin ; else lower = bin ; } indices_data [ i ] = bin ; } } else { #pragma  omp parallel for schedule(static, THREAD_CHUNK_SIZE) for ( npy_intp i = 0 ; i < x_length ; ++ i ) { auto const tmp = x_data [ i ] ; int64_t lower = 0 ; int64_t upper = bins_length ; int64_t bin = 0 ; while ( true ) { bin = lower + ( upper - lower ) / 2 ; if ( bin == 0 || bin == bins_length || ( bins_data [ bin - 1 ] >= tmp && bins_data [ bin ] < tmp ) ) break ; if ( bins_data [ bin ] < tmp ) upper = bin ; else lower = bin + 1 ; } indices_data [ i ] = bin ; } } goto success ; success : Py_DECREF ( x_shape ) ; Py_DECREF ( bins_np ) ; Py_DECREF ( x_np ) ; return Py_BuildValue ( " " , indices ) ; fail : Py_XDECREF ( x_shape ) ; Py_XDECREF ( bins_np ) ; Py_XDECREF ( x_np ) ; Py_XDECREF ( indices ) ; return NULL ; }
@brief Return the indices of the bins to which each value in input array belongs.
@brief Return the indices of the bins to which each value in input array belongs .
/* input */ /* output */ /* helpers */ /* argument parsing */ /* obtain ndarray behind `x` */ /* obtain ndarray behind `bins` */ /* allocate output arrays */ /* wörkwörk */