Translating python function with boolean indexing within loop to cython

Below I’ll write a function in pure python that I want to Cythonize.

def do_stuff(M_i, C_i):
    return M_i.dot(C_i).dot(M_i)

def foo(M, C):
    '''
    M : np.array
        N x J matrix
    C : np.array
        J x J matrix
    '''

    N = M.shape[0]

    tot = 0

    for i in range(N):
        nonmiss = ~np.isnan(M[i, :])
        M_i = M[i, nonmiss] # select non empty components of M
        C_i = C[nonmiss, :] # select corresponding components of C
        C_i = C_i[:, nonmiss] # select corresponding components of C

        tot = tot + do_stuff(M_i, C_i)

    return tot

Suppose that I know how to Cythonize the function do_stuff. The actual do_stuff function I’m interested in is more complicated than the one above but I thought I’d provide an example. The real do_stuff function computes determinants and takes inverses, in addition to matrix multiplication.

My main problem has to do with creating the M_i and C_i subvector and submatrix. I’m not sure I can do the same sort of boolean indexing in Cython. And if I can, I don’t know how. But I can get started with the bit of Cython I know.

def foo_c(double[:, ::1] M, double[:, ::1] C):

    cdef int N = M.shape[0]
    cdef double tot = 0
    ...

    for i in range(N):
        ...
        tot = tot + do_stuff_c(M_i, C_i)

    return tot

Answer

You probably won’t gain too much speed here, since boolean indexing in Numpy is implemented in C anyway, so should be fairly fast. The main thing you’d avoid is creating some unnecessary intermediates (which involves a memory allocation so is potentially slow).

What you want to do is create temporary arrays for M_i and C_i that are the largest possible size (i.e. J or JxJ). As you iterate through isnan(M_I) you keep track of how many values you’ve actually stored. Then at the end you trim M_i and C_i to only the part that you’ve used:

Untested code:

for i in range(N):
    filled_count_j = 0
    M_i = np.empty((M.shape[1],))
    C_i = np.empty((M.shape[1],M.shape[1]))
    for j in range(M.shape[1]):
        if ~isnan(M[i,j]):
            M_i[filled_count] = M[i,j]

            filled_count_k = 0
            for k in range(M.shape[1]):
                if ~isnan(M[i,k]):
                    C_i[filled_count_j,filled_count_k] = C[j,k]
                    filled_count_k += 1
            filled_count_j += 1
     M_i = M_i[:filled_count_j]
     C_i = C_i[:filled_count_j,:filled_count_j]