spherical spline interpolation

sspline.py

import numpy as np
from numpy.polynomial.legendre import legval
from numpy.linalg import solve


def calc_spline(cosang, stiffnes=4, num_lterms=50):
    """Calculate spherical spline between points on a sphere.
    
    Parameters
    ----------
    cosang : array-like | float
        cosine of angles between pairs of points on a spherical surface. This
        is equivalent to the dot product of unit vectors.
    stiffnes : float
        stiffnes of the spline
    num_lterms : int
        number of Legendre terms to evaluate
    """
    factors = [(2 * n + 1) / (n ** stiffnes * (n + 1) ** stiffnes * 4 * np.pi)
               for n in range(1, num_lterms + 1)]
    return legval(cosang, factors)


def interpolate(data, datapos, targetpos, stiffnes=4, num_lterms=50):
    """Spherical spline interpolation of time series.
    
    Parameters
    ----------
    data : array-like, shape = [n_epochs, n_channels, n_samples]
        Time series data
    datapos : array-like, shape = [n_channels, 3]
        3D positions of the data channels. Each position vector is assumed to
        lie on the unit sphere (must me normalized to 1).
    targetpos : array-like, shape = [n_channels, 3]
        3D positions of interpolation points. Each position vector is assumed
        to lie on the unit sphere (must me normalized to 1).
    stiffnes : float
        stiffnes of the spline
    num_lterms : int
        number of Legendre terms to evaluate
    """
    
    data = np.asarray(data)
    datapos = np.asarray(datapos)
    targetpos = np.asarray(targetpos)
    
    _, n_channels, n_samples = data.shape
    
    source_cosangles = np.dot(datapos, datapos.T)
    source_splines = calc_spline(source_cosangles, stiffnes, num_lterms)
    
    source_splines = np.ones((1 + n_channels, 1 + n_channels))
    source_splines[-1, 0] = 0
    source_splines[0:-1, 1:] = calc_spline(source_cosangles, stiffnes, num_lterms)
        
    transfer_cosangles = np.dot(targetpos, datapos.T)
    transfer_splines = calc_spline(transfer_cosangles, stiffnes, num_lterms)

    output = []    
    for epoch in data:
        z = np.concatenate((epoch, np.zeros((1, n_samples))))
        coefficients = solve(source_splines, z)
        output.append(np.dot(transfer_splines, coefficients[1:, :]) + coefficients[0, :])
    return np.array(output)
        
        
if __name__ == '__main__':
    data = np.random.randn(1, 6, 10)
    pos = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1], [-1, 0, 0], [0, -1, 0], [0, 0, -1]])
    
    out = interpolate(data, pos, pos)
    
    assert(np.allclose(data, out))