How to use a templated function in C++ for two types in Cython?

Question

I have a templated C++ function of which I would like to be able to use both types. As Python does not support overloading, I am a little stuck how to solve this. I have a .pyx as shown below. How can I use the C++ function in both float and double?

import cython
import numpy as np
cimport numpy as np

# declare the interface to the C code
cdef extern from "diff_cpp.cpp" namespace "diff":
    cdef void diff_cpp[float] (float* at, const float* a, const float visc,
                               const float dxidxi, const float dyidyi, const float dzidzi,
                               const int itot, const int jtot, const int ktot)

cdef extern from "diff_cpp.cpp" namespace "diff":
    cdef void diff_cpp[double] (double* at, const double* a, const double visc,
                                const double dxidxi, const double dyidyi, const double dzidzi,
                                const int itot, const int jtot, const int ktot)

@cython.boundscheck(False)
@cython.wraparound(False)
def diff(np.ndarray[double, ndim=3, mode="c"] at not None,
         np.ndarray[double, ndim=3, mode="c"] a not None,
         double visc, double dxidxi, double dyidyi, double dzidzi):
    cdef int ktot, jtot, itot
    ktot, jtot, itot = at.shape[0], at.shape[1], at.shape[2]
    diff_cpp[double](&at[0,0,0], &a[0,0,0], visc, dxidxi, dyidyi, dzidzi, itot, jtot, ktot)
    return None

@cython.boundscheck(False)
@cython.wraparound(False)
def diff_f(np.ndarray[float, ndim=3, mode="c"] at not None,
           np.ndarray[float, ndim=3, mode="c"] a not None,
           float visc, float dxidxi, float dyidyi, float dzidzi):
    cdef int ktot, jtot, itot
    ktot, jtot, itot = at.shape[0], at.shape[1], at.shape[2]
    diff_cpp[float](&at[0,0,0], &a[0,0,0], visc, dxidxi, dyidyi, dzidzi, itot, jtot, ktot)
    return None

UPDATE WITH SOLUTION

The answer of @oz1 provided the correct way of doing this. This the code that works, for those who are interested in the solution to this particular problem.

import cython
import numpy as np
cimport numpy as np

# declare the interface to the C code
cdef extern from "diff_cpp.cpp" namespace "diff":
    cdef void diff_cpp[T](T* at, const T* a, const T visc,
                          const T dxidxi, const T dyidyi, const T dzidzi,
                          const int itot, const int jtot, const int ktot)

ctypedef fused float_t:
    cython.float
    cython.double

@cython.boundscheck(False)
@cython.wraparound(False)
def diff(np.ndarray[float_t, ndim=3, mode="c"] at not None,
         np.ndarray[float_t, ndim=3, mode="c"] a not None,
         float_t visc, float_t dxidxi, float_t dyidyi, float_t dzidzi):
    cdef int ktot, jtot, itot
    ktot, jtot, itot = at.shape[0], at.shape[1], at.shape[2]
    diff_cpp(&at[0,0,0], &a[0,0,0], visc, dxidxi, dyidyi, dzidzi, itot, jtot, ktot)
    return None

user_6610031 · Answer

Two notes:

Cython supports c++ template(http://docs.cython.org/en/latest/src/userguide/wrapping_CPlusPlus.html)
Cython has fused types(http://docs.cython.org/en/latest/src/userguide/fusedtypes.html)

An example:

// lib.cpp
template
T arr_sum(T *arr, size_t size)
{
    T temp=0;
    for (size_t i=0; i != size; ++i){
        temp += arr[i];
    }
    return temp;
}

# lib_wrapper.pyx
cimport cython

ctypedef fused  float_t:
    cython.float
    cython.double

cdef extern from "lib.cpp" nogil:
    T arr_sum[T](T *arr, size_t size)

def py_arr_sum(float_t[:] arr not None):
    print(sizeof(arr[0]))  # check the element size
    return arr_sum(&arr[0], arr.shape[0])

# use.py
import numpy as np
from lib_wrapper import py_arr_sum

print(py_arr_sum(np.array([1,2,3], dtype=np.float32)))
print(py_arr_sum(np.array([1,2,3], dtype=np.float64)))
print(py_arr_sum(np.array([1,2,3], dtype=np.int32)))  # oops

How to use a templated function in C++ for two types in Cython?

1 Answers