summator.pyx

#cython: language_level=3, boundscheck=False, wraparound=False, cdivision=True
# -*- coding: utf-8 -*-
"""
This is the variogram estimater, implemented in cython.
"""

import numpy as np

cimport cython
from cython.parallel import prange
from libc.math cimport sin, cos
cimport numpy as np


DTYPE = np.double
ctypedef np.double_t DTYPE_t


def summate_unstruct(
    const double[:,:] cov_samples,
    const double[:] z_1,
    const double[:] z_2,
    const double[:,:] pos
    ):
    cdef int i, j, d, X_len, N
    cdef double phase
    cdef int dim
    dim = pos.shape[0]

    X_len = pos.shape[1]
    N = cov_samples.shape[1]

    cdef double[:] summed_modes = np.zeros(X_len, dtype=DTYPE)

    for i in prange(X_len, nogil=True):
        for j in range(N):
            phase = 0.
            for d in range(dim):
                phase += cov_samples[d,j] * pos[d,i]
            summed_modes[i] += z_1[j] * cos(phase) + z_2[j] * sin(phase)

    return np.asarray(summed_modes)

def summate_struct(
    const double[:,:] cov_samples,
    const double[:] z_1,
    const double[:] z_2,
    const double[:] x,
    const double[:] y=None,
    const double[:] z=None,
):
    if y == None and z == None:
        return summate_struct_1d(cov_samples, z_1, z_2, x)
    elif z == None:
        return summate_struct_2d(cov_samples, z_1, z_2, x, y)
    else:
        return summate_struct_3d(cov_samples, z_1, z_2, x, y, z)

def summate_struct_1d(
    const double[:,:] cov_samples,
    const double[:] z_1,
    const double[:] z_2,
    const double[:] x,
    ):

    cdef int i, j, X_len, N
    cdef double phase

    X_len = x.shape[0]
    N = cov_samples.shape[1]

    cdef double[:] summed_modes = np.zeros(X_len, dtype=DTYPE)

    for i in prange(X_len, nogil=True):
        for j in range(N):
            phase = cov_samples[0,j] * x[i]
            summed_modes[i] += z_1[j] * cos(phase) + z_2[j] * sin(phase)

    return np.asarray(summed_modes)

def summate_struct_2d(
    const double[:,:] cov_samples,
    const double[:] z_1,
    const double[:] z_2,
    const double[:] x,
    const double[:] y,
    ):
    cdef int i, j, k, X_len, Y_len, N
    cdef double phase

    X_len = x.shape[0]
    Y_len = y.shape[0]
    N = cov_samples.shape[1]

    cdef double[:,:] summed_modes = np.zeros((X_len, Y_len), dtype=DTYPE)

    for i in prange(X_len, nogil=True):
        for j in range(Y_len):
            for k in range(N):
                phase = cov_samples[0,k] * x[i] + cov_samples[1,k] * y[j]
                summed_modes[i,j] += z_1[k] * cos(phase) + z_2[k] * sin(phase)

    return np.asarray(summed_modes)

def summate_struct_3d(
    const double[:,:] cov_samples,
    const double[:] z_1,
    const double[:] z_2,
    const double[:] x,
    const double[:] y,
    const double[:] z,
    ):
    cdef int i, j, k, l, X_len, Y_len, Z_len, N
    cdef double phase

    X_len = x.shape[0]
    Y_len = y.shape[0]
    Z_len = z.shape[0]
    N = cov_samples.shape[1]

    cdef double[:,:,:] summed_modes = np.zeros((X_len, Y_len, Z_len), dtype=DTYPE)

    for i in prange(X_len, nogil=True):
        for j in range(Y_len):
            for k in range(Z_len):
                for l in range(N):
                    phase = (
                        cov_samples[0,l] * x[i] +
                        cov_samples[1,l] * y[j] +
                        cov_samples[2,l] * z[k]
                    )
                    summed_modes[i,j,k] += z_1[l] * cos(phase) + z_2[l] * sin(phase)

    return np.asarray(summed_modes)

cdef (double) abs_square(const double[:] vec) nogil:
    cdef int i
    cdef double r = 0.

    for i in range(vec.shape[0]):
        r += vec[i]**2

    return r

