windowize.pyx

# cython windowize.pyx
# gcc -shared -pthread -fPIC -fwrapv -O2 -Wall -fno-strict-aliasing -I/usr/include/python2.7 -o windowize.so windowize.c

import numpy as np
cimport numpy as np
cimport cython

DTYPE = np.float64
ctypedef np.float64_t dtype_t

@cython.boundscheck(False)
cpdef np.ndarray[dtype_t, ndim=2] windowize(np.ndarray[dtype_t, ndim=3] vol, int window_size):
    """
    vol: numpy ndarray, dtype `numpy.float64`, shape `(nx,ny,nz)`
        The volume to be windowized.

    window_size: int
        Each window will be size, `(window_size, window_size, window_size)`.

    returns: numpy ndarray, dtype `numpy.float64`, shape `(n_windows, window_size**3)`.
        Each row in the returned matrix is a flattened window. `n_windows` is found by taking the product of `(nx-window_size+`)*(ny-window_size+1)*(nz-window_size+1)`.
        
    Example:
        >>> import numpy as np
        >>> import windowize as wd
        >>> V = np.random.randn(50,50,50).astype(np.float64)
        >>> windows = wd.windowize(vol=V, window_size=5)
    """
    assert vol.dtype == DTYPE, "Input volume must be dtype, %s" % str(DTYPE)
    assert window_size % 2 == 1, "Window size must be odd."

    cdef int w = int(np.floor(window_size / 2.))

    cdef int n_windows = 1
    cdef int nx = vol.shape[0]
    cdef int ny = vol.shape[1]
    cdef int nz = vol.shape[2]
    n_windows *= int(nx-window_size+1)
    n_windows *= int(ny-window_size+1)
    n_windows *= int(nz-window_size+1)

    cdef np.ndarray[dtype_t, ndim=2] windows = np.empty((n_windows, window_size**3), dtype=DTYPE)

    cdef int i,j,k, ii,jj,kk
    cdef n = 0
    cdef l = 0

    for i in range(w,nx-w):
        for j in range(w,ny-w):
            for k in range(w,nz-w):
                l = 0
                for ii in range(i-w,i+w+1):
                    for jj in range(j-w,j+w+1):
                        for kk in range(k-w,k+w+1):
                            windows[n,l] = vol[ii,jj,kk]
                            l+=1
                n+=1

    return windows