use_gpu = True # Set to False to use CPU
dev = str(0)

import os
if use_gpu:
    try:
        import cupy as xp
        from cupyx.scipy.interpolate import Akima1DInterpolator as stock_akima
        os.environ["OMP_NUM_THREADS"] = str(1)
        os.system("OMP_NUM_THREADS=1")
        os.system("CUDA_VISIBLE_DEVICES=" + dev)
        os.environ["CUDA_VISIBLE_DEVICES"] = dev
        os.system("echo $CUDA_VISIBLE_DEVICES")

        os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=1'
    except:
        print('No GPUs available, using CPUs')
        import numpy as np
        from scipy.interpolate import Akima1DInterpolator as stock_akima

else:
    import numpy as xp
    from scipy.interpolate import Akima1DInterpolator as stock_akima

import numpy as np
import matplotlib.pyplot as plt
from cudakima import AkimaInterpolant1D

import scipy

0

%load_ext autoreload
%autoreload 2

# just for plotting purposes
def to_cpu(x):
    try:
        return x.get()
    except:
        return x

Getting started

cubic_interp = AkimaInterpolant1D(sanitize=True, use_gpu=use_gpu, threadsperblock=512, order='cubic')
linear_interp = AkimaInterpolant1D(sanitize=True, use_gpu=use_gpu, threadsperblock=512, order='linear')

Arrays of equal shape

# Generate some random data
xmin, xmax = 0, 10 # X Domain
ymin, ymax = -100, 100 # Y Domain
npoints = 10 # Number of points
nrealizations = 100 # Number of realizations

x = xp.random.random(size=(nrealizations, npoints)) * (xmax - xmin) + xmin
# Make sure the first and last points are the same
x[:, 0] = xmin
x[:, -1] = xmax
y = xp.random.random(size=(nrealizations, npoints)) * (ymax - ymin) + ymin


# Define the x values to interpolate
xinterp = xp.linspace(xmin, xmax, 1000)

# Generate some random data
xmin, xmax = -4, np.log10(2.9e-2) # X Domain
ymin, ymax = -100, 100 # Y Domain
npoints = 10 # Number of points
nrealizations = 100 # Number of realizations

x = xp.random.random(size=(nrealizations, npoints)) * (xmax - xmin) + xmin
# Make sure the first and last points are the same
x[:, 0] = xmin
x[:, -1] = xmax
y = xp.random.random(size=(nrealizations, npoints)) * (ymax - ymin) + ymin

y = scipy.stats.truncnorm(loc=0, scale=0.1, a=-3, b=3).rvs(size=(nrealizations, npoints))


# Define the x values to interpolate
xinterp = xp.linspace(xmin, xmax, 1000)

linear_result = linear_interp(xinterp, x, y)
cubic_result = cubic_interp(xinterp, x, y)

realization = 4

for result, label in zip([linear_result, cubic_result], ['Linear Akima', 'Cubic Akima']):
    plt.plot(to_cpu(xinterp), 10**to_cpu(result)[realization], label=label)

plt.scatter(to_cpu(x)[realization], 10**to_cpu(y)[realization], marker='x', c='k', label='Knots')



plt.legend()
plt.show()

Arrays of different shapes

You can combine arrays of different shapes padding them with xp.nan and stacking them together

# generate random arrays as before, but fill some of them with nans
x = xp.random.random(size=(nrealizations, npoints)) * (xmax - xmin) + xmin
# Make sure the first and last points are the same
x[:, 0] = xmin
x[:, -4] = xmax
y = xp.random.random(size=(nrealizations, npoints)) * (ymax - ymin) + ymin

x[0, -3:] = xp.nan
x[1, -2:] = xp.nan
x[2, -1:] = xp.nan

y[0, -3:] = xp.nan
y[1, -2:] = xp.nan
y[2, -1:] = xp.nan

linear_result = linear_interp(xinterp, x, y)
cubic_result = cubic_interp(xinterp, x, y)