def summate_incompr_unstruct(
    const double[:,:] cov_samples,
    const double[:] z_1,
    const double[:] z_2,
    const double[:,:] pos
    ):
    cdef int i, j, d, X_len, N
    cdef double phase
    cdef int dim
    cdef double k_2
    dim = pos.shape[0]

    cdef double[:] e1 = np.zeros(dim, dtype=DTYPE)
    e1[0] = 1.
    cdef double[:] proj = np.empty(dim, dtype=DTYPE)

    X_len = pos.shape[1]
    N = cov_samples.shape[1]

    cdef double[:,:] summed_modes = np.zeros((dim, X_len), dtype=DTYPE)

    for i in prange(X_len, nogil=True):
        for j in range(N):
            k_2 = abs_square(cov_samples[:,j])
            phase = 0.
            for d in range(dim):
                phase += cov_samples[d,j] * pos[d,i]
            for d in range(dim):
                proj[d] = e1[d] - cov_samples[d,j] * cov_samples[0,j] / k_2
                summed_modes[d,i] += proj[d] * (z_1[j] * cos(phase) + z_2[j] * sin(phase))

    return np.asarray(summed_modes)

def summate_incompr_struct(
    const double[:,:] cov_samples,
    const double[:] z_1,
    const double[:] z_2,
    const double[:] x,
    const double[:] y=None,
    const double[:] z=None,
):
    if z == None:
        return summate_incompr_struct_2d(cov_samples, z_1, z_2, x, y)
    else:
        return summate_incompr_struct_3d(cov_samples, z_1, z_2, x, y, z)

def summate_incompr_struct_2d(
    const double[:,:] cov_samples,
    const double[:] z_1,
    const double[:] z_2,
    const double[:] x,
    const double[:] y,
    ):
    cdef int i, j, k, d, X_len, Y_len, N
    cdef double phase
    cdef int dim = 2
    cdef double k_2

    cdef double[:] e1 = np.zeros(dim, dtype=DTYPE)
    e1[0] = 1.
    cdef double[:] proj = np.empty(dim, dtype=DTYPE)

    X_len = x.shape[0]
    Y_len = y.shape[0]
    N = cov_samples.shape[1]

    cdef double[:,:,:] summed_modes = np.zeros((dim, X_len, Y_len), dtype=DTYPE)

    for i in prange(X_len, nogil=True):
        for j in range(Y_len):
            for k in range(N):
                k_2 = abs_square(cov_samples[:,k])
                phase = cov_samples[0,k] * x[i] + cov_samples[1,k] * y[j]
                for d in range(dim):
                    proj[d] = e1[d] - cov_samples[d,k] * cov_samples[0,k] / k_2
                    summed_modes[d,i,j] += proj[d] * (z_1[k] * cos(phase) + z_2[k] * sin(phase))

    return np.asarray(summed_modes)

def summate_incompr_struct_3d(
    const double[:,:] cov_samples,
    const double[:] z_1,
    const double[:] z_2,
    const double[:] x,
    const double[:] y,
    const double[:] z,
    ):
    cdef int i, j, k, l, d, X_len, Y_len, Z_len, N
    cdef double phase
    cdef int dim = 3
    cdef double k_2

    cdef double[:] e1 = np.zeros(dim, dtype=DTYPE)
    e1[0] = 1.
    cdef double[:] proj = np.empty(dim, dtype=DTYPE)

    X_len = x.shape[0]
    Y_len = y.shape[0]
    Z_len = z.shape[0]
    N = cov_samples.shape[1]

    cdef double[:,:,:,:] summed_modes = np.zeros((dim, X_len, Y_len, Z_len), dtype=DTYPE)

    for i in prange(X_len, nogil=True):
        for j in range(Y_len):
            for k in range(Z_len):
                for l in range(N):
                    k_2 = abs_square(cov_samples[:,l])
                    phase = (
                        cov_samples[0,l] * x[i] +
                        cov_samples[1,l] * y[j] +
                        cov_samples[2,l] * z[k]
                    )
                    for d in range(dim):
                        proj[d] = e1[d] - cov_samples[d,l] * cov_samples[0,l] / k_2
                        summed_modes[d,i,j,k] += proj[d] * (z_1[l] * cos(phase) + z_2[l] * sin(phase))

    return np.asarray(summed_modes)