Add FFT evaluation using TrigPolys

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"

docs/src/index.md

@@ -46,6 +46,7 @@ Legendre
 ChebyshevFirstKind
 ChebyshevSecondKind
 Trigonometric
+TrigEvalMatrix
 ```
 
 ## Additional basis

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)

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)`.

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]