Pastebin launched a little side project called VERYVIRAL.com, check it out ;-) Want more features on Pastebin? Sign Up, it's FREE!
Guest

Numpy find first/last function for 1-D arrays

By: dpitch40 on Jul 25th, 2012  |  syntax: C  |  size: 7.91 KB  |  views: 247  |  expires: Never
download  |  raw  |  embed  |  report abuse  |  print
Text below is selected. Please press Ctrl+C to copy to your clipboard. (⌘+C on Mac)
  1. #include "Python.h"
  2. #include "Numeric/arrayobject.h"
  3.  
  4. #define SEARCH_LOOP(OP, INC) { \
  5. for (i = start_addr; i != end_addr; INC) { \
  6.     if (*i OP value) { \
  7.         return (void*)i; \
  8.     } \
  9. } \
  10. return (void*)0; }
  11.  
  12. #define SWITCH_ON_OP(TYPE, INC) { \
  13. TYPE* i; \
  14. switch (op) { \
  15.     case EQ: \
  16.         SEARCH_LOOP(==, INC) \
  17.     case GT: \
  18.         SEARCH_LOOP(>, INC) \
  19.     case GTE: \
  20.         SEARCH_LOOP(>=, INC) \
  21.     case LT: \
  22.         SEARCH_LOOP(<, INC) \
  23.     case LTE: \
  24.         SEARCH_LOOP(<=, INC) \
  25.     case NEQ: \
  26.         SEARCH_LOOP(!=, INC) \
  27.     default: \
  28.         return (void*)0; \
  29.     } \
  30. }
  31.  
  32. #define SWITCH_ON_STRIDE(TYPE) { \
  33.     if (stride > 0) \
  34.     { \
  35.         SWITCH_ON_OP(TYPE, i++) \
  36.     } \
  37.     else \
  38.     { \
  39.         SWITCH_ON_OP(TYPE, i--) \
  40.     } \
  41. }
  42.  
  43. enum comp_operators {
  44.     EQ,
  45.     GT,
  46.     GTE,
  47.     LT,
  48.     LTE,
  49.     NEQ
  50. };
  51.  
  52. static void* array_find_byte(PyObject* searchval, void* start_addr, void* end_addr, int stride, int op) {
  53.     char value = (char)(PyInt_AsLong(searchval));
  54.     SWITCH_ON_STRIDE(char)
  55. }
  56.  
  57. static void* array_find_short(PyObject* searchval, void* start_addr, void* end_addr, int stride, int op) {
  58.     short value = (short)(PyInt_AsLong(searchval));
  59.     SWITCH_ON_STRIDE(short)
  60. }
  61.  
  62. static void* array_find_long(PyObject* searchval, void* start_addr, void* end_addr, int stride, int op) {
  63.     long value = PyInt_AsLong(searchval);
  64.     SWITCH_ON_STRIDE(long)
  65. }
  66.  
  67. static void* array_find_longlong(PyObject* searchval, void* start_addr, void* end_addr, int stride, int op) {
  68.     long long value = PyLong_AsLongLong(searchval);
  69.     SWITCH_ON_STRIDE(long long)
  70. }
  71.  
  72. static void* array_find_ubyte(PyObject* searchval, void* start_addr, void* end_addr, int stride, int op) {
  73.     unsigned char value = (unsigned char)(PyInt_AsUnsignedLongMask(searchval));
  74.     SWITCH_ON_STRIDE(unsigned char)
  75. }
  76.  
  77. static void* array_find_ushort(PyObject* searchval, void* start_addr, void* end_addr, int stride, int op) {
  78.     unsigned short value = (unsigned short)(PyInt_AsUnsignedLongMask(searchval));
  79.     SWITCH_ON_STRIDE(unsigned short)
  80. }
  81.  
  82. static void* array_find_ulong(PyObject* searchval, void* start_addr, void* end_addr, int stride, int op) {
  83.     unsigned long value = PyInt_AsUnsignedLongMask(searchval);
  84.     SWITCH_ON_STRIDE(unsigned long)
  85. }
  86.  
  87. static void* array_find_ulonglong(PyObject* searchval, void* start_addr, void* end_addr, int stride, int op) {
  88.     unsigned long long value = PyLong_AsUnsignedLongLong(searchval);
  89.     SWITCH_ON_STRIDE(unsigned long long)
  90. }
  91.  
  92. static void* array_find_float(PyObject* searchval, void* start_addr, void* end_addr, int stride, int op) {
  93.     float value = (float)(PyFloat_AsDouble(searchval));
  94.     SWITCH_ON_STRIDE(float)
  95. }
  96.  
  97. static void* array_find_double(PyObject* searchval, void* start_addr, void* end_addr, int stride, int op) {
  98.     double value = PyFloat_AsDouble(searchval);
  99.     SWITCH_ON_STRIDE(double)
  100. }
  101.  
  102. static void* array_find_eq_string(char* value, void* start_addr, void* end_addr, int stride, int elsize)
  103. {
  104.     char dest[elsize+1];
  105.     dest[elsize] = '\0';
  106.     void* i;
  107.     for (i = start_addr; i != end_addr; i += stride)
  108.     {
  109.         memcpy((void*)dest, i, elsize);
  110.         if (strcmp(dest, value) == 0)
  111.         {
  112.             return i;
  113.         }
  114.     }
  115.     return 0;
  116. }
  117.  
  118. static void* array_find_neq_string(char* value, void* start_addr, void* end_addr, int stride, int elsize)
  119. {
  120.     char dest[elsize+1];
  121.     dest[elsize] = '\0';
  122.     void* i;
  123.     for (i = start_addr; i != end_addr; i += stride)
  124.     {
  125.         memcpy((void*)dest, i, elsize);
  126.         if (strcmp(dest, value) != 0)
  127.         {
  128.             return i;
  129.         }
  130.     }
  131.     return 0;
  132. }
  133.  
  134. static void* array_find_string(PyObject* searchval, void* start_addr, void* end_addr, int stride, int op, int elsize) {
  135.     char* value = PyString_AsString(searchval);
  136.     if (op == NEQ)
  137.     {
  138.         return array_find_neq_string(value, start_addr, end_addr, stride, elsize);
  139.     }
  140.     else
  141.     {
  142.         return array_find_eq_string(value, start_addr, end_addr, stride, elsize);
  143.     }
  144. }
  145.  
  146. static PyObject* multi_find(PyObject* self, PyObject* args, short direction)
  147. {
  148.     PyArrayObject* arr;
  149.     PyObject* value;
  150.     short eqtest;
  151.     int start, end;
  152.     if (!PyArg_ParseTuple(args, "O!Ohii", &PyArray_Type, &arr, &value, &eqtest, &start, &end))
  153.     {
  154.         return NULL;
  155.     }
  156.     char type = arr->descr->type;
  157.     int elsize = arr->descr->elsize;
  158.     int stride = arr->strides[0];
  159.  
  160.     void* base_addr = PyArray_DATA(arr);
  161.     void* start_addr;
  162.     void* end_addr;
  163.     int dirstride = stride;
  164.     void* temp;
  165.  
  166.     if (stride > 0)
  167.     {
  168.         start_addr = base_addr + stride * start;
  169.         end_addr = base_addr + stride * end;
  170.     }
  171.     else
  172.     {
  173.         start_addr = base_addr + stride * start;
  174.         end_addr = base_addr + stride * end;
  175.     }
  176.  
  177.     if (!direction)
  178.     {
  179.         temp = start_addr;
  180.         //Shift by one place so we read the same set of data backwards
  181.         start_addr = end_addr - stride;
  182.         end_addr = temp - stride;
  183.         dirstride = -1 * stride;
  184.     }
  185.  
  186.     void* dest_addr;
  187.     switch (type)
  188.     {
  189.         case 'b':
  190.             dest_addr = array_find_byte(value, start_addr, end_addr, dirstride, eqtest);
  191.             break;
  192.         case 'h':
  193.             dest_addr = array_find_short(value, start_addr, end_addr, dirstride, eqtest);
  194.             break;
  195.         case 'i':
  196.             dest_addr = array_find_long(value, start_addr, end_addr, dirstride, eqtest);
  197.             break;
  198.         case 'l':
  199.             dest_addr = array_find_long(value, start_addr, end_addr, dirstride, eqtest);
  200.             break;
  201.         case 'q':
  202.             dest_addr = array_find_longlong(value, start_addr, end_addr, dirstride, eqtest);
  203.             break;
  204.         case 'B':
  205.             dest_addr = array_find_ubyte(value, start_addr, end_addr, dirstride, eqtest);
  206.             break;
  207.         case 'H':
  208.             dest_addr = array_find_ushort(value, start_addr, end_addr, dirstride, eqtest);
  209.             break;
  210.         case 'I':
  211.             dest_addr = array_find_ulong(value, start_addr, end_addr, dirstride, eqtest);
  212.             break;
  213.         case 'L':
  214.             dest_addr = array_find_ulong(value, start_addr, end_addr, dirstride, eqtest);
  215.             break;
  216.         case 'Q':
  217.             dest_addr = array_find_ulonglong(value, start_addr, end_addr, dirstride, eqtest);
  218.             break;
  219.         case 'f':
  220.             dest_addr = array_find_float(value, start_addr, end_addr, dirstride, eqtest);
  221.             break;
  222.         case 'd':
  223.             dest_addr = array_find_double(value, start_addr, end_addr, dirstride, eqtest);
  224.             break;
  225.         case 'S':
  226.             dest_addr = array_find_string(value, start_addr, end_addr, dirstride, eqtest, elsize);
  227.             break;
  228.         default:
  229.             dest_addr = 0;
  230.     }
  231.  
  232.     // printf("%s\t%s\t%d\t%d\t%p\t%p\t%p\t%p\n", &(arr->descr->kind), &type, elsize, dirstride,
  233.     //                                            base_addr, start_addr, end_addr, dest_addr);
  234.  
  235.     int dest_index;
  236.     if (dest_addr > 0)
  237.     {
  238.         dest_index = (dest_addr - base_addr) / stride;
  239.     }
  240.     else
  241.     {
  242.         dest_index = -1;
  243.     }
  244.     return Py_BuildValue("i", dest_index);
  245. }
  246.  
  247. static char py_find_doc[] = "Finds the first index of the occurrence of a value in an array.";
  248. static PyObject* py_find(PyObject* self, PyObject* args)
  249. {
  250.     return multi_find(self, args, 1);
  251. }
  252.  
  253. static char py_rfind_doc[] = "Finds the last index of the occurrence of a value in an array.";
  254. static PyObject* py_rfind(PyObject* self, PyObject* args)
  255. {
  256.     return multi_find(self, args, 0);
  257. }
  258.  
  259. static PyMethodDef _npfindmethods[] = {
  260.     {"np_find", py_find, METH_VARARGS, py_find_doc},
  261.     {"np_rfind", py_rfind, METH_VARARGS, py_rfind_doc},
  262.     {NULL, NULL, 0, NULL}
  263. };
  264.  
  265. void init_npfind(void)
  266. {
  267.     Py_InitModule("_npfind", _npfindmethods);
  268.     import_array();
  269. }