diff --git a/Project.toml b/Project.toml index d716eed..9bcf699 100644 --- a/Project.toml +++ b/Project.toml @@ -8,9 +8,11 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" MultivariatePolynomials = "102ac46a-7ee4-5c85-9060-abc95bfdeaa3" MutableArithmetics = "d8a4904e-b15c-11e9-3269-09a3773c0cb0" StarAlgebras = "0c0c59c1-dc5f-42e9-9a8b-b5dc384a6cd1" +TrigPolys = "bbdedc48-cb31-4a37-9fe3-b015aecc8dd3" [compat] MultivariatePolynomials = "0.5.16" MutableArithmetics = "1" StarAlgebras = "0.3" +TrigPolys = "0.2" julia = "1.10" diff --git a/docs/src/index.md b/docs/src/index.md index e3d1762..a4225a2 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -46,6 +46,7 @@ Legendre ChebyshevFirstKind ChebyshevSecondKind Trigonometric +TrigEvalMatrix ``` ## Additional basis diff --git a/src/lagrange.jl b/src/lagrange.jl index b756928..cb4abc9 100644 --- a/src/lagrange.jl +++ b/src/lagrange.jl @@ -133,6 +133,16 @@ function transformation_to(basis::SubBasis, lag::LagrangeBasis{T}) where {T} return V end +function transformation_to( + basis::SubBasis{Trigonometric}, + lag::LagrangeBasis{T}, +) where {T} + # FFT-friendly lazy matrix. The hot-path `*`/`mul!` use TrigPolys.evaluate / + # evaluateT in `O(d log d)` instead of the dense `O(n d)` Vandermonde path + # above. + return TrigEvalMatrix{T}(lag.points, length(basis)) +end + # Heuristic taken from Hypatia function num_samples(sample_factor, dim) if iszero(sample_factor) diff --git a/src/trigonometric.jl b/src/trigonometric.jl index 8624f00..9667579 100644 --- a/src/trigonometric.jl +++ b/src/trigonometric.jl @@ -94,3 +94,205 @@ end # FIXME The cos part is, like Chebysev, maybe the sin part too ? We should do better here, this is just a stopgap even_odd_separated(::Type{Trigonometric}) = false + +import TrigPolys + +""" + TrigEvalMatrix{T,P<:AbstractVector} <: AbstractMatrix{T} + +Change-of-basis matrix from a univariate `SubBasis{Trigonometric}` (rows +indexing samples / columns indexing the canonical trig basis +`[1, cos(ω), sin(ω), cos(2ω), sin(2ω), ...]` in `MultivariateBases`' +interleaved order) to a `LagrangeBasis` whose nodes are `points`. + +`mul!`/`Base.*` use `TrigPolys.evaluate` / `TrigPolys.evaluateT`, i.e. FFTs +on the canonical TrigPolys uniform grid. + +TODO: this currently ignores `points` for the fast path — we assume they +are the `2n+1` equispaced grid points TrigPolys uses internally. Calling +`mul!` on points that do not match that grid will give silently wrong +results. A proper implementation should detect mismatch and fall back to +the dense Vandermonde matrix built via `recurrence_eval`. +""" +struct TrigEvalMatrix{T,P<:AbstractVector} <: AbstractMatrix{T} + points::P + n_coef::Int + # Dense materialization computed at construction. Building it costs + # `O(n d)` (Chebyshev-style recurrence per row), the same total work as + # one BLAS matmul would do on a stored matrix. Holding it makes per-element + # `getindex` `O(1)` — required so the downstream `LRO.Factorization`'s + # `dot` chain (which iterates `factor[k]` for `factor::SubArray{T,1,<:TrigEvalMatrix}`) + # stays cheap. `mul!` on `AbstractVector{<:Real}` / `AbstractMatrix` columns + # still routes through `TrigPolys.evaluate`/`evaluateT`, so when downstream + # code dispatches at the parent-matrix level (batched FFT in BM, see TODO + # above) the FFT path remains live. + dense::Matrix{T} +end + +function _materialize_trig_eval( + ::Type{T}, + points::AbstractVector, + n_coef::Integer, +) where {T} + n_pts = length(points) + out = Matrix{T}(undef, n_pts, n_coef) + for i in 1:n_pts + # `LagrangeBasis` stores each point as a 1-element `AbstractVector` + # (univariate); the basis-recurrence expects the scalar value. + val = T(only(points[i])) + if n_coef >= 1 + out[i, 1] = one(T) + end + if n_coef >= 2 + out[i, 2] = degree_one_univariate_polynomial(Trigonometric, val) + end + for d in 2:(n_coef-1) + out[i, d+1] = + recurrence_eval(Trigonometric, view(out, i, 1:d), val, d) + end + end + return out +end + +function TrigEvalMatrix{T}( + points::P, + n_coef::Integer, +) where {T,P<:AbstractVector} + return TrigEvalMatrix{T,P}( + points, + n_coef, + _materialize_trig_eval(T, points, n_coef), + ) +end + +Base.size(M::TrigEvalMatrix) = (length(M.points), M.n_coef) + +# Forward the strided-array interface to the underlying dense storage so that +# `view(M, j, :)` is a strided `SubArray` that BLAS' `gemv!` accepts. Without +# this, downstream `LinearAlgebra.dot(::Factorization{T,<:SubArray{T,1,<:TrigEvalMatrix}}, ...)` +# in `LowRankOpt.BurerMonteiro` would fall back to the generic Julia loop with +# allocations per dot — orders of magnitude slower than BLAS. +Base.IndexStyle(::Type{<:TrigEvalMatrix}) = IndexLinear() +Base.@propagate_inbounds Base.getindex(M::TrigEvalMatrix, i::Int) = M.dense[i] +Base.@propagate_inbounds Base.getindex( + M::TrigEvalMatrix, + i::Integer, + j::Integer, +) = M.dense[i, j] +Base.strides(M::TrigEvalMatrix) = strides(M.dense) +function Base.unsafe_convert(::Type{Ptr{T}}, M::TrigEvalMatrix{T}) where {T} + return Base.unsafe_convert(Ptr{T}, M.dense) +end +Base.elsize(::Type{<:TrigEvalMatrix{T}}) where {T} = sizeof(T) + +# Reorder a coefficient vector laid out in MultivariateBases' interleaved order +# `[a0, c1, s1, c2, s2, ...]` into TrigPolys' layout `[a0, c1, ..., cn, s1, ..., sn]`. +function _trig_coefs_to_trigpoly!(out::AbstractVector, x::AbstractVector) + n_x = length(x) + n_out = length(out) + @assert isodd(n_x) + @assert isodd(n_out) + d = div(n_x - 1, 2) + n = div(n_out - 1, 2) + @assert n >= d + out[1] = x[1] + @inbounds for k in 1:d + out[1+k] = x[2k] + out[1+n+k] = x[2k+1] + end + return out +end + +function _trigpoly_to_trig_coefs!(out::AbstractVector, a::AbstractVector) + n_out = length(out) + n_a = length(a) + @assert isodd(n_out) + @assert isodd(n_a) + d = div(n_out - 1, 2) + n = div(n_a - 1, 2) + @assert n >= d + out[1] = a[1] + @inbounds for k in 1:d + out[2k] = a[1+k] + out[2k+1] = a[1+n+k] + end + return out +end + +function _trigpoly_for_eval(M::TrigEvalMatrix{T}, x::AbstractVector) where {T} + n_coef = M.n_coef + @assert length(x) == n_coef + @assert isodd(n_coef) + grid_d = div(length(M.points) - 1, 2) + n = max(div(n_coef - 1, 2), grid_d) + a = zeros(T, 2n + 1) + # `a` is laid out as `[a0, c1..cn, s1..sn]`; padding for the high + # frequencies stays at zero (mirroring `vectorized_pad_to`). + _trig_coefs_to_trigpoly!(a, x) + return TrigPolys.TrigPoly(a) +end + +function LinearAlgebra.mul!( + y::AbstractVector{<:Real}, + M::TrigEvalMatrix{T}, + x::AbstractVector{<:Real}, +) where {T} + p = _trigpoly_for_eval(M, x) + copyto!(y, TrigPolys.evaluate(p)) + return y +end + +function Base.:*(M::TrigEvalMatrix{T}, x::AbstractVector{<:Real}) where {T} + y = Vector{T}(undef, size(M, 1)) + return LinearAlgebra.mul!(y, M, x) +end + +# Fallback when `x` carries symbolic / MOI-style entries (e.g. during +# `SA.coeffs(::AlgebraElement{<:MOI.ScalarAffineFunction}, ...)` in the SOS +# constraint bridge). The FFT pad/unpad path needs numeric storage, so we +# materialize a dense matrix and hand off to the generic `*`. +function Base.:*(M::TrigEvalMatrix, x::AbstractVector) + return Matrix(M) * x +end + +function LinearAlgebra.mul!( + Y::AbstractMatrix, + M::TrigEvalMatrix{T}, + X::AbstractMatrix, +) where {T} + @assert size(Y, 1) == size(M, 1) + @assert size(M, 2) == size(X, 1) + for j in axes(X, 2) + @views LinearAlgebra.mul!(Y[:, j], M, X[:, j]) + end + return Y +end + +# Adjoint operator: `evaluateT` is the adjoint of `evaluate`. +function LinearAlgebra.mul!( + x::AbstractVector, + Madj::LinearAlgebra.Adjoint{<:Any,<:TrigEvalMatrix{T}}, + y::AbstractVector, +) where {T} + M = parent(Madj) + @assert length(x) == M.n_coef + @assert length(y) == size(M, 1) + # `a` is laid out as `[a0, c1..cn, s1..sn]`; `_trigpoly_to_trig_coefs!` + # unpads it into MB's interleaved layout. + _trigpoly_to_trig_coefs!(x, TrigPolys.evaluateT(y)) + return x +end + +function Base.:*( + Madj::LinearAlgebra.Adjoint{<:Any,<:TrigEvalMatrix{T}}, + y::AbstractVector, +) where {T} + x = Vector{T}(undef, size(Madj, 1)) + return LinearAlgebra.mul!(x, Madj, y) +end + +# Override `transformation_to(::SubBasis{Trigonometric}, ::LagrangeBasis)` so +# that the bridge layer carries an FFT-backed `AbstractMatrix` instead of the +# dense Vandermonde matrix built by the generic `transformation_to` in +# `lagrange.jl`. The new lazy matrix multiplies via TrigPolys' FFT, so the +# downstream `LowRankOpt.BurerMonteiro` `mul!` calls run in `O(d log d)`. diff --git a/test/Project.toml b/test/Project.toml index a3623d8..fc6bf93 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -7,6 +7,7 @@ MutableArithmetics = "d8a4904e-b15c-11e9-3269-09a3773c0cb0" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" StarAlgebras = "0c0c59c1-dc5f-42e9-9a8b-b5dc384a6cd1" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" +TrigPolys = "bbdedc48-cb31-4a37-9fe3-b015aecc8dd3" TypedPolynomials = "afbbf031-7a57-5f58-a1b9-b774a0fad08d" [compat] diff --git a/test/trigeval.jl b/test/trigeval.jl new file mode 100644 index 0000000..43132b7 --- /dev/null +++ b/test/trigeval.jl @@ -0,0 +1,242 @@ +module TestTrigEval + +using Test +using LinearAlgebra +using DynamicPolynomials +import StarAlgebras as SA +import MultivariateBases as MB +import TrigPolys + +# A `2n+1`-point grid that matches `TrigPolys.evaluate`'s canonical layout: +# `θ_k = 2π k / (2n+1)` for `k = 0, …, 2n`. The univariate `Trigonometric` +# basis is parameterised by `value = cos θ`, so we wrap each `cos θ_k` in a +# 1-element vector to mimic how `LagrangeBasis` stores multivariate points. +_grid_points(n::Int) = [Float64[cos(2π * k / (2n + 1))] for k in 0:(2n)] + +_grid_thetas(n::Int) = [2π * k / (2n + 1) for k in 0:(2n)] + +# Hand-built reference: the trig polynomial +# `p(θ) = a₀ + Σₖ (c_k cos(kθ) + s_k sin(kθ))` evaluated at each grid `θ_k`. +# Coefficients are in `MultivariateBases`' interleaved order +# `[a₀, c₁, s₁, c₂, s₂, …, c_n, s_n]`. +function _eval_reference(coefs::AbstractVector, thetas::AbstractVector) + n_coef = length(coefs) + @assert isodd(n_coef) + d = div(n_coef - 1, 2) + out = zeros(eltype(coefs), length(thetas)) + for (i, θ) in pairs(thetas) + out[i] = coefs[1] + for k in 1:d + out[i] += coefs[2k] * cos(k * θ) + coefs[2k+1] * sin(k * θ) + end + end + return out +end + +function test_construction_and_size() + n = 2 + n_coef = 2n + 1 + pts = _grid_points(n) + M = MB.TrigEvalMatrix{Float64}(pts, n_coef) + @test M isa MB.TrigEvalMatrix{Float64} + @test eltype(M) === Float64 + @test size(M) == (length(pts), n_coef) + @test size(M, 1) == length(pts) + @test size(M, 2) == n_coef +end + +function test_dense_materialization() + # The `.dense` field is built at construction; `M[i, j]` and `Matrix(M)` + # must forward to it, and a manual recurrence on the points must give + # the same answer. + n = 3 + n_coef = 2n + 1 + pts = _grid_points(n) + M = MB.TrigEvalMatrix{Float64}(pts, n_coef) + den = Matrix(M) + @test den == M.dense + @test size(den) == size(M) + for i in 1:size(M, 1), j in 1:size(M, 2) + @test M[i, j] == den[i, j] + end + # Constant column is 1; first non-constant column equals each point value + # (degree-1 univariate polynomial is `value`). + @test all(M[:, 1] .== 1.0) + @test M[:, 2] == [only(p) for p in pts] +end + +function test_index_style_and_strides() + # Strided forwarding is required so `view(M, j, :)` plus + # `LinearAlgebra.dot` dispatches to BLAS rather than the generic loop. + n = 4 + n_coef = 2n + 1 + pts = _grid_points(n) + M = MB.TrigEvalMatrix{Float64}(pts, n_coef) + @test IndexStyle(M) == IndexLinear() + @test IndexStyle(typeof(M)) == IndexLinear() + @test strides(M) == strides(M.dense) + @test Base.elsize(typeof(M)) == sizeof(Float64) + # `unsafe_convert` must point at the same memory as the underlying dense. + @test Base.unsafe_convert(Ptr{Float64}, M) === + Base.unsafe_convert(Ptr{Float64}, M.dense) + # Linear indexing reaches the same elements as cartesian. + @test M[1] == M[1, 1] + @test M[length(pts)+1] == M[1, 2] +end + +function test_fft_mul_matches_reference() + # The FFT path (`M * x` for `x::AbstractVector{<:Real}`) routes through + # `TrigPolys.evaluate` after reordering MB's interleaved layout into + # TrigPolys' separated layout. On the canonical grid it must match the + # hand-evaluated trig polynomial. + n = 5 + n_coef = 2n + 1 + pts = _grid_points(n) + thetas = _grid_thetas(n) + M = MB.TrigEvalMatrix{Float64}(pts, n_coef) + coefs = [0.7, -0.4, 0.2, 1.1, -0.3, 0.05, 0.9, -1.2, 0.6, 0.1, -0.5] + @assert length(coefs) == n_coef + expected = _eval_reference(coefs, thetas) + y_alloc = M * coefs + @test y_alloc ≈ expected + y_inplace = similar(y_alloc) + LinearAlgebra.mul!(y_inplace, M, coefs) + @test y_inplace ≈ expected +end + +function test_fft_mul_matrix_rhs() + # Matrix-RHS `mul!` iterates columns; each column must match `M * col`. + n = 4 + n_coef = 2n + 1 + pts = _grid_points(n) + M = MB.TrigEvalMatrix{Float64}(pts, n_coef) + k = 3 + X = randn(n_coef, k) + Y = Matrix{Float64}(undef, size(M, 1), k) + LinearAlgebra.mul!(Y, M, X) + for j in 1:k + @test Y[:, j] ≈ M * X[:, j] + end +end + +function test_adjoint_mul_is_real_adjoint() + # `mul!(x, M', y)` routes through `TrigPolys.evaluateT`. We don't have a + # closed form for the unpadded adjoint, so verify the inner-product + # identity `⟨Mx, y⟩ == ⟨x, M' y⟩` directly. + n = 4 + n_coef = 2n + 1 + pts = _grid_points(n) + M = MB.TrigEvalMatrix{Float64}(pts, n_coef) + x = randn(n_coef) + y = randn(size(M, 1)) + Mx = M * x + Mty = M' * y + @test Mty isa Vector{Float64} + @test length(Mty) == n_coef + @test dot(Mx, y) ≈ dot(x, Mty) + # In-place adjoint + Mty_ip = similar(Mty) + LinearAlgebra.mul!(Mty_ip, M', y) + @test Mty_ip ≈ Mty +end + +# A minimal symbolic-coefficient type. Mirrors what MOI's +# `ScalarAffineFunction` looks like to `Matrix(M) * x`: an opaque element +# closed under `Float64 * _` and `_ + _`. The FFT path is gated on +# `<:Real`, so this exercises the dense fallback. +struct Symb + label::Symbol +end +Base.zero(::Type{Symb}) = Symb(:zero) +Base.zero(x::Symb) = zero(Symb) +Base.iszero(::Symb) = false +Base.:+(::Symb, ::Symb) = Symb(:sum) +Base.:*(::Float64, ::Symb) = Symb(:scaled) +Base.:*(::Symb, ::Float64) = Symb(:scaled) +Base.promote_rule(::Type{Symb}, ::Type{Float64}) = Symb + +function test_symbolic_fallback_uses_dense() + # When the RHS isn't `<:Real`, the generic `*` overload should + # materialize the dense matrix and hand off to `Matrix * x`. This is + # what the SumOfSquares SOS bridge relies on when the constraint + # function coefficients are `MOI.ScalarAffineFunction`. + n = 2 + n_coef = 2n + 1 + pts = _grid_points(n) + M = MB.TrigEvalMatrix{Float64}(pts, n_coef) + x = Symb[Symb(Symbol("c$i")) for i in 1:n_coef] + y = M * x + @test y isa AbstractVector{Symb} + @test length(y) == size(M, 1) + # Each output is a sum (across the row) of scaled symbolic terms. + @test all(yi -> yi isa Symb, y) +end + +function test_transformation_to_dispatch() + # The new `transformation_to(::SubBasis{Trigonometric}, ::LagrangeBasis)` + # method should fire and return a `TrigEvalMatrix` — NOT the generic + # dense `Matrix{T}` built by the fallback at `lagrange.jl:119`. + @polyvar t + n = 3 + monos = monomials(t, 0:(2n)) + sub = MB.SubBasis{MB.Trigonometric}(monos) + pts = _grid_points(n) + lag = MB.LagrangeBasis([t], pts) + U = MB.transformation_to(sub, lag) + @test U isa MB.TrigEvalMatrix + @test size(U) == (length(lag), length(sub)) +end + +function test_transformation_to_promote_op() + # `Base.promote_op(MB.transformation_to, G, B)` must resolve the matrix + # type at compile time so SumOfSquares' `LowRankBridge`'s + # `added_constrained_variable_types` can declare a concrete factor type. + @polyvar t + n = 2 + monos = monomials(t, 0:(2n)) + sub = MB.SubBasis{MB.Trigonometric}(monos) + pts = _grid_points(n) + lag = MB.LagrangeBasis([t], pts) + MT = Base.promote_op(MB.transformation_to, typeof(sub), typeof(lag)) + @test MT <: MB.TrigEvalMatrix{Float64} + @test MT !== Any + # Same query for a non-Trig basis should drop to the generic + # `Matrix{Float64}` builder. + sub_mono = MB.SubBasis{MB.Monomial}(monos) + MT_mono = + Base.promote_op(MB.transformation_to, typeof(sub_mono), typeof(lag)) + @test MT_mono <: AbstractMatrix{Float64} + @test !(MT_mono <: MB.TrigEvalMatrix) +end + +function test_row_view_is_strided_blas_path() + # `view(M, j, :)` should be a strided `SubArray`. The whole point of + # `IndexStyle = IndexLinear` + `strides` + `unsafe_convert` is that + # downstream `dot(::SubArray, ::AbstractVector)` reaches BLAS rather than + # falling back to the generic Julia loop. + n = 4 + n_coef = 2n + 1 + pts = _grid_points(n) + M = MB.TrigEvalMatrix{Float64}(pts, n_coef) + row = view(M, 2, :) + @test row isa SubArray + @test parent(row) === M + @test Vector(row) == M.dense[2, :] + # Inner product equals the dense reference. + v = randn(n_coef) + @test dot(row, v) ≈ dot(M.dense[2, :], v) +end + +function runtests() + for name in names(@__MODULE__; all = true) + if startswith("$name", "test_") + @testset "$(name)" begin + getfield(@__MODULE__, name)() + end + end + end +end + +end + +TestTrigEval.runtests()