binorm.pyx

# cython: boundscheck=False, wraparound=False
# distutils: extra_compile_args = -fopenmp
# distutils: extra_link_args = -fopenmp

from cython.parallel import prange
from libc.math cimport exp, sqrt
import numpy as np

cdef double pi = np.pi
cdef double tau = np.pi * 2


def binorm(double[:] a, double[:] b, double[:] x, double[:] y,
           double[:] xvar, double[:] yvar, double[:] xyvar, 
           int nthread=1):
    """binorm(a, b, x, y, xvar, yvar, xyvar, nthread=1)

    Return the probability density of data points in given model.
    The model is described by points (x, y, xvar, yvar, xyvar).

    a, b :
        data points
    x, y, xvar, yvar, xyvar :
        model points and weights
    nthread :
        number of threads
    """
    cdef:
        int m, n, i, j
        double aj, bj, sx, sy, sxy, det, dx, dy
        double p, psum
        double[:] out

    assert x.size == y.size == xvar.size == yvar.size == xyvar.size
    assert a.size == b.size

    m, n = x.size, a.size
    out = np.empty(n, dtype='double')
    
    for j in prange(n, nogil=True, num_threads=nthread):
        aj, bj = a[j], b[j]
        psum = 0.
        for i in range(m):
            dx, dy = aj - x[i], bj - y[i]
            sx, sy, sxy = xvar[i], yvar[i], xyvar[i]
            det = sx * sy - sxy * sxy
            p = (2 * dx * dy * sxy - dx * dx * sy - dy * dy * sx) / 2. / det
            psum = psum + exp(p) / sqrt(det)
        out[j] = psum / tau  / m
    return out.base


def binorm_alt(double[:] x, double[:] y, 
           double[:] xvar, double[:] yvar, double[:] xyvar, 
           double[:] a, double[:] b, double[:] w,
           int nthread=1):
    """binorm(x, y, xvar, yvar, xyvar, a, b, w, nthread=1)

    Return the probability density of data points in given model.
    The model is described by points (a, b, w).

    a, b, w :
        model points and weights
    x, y, xvar, yvar, xyvar :
        data points
    nthread :
        number of threads
    """
    cdef:
        int m, n, i, j
        double xi, yi, sx, sy, sxy, det, dx, dy
        double p, psum, wsum
        double[:] out

    assert x.size == y.size == xvar.size == yvar.size == xyvar.size
    assert a.size == b.size == w.size

    m, n = x.size, a.size
    out = np.empty(m, dtype='double')
    
    wsum = 0
    for j in range(n):
        wsum += w[j]

    # for i in range(m):
    for i in prange(m, nogil=True, num_threads=nthread):
        xi, yi = x[i], y[i]
        sx, sy, sxy = xvar[i], yvar[i], xyvar[i]
        det = sx * sy - sxy * sxy
        psum = 0.
        for j in range(n):
            dx, dy = a[j] - xi, b[j] - yi
            p = (2 * dx * dy * sxy - dx * dx * sy - dy * dy * sx) / 2. / det
            psum = psum + exp(p) * w[j]
        out[i] = psum / tau / sqrt(det) / wsum
    return out.base

binorm_weighted.py

code = '''
from cython.parallel import prange
import numpy as np
from libc.math cimport exp, pi

def binorm_weighted(double[:] xi, double[:] yi, double[:] x, double[:] y, double[:] s, double[:] w, int nproc=1):
    """
    xi, yi: where to calculate the density
    x, y: points who define the density field
    s: Gaussian kernel size
    w: weight of the points
    """
    cdef:
        int ni = len(xi), n = len(x)
        int i, j
        double[:] zi = np.zeros(ni, dtype='f8')

    for i in prange(ni, num_threads=nproc, nogil=True):
        for j in range(n):
            zi[i] += (
                w[j] / (2 * pi * s[j]**2) *
                exp(-0.5 * ((xi[i] - x[j]) ** 2 + (yi[i] - y[j])**2) / s[j]**2)
            )
    return zi.base
'''
binorm_weighted = cythonmagic(code, fast_indexing=True).binorm_weighted

It seems the numba version is slightly faster than Cython.

from numba import vectorize

@vectorize(["float64(float64, float64, float64, float64, float64, float64, float64)"])
def binorm(xi, yi, x, y, xvar, yvar, xyvar):
    X, Y = xi - x, yi - y
    det = xvar * yvar - xyvar**2
    p = (np.exp(-(X**2 * yvar + Y**2 * xvar - 2 * X * Y * xyvar) / (2 * det)) / (2 * np.pi * np.sqrt(det)))
    return p

def binorm_grid(xi, yi, x, y, xvar, yvar, xyvar):
    xi, yi = xi[..., None, None], yi[..., None, :, None]
    p = binorm(xi, yi, x, y, xvar, yvar, xyvar)
    return p.mean(-1)