summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
Diffstat
-rw-r--r--numpy/add_newdocs.py19
-rw-r--r--numpy/lib/src/_compiled_base.c96
-rw-r--r--numpy/lib/tests/test_function_base.py27
3 files changed, 122 insertions, 20 deletions
diff --git a/numpy/add_newdocs.py b/numpy/add_newdocs.py
index d6bd881cc..6c6488547 100644
--- a/numpy/add_newdocs.py
+++ b/numpy/add_newdocs.py
@@ -4865,14 +4865,17 @@ add_newdoc('numpy.core.umath', 'seterrobj',
add_newdoc('numpy.lib._compiled_base', 'digitize',
"""
- digitize(x, bins)
+ digitize(x, bins, right=False)
Return the indices of the bins to which each value in input array belongs.
Each index ``i`` returned is such that ``bins[i-1] <= x < bins[i]`` if
`bins` is monotonically increasing, or ``bins[i-1] > x >= bins[i]`` if
`bins` is monotonically decreasing. If values in `x` are beyond the
- bounds of `bins`, 0 or ``len(bins)`` is returned as appropriate.
+ bounds of `bins`, 0 or ``len(bins)`` is returned as appropriate. If right
+ is True, then the right bin is closed so that the index ``i`` is such
+ that ``bins[i-1] < x <= bins[i]`` or bins[i-1] >= x > bins[i]`` if `bins`
+ is monotonically increasing or decreasing, respectively.
Parameters
----------
@@ -4880,6 +4883,12 @@ add_newdoc('numpy.lib._compiled_base', 'digitize',
Input array to be binned. It has to be 1-dimensional.
bins : array_like
Array of bins. It has to be 1-dimensional and monotonic.
+ right : bool, optional
+ Indicating whether the intervals include the right or the left bin
+ edge. Default behavior is (right==False) indicating that the interval
+ does not include the right edge. The left bin and is open in this
+ case. Ie., bins[i-1] <= x < bins[i] is the default behavior for
+ monotonically increasing bins.
Returns
-------
@@ -4918,6 +4927,12 @@ add_newdoc('numpy.lib._compiled_base', 'digitize',
2.5 <= 3.0 < 4.0
1.0 <= 1.6 < 2.5
+ >>> x = np.array([1.2, 10.0, 12.4, 15.5, 20.])
+ >>> bins = np.array([0,5,10,15,20])
+ >>> np.digitize(x,bins,right=True)
+ array([1, 2, 3, 4, 4])
+ >>> np.digitize(x,bins,right=False)
+ array([1, 3, 3, 4, 5])
""")
add_newdoc('numpy.lib._compiled_base', 'bincount',
diff --git a/numpy/lib/src/_compiled_base.c b/numpy/lib/src/_compiled_base.c
index d09e5f58c..61935a39b 100644
--- a/numpy/lib/src/_compiled_base.c
+++ b/numpy/lib/src/_compiled_base.c
@@ -33,6 +33,32 @@ decr_slot_(double x, double * bins, npy_intp lbins)
return 0;
}
+static npy_intp
+incr_slot_right_(double x, double *bins, npy_intp lbins)
+{
+ npy_intp i;
+
+ for ( i = 0; i < lbins; i ++ ) {
+ if ( x <= bins [i] ) {
+ return i;
+ }
+ }
+ return lbins;
+}
+
+static npy_intp
+decr_slot_right_(double x, double * bins, npy_intp lbins)
+{
+ npy_intp i;
+
+ for ( i = lbins - 1; i >= 0; i -- ) {
+ if (x <= bins [i]) {
+ return i + 1;
+ }
+ }
+ return 0;
+}
+
/**
* Returns -1 if the array is monotonic decreasing,
* +1 if the array is monotonic increasing,
@@ -219,11 +245,13 @@ fail:
/*
- * digitize (x, bins) returns an array of python integers the same
+ * digitize (x, bins, right=False) returns an array of python integers the same
* length of x. The values i returned are such that bins [i - 1] <= x <
* bins [i] if bins is monotonically increasing, or bins [i - 1] > x >=
* bins [i] if bins is monotonically decreasing. Beyond the bounds of
* bins, returns either i = 0 or i = len (bins) as appropriate.
+ * if right == True the comparison is bins [i - 1] < x <= bins[i]
+ * or bins [i - 1] >= x > bins[i]
*/
static PyObject *
arr_digitize(PyObject *NPY_UNUSED(self), PyObject *args, PyObject *kwds)
@@ -233,13 +261,15 @@ arr_digitize(PyObject *NPY_UNUSED(self), PyObject *args, PyObject *kwds)
PyArrayObject *ax = NULL, *abins = NULL, *aret = NULL;
double *dx, *dbins;
npy_intp lbins, lx; /* lengths */
+ npy_intp right = 0; /* whether right or left is inclusive */
npy_intp *iret;
int m, i;
- static char *kwlist[] = {"x", "bins", NULL};
+ static char *kwlist[] = {"x", "bins", "right", NULL};
PyArray_Descr *type;
char bins_non_monotonic = 0;
- if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO", kwlist, &ox, &obins)) {
+ if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO|i", kwlist, &ox, &obins,
+ &right)) {
goto fail;
}
type = PyArray_DescrFromType(NPY_DOUBLE);
@@ -272,30 +302,64 @@ arr_digitize(PyObject *NPY_UNUSED(self), PyObject *args, PyObject *kwds)
}
NPY_BEGIN_ALLOW_THREADS;
if (lbins == 1) {
- for (i = 0; i < lx; i++) {
- if (dx [i] >= dbins[0]) {
- iret[i] = 1;
+ if (right == 0) {
+ for (i = 0; i < lx; i++) {
+ if (dx [i] >= dbins[0]) {
+ iret[i] = 1;
+ }
+ else {
+ iret[i] = 0;
+ }
}
- else {
- iret[i] = 0;
+ }
+ else {
+ for (i = 0; i < lx; i++) {
+ if (dx [i] > dbins[0]) {
+ iret[i] = 1;
+ }
+ else {
+ iret[i] = 0;
+ }
}
+
}
}
else {
m = check_array_monotonic(dbins, lbins);
- if ( m == -1 ) {
- for ( i = 0; i < lx; i ++ ) {
- iret [i] = decr_slot_ ((double)dx[i], dbins, lbins);
+ if (right == 0) {
+ if ( m == -1 ) {
+ for ( i = 0; i < lx; i ++ ) {
+ iret [i] = decr_slot_ ((double)dx[i], dbins, lbins);
+ }
}
- }
- else if ( m == 1 ) {
- for ( i = 0; i < lx; i ++ ) {
- iret [i] = incr_slot_ ((double)dx[i], dbins, lbins);
+ else if ( m == 1 ) {
+ for ( i = 0; i < lx; i ++ ) {
+ iret [i] = incr_slot_ ((double)dx[i], dbins, lbins);
+ }
+ }
+ else {
+ /* defer PyErr_SetString until after NPY_END_ALLOW_THREADS */
+ bins_non_monotonic = 1;
}
}
else {
+ if ( m == -1 ) {
+ for ( i = 0; i < lx; i ++ ) {
+ iret [i] = decr_slot_right_ ((double)dx[i], dbins,
+ lbins);
+ }
+ }
+ else if ( m == 1 ) {
+ for ( i = 0; i < lx; i ++ ) {
+ iret [i] = incr_slot_right_ ((double)dx[i], dbins,
+ lbins);
+ }
+ }
+ else {
/* defer PyErr_SetString until after NPY_END_ALLOW_THREADS */
- bins_non_monotonic = 1;
+ bins_non_monotonic = 1;
+ }
+
}
}
NPY_END_ALLOW_THREADS;
diff --git a/numpy/lib/tests/test_function_base.py b/numpy/lib/tests/test_function_base.py
index 910fd5283..d3d3de5e8 100644
--- a/numpy/lib/tests/test_function_base.py
+++ b/numpy/lib/tests/test_function_base.py
@@ -191,10 +191,10 @@ class TestCumsum(TestCase):
tgt = np.array([1, 3, 13, 24, 30, 35, 39], ctype)
assert_array_equal(np.cumsum(a, axis=0), tgt)
-
+
tgt = np.array([[1, 2, 3, 4], [6, 8, 10, 13], [16, 11, 14, 18]], ctype)
assert_array_equal(np.cumsum(a2, axis=0), tgt)
-
+
tgt = np.array([[1, 3, 6, 10], [5, 11, 18, 27], [10, 13, 17, 22]], ctype)
assert_array_equal(np.cumsum(a2, axis=1), tgt)
@@ -429,6 +429,29 @@ class TestDigitize(TestCase):
bin = np.linspace(x.min(), x.max(), 10)
assert_(np.all(digitize(x, bin) != 0))
+ def test_right_basic(self):
+ x = [1,5,4,10,8,11,0]
+ bins = [1,5,10]
+ default_answer = [1,2,1,3,2,3,0]
+ assert_array_equal(digitize(x, bins), default_answer)
+ right_answer = [0,1,1,2,2,3,0]
+ assert_array_equal(digitize(x, bins, True), right_answer)
+
+ def test_right_open(self):
+ x = np.arange(-6, 5)
+ bins = np.arange(-6, 4)
+ assert_array_equal(digitize(x,bins,True), np.arange(11))
+
+ def test_right_open_reverse(self):
+ x = np.arange(5, -6, -1)
+ bins = np.arange(4, -6, -1)
+ assert_array_equal(digitize(x, bins, True), np.arange(11))
+
+ def test_right_open_random(self):
+ x = rand(10)
+ bins = np.linspace(x.min(), x.max(), 10)
+ assert_(all(digitize(x, bins, True) != 10))
+
class TestUnwrap(TestCase):
def test_simple(self):
View Lorry Controller Status