C-API register a ufunc loop against builtin ufunc-CodePudding

I am looking to build a numpy extension module which declares a new structured dtype and provides the necessary inner ufunc loops to allow built-in math operations to be performed on it.

Following the guide on creating your own ufunc, I have been able to:

Create a new PyUFuncGenericFunction to mock a built-in ufunc:

PyUFuncGenericFunction *add_uncertain = PyUFunc_FromFuncAndData(
    NULL,
    NULL,
    NULL,
    0,
    2,
    1,
    PyUFunc_None,
    "add_uncertain",
    "adds arrays of uncertain values",
    0);

Declare the data-type, as:

PyObject *uncertain_double_dtype_dict = Py_BuildValue(
    "[(s, s), (s, s)]",
    "nominal", "f8", "uncertainty", "f8");

Write an implementation of this for a specific dtype:

static PyObject *
add_uncertain_double(char **args, npy_intp *dimensions, npy_intp *steps, void *data)
{
    char *in_arr_a = args[0];
    char *in_arr_b = args[1];
    char *out_arr = args[2];

    npy_intp in_a_step = steps[0];
    npy_intp in_b_step = steps[1];
    npy_intp out_step = steps[2];

    for (npy_intp i = 0; i < dimensions[0]; i  )
    {
        double *in_a = (double *)in_arr_a;
        double *in_b = (double *)in_arr_b;
        double *out = (double *)out_arr;

        out[0] = in_a[0]   in_b[0];
        out[1] = hypot(in_a[1], in_b[1]);

        in_arr_a  = in_a_step;
        out_arr  = out_step;
    }
};

PyUFunc_RegisterLoopForDescr(
    add_uncertain,
    uncertain_double_dtype,
    &add_uncertain_double,
    add_uncertain_double_dtypes,
    NULL);

When exported as a module this works as intended, but requires the user to call the function via the package exported ufunc (i.e. my_package.add_uncertain); I would rather this implementation was available via numpy.add.

The language used in the registering a ufunc loop section of the numpy C-API docs seems to suggest that I should be able to register the ufunc loop against built-in ufuncs. To do so I believe I should pass the built-in PyUFuncGenericFunction to PyUFunc_RegisterLoopForDescr.

I would be very grateful to know if I am on the right track with this, and if so where I should look for the built-in PyUFuncGenericFunction

CodePudding user response：

The solution, as contained within the numpy rational type test, was to import numpy with PyImport_Import then get the add ufunc using PyObject_GetAttrString, at which point the new inner loop could be registered:

  numpy_str = PyUnicode_FromString("numpy");
  if (!numpy_str) 
      return NULL;
  numpy = PyImport_Import(numpy_str);
  Py_DECREF(numpy_str);
  if (!numpy)
      return NULL;

  PyUFuncObject *add_ufunc = PyObject_GetAttrString(numpy, "add");
  if (!add_ufunc)
      return NULL;

  PyUFunc_RegisterLoopForDescr(
      add_ufunc,
      uncertain_double_dtype,
      &add_uncertain_double,
      add_uncertain_double_dtypes,
      NULL);