为 cython 类中的函数定义参数和 cython 中的快速积分计算

Question

我是 cython 的新手，正在尝试将 python 类转换为 cython。我不知道应该如何在 instance Da 中定义参数 z，因为它可以同时处理 numpy.array 或单个 float号码。

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)

我也想知道我是否将所有参数正确转换为 cython 参数？我想改变我原来的python代码来提高计算速度。我认为我的代码中的瓶颈之一应该是integrate.quad。 cython 中的这个函数是否有任何替代品，有助于加快我的代码的性能？

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]

如果我想将一个数组传递给halo_positions 实例，这样做是否正确？

Answer 1

如果您的类被定义为cdef，那么它只能在 Cython（而不是 Python）中访问，因此在类方法中使用 cpdef 和 def 是不必要且效率不高的。您可以将它们全部转换为cdef。

当您告诉z 是double 时，它只会接受double。如果你希望这个参数是两种不同的类型，你应该保持它的类型未声明，但是当z是ndarray时，这将直接影响循环性能。

或者，您可以使用double * 并传递它的大小，当大小为1 时，它是一个双精度数，当大小是>1 一个数组时。函数是：

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])