Defining the arguments for the functions inside the cython class and fast integral calculation in cython

Question

I am a newbie with cython and trying to convert a python class to cython. I don't know how I should define argument z in instance Da, in the way that it can deal with both numpy.array or just a single float number.

cdef class Cosmology(object):
    cdef double omega_m, omega_lam, omega_c  

    def __init__(self,double omega_m=0.3,double omega_lam=0.7):
        self.omega_m = omega_m
        self.omega_lam = omega_lam
        self.omega_c = (1. - omega_m - omega_lam)


    cpdef double a(self, double z):
        cdef double a
        return 1./(1+z)

    cpdef double E(self, double a):
        cdef double E
        return (self.omega_m*a**(-3) + self.omega_c*a**(-2) + self.omega_lam)**0.5

    cpdef double __angKernel(self, double x):
        cdef __angKernel:
        """Integration kernel"""
        return self.E(x**-1)**-1

    cpdef double Da(self, double z, double z_ref=0):
        cdef double Da
        if isinstance(z, np.ndarray):
            da = np.zeros_like(z)
            for i in range(len(da)):
                da[i] = self.Da(z[i], z_ref)
            return da
        else:
            if z < 0:
                raise ValueError("Redshift z must not be negative")
            if z < z_ref:
                raise ValueError("Redshift z must not be smaller than the reference redshift")

            d = integrate.quad(self.__angKernel, z_ref+1, z+1,epsrel=1.e-6, epsabs=1.e-12)
            rk = (abs(self.omega_c))**0.5
            if (rk*d[0] > 0.01):
                if self.omega_c > 0:
                    d[0] = sinh(rk*d[0])/rk
                if self.omega_c < 0:
                    d[0] = sin(rk*d[0])/rk
            return d[0]/(1+z)

I also wonder whether I convert all the arguments correctly into cython argument? I want to change my original python code to improve the speed of calculation. One of the bottleneck in my code I reckon, should be integrate.quad. Is there any substitution for this function in cython which helps to speed up the performance of my code?

cdef class halo_positions(object):
     cdef double x = None
     cdef double y = None
     def __init__(self,numpy.ndarray[double, ndim=1] positions):
         self.x = positions[0]
         self.y = positions[1]

And if I want to pass an array to halo_positions instance is it a right way to do it?

Answer 1

If your class is defined as cdef it will be accessible only in Cython (not in Python) making it unnecessary and not efficient to use cpdef and def for the class methods. You can convert them all to cdef.

When you tell that z is double, it will accept only a double. If you want this argument to be of two different types, you should keep its type undeclared, but this will directly affect the loop performance when z is a ndarray.

Alternatively you could use double * and pass the size of it, when the size is 1 it is a double, when the size is >1 an array. The function would be:

cdef double Da(self, int size, double *z, double z_ref=0):
    if size>1:
        da = np.zeros(size)
        for i in range(size):
            da[i] = self.Da(1, &z[i], z_ref)
        return da
    else:
        if z[0] < 0:
            raise ValueError("Redshift z must not be negative")
        if z[0] < z_ref:
            raise ValueError("Redshift z must not be smaller than the reference redshift")

        d = integrate.quad(self.__angKernel, z_ref+1, z[0]+1,
                           epsrel=1.e-6, epsabs=1.e-12)
        rk = (abs(self.omega_c))**0.5
        if (rk*d[0] > 0.01):
            if self.omega_c > 0:
                d[0] = sinh(rk*d[0])/rk
            if self.omega_c < 0:
                d[0] = sin(rk*d[0])/rk
        return d[0]/(1+z[0])