Skip to content

Move out recurrence code #203

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 8 commits into from
Oct 24, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion Project.toml
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
name = "ClassicalOrthogonalPolynomials"
uuid = "b30e2e7b-c4ee-47da-9d5f-2c5c27239acd"
authors = ["Sheehan Olver <[email protected]>"]
version = "0.13.6"
version = "0.13.7"

[deps]
ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
Expand All @@ -22,6 +22,8 @@ LazyArrays = "5078a376-72f3-5289-bfd5-ec5146d43c02"
LazyBandedMatrices = "d7e5e226-e90b-4449-9968-0f923699bf6f"
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
QuasiArrays = "c4ea9172-b204-11e9-377d-29865faadc5c"
RecurrenceRelationships = "807425ed-42ea-44d6-a357-6771516d7b2c"
RecurrenceRelationshipArrays = "b889d2dc-af3c-4820-88a8-238fa91d3518"
SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"

[compat]
Expand All @@ -42,6 +44,8 @@ IntervalSets = "0.7"
LazyArrays = "2.2"
LazyBandedMatrices = "0.10"
QuasiArrays = "0.11"
RecurrenceRelationships = "0.1"
RecurrenceRelationshipArrays = "0.1"
SpecialFunctions = "1.0, 2"
julia = "1.10"

Expand Down
3 changes: 0 additions & 3 deletions docs/src/index.md
Original file line number Diff line number Diff line change
Expand Up @@ -180,9 +180,6 @@ ClassicalOrthogonalPolynomials.normalizationconstant
ClassicalOrthogonalPolynomials.OrthogonalPolynomialRatio
```
```@docs
ClassicalOrthogonalPolynomials.Clenshaw
```
```@docs
ClassicalOrthogonalPolynomials.singularities
```
```@docs
Expand Down
11 changes: 6 additions & 5 deletions src/ClassicalOrthogonalPolynomials.jl
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ using IntervalSets: UnitRange
using ContinuumArrays, QuasiArrays, LazyArrays, FillArrays, BandedMatrices, BlockArrays,
IntervalSets, DomainSets, ArrayLayouts, SpecialFunctions,
InfiniteLinearAlgebra, InfiniteArrays, LinearAlgebra, FastGaussQuadrature, FastTransforms, FFTW,
LazyBandedMatrices, HypergeometricFunctions
LazyBandedMatrices, HypergeometricFunctions, RecurrenceRelationships

import Base: @_inline_meta, axes, getindex, unsafe_getindex, convert, prod, *, /, \, +, -,
IndexStyle, IndexLinear, ==, OneTo, tail, similar, copyto!, copy, setindex,
Expand All @@ -18,7 +18,7 @@ import LazyArrays: MemoryLayout, Applied, ApplyStyle, flatten, _flatten, adjoint
_mul_arguments, CachedVector, CachedMatrix, LazyVector, LazyMatrix, axpy!, AbstractLazyLayout, BroadcastLayout,
AbstractCachedVector, AbstractCachedMatrix, paddeddata, cache_filldata!,
simplifiable, PaddedArray, converteltype, simplify
import ArrayLayouts: MatMulVecAdd, materialize!, _fill_lmul!, sublayout, sub_materialize, lmul!, ldiv!, ldiv, transposelayout, triangulardata,
import ArrayLayouts: MatMulVecAdd, materialize!, sublayout, sub_materialize, lmul!, ldiv!, ldiv, transposelayout, triangulardata,
subdiagonaldata, diagonaldata, supdiagonaldata, mul, rowsupport, colsupport
import LazyBandedMatrices: SymTridiagonal, Bidiagonal, Tridiagonal, unitblocks, BlockRange1, AbstractLazyBandedLayout
import LinearAlgebra: pinv, factorize, qr, adjoint, transpose, dot, mul!, reflectorApply!
Expand All @@ -39,9 +39,10 @@ import ContinuumArrays: Basis, Weight, basis_axes, @simplify, Identity, Abstract
AffineQuasiVector, AffineMap, AbstractWeightLayout, AbstractWeightedBasisLayout, WeightedBasisLayout, WeightedBasisLayouts, demap, AbstractBasisLayout, BasisLayout,
checkpoints, weight, unweighted, MappedBasisLayouts, sum_layout, invmap, plan_ldiv, layout_broadcasted, MappedBasisLayout, SubBasisLayout, broadcastbasis_layout,
plan_grid_transform, plan_transform, MAX_PLOT_POINTS, MulPlan, grammatrix, AdjointBasisLayout, grammatrix_layout, plan_transform_layout, _cumsum
import FastTransforms: Λ, forwardrecurrence, forwardrecurrence!, _forwardrecurrence!, clenshaw, clenshaw!,
_forwardrecurrence_next, _clenshaw_next, check_clenshaw_recurrences, ChebyshevGrid, chebyshevpoints, Plan, ScaledPlan, th_cheb2leg

import FastTransforms: Λ, ChebyshevGrid, chebyshevpoints, Plan, ScaledPlan, th_cheb2leg
import RecurrenceRelationships: forwardrecurrence, forwardrecurrence!, clenshaw, clenshaw!,
check_clenshaw_recurrences
import RecurrenceRelationshipArrays: initiateforwardrecurrence, Clenshaw
import FastGaussQuadrature: jacobimoment

import BlockArrays: blockedrange, _BlockedUnitRange, unblock, _BlockArray, block, blockindex, BlockSlice, blockvec
Expand Down
247 changes: 7 additions & 240 deletions src/clenshaw.jl
Original file line number Diff line number Diff line change
Expand Up @@ -2,16 +2,6 @@
# Assume 1 normalization
_p0(A) = one(eltype(A))

function initiateforwardrecurrence(N, A, B, C, x, μ)
T = promote_type(eltype(A), eltype(B), eltype(C), typeof(x))
p0 = convert(T, μ)
N == 0 && return zero(T), p0
p1 = convert(T, muladd(A[1],x,B[1])*p0)
@inbounds for n = 2:N
p1,p0 = _forwardrecurrence_next(n, A, B, C, x, p0, p1),p1
end
p0,p1
end

for (get, vie) in ((:getindex, :view), (:(Base.unsafe_getindex), :(Base.unsafe_view)))
@eval begin
Expand Down Expand Up @@ -39,7 +29,7 @@ function forwardrecurrence_copyto!(dest::AbstractMatrix, V)
Ã,B̃,C̃ = A[shift:∞],B[shift:∞],C[shift:∞]
for (k,x) = enumerate(xr)
p0, p1 = initiateforwardrecurrence(shift, A, B, C, x, _p0(P))
_forwardrecurrence!(view(dest,k,:), Ã, B̃, C̃, x, p0, p1)
forwardrecurrence!(view(dest,k,:), Ã, B̃, C̃, x, p0, p1)
end
dest
end
Expand All @@ -54,7 +44,7 @@ function copyto!(dest::AbstractVector, V::SubArray{<:Any,1,<:OrthogonalPolynomia
shift = first(jr)
Ã,B̃,C̃ = A[shift:∞],B[shift:∞],C[shift:∞]
p0, p1 = initiateforwardrecurrence(shift, A, B, C, x, _p0(P))
_forwardrecurrence!(dest, Ã, B̃, C̃, x, p0, p1)
forwardrecurrence!(dest, Ã, B̃, C̃, x, p0, p1)
dest
end

Expand Down Expand Up @@ -109,228 +99,6 @@ end
Base.@propagate_inbounds getindex(f::Mul{<:WeightedOPLayout,<:AbstractPaddedLayout}, x::Number, j...) =
weight(f.A)[x] * (unweighted(f.A) * f.B)[x, j...]

###
# Operator clenshaw
###


Base.@propagate_inbounds function _clenshaw_next!(n, A::AbstractFill, ::Zeros, C::Ones, x::AbstractMatrix, c, bn1::AbstractMatrix{T}, bn2::AbstractMatrix{T}) where T
muladd!(getindex_value(A), x, bn1, -one(T), bn2)
view(bn2,band(0)) .+= c[n]
bn2
end

Base.@propagate_inbounds function _clenshaw_next!(n, A::AbstractVector, ::Zeros, C::AbstractVector, x::AbstractMatrix, c, bn1::AbstractMatrix{T}, bn2::AbstractMatrix{T}) where T
muladd!(A[n], x, bn1, -C[n+1], bn2)
view(bn2,band(0)) .+= c[n]
bn2
end

Base.@propagate_inbounds function _clenshaw_next!(n, A::AbstractVector, B::AbstractVector, C::AbstractVector, x::AbstractMatrix, c, bn1::AbstractMatrix{T}, bn2::AbstractMatrix{T}) where T
# bn2 .= B[n] .* bn1 .- C[n+1] .* bn2
lmul!(-C[n+1], bn2)
LinearAlgebra.axpy!(B[n], bn1, bn2)
muladd!(A[n], x, bn1, one(T), bn2)
view(bn2,band(0)) .+= c[n]
bn2
end

# Operator * f Clenshaw
Base.@propagate_inbounds function _clenshaw_next!(n, A::AbstractFill, ::Zeros, C::Ones, X::AbstractMatrix, c, f::AbstractVector, bn1::AbstractVector{T}, bn2::AbstractVector{T}) where T
muladd!(getindex_value(A), X, bn1, -one(T), bn2)
bn2 .+= c[n] .* f
bn2
end

Base.@propagate_inbounds function _clenshaw_next!(n, A, ::Zeros, C, X::AbstractMatrix, c, f::AbstractVector, bn1::AbstractVector{T}, bn2::AbstractVector{T}) where T
muladd!(A[n], X, bn1, -C[n+1], bn2)
bn2 .+= c[n] .* f
bn2
end

Base.@propagate_inbounds function _clenshaw_next!(n, A, B, C, X::AbstractMatrix, c, f::AbstractVector, bn1::AbstractVector{T}, bn2::AbstractVector{T}) where T
bn2 .= B[n] .* bn1 .- C[n+1] .* bn2 .+ c[n] .* f
muladd!(A[n], X, bn1, one(T), bn2)
bn2
end

# allow special casing first arg, for ChebyshevT in ClassicalOrthogonalPolynomials
Base.@propagate_inbounds function _clenshaw_first!(A, ::Zeros, C, X, c, bn1, bn2)
muladd!(A[1], X, bn1, -C[2], bn2)
view(bn2,band(0)) .+= c[1]
bn2
end

Base.@propagate_inbounds function _clenshaw_first!(A, B, C, X, c, bn1, bn2)
lmul!(-C[2], bn2)
LinearAlgebra.axpy!(B[1], bn1, bn2)
muladd!(A[1], X, bn1, one(eltype(bn2)), bn2)
view(bn2,band(0)) .+= c[1]
bn2
end

Base.@propagate_inbounds function _clenshaw_first!(A, ::Zeros, C, X, c, f::AbstractVector, bn1, bn2)
muladd!(A[1], X, bn1, -C[2], bn2)
bn2 .+= c[1] .* f
bn2
end

Base.@propagate_inbounds function _clenshaw_first!(A, B, C, X, c, f::AbstractVector, bn1, bn2)
bn2 .= B[1] .* bn1 .- C[2] .* bn2 .+ c[1] .* f
muladd!(A[1], X, bn1, one(eltype(bn2)), bn2)
bn2
end

_clenshaw_op(::AbstractBandedLayout, Z, N) = BandedMatrix(Z, (N-1,N-1))

function clenshaw(c::AbstractVector, A::AbstractVector, B::AbstractVector, C::AbstractVector, X::AbstractMatrix)
N = length(c)
T = promote_type(eltype(c),eltype(A),eltype(B),eltype(C),eltype(X))
@boundscheck check_clenshaw_recurrences(N, A, B, C)
m = size(X,1)
m == size(X,2) || throw(DimensionMismatch("X must be square"))
N == 0 && return zero(T)
bn2 = _clenshaw_op(MemoryLayout(X), Zeros{T}(m, m), N)
bn1 = _clenshaw_op(MemoryLayout(X), c[N]*Eye{T}(m), N)
_clenshaw_op!(c, A, B, C, X, bn1, bn2)
end

function clenshaw(c::AbstractVector, A::AbstractVector, B::AbstractVector, C::AbstractVector, X::AbstractMatrix, f::AbstractVector)
N = length(c)
T = promote_type(eltype(c),eltype(A),eltype(B),eltype(C),eltype(X))
@boundscheck check_clenshaw_recurrences(N, A, B, C)
m = size(X,1)
m == size(X,2) || throw(DimensionMismatch("X must be square"))
m == length(f) || throw(DimensionMismatch("Dimensions must match"))
N == 0 && return [zero(T)]
bn2 = zeros(T,m)
bn1 = Vector{T}(undef,m)
bn1 .= c[N] .* f
_clenshaw_op!(c, A, B, C, X, f, bn1, bn2)
end

function _clenshaw_op!(c, A, B, C, X, bn1, bn2)
N = length(c)
N == 1 && return bn1
@inbounds begin
for n = N-1:-1:2
bn1,bn2 = _clenshaw_next!(n, A, B, C, X, c, bn1, bn2),bn1
end
bn1 = _clenshaw_first!(A, B, C, X, c, bn1, bn2)
end
bn1
end

function _clenshaw_op!(c, A, B, C, X, f::AbstractVector, bn1, bn2)
N = length(c)
N == 1 && return bn1
@inbounds begin
for n = N-1:-1:2
bn1,bn2 = _clenshaw_next!(n, A, B, C, X, c, f, bn1, bn2),bn1
end
bn1 = _clenshaw_first!(A, B, C, X, c, f, bn1, bn2)
end
bn1
end



"""
Clenshaw(a, X)

represents the operator `a(X)` where a is a polynomial.
Here `a` is to stored as a quasi-vector.
"""
struct Clenshaw{T, Coefs<:AbstractVector, AA<:AbstractVector, BB<:AbstractVector, CC<:AbstractVector, Jac<:AbstractMatrix} <: AbstractBandedMatrix{T}
c::Coefs
A::AA
B::BB
C::CC
X::Jac
p0::T
end

Clenshaw(c::AbstractVector{T}, A::AbstractVector, B::AbstractVector, C::AbstractVector, X::AbstractMatrix{T}, p0::T) where T =
Clenshaw{T,typeof(c),typeof(A),typeof(B),typeof(C),typeof(X)}(c, A, B, C, X, p0)

Clenshaw(c::Number, A, B, C, X, p) = Clenshaw([c], A, B, C, X, p)

function Clenshaw(a::AbstractQuasiVector, X::AbstractQuasiMatrix)
P,c = arguments(a)
Clenshaw(paddeddata(c), recurrencecoefficients(P)..., jacobimatrix(X), _p0(P))
end

copy(M::Clenshaw) = M
size(M::Clenshaw) = size(M.X)
axes(M::Clenshaw) = axes(M.X)
bandwidths(M::Clenshaw) = (length(M.c)-1,length(M.c)-1)

Base.array_summary(io::IO, C::Clenshaw{T}, inds::Tuple{Vararg{OneToInf{Int}}}) where T =
print(io, Base.dims2string(length.(inds)), " Clenshaw{$T} with $(length(C.c)) degree polynomial")

struct ClenshawLayout <: AbstractLazyBandedLayout end
MemoryLayout(::Type{<:Clenshaw}) = ClenshawLayout()
sublayout(::ClenshawLayout, ::Type{<:NTuple{2,AbstractUnitRange{Int}}}) = ClenshawLayout()
sublayout(::ClenshawLayout, ::Type{<:Tuple{AbstractUnitRange{Int},Union{Slice,AbstractInfUnitRange{Int}}}}) = LazyBandedLayout()
sublayout(::ClenshawLayout, ::Type{<:Tuple{Union{Slice,AbstractInfUnitRange{Int}},AbstractUnitRange{Int}}}) = LazyBandedLayout()
sublayout(::ClenshawLayout, ::Type{<:Tuple{Union{Slice,AbstractInfUnitRange{Int}},Union{Slice,AbstractInfUnitRange{Int}}}}) = LazyBandedLayout()
sub_materialize(::ClenshawLayout, V) = BandedMatrix(V)

function _BandedMatrix(::ClenshawLayout, V::SubArray{<:Any,2})
M = parent(V)
kr,jr = parentindices(V)
b = bandwidth(M,1)
jkr = max(1,min(first(jr),first(kr))-b÷2):max(last(jr),last(kr))+b÷2
# relationship between jkr and kr, jr
kr2,jr2 = kr.-first(jkr).+1,jr.-first(jkr).+1
lmul!(M.p0, clenshaw(M.c, M.A, M.B, M.C, M.X[jkr, jkr])[kr2,jr2])
end

function getindex(M::Clenshaw{T}, kr::AbstractUnitRange, j::Integer) where T
b = bandwidth(M,1)
jkr = max(1,min(j,first(kr))-b÷2):max(j,last(kr))+b÷2
# relationship between jkr and kr, jr
kr2,j2 = kr.-first(jkr).+1,j-first(jkr)+1
f = [Zeros{T}(j2-1); one(T); Zeros{T}(length(jkr)-j2)]
lmul!(M.p0, clenshaw(M.c, M.A, M.B, M.C, M.X[jkr, jkr], f)[kr2])
end

getindex(M::Clenshaw, k::Int, j::Int) = M[k:k,j][1]

function getindex(S::Symmetric{T,<:Clenshaw}, k::Integer, jr::AbstractUnitRange) where T
m = max(jr.start,jr.stop,k)
return Symmetric(getindex(S.data,1:m,1:m),Symbol(S.uplo))[k,jr]
end
function getindex(S::Symmetric{T,<:Clenshaw}, kr::AbstractUnitRange, j::Integer) where T
m = max(kr.start,kr.stop,j)
return Symmetric(getindex(S.data,1:m,1:m),Symbol(S.uplo))[kr,j]
end
function getindex(S::Symmetric{T,<:Clenshaw}, kr::AbstractUnitRange, jr::AbstractUnitRange) where T
m = max(kr.start,jr.start,kr.stop,jr.stop)
return Symmetric(getindex(S.data,1:m,1:m),Symbol(S.uplo))[kr,jr]
end

transposelayout(M::ClenshawLayout) = LazyBandedMatrices.LazyBandedLayout()
# TODO: generalise for layout, use Base.PermutedDimsArray
Base.permutedims(M::Clenshaw{<:Number}) = transpose(M)


function materialize!(M::MatMulVecAdd{<:ClenshawLayout,<:AbstractPaddedLayout,<:AbstractPaddedLayout})
α,A,x,β,y = M.α,M.A,M.B,M.β,M.C
length(y) == size(A,1) || throw(DimensionMismatch("Dimensions must match"))
length(x) == size(A,2) || throw(DimensionMismatch("Dimensions must match"))
x̃ = paddeddata(x);
m = length(x̃)
b = bandwidth(A,1)
jkr=1:m+b
p = [x̃; zeros(eltype(x̃),length(jkr)-m)];
Ax = lmul!(A.p0, clenshaw(A.c, A.A, A.B, A.C, A.X[jkr, jkr], p))
_fill_lmul!(β,y)
resizedata!(y, last(jkr))
v = view(paddeddata(y),jkr)
LinearAlgebra.axpy!(α, Ax, v)
y
end

# TODO: generalise this to be trait based
function layout_broadcasted(::Tuple{ExpansionLayout{<:AbstractOPLayout},AbstractOPLayout}, ::typeof(*), a, P)
Expand All @@ -357,9 +125,8 @@ function _broadcasted_layout_broadcasted_mul(::Tuple{AbstractWeightLayout,Polyno
a .* P
end


##
# Banded dot is slow
###

LinearAlgebra.dot(x::AbstractVector, A::Clenshaw, y::AbstractVector) = dot(x, mul(A, y))
# constructor for Clenshaw
function Clenshaw(a::AbstractQuasiVector, X::AbstractQuasiMatrix)
P,c = arguments(a)
Clenshaw(paddeddata(c), recurrencecoefficients(P)..., jacobimatrix(X), _p0(P))
end