Cython: double or array as input

user_14104321 0
Viewed 148

Let's consider the following simple example:

def add(double a, double b):
    cdef double c
    c = a + b
    return c

How can I define it such that the inputs a and b (and the outputc as a consequence) might be double or np.array?

Thanks a lot!

cython python-3.x
1 Answers

For 1D NumPy arrays:

def add(np.ndarray[np.float64_t, ndim=1] a, np.ndarray[np.float64_t, ndim=1] b):
    cdef np.ndarray [np.float64_t, ndim=1] c
    c = a + b
    return c

For 2D NumPy arrays:

    def add(np.ndarray[np.float64_t, ndim=2] a, np.ndarray[np.float64_t, ndim=2] b):
    cdef np.ndarray [np.float64_t, ndim=2] c
    c = a + b
    return c

You can see that ndim is the number of dimensions. np.float64_t is the data type for the array elements, np.float64. np.ndarray defines the data type as a NumPy array.

For arrays of double, you can look into memoryviews, but it might be not what you want. It is for fast indexing, not arithmetic operations.

user_143923080
Related