realization = 3

for result, label in zip([linear_result, cubic_result], ['Linear Akima', 'Cubic Akima']):
    plt.plot(to_cpu(xinterp), to_cpu(result)[realization], label=label)

plt.scatter(to_cpu(x)[realization], to_cpu(y)[realization], marker='x', c='k', label='Knots')

plt.legend()
plt.show()

arrays containing less than 4 points

The boundary conditions currently implemented require at least 4 points in the interpolation grid. If there are less then 4 non-NaN values, the code falls back on linear interpolation between them.

# generate random arrays as before, but fill some of them with nans
x = xp.random.random(size=(nrealizations, npoints)) * (xmax - xmin) + xmin
# Make sure the first and last points are the same
x[:, 0] = xmin
x[:, 1] = xmax
y = xp.random.random(size=(nrealizations, npoints)) * (ymax - ymin) + ymin

# introduce a case with only 2 grid points and one with 3 grid points

x[0, 2:] = xp.nan
x[1, 3:] = xp.nan
x[2, 4:] = xp.nan

y[0, 2:] = xp.nan
y[1, 3:] = xp.nan
y[2, 4:] = xp.nan

linear_result = linear_interp(xinterp, x, y)
cubic_result = cubic_interp(xinterp, x, y)

fig = plt.figure(figsize=(15, 5))

titles = ['2 points, linear', '3 points, linear', '4+ points, Akima']

for realization, title in zip(range(3), titles):

    plt.subplot(1, 3, realization + 1)
    for result, label in zip([linear_result, cubic_result], ['Linear Akima', 'Cubic Akima']):
        plt.plot(to_cpu(xinterp), to_cpu(result)[realization], label=label)
    plt.scatter(to_cpu(x)[realization], to_cpu(y)[realization], marker='x', c='k', label='Knots')
    plt.title(title)

    plt.legend()
plt.show()

Comparison with Cubic Splines

from scipy.interpolate import CubicSpline

realization = 5 # Choose a realization without NaNs

xhere = to_cpu(x[realization])
yhere = to_cpu(y[realization])
resulthere = to_cpu(cubic_result[realization])

#sort the arrays
idx = np.argsort(xhere)
xhere = xhere[idx]
yhere = yhere[idx]

cs = CubicSpline(xhere, yhere)
resultcs = cs(to_cpu(xinterp))

plt.plot(to_cpu(xinterp), to_cpu(result)[realization], c='k', ls = '--', lw=1, label='Akima Splines')
plt.plot(to_cpu(xinterp), resultcs, c='r', ls = '--', lw=1, label='Cubic Splines')
plt.scatter(to_cpu(x)[realization], to_cpu(y)[realization], marker='x', c='k', label='Knots')

plt.legend()
plt.show()

Akima splines show smoother behaviour with respect to Cubic Splines

Timing against the stock implementation in Scipy / Cupy

# Generate some random data
xmin, xmax = 0, 10 # X Domain
ymin, ymax = -100, 100 # Y Domain
npoints = 10 # Number of points
nrealizations = 100 # Number of realizations

x = xp.random.random(size=(nrealizations, npoints)) * (xmax - xmin) + xmin
# Make sure the first and last points are the same
x[:, 0] = xmin
x[:, -1] = xmax
y = xp.random.random(size=(nrealizations, npoints)) * (ymax - ymin) + ymin

# Sort the array, cupy and scipy demand sorted x values
x = xp.sort(x, axis=1)

# Define the x values to interpolate
xinterp = xp.linspace(xmin, xmax, 1000)

def stock_interpolator(x, x_all, y_all):

    result = xp.empty((len(x_all), len(x)))
    for i, (x_here, y_here) in enumerate(zip(x_all, y_all)):
        interpolator = stock_akima(x_here, y_here)
        result[i] = interpolator(x)
    return result

cudakima_interpolator = AkimaInterpolant1D(sanitize=False, use_gpu=use_gpu, threadsperblock=128)

