Generate M-spline and I-spline functions in Python

splines.py

"""
Python code to generate M-splines and I-splines.

References
----------
Ramsay, J. O. (1988). Monotone regression splines in action.
Statistical science, 3(4), 425-441.
"""

import numpy as np
import matplotlib.pyplot as plt

def mspline_grid(order, num_basis_funcs, nt):
    """
    Generate a set of M-spline basis functions with evenly
    spaced knots.

    Parameters
    ----------
    order : int
        Order parameter of the splines.
    num_basis_funcs : int
        Number of desired basis functions. Note that we
        require num_basis_funcs >= order.
    nt : int
        Number of points to evaluate the basis functions.

    Returns
    -------
    spine_basis : array
        Matrix with shape (num_basis_funcs, nt), holding the
        desired spline basis functions.
    """

    # Determine number of interior knots.
    num_interior_knots = num_basis_funcs - order
    
    if num_interior_knots < 0:
        raise ValueError(
            "Spline `order` parameter cannot be larger "
            "than `num_basis_funcs` parameter."
        )

    # Fine grid of numerically evaluated points.
    x = np.linspace(0, 1 - 1e-6, nt)

    # Set of spline knots. We need to add extra knots to
    # the end to handle boundary conditions for higher-order
    # spline bases. See Ramsay (1988) cited above.
    #
    # Note - this is poorly explained on most corners of the
    # internet that I've found.
    knots = np.concatenate((
        np.zeros(order - 1),
        np.linspace(0, 1, num_interior_knots + 2),
        np.ones(order - 1),
    ))

    # Evaluate and stack each basis function.
    return np.row_stack(
        [mspline(x, order, i, knots) for i in range(num_basis_funcs)]
    )

def mspline(x, k, i, T):
    """
    Compute M-spline basis function `i` at points `x` for a spline
    basis of order-`k` with knots `T`.

    Parameters
    ----------
    x : array
        Vector holding points to evaluate the spline.
    """

    # Boundary conditions.
    if (T[i + k] - T[i]) < 1e-6:
        return np.zeros_like(x)

    # Special base case of first-order spline basis.
    elif k == 1:
        v = np.zeros_like(x)
        v[(x >= T[i]) & (x < T[i + 1])] = 1 / (T[i + 1] - T[i])
        return v

    # General case, defined recursively
    else:
        return k * (
            (x - T[i]) * mspline(x, k - 1, i, T)
            + (T[i + k] - x) * mspline(x, k - 1, i + 1, T)
        ) / ((k-1) * (T[i + k] - T[i]))

    
def ispline_grid(order, num_basis_funcs, nt):
    """
    Generate a set of I-spline basis functions with evenly
    spaced knots.

    Parameters
    ----------
    order : int
        Order parameter of the splines.
    num_basis_funcs : int
        Number of desired basis functions. Note that we
        require num_basis_funcs >= order.
    nt : int
        Number of points to evaluate the basis functions.

    Returns
    -------
    spine_basis : array
        Matrix with shape (num_basis_funcs, nt), holding the
        desired spline basis functions.
    """

    # Determine number of interior knots.
    num_interior_knots = num_basis_funcs - order
    
    if num_interior_knots < 0:
        raise ValueError(
            "Spline `order` parameter cannot be larger "
            "than `num_basis_funcs` parameter."
        )

    # Fine grid of numerically evaluated points.
    x = np.linspace(0, 1 - 1e-6, nt)

    # Set of spline knots. We need to add extra knots to
    # the end to handle boundary conditions for higher-order
    # spline bases. See Ramsay (1988) cited above.
    #
    # Note - this is poorly explained on most corners of the
    # internet that I've found.
    knots = np.concatenate((
        np.zeros(order),
        np.linspace(0, 1, num_interior_knots + 2),
        np.ones(order),
    ))

    # Evaluate and stack each basis function.
    return np.row_stack(
        [ispline(x, order, i, knots) for i in range(num_basis_funcs)]
    )

def ispline(x, k, i, T):
    """
    Compute I-spline basis function `i` at points `x` for a spline
    basis of order-`k` with knots `T`.

    Parameters
    ----------
    x : array
        Vector holding points to evaluate the spline.
    """

    # index j where t_j ≤ x < t_{j+1}
    i += 1
    j = np.array([np.argwhere(T <= xi)[-1, 0] if T[0] <= xi else -1 for xi in x])

    # Boundary conditions
    v = np.zeros_like(x)
    v[j-k+1 > i] = 1
    
    # General case, defined by M-spline
    j_valid = j[(j-k+1 <= i) & (i <= j)]
    x_valid = x[(j-k+1 <= i) & (i <= j)]
    
    # (T[i+k] - t[i-1])  ~  (T[j+k] - t[j-1])
    Tms = [T[(i+k+1):(j_+k+1+1)] - T[i:(j_+1)] for j_ in j_valid]
    
    # M_i(x|k+1,t)  ~  M_j(x|k+1,t)
    Mms = [np.array([mspline(x_, k+1, m, T) for m in range(i, j_+1)]) \
           for x_, j_ in zip(x_valid, j_valid)]
    
    v[(j-k+1 <= i) & (i <= j)] = [(Tm * Mm).sum()/(k+1) for Tm, Mm in zip(Tms, Mms)]
        
    return v

# Test code
if __name__ == "__main__":
    # Plot basis funcs, varying the order parameter.
    fig, axes = plt.subplots(1, 3, figsize=(10, 3))
    for k, ax in enumerate(axes):
        order = k + 1
        # I-spline
        ax.plot(ispline_grid(order, 7, 1000).T)
        # # M-spline
        # ax.plot(mspline_grid(order, 7, 1000).T)
        ax.set_yticks([])
        ax.set_title(f"order-{order}")
    plt.suptitle("I-spline")
    fig.tight_layout()
    plt.show()