Implementation of B-splines and derivatives in Stan

splines.stan

/* splines.stan
 * Include this in the functions block to get splines in Stan
 *
 * Construct B-spline basis functions; taken from Kharratzadeh's example.
 * see https://mc-stan.org/users/documentation/case-studies/splines_in_stan.html
 * N.B. recursive functions need a forward declaration.
 */
vector build_b_spline(vector t, vector knots, int ind, int order);
vector build_b_spline(vector t, vector knots, int ind, int order) {
    int n = num_elements(t); int k = num_elements(knots);
    vector[n] b_spline;
    vector[n] w1 = rep_vector(0, n);
    vector[n] w2 = rep_vector(0, n);
    if ( order == 1 ) {
        for ( i in 1:n ) {
            b_spline[i] = (knots[ind] <= t[i]) && (t[i] < knots[ind+1]);
        }
    } else { // order > 1
        if ( knots[ind] != knots[ind+order-1] ) {
            w1 = (t - rep_vector(knots[ind], n)) / (knots[ind+order-1] - knots[ind]);
        }
        if ( knots[ind+1] != knots[ind+order] ) {
            w2 = (rep_vector(knots[ind+order], n) - t) / (knots[ind+order] - knots[ind+1]);
        }
        b_spline = w1 .* build_b_spline(t, knots, ind, order-1) 
                 + w2 .* build_b_spline(t, knots, ind+1, order-1);
    }
    return b_spline;
}

/* Derivative of the B-spline basis.
 * B_{k,i}'(t) = (k-1) / (tau_{i+k-1} - tau_i}) * B_{i, k-1}(t) 
 *    - (k-1) / (tau_{i+k} - tau_{i+1}) * B_{i+1, k-1}(t)
 */
vector build_derivative_b_spline(vector t, vector knots, int ind, int order) {
    int n = num_elements(t);
    vector[n] deriv_b_spline;
    if ( order == 1 ) {
        // piece-wise constant, hence 0 derivative
        deriv_b_spline = rep_vector(0, n);
    } else { // order > 1
        real w1 = 0; real w2 = 0;
        if ( knots[ind] != knots[ind+order-1] ) {
            w1 = (order-1) / (knots[ind+order-1] - knots[ind]);
        }
        if ( knots[ind+1] != knots[ind+order] ) {
            w2 = (order-1) / (knots[ind+order] - knots[ind+1]);
        }
        deriv_b_spline = w1 * build_b_spline(t, knots, ind, order-1) 
                       - w2 * build_b_spline(t, knots, ind+1, order-1);
    }
    return deriv_b_spline;
}