%%timeit
# time stock scipy / cupy
result_stock = stock_interpolator(xinterp, x, y)

236 ms ± 2.21 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%%timeit
#time cudakima
result_cudakima = cudakima_interpolator(xinterp, x, y)

2.83 ms ± 13.9 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

result_cudakima = to_cpu(cudakima_interpolator(xinterp, x, y))
result_stock = to_cpu(stock_interpolator(xinterp, x, y))

realization = np.random.randint(0, nrealizations)
fig = plt.figure(figsize=(12, 6))

plt.subplot(121)
plt.plot(to_cpu(xinterp), result_cudakima[realization], c='k', ls = '-', lw=2, label='Cudakima')
plt.plot(to_cpu(xinterp), result_stock[realization], c='red', ls = '--', lw=2, label='Stock scipy/cupy', alpha=0.5)
plt.scatter(to_cpu(x)[realization], to_cpu(y)[realization], marker='x', c='k', label='Knots')
plt.legend()
plt.subplot(122)
plt.plot(to_cpu(xinterp), result_cudakima[realization] - result_stock[realization], c='k', ls = '-', lw=2, label='difference')
plt.legend()
plt.show()

Multidimensional new arrays

from cudakima import AkimaInterpolant1DMultiDim

linear_multidim_interp = AkimaInterpolant1DMultiDim(sanitize=False, use_gpu=use_gpu, threadsperblock=512, order='linear')
cubic_multidim_interp = AkimaInterpolant1DMultiDim(sanitize=False, use_gpu=use_gpu, threadsperblock=512, order='cubic')

# generate some data
batch_shape = (100, 5)
max_points_per_group = 100
full_shape = batch_shape + (max_points_per_group,)

n_interp = 10000

x_data = xp.full(full_shape, xp.nan)
y_data = xp.full(full_shape, xp.nan)
x_new_multidim = xp.full(batch_shape + (n_interp,), xp.nan)

for idx in xp.ndindex(batch_shape):
    # Random group size (at least 4 for Akima)
    group_size = int(xp.random.randint(4, max_points_per_group + 1))

    # Create smooth test function with some noise
    tmp = xp.random.uniform(0, 10, group_size)
    x_group = xp.sort(tmp)
    y_group = xp.sin(x_group) + 0.1 * x_group + 0.05 * xp.random.randn(group_size)

    x_data[idx][:group_size] = x_group
    y_data[idx][:group_size] = y_group

    # Create interpolation points for this group (within data range)
    x_min, x_max = x_group[0], x_group[-1]
    x_new_multidim[idx] = xp.linspace(x_min, x_max, n_interp)

for _ in [x_data, y_data, x_new_multidim]:
    print(_.shape)

(100, 5, 100)
(100, 5, 100)
(100, 5, 10000)

linear_result = linear_multidim_interp(x_new_multidim, x_data, y_data)
cubic_result = cubic_multidim_interp(x_new_multidim, x_data, y_data)

%%timeit

_ = cubic_multidim_interp(x_new_multidim, x_data, y_data)

2.89 ms ± 14.1 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

# visualize a few results
N1 = 3
N2 = 3
Nsubplots = 101 + N2 * 10

for i in range(N1):
    plt.figure(figsize=(15, 10))
    for j in range(N2):
        plt.subplot(Nsubplots + j)
        idx = (i, j)
        for result, label in zip([linear_result, cubic_result], ['Linear Akima', 'Cubic Akima']):
            plt.plot(to_cpu(x_new_multidim[idx]), to_cpu(result[idx]), ls = '-', lw=1, label=label)
        valid = ~xp.isnan(x_data[idx]) & ~xp.isnan(y_data[idx])
        plt.scatter(to_cpu(x_data[idx][valid]), to_cpu(y_data[idx][valid]), marker='x', c='k', label='Knots')
        plt.title(f'Indices {idx}, {valid.sum()} points')
        if i == 0 and j == 0:
            plt.legend()