syrte
10/26/2017 - 4:04 PM

sum_tidal.pyx

%%cython
import numpy as np
cdef extern from "<math.h>" nogil:
    float sqrtf(float)

def sum_tidal(int[:] ix, float[:, :] rs, int length=0):
    cdef:
        int h, i, j, k
        int n = ix.size, m = max(np.max(ix) + 1, length)
        float r1, r2, r3, r5
        float[:, :, :] res = np.zeros((m, 3, 3), dtype='f4')

    for i in range(n):
        h = ix[i]
        r2 = 0
        for j in range(3):
            r2 += rs[i, j] * rs[i, j]
        r1 = sqrtf(r2)
        r3 = r1 * r2
        r5 = r2 * r3
        for j in range(3):
            for k in range(3):
                if j == k:
                    res[h, j, k] += 3 * rs[i, j] * rs[i, k] / r5 - 1 / r3
                else:
                    res[h, j, k] += 3 * rs[i, j] * rs[i, k] / r5
    return res.base