Skip to content

Commit

Permalink
Add getindex and setindex! for StepRanges
Browse files Browse the repository at this point in the history
  • Loading branch information
@lbenet
lbenet committed Jun 2, 2019
1 parent f1d6e25 commit 97de6be
Show file tree
Hide file tree
Showing 3 changed files with 81 additions and 21 deletions.
60 changes: 47 additions & 13 deletions src/auxiliary.jl
View file
Expand Up @@ -74,21 +74,30 @@ Return the coefficient of order `n::Int` of a `a::Taylor1` polynomial.
getcoeff(a::Taylor1, n::Int) = (@assert 0 ≤ n ≤ a.order; return a[n])

getindex(a::Taylor1, n::Int) = a.coeffs[n+1]
getindex(a::Taylor1, u::UnitRange) = view(a.coeffs, u .+ 1 )
getindex(a::Taylor1, u::UnitRange{Int}) = view(a.coeffs, u .+ 1 )
getindex(a::Taylor1, c::Colon) = view(a.coeffs, c)
getindex(a::Taylor1{T}, u::StepRange{Int,Int}) where {T<:Number} =
view(a.coeffs, u[:] .+ 1)

setindex!(a::Taylor1{T}, x::T, n::Int) where {T<:Number} = a.coeffs[n+1] = x
setindex!(a::Taylor1{T}, x::T, u::UnitRange) where {T<:Number} =
setindex!(a::Taylor1{T}, x::T, u::UnitRange{Int}) where {T<:Number} =
a.coeffs[u .+ 1] .= x
function setindex!(a::Taylor1{T}, x::Array{T,1}, u::UnitRange) where {T<:Number}
# a.coeffs[u .+ 1] .= x
function setindex!(a::Taylor1{T}, x::Array{T,1}, u::UnitRange{Int}) where {T<:Number}
@assert length(u) == length(x)
for ind in eachindex(x)
a.coeffs[u[ind]+1] = x[ind]
end
end
setindex!(a::Taylor1{T}, x::T, c::Colon) where {T<:Number} = a.coeffs[c] .= x
setindex!(a::Taylor1{T}, x::Array{T,1}, c::Colon) where {T<:Number} = a.coeffs[c] .= x
setindex!(a::Taylor1{T}, x::T, u::StepRange{Int,Int}) where {T<:Number} =
a.coeffs[u[:] .+ 1] .= x
function setindex!(a::Taylor1{T}, x::Array{T,1}, u::StepRange{Int,Int}) where {T<:Number}
@assert length(u) == length(x)
for ind in eachindex(x)
a.coeffs[u[ind]+1] = x[ind]
end
end


"""
Expand All @@ -107,19 +116,25 @@ end
getcoeff(a::HomogeneousPolynomial, v::Array{Int,1}) = getcoeff(a, (v...,))

getindex(a::HomogeneousPolynomial, n::Int) = a.coeffs[n]
getindex(a::HomogeneousPolynomial, n::UnitRange) = view(a.coeffs, n)
getindex(a::HomogeneousPolynomial, n::UnitRange{Int}) = view(a.coeffs, n)
getindex(a::HomogeneousPolynomial, c::Colon) = view(a.coeffs, c)
getindex(a::HomogeneousPolynomial, u::StepRange{Int,Int}) where {T<:Number} =
view(a.coeffs, u[:])

setindex!(a::HomogeneousPolynomial{T}, x::T, n::Int) where {T<:Number} =
a.coeffs[n] = x
setindex!(a::HomogeneousPolynomial{T}, x::T, n::UnitRange) where {T<:Number} =
setindex!(a::HomogeneousPolynomial{T}, x::T, n::UnitRange{Int}) where {T<:Number} =
a.coeffs[n] .= x
setindex!(a::HomogeneousPolynomial{T}, x::Array{T,1}, n::UnitRange) where {T<:Number} =
setindex!(a::HomogeneousPolynomial{T}, x::Array{T,1}, n::UnitRange{Int}) where {T<:Number} =
a.coeffs[n] .= x
setindex!(a::HomogeneousPolynomial{T}, x::T, c::Colon) where {T<:Number} =
a.coeffs[c] .= x
setindex!(a::HomogeneousPolynomial{T}, x::Array{T,1}, c::Colon) where {T<:Number} =
a.coeffs[c] = x
setindex!(a::HomogeneousPolynomial{T}, x::T, u::StepRange{Int,Int}) where {T<:Number} =
a.coeffs[u[:]] .= x
setindex!(a::HomogeneousPolynomial{T}, x::Array{T,1}, u::StepRange{Int,Int}) where {T<:Number} =
a.coeffs[u[:]] .= x[:]


"""
Expand All @@ -137,8 +152,10 @@ end
getcoeff(a::TaylorN, v::Array{Int,1}) = getcoeff(a, (v...,))

getindex(a::TaylorN, n::Int) = a.coeffs[n+1]
getindex(a::TaylorN, u::UnitRange) = view(a.coeffs, u .+ 1)
getindex(a::TaylorN, u::UnitRange{Int}) = view(a.coeffs, u .+ 1)
getindex(a::TaylorN, c::Colon) = view(a.coeffs, c)
getindex(a::TaylorN, u::StepRange{Int,Int}) where {T<:Number} =
view(a.coeffs, u[:] .+ 1)

function setindex!(a::TaylorN{T}, x::HomogeneousPolynomial{T}, n::Int) where
{T<:Number}
Expand All @@ -147,21 +164,19 @@ function setindex!(a::TaylorN{T}, x::HomogeneousPolynomial{T}, n::Int) where
end
setindex!(a::TaylorN{T}, x::T, n::Int) where {T<:Number} =
a.coeffs[n+1] = HomogeneousPolynomial(x, n)
function setindex!(a::TaylorN{T}, x::T, u::UnitRange) where {T<:Number}
function setindex!(a::TaylorN{T}, x::T, u::UnitRange{Int}) where {T<:Number}
for ind in u
a[ind] = x
end
a[u]
end
function setindex!(a::TaylorN{T}, x::Array{HomogeneousPolynomial{T},1}, u::UnitRange) where {T<:Number}
# a[u[:]] .= x[:]
function setindex!(a::TaylorN{T}, x::Array{HomogeneousPolynomial{T},1}, u::UnitRange{Int}) where {T<:Number}
@assert length(u) == length(x)
for ind in eachindex(x)
a[u[ind]] = x[ind]
end
end
function setindex!(a::TaylorN{T}, x::Array{T,1}, u::UnitRange) where {T<:Number}
# a[u] .= x
function setindex!(a::TaylorN{T}, x::Array{T,1}, u::UnitRange{Int}) where {T<:Number}
@assert length(u) == length(x)
for ind in eachindex(x)
a[u[ind]] = x[ind]
Expand All @@ -173,6 +188,25 @@ setindex!(a::TaylorN{T}, x::Array{HomogeneousPolynomial{T},1}, ::Colon) where
{T<:Number} = (a[0:end] = x; a[:])
setindex!(a::TaylorN{T}, x::Array{T,1}, ::Colon) where {T<:Number} =
(a[0:end] = x; a[:])
function setindex!(a::TaylorN{T}, x::T, u::StepRange{Int,Int}) where {T<:Number}
for ind in u
a[ind] = x
end
a[u]
end
function setindex!(a::TaylorN{T}, x::Array{HomogeneousPolynomial{T},1}, u::StepRange{Int,Int}) where {T<:Number}
# a[u[:]] .= x[:]
@assert length(u) == length(x)
for ind in eachindex(x)
a[u[ind]] = x[ind]
end
end
function setindex!(a::TaylorN{T}, x::Array{T,1}, u::StepRange{Int,Int}) where {T<:Number}
@assert length(u) == length(x)
for ind in eachindex(x)
a[u[ind]] = x[ind]
end
end


## eltype, length, get_order ##
Expand Down
21 changes: 18 additions & 3 deletions test/manyvariables.jl
View file
Expand Up @@ -71,10 +71,16 @@ eeuler = Base.MathConstants.e
@test typeof(TaylorSeries.resize_coeffsHP!(v,2)) == Nothing
@test v == [1,2,0]
@test_throws AssertionError TaylorSeries.resize_coeffsHP!(v,1)
HomogeneousPolynomial(v)[3] = 3
hpol_v = HomogeneousPolynomial(v)
hpol_v[3] = 3
@test v == [1,2,3]
HomogeneousPolynomial(v)[1:3] = 3
hpol_v[1:3] = 3
@test v == [3,3,3]
hpol_v[1:2:2] = 0
@test v == [0,3,3]
hpol_v[1:1:2] = [1,2]
@test all(hpol_v[1:1:2] .== [1,2])
@test v == [1,2,3]

xH = HomogeneousPolynomial([1,0])
yH = HomogeneousPolynomial([0,1],1)
Expand Down Expand Up @@ -241,6 +247,9 @@ eeuler = Base.MathConstants.e
rv = a[end][1:end]
@test a[end][1:end] == rv
@test a[end][1:end] != b[end][1:end]
@test a[0:2:end] == a.coeffs[1:2:end]
a[0:1:end] = 0.0
@test a == zero(a)

hp = HomogeneousPolynomial(1)^8
rv1 = rand( length(hp) )
Expand All @@ -257,6 +266,7 @@ eeuler = Base.MathConstants.e
@test hp[end-1:end] == rv1[end-1:end]
hp[:] = 0.0
@test hp[:] == zero(rv1)
@test all(hp[end-1:1:end] .== 0.0)

pol = sin(xT+yT*xT)+yT^2-(1-xT)^3
q = deepcopy(pol)
Expand Down Expand Up @@ -295,7 +305,7 @@ eeuler = Base.MathConstants.e
q[:] = pol.coeffs
q[1:end-1] = zeros(q.order+1)[2:end-1]
@test q != pol
@test q[:] != pol[:]
@test all(q[1:1:end-1] .== 0.0)
@test q[1:end-1] == zeros(q.order+1)[2:end-1]
@test q[0] == pol[0]
@test q[end] == pol[end]
Expand All @@ -305,6 +315,11 @@ eeuler = Base.MathConstants.e
@test q[0:1] == pol[0:1]
@test q[2:end-2] == pol2[2:end-2]
@test q[end-1:end] == pol[end-1:end]
@test q[2:2:end-2] == pol2[2:2:end-2]
@test q[end-1:1:end] == pol[end-1:1:end]
q[end-2:2:end] = [0.0, 0.0]
@test q[end-2] == 0.0
@test_throws AssertionError q[end-2:2:end] = [0.0, 0.0, 0.0]

@test_throws AssertionError yT^(-2)
@test_throws AssertionError yT^(-2.0)
Expand Down
21 changes: 16 additions & 5 deletions test/onevariable.jl
View file
Expand Up @@ -43,14 +43,20 @@ eeuler = Base.MathConstants.e
TaylorSeries.resize_coeffs1!(v,3)
setindex!(Taylor1(v),3,2)
@test v == [1,0,3,0]
@test Taylor1(v)[:] == [1,0,3,0]
@test Taylor1(v)[:] == Taylor1(v).coeffs[:]
setindex!(Taylor1(v),0,0:2)
pol_int = Taylor1(v)
@test pol_int[:] == [1,0,3,0]
@test pol_int[:] == pol_int.coeffs[:]
@test pol_int[1:2:3] == pol_int.coeffs[2:2:4]
setindex!(pol_int,0,0:2)
@test v == zero(v)
setindex!(Taylor1(v),1,:)
setindex!(pol_int,1,:)
@test v == ones(Int, 4)
Taylor1(v)[:] = 0
pol_int[:] = 0
@test v == zero(v)
pol_int[0:2:end] = 2
@test all(v[1:2:end] .== 2)
pol_int[0:2:3] = [0, 1]
@test all(v[1:2:3] .== [0, 1])
rv = [rand(0:3) for i in 1:4]
@test Taylor1(rv)[:] == rv
y = sin(Taylor1(16))
Expand Down Expand Up @@ -79,6 +85,11 @@ eeuler = Base.MathConstants.e
rv = rand.(length(y))
y[:] = rv
@test y[:] == rv
y[0:4:end] = 1.0
@test all(y.coeffs[1:4:end] .== 1.0)
y[0:2:8] = rv[1:5]
@test y.coeffs[1:2:9] == rv[1:5]
@test_throws AssertionError y[0:2:3] = rv

@test Taylor1([0,1,0,0]) == Taylor1(3)
@test getcoeff(Taylor1(Complex{Float64},3),1) == complex(1.0,0.0)
Expand Down

0 comments on commit 97de6be

Please sign in to comment.