|
| 1 | +""" |
| 2 | + MutableSparseVectorPolynomial{B,T,X} |
| 3 | +
|
| 4 | +This polynomial type uses an `SparseVector{T,Int}` to store the coefficients of a polynomial relative to the basis `B` with indeterminate `X`. |
| 5 | +The type `T` should have `zero(T)` defined. |
| 6 | +
|
| 7 | +
|
| 8 | +""" |
| 9 | +struct MutableSparseVectorPolynomial{B,T,X} <: AbstractUnivariatePolynomial{B, T,X} |
| 10 | + coeffs::SparseVector{T, Int} |
| 11 | + function MutableSparseVectorPolynomial{B,T,X}(cs::SparseVector{S,Int}, order::Int=0) where {B,T,S,X} |
| 12 | + new{B,T,Symbol(X)}(cs) |
| 13 | + end |
| 14 | +end |
| 15 | + |
| 16 | +MutableSparseVectorPolynomial{B,T,X}(check::Val{:false}, coeffs::SparseVector{Int,S}) where {B,T,S,X} = |
| 17 | + MutableSparseVectorPolynomial{B,T,X}(coeffs) |
| 18 | +MutableSparseVectorPolynomial{B,T,X}(checked::Val{:true}, coeffs::SparseVector{Int,T}) where {B,T,X<:Symbol} = |
| 19 | + MutableSparseVectorPolynomial{B,T,X}(coeffs) |
| 20 | + |
| 21 | +# --- |
| 22 | +function MutableSparseVectorPolynomial{B,T}(coeffs::SparseVector{S,Int}, var::SymbolLike=Var(:x)) where {B,T,S} |
| 23 | + MutableSparseVectorPolynomial{B,T,Symbol(var)}(coeffs) |
| 24 | +end |
| 25 | + |
| 26 | +function MutableSparseVectorPolynomial{B}(cs::SparseVector{T,Int}, var::SymbolLike=Var(:x)) where {B,T} |
| 27 | + MutableSparseVectorPolynomial{B,T,Symbol(var)}(cs) |
| 28 | +end |
| 29 | + |
| 30 | +# From a Dictionary |
| 31 | +function MutableSparseVectorPolynomial{B,X}(cs::AbstractDict{Int, T}) where {B,T,X} |
| 32 | + N = maximum(keys(cs)) + 1 |
| 33 | + v = SparseVector(N, 1 .+ keys(cs), collect(values(cs))) |
| 34 | + MutableSparseVectorPolynomial{B,T,X}(v) |
| 35 | +end |
| 36 | + |
| 37 | +function MutableSparseVectorPolynomial{B}(cs::AbstractDict{Int, T}, var::SymbolLike=Var(:x)) where {B,T} |
| 38 | + MutableSparseVectorPolynomial{B,Symbol(var)}(cs) |
| 39 | +end |
| 40 | + |
| 41 | + |
| 42 | +# abstract vector has order/symbol |
| 43 | +function MutableSparseVectorPolynomial{B,T,X}(coeffs::AbstractVector{S}, order::Int=0) where {B,T,S,X} |
| 44 | + if Base.has_offset_axes(coeffs) |
| 45 | + @warn "ignoring the axis offset of the coefficient vector" |
| 46 | + coeffs = parent(coeffs) |
| 47 | + end |
| 48 | + |
| 49 | + MutableSparseVectorPolynomial{B,T,X}(convert(SparseVector, coeffs)) |
| 50 | +end |
| 51 | + |
| 52 | + |
| 53 | +# # cs iterable of pairs; ensuring tight value of T |
| 54 | +# function MutableSparseVectorPolynomial{B}(cs::Tuple, var::SymbolLike=:x) where {B} |
| 55 | +# isempty(cs) && throw(ArgumentError("No type attached")) |
| 56 | +# X = Var(var) |
| 57 | +# if length(cs) == 1 |
| 58 | +# c = only(cs) |
| 59 | +# d = Dict(first(c) => last(c)) |
| 60 | +# T = eltype(last(c)) |
| 61 | +# return MutableSparseVectorPolynomial{B,T,X}(d) |
| 62 | +# else |
| 63 | +# c₁, c... = cs |
| 64 | +# T = typeof(last(c₁)) |
| 65 | +# for b ∈ c |
| 66 | +# T = promote_type(T, typeof(b)) |
| 67 | +# end |
| 68 | +# ks = 0:length(cs)-1 |
| 69 | +# vs = cs |
| 70 | +# d = Dict{Int,T}(Base.Generator(=>, ks, vs)) |
| 71 | +# return MutableSparseVectorPolynomial{B,T,X}(d) |
| 72 | +# end |
| 73 | +# end |
| 74 | + |
| 75 | +constructorof(::Type{<:MutableSparseVectorPolynomial{B}}) where {B <: AbstractBasis} = MutableSparseVectorPolynomial{B} |
| 76 | +@poly_register MutableSparseVectorPolynomial |
| 77 | + |
| 78 | +function Base.map(fn, p::P, args...) where {B,T,X, P<:MutableSparseVectorPolynomial{B,T,X}} |
| 79 | + xs = map(fn, p.coeffs) |
| 80 | + R = eltype(xs) |
| 81 | + return MutableSparseVectorPolynomial{B, R, X}(xs) |
| 82 | +end |
| 83 | + |
| 84 | +function Base.map!(fn, q::Q, p::P, args...) where {B,T,X, P<:MutableSparseVectorPolynomial{B,T,X},S,Q<:MutableSparseVectorPolynomial{B,S,X}} |
| 85 | + map!(fn, p.coeffs, p.coeffs) |
| 86 | + nothing |
| 87 | +end |
| 88 | + |
| 89 | +## --- |
| 90 | +Base.collect(p::MutableSparseVectorPolynomial) = collect(p.coeffs) |
| 91 | +Base.collect(::Type{T}, p::MutableSparseVectorPolynomial) where {T} = collect(T, p.coeffs) |
| 92 | +minimumexponent(::Type{<:MutableSparseVectorPolynomial}) = 0 |
| 93 | + |
| 94 | +Base.length(p::MutableSparseVectorPolynomial) = length(p.coeffs) |
| 95 | + |
| 96 | +function degree(p::MutableSparseVectorPolynomial) |
| 97 | + idx = findall(!iszero, p.coeffs) |
| 98 | + isempty(idx) && return -1 |
| 99 | + n = maximum(idx) |
| 100 | + n - 1 |
| 101 | +end |
| 102 | + |
| 103 | +Base.copy(p::MutableSparseVectorPolynomial{B,T,X}) where {B,T,X} = MutableSparseVectorPolynomial{B,T,X}(copy(p.coeffs)) |
| 104 | + |
| 105 | +function Base.convert(::Type{MutableSparseVectorPolynomial{B,T,X}}, p::MutableSparseVectorPolynomial{B,S,X}) where {B,T,S,X} |
| 106 | + cs = convert(SparseVector{T,Int}, p.coeffs) |
| 107 | + MutableSparseVectorPolynomial{B,T,X}(cs) |
| 108 | +end |
| 109 | + |
| 110 | +function Base.:(==)(p1::P, p2::P) where {P <: MutableSparseVectorPolynomial} |
| 111 | + iszero(p1) && iszero(p2) && return true |
| 112 | + |
| 113 | + ks1 = findall(!iszero, p1.coeffs) |
| 114 | + ks2 = findall(!iszero, p2.coeffs) |
| 115 | + length(ks1) == length(ks2) || return false |
| 116 | + idx = sortperm(ks1) |
| 117 | + for i ∈ idx |
| 118 | + ks1[i] == ks2[i] || return false |
| 119 | + p1.coeffs[ks1[i]] == p2.coeffs[ks2[i]] || return false |
| 120 | + end |
| 121 | + |
| 122 | + return true |
| 123 | + # # eachindex(p1) == eachindex(p2) || return false |
| 124 | + # # coeffs(p1) == coeffs(p2), but non-allocating |
| 125 | + # p1val = (p1[i] for i in eachindex(p1)) |
| 126 | + # p2val = (p2[i] for i in eachindex(p2)) |
| 127 | + # all(((a,b),) -> a == b, zip(p1val, p2val)) |
| 128 | +end |
| 129 | + |
| 130 | +# --- |
| 131 | + |
| 132 | +Base.firstindex(p::MutableSparseVectorPolynomial) = 0 |
| 133 | +function Base.lastindex(p::MutableSparseVectorPolynomial) |
| 134 | + isempty(p.coeffs) && return 0 |
| 135 | + maximum(keys(p.coeffs)) |
| 136 | +end |
| 137 | + |
| 138 | +function Base.getindex(p::MutableSparseVectorPolynomial{B,T,X}, i::Int) where {B,T,X} |
| 139 | + get(p.coeffs, i + 1, zero(T)) |
| 140 | +end |
| 141 | + |
| 142 | +# errors if extending |
| 143 | +function Base.setindex!(p::MutableSparseVectorPolynomial{B,T,X}, value, i::Int) where {B,T,X} |
| 144 | + p.coeffs[i+1] = value |
| 145 | +end |
| 146 | + |
| 147 | + |
| 148 | +function Base.pairs(p::MutableSparseVectorPolynomial) |
| 149 | + ks, vs = findnz(p.coeffs) |
| 150 | + idx = sortperm(ks) # guarantee order here |
| 151 | + Base.Generator(=>, ks[idx] .- 1, vs) |
| 152 | +end |
| 153 | +Base.keys(p::MutableSparseVectorPolynomial) = Base.Generator(first, pairs(p)) |
| 154 | +Base.values(p::MutableSparseVectorPolynomial) = Base.Generator(last, pairs(p)) |
| 155 | + |
| 156 | +basis(P::Type{<:MutableSparseVectorPolynomial{B, T, X}}, i::Int) where {B,T,X} = P(SparseVector(1+i, [i+1], [1])) |
| 157 | + |
| 158 | +# return coeffs as a vector |
| 159 | +function coeffs(p::MutableSparseVectorPolynomial{B,T}) where {B,T} |
| 160 | + d = degree(p) |
| 161 | + ps = p.coeffs |
| 162 | + [ps[i] for i ∈ 1:(d+1)] |
| 163 | +end |
| 164 | + |
| 165 | + |
| 166 | +hasnan(p::MutableSparseVectorPolynomial) = any(hasnan, values(p.coeffs))::Bool |
| 167 | + |
| 168 | + |
| 169 | +offset(p::MutableSparseVectorPolynomial) = 1 |
| 170 | + |
| 171 | +function keys_union(p::MutableSparseVectorPolynomial, q::MutableSparseVectorPolynomial) |
| 172 | + # IterTools.distinct(Base.Iterators.flatten((keys(p), keys(q)))) may allocate less |
| 173 | + unique(Base.Iterators.flatten((keys(p), keys(q)))) |
| 174 | +end |
| 175 | + |
| 176 | + |
| 177 | + |
| 178 | +## --- |
| 179 | + |
| 180 | +chop_exact_zeros!(d::SparseVector{T, Int}) where {T} = d |
| 181 | + |
| 182 | + |
| 183 | +function _truncate!(v::SparseVector{T,X}; |
| 184 | + rtol::Real = Base.rtoldefault(real(T)), |
| 185 | + atol::Real = 0) where {T,X} |
| 186 | + isempty(v) && return v |
| 187 | + δ = something(rtol,0) |
| 188 | + ϵ = something(atol,0) |
| 189 | + τ = max(ϵ, norm(values(v),2) * δ) |
| 190 | + for (i,pᵢ) ∈ pairs(v) |
| 191 | + abs(pᵢ) ≤ τ && (v[i] = zero(T)) |
| 192 | + end |
| 193 | + v |
| 194 | +end |
| 195 | + |
| 196 | + |
| 197 | +chop!(p::MutableSparseVectorPolynomial; kwargs...) = (chop!(p.coeffs; kwargs...); p) |
| 198 | +function chop!(d::SparseVector{T, Int}; atol=nothing, rtol=nothing) where {T} |
| 199 | + isempty(d) && return d |
| 200 | + δ = something(rtol,0) |
| 201 | + ϵ = something(atol,0) |
| 202 | + τ = max(ϵ, norm(values(d),2) * δ) |
| 203 | + for (i, pᵢ) ∈ Base.Iterators.reverse(pairs(d)) |
| 204 | + abs(pᵢ) ≥ τ && break |
| 205 | + d[i] = zero(T) |
| 206 | + end |
| 207 | + d |
| 208 | +end |
| 209 | + |
| 210 | +## --- |
| 211 | + |
| 212 | +_zeros(::Type{MutableSparseVectorPolynomial{B,T,X}}, z::S, N) where {B,T,X,S} = zeros(T, N) |
| 213 | + |
| 214 | +Base.zero(::Type{MutableSparseVectorPolynomial{B,T,X}}) where {B,T,X} = MutableSparseVectorPolynomial{B,T,X}(spzeros(T,0)) |
| 215 | + |
| 216 | +## --- |
| 217 | + |
| 218 | +function isconstant(p::MutableSparseVectorPolynomial) |
| 219 | + degree(p) <= 0 |
| 220 | +end |
| 221 | + |
| 222 | +Base.:+(p::MutableSparseVectorPolynomial{B,T,X}, q::MutableSparseVectorPolynomial{B,S,X}) where{B,X,T,S} = |
| 223 | + _sparse_vector_combine(+, p, q) |
| 224 | +Base.:-(p::MutableSparseVectorPolynomial{B,T,X}, q::MutableSparseVectorPolynomial{B,S,X}) where{B,X,T,S} = |
| 225 | + _sparse_vector_combine(-, p, q) |
| 226 | + |
| 227 | +# embed into bigger vector |
| 228 | +function _embed(v::SparseVector{T, Int}, l) where {T} |
| 229 | + l == length(v) && return v |
| 230 | + ks,vs = findnz(v) |
| 231 | + SparseVector(l, ks, vs) |
| 232 | +end |
| 233 | + |
| 234 | + |
| 235 | +function _sparse_vector_combine(op, p::MutableSparseVectorPolynomial{B,T,X}, q::MutableSparseVectorPolynomial{B,S,X}) where{B,X,T,S} |
| 236 | + R = promote_type(T,S) |
| 237 | + ps, qs = p.coeffs, q.coeffs |
| 238 | + m = max(length(ps), length(qs)) |
| 239 | + ps′, qs′ = _embed(ps, m), _embed(qs, m) |
| 240 | + cs = op(ps′, qs′) |
| 241 | + MutableSparseVectorPolynomial{B,R,X}(cs) |
| 242 | +end |
0 commit comments