Allow smoothing splines with vector data

Author: jipolanco

src/SplineInterpolations/smoothing.jl

@@ -50,9 +50,9 @@ function fit(
     λ ≥ 0 || throw(DomainError(λ, "the smoothing parameter λ must be non-negative"))
     eachindex(xs) == eachindex(ys) || throw(DimensionMismatch("x and y vectors must have the same length"))
     N = length(xs)
-    cs = similar(xs)
+    cs = similar(ys)
 
-    T = eltype(cs)
+    T = eltype(xs)
 
     # Create natural cubic B-spline basis with knots = input points
     B = BSplineBasis(order, copy(xs))
@@ -98,15 +98,24 @@ function fit(
     # ldiv!(cs, F, z)
 
     # Construct RHS trying to reduce allocations
-    zs = copy(ys)
+    Y = eltype(ys)
+    if ndims(Y) == 0  # scalar data -- this includes float types (ndims(Float64) == 0)
+        cs_lin = cs
+        zs = copy(ys)
+    else  # vector or multidimensional data, for example for SVector values
+        Z = eltype(eltype(ys))  # for example, if eltype(ys) <: SVector{D,Z}
+        @assert Z <: Number
+        cs_lin = reinterpret(reshape, Z, cs)'
+        zs = copy(reinterpret(reshape, Z, ys)')  # dimensions (N, D)
+    end
     if weights !== nothing
         eachindex(weights) == eachindex(xs) || throw(DimensionMismatch("the `weights` vector must have the same length as the data"))
         lmul!(Diagonal(weights), zs)  # zs = W * ys
     end
-    mul!(cs, A', zs)  # cs = A' * (W * ys)
-    lmul!(2, cs)      # cs = 2 * A' * (W * ys)
+    mul!(cs_lin, A', zs)  # cs = A' * (W * ys)
+    lmul!(2, cs_lin)      # cs = 2 * A' * (W * ys)
 
-    ldiv!(F, cs)      # solve linear system
+    ldiv!(F, cs_lin)      # solve linear system
 
     Spline(R, cs)
 end
@@ -120,9 +129,9 @@ function fit(
     λ ≥ 0 || throw(DomainError(λ, "the smoothing parameter λ must be non-negative"))
     eachindex(xs) == eachindex(ys) || throw(DimensionMismatch("x and y vectors must have the same length"))
     N = length(xs)
-    cs = similar(xs)
+    cs = similar(ys)
 
-    T = eltype(cs)
+    T = eltype(xs)
     ts_in = make_knots(xs, order_in, bc)
     R = PeriodicBSplineBasis(order_in, ts_in)
     k = order(R)  # = 4
@@ -170,15 +179,25 @@ function fit(
     F = cholesky(M)      # factorise matrix (assuming posdef)
 
     # Construct RHS trying to reduce allocations
-    zs = copy(ys)
+    Y = eltype(ys)
+    if ndims(Y) == 0  # scalar data
+        cs_lin = cs
+        zs = copy(ys)
+    else  # multidimensional data
+        Z = eltype(eltype(ys))  # for example, if eltype(ys) <: SVector{D,Z}
+        @assert Z <: Number
+        cs_lin = reinterpret(reshape, Z, cs)'
+        zs = copy(reinterpret(reshape, Z, ys)')  # dimensions (N, D)
+    end
     if weights !== nothing
         eachindex(weights) == eachindex(xs) || throw(DimensionMismatch("the `weights` vector must have the same length as the data"))
         lmul!(Diagonal(weights), zs)  # zs = W * ys
     end
-    mul!(cs, A', zs)  # cs = A' * (W * ys)
-    lmul!(2, cs)      # cs = 2 * A' * (W * ys)
+    @show summary(cs_lin) summary(A) summary(zs)
+    mul!(cs_lin, A', zs)  # cs = A' * (W * ys)
+    lmul!(2, cs_lin)      # cs = 2 * A' * (W * ys)
 
-    cs .= F \ cs  # solve linear system (allocates intermediate array)
+    cs_lin .= F \ cs_lin  # solve linear system (allocates intermediate array)
 
     Spline(R, cs)
 end

test/smoothing.jl

@@ -2,6 +2,7 @@
 
 using BSplineKit
 using QuadGK: quadgk
+using StaticArrays
 using ReverseDiff
 using Test
 
@@ -49,7 +50,8 @@ function smoothing_objective(cs, R::AbstractBSplineBasis, xs, ys; weights = noth
     loss
 end
 
-function check_zero_gradient(S::Spline, xs, ys; weights = nothing, λ)
+# Scalar data
+function _check_zero_gradient(::Type{T}, S::Spline, xs, ys; weights = nothing, λ, rtol = 1e-12) where {T <: Real}
     R = basis(S)  # usually a RecombinedBSplineBasis
     cs = parent(coefficients(S))  # `parent` is useful if this is a PeriodicBSplineBasis
 
@@ -67,11 +69,27 @@ function check_zero_gradient(S::Spline, xs, ys; weights = nothing, λ)
     # f ~ Y² and therefore ∂f/∂cⱼ ~ Y. So we compare it with the sum of |y_i|².
     reference = sum(abs2, ys)
     err = sum(abs2, ∇f)
-    @test err / reference < 1e-12
+    @test err / reference < rtol
 
     nothing
 end
 
+# Vector data (e.g. parametric splines)
+# Currently all components are separately smoothed, so we verify them as separate scalar functions.
+function _check_zero_gradient(::Type{SVector{N, T}}, S::Spline, xs, ys; kws...) where {N, T}
+    R = basis(S)
+    cs = coefficients(S)
+    for i in 1:N
+        Si = Spline(R, getindex.(cs, i))
+        check_zero_gradient(Si, xs, getindex.(ys, i); kws...)
+    end
+    nothing
+end
+
+function check_zero_gradient(S::Spline, xs, ys; kws...)
+    _check_zero_gradient(eltype(ys), S, xs, ys; kws...)
+end
+
 # Returns the integral of |S''(x)| (the "curvature") over the whole spline.
 function total_curvature(S::Spline)
     ts = knots(S)
@@ -148,4 +166,27 @@ end
         Sw = fit(xs, ys, λ; weights)
         check_zero_gradient(Sw, xs, ys; λ, weights)
     end
+
+    @testset "Parametric" begin
+        λ = 1e-2
+        N = 100
+        ts = range(0, 2π; length = N + 1)[1:N]
+        vs = [0.1 * SVector(cos(t), sin(t)) .+ 0.01 * sin(10 * t) for t in ts]
+        S_nat = @inferred fit(ts, vs, λ, Natural())
+        S_per = @inferred fit(ts, vs, λ, Periodic(2π))
+
+        @testset "Natural" check_zero_gradient(S_nat, ts, vs; λ)
+        @testset "Periodic" check_zero_gradient(S_per, ts, vs; λ)
+
+        @testset "With weights" begin
+            weights = fill!(similar(ts), 1)
+            weights[3] = 1000
+
+            S_nat = @inferred fit(ts, vs, λ, Natural(); weights)
+            S_per = @inferred fit(ts, vs, λ, Periodic(2π); weights)
+
+            @testset "Natural" check_zero_gradient(S_nat, ts, vs; λ, weights)
+            @testset "Periodic" check_zero_gradient(S_per, ts, vs; λ, weights)
+        end
+    end
 end