Skip to content

Commit

Permalink
before switching to mvector
Browse files Browse the repository at this point in the history
  • Loading branch information
@mattsignorelli
mattsignorelli committed Feb 8, 2025
1 parent 6f140b2 commit 09a139a
Show file tree
Hide file tree
Showing 3 changed files with 113 additions and 98 deletions.
89 changes: 45 additions & 44 deletions src/map/ctors.jl
View file
Original file line number Diff line number Diff line change
@@ -1,65 +1,66 @@
function copy!(m::Union{TaylorMap,VectorField}, m1::Union{TaylorMap,VectorField})
checkstates(m, m1)
if m1 isa TaylorMap && m isa TaylorMap
m.x0 = m1.x0
end
NV = nvars(m)
foreach((xi, x1i)->TI.copy!(xi, x1i), view(m.x, 1:NV), m1.x)
if !isnothing(m.q)
foreach((qi, q1i)->copy!(qi, q1i), m.q, m1.q)
end
if !isnothing(m.s) && m1 isa TaylorMap
m.s .= m1.s
end
return m
end


"""
zero(m::TaylorMap)
Creates a zero `m` with the same properties as `m` including TPSA definiton,
spin, and stochasticity.
"""
function zero(m::TaylorMap)
return _zero_map(typeof(m), getdef(m), m)
end

for t = (:DAMap, :TPSAMap)
@eval begin

# Explicit type specification
# Def change would be static (in type)
# This one will probably not be used
function $t{X0,X,Q,S}(m::Union{TaylorMap,Nothing}=nothing) where {X0,X,Q,S}
def = getdef(eltype(X))
# Lowest-level, internal
function _zero_map(::Type{$t{X0,X,Q,S}}, def::AbstractTPSADef, reuse::Union{TaylorMap,Nothing}) where {X0,X,Q,S}
out_x0 = init_x0(X0, def)
out_x = init_x(X, def, m)
out_x = init_x(X, def, reuse)
out_q = init_q(Q, def)
out_s = init_s(S, def)

out_m = $t(out_x0, out_x, out_q, out_s)
return out_m
end

function zero(::Type{$t{X0,X,Q,S}}) where {X0,X,Q,S}
return _zero_map($t{X0,X,Q,S}, getdef(eltype(X)), nothing)
end

# Explicit type specification
# Def change would be static (in type)
function $t{X0,X,Q,S}(m::Union{TaylorMap,Nothing}=nothing) where {X0,X,Q,S}
out_m = _zero_map($t{X0,X,Q,S}, getdef(eltype(X)), m)
if !isnothing(m)
nvars(def) == nvars(m) || error("Number of variables in TPSAs for `m` and output map disagree!")
nparams(def) == nparams(m) || error("Number of parameters in TPSAs for `m` and output map disagree!")
out_m.x0 = m.x0
foreach((out_xi, xi)->TI.copy!(out_xi, xi), view(out_m.x, 1:nvars(def)), m.x)
if !isnothing(out_m.q)
foreach((out_qi, Qi)->copy!(out_qi, Qi), out_m.q, m.q)
end
if !isnothing(out_m.s)
out_m.s .= m.s
end
copy!(out_m, m)
end


return $t(out_x0, out_x, out_q, out_s)
return out_m
end

# Copy ctor including optional TPSA def change
function $t(m::TaylorMap; def::Union{AbstractTPSADef,Nothing}=nothing)
if isnothing(def)
def = getdef(m)
else
nvars(def) == nvars(m) || error("Number of variables in TPSAs for `m` and `def` disagree!")
nparams(def) == nparams(m) || error("Number of parameters in TPSAs for `m` and `def` disagree!")
end
function $t(m::TaylorMap; def::AbstractTPSADef=getdef(m))
W = TI.numtype(eltype(m.x0))
X0 = promote_x0_type(typeof(m.x0), W)
X = promote_x_type(typeof(m.x), TI.init_tps_type(W, def))
Q = promote_q_type(typeof(m.q), TI.init_tps_type(W, def))
S = promote_s_type(typeof(m.s), W)

out_x0 = init_x0(X0, def)
out_x = init_x(X, def, m)
out_q = init_q(Q, def)
out_s = init_s(S, def)

out_m = $t(out_x0, out_x, out_q, out_s)

out_m.x0 .= m.x0
foreach((out_xi, xi)->TI.copy!(out_xi, xi), view(out_m.x, 1:nvars(def)), m.x)
if !isnothing(out_m.q)
foreach((out_qi, Qi)->copy!(out_qi, Qi), out_m.q, m.q)
end
if !isnothing(out_m.s)
out_m.s .= m.s
end

out_m = _zero_map($t{X0,X,Q,S}, def, m)
copy!()
return out_m
end

Expand Down
10 changes: 9 additions & 1 deletion src/sanity.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,15 @@ and the checks should be optimized away in the JIT compilation if valid.
@inline checkspin(stuff...) = all(x->isnothing(x.Q), stuff) || all(x->!isnothing(x.Q), stuff) || error("Atleast one map/vector field includes spin while others do not")
@inline checkstochastic(maps::TaylorMap...) = all(x->isnothing(x.E), maps) || all(x->!isnothing(x.E), maps) || error("Atleast one map includes stochasticity while others do not")

@inline checkstates(stuff...) = checkspin(filter(x->(x isa Union{TaylorMap,VectorField}), stuff)...) && checkstochastic(filter(x->(x isa TaylorMap), stuff)...)
@inline checkstates(stuff...) = checktpsas(stuff...) && checkspin(filter(x->(x isa Union{TaylorMap,VectorField}), stuff)...) && checkstochastic(filter(x->(x isa TaylorMap), stuff)...)
@inline function checktpsas(stuff...)
NV = nvars(first(stuff))
NN = ndiffs(first(stuff))
return all(stuff) do m
nvars(m) == NV || error("Atleast one map/vector field has TPSA with disagreeing number of variables")
ndiffs(m) == NN || error("Atleast one map/vector field has TPSA with disagreeing number of parameters")
end
end

# checkinplace is preferred to using the "where {S,..}.." syntax as this
# gives descriptive errors rather than just "function not found"
Expand Down
112 changes: 59 additions & 53 deletions src/types.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -108,6 +108,65 @@ end
TI.is_tps_type(eltype(X)) isa TI.IsTPSType || error("Orbital ray element type must be a truncated power series type supported by `TPSAInterface.jl`!")
Q == Nothing || eltype(Q) == eltype(X) || error("Quaternion number type $(eltype(Q)) must be $(eltype(X)) (equal to orbital ray)")
end


# =================================================================================== #
# Field initialization functions.
# Note that even when a TPSA definition is inferrable from the types (X0, X, ...),
# we still explicitly pass the def in case it is not. This is redundant, however ensures
# flexibility for different TPSA packages/modes.

# These may be overrided by external array packages.
function init_x0(::Type{X0_0}, ::Type{W}, def::AbstractTPSADef) where {X0_0,W}
nv = nvars(def)
x0 = similar(X0_0, W, nv)
x0 .= 0
return x0
end

function init_x(::Type{X_0}, ::Type{W}, def::AbstractTPSADef, reuse::Union{Nothing,TaylorMap}=nothing) where {X_0,W}
nv = nvars(def)
nn = ndiffs(def)
x = similar(X_0, nn)
for i in 1:nv
x[i] = TI.init_tps(TI.numtype(eltype(X_0)), def)
end
# use same parameters if use isa TaylorMap and eltype(x) == eltype(reuse.x)
if reuse isa TaylorMap && eltype(x) == eltype(reuse.x)
x[nv+1:nn] .= view(reuse.x, nv+1:nn)
else # allocate
for i in nv+1:nn
x[i] = TI.init_tps(TI.numtype(eltype(X_0)), def)
TI.seti!(x[i], 1, i)
end
end
return x
end

function init_q(::Type{Q}, def::AbstractTPSADef) where {Q}
if Q != Nothing
q0 = TI.init_tps(TI.numtype(eltype(Q)), def)
q1 = TI.init_tps(TI.numtype(eltype(Q)), def)
q2 = TI.init_tps(TI.numtype(eltype(Q)), def)
q3 = TI.init_tps(TI.numtype(eltype(Q)), def)
q = Quaternion(q0,q1,q2,q3)
else
q = nothing
end
return q
end

function init_s(::Type{S}, def::AbstractTPSADef) where {S}
if S != Nothing
nv = nvars(def)
s = similar(S, nv, nv)
s .= 0
else
s = nothing
end
return s
end



# =================================================================================== #
Expand Down Expand Up @@ -178,59 +237,6 @@ function real(::Type{VectorField{X,Q}}) where {X,Q}
end


# =================================================================================== #
# Field initialization functions.

# These may be overrided by external array packages.
function init_x0(::Type{X0}, def::AbstractTPSADef) where {X0}
nv = nvars(def)
x0 = similar(X0, nv)
x0 .= 0
return x0
end

function init_x(::Type{X}, def::AbstractTPSADef, reuse::Union{Nothing,TaylorMap}=nothing) where {X}
nv = nvars(def)
nn = ndiffs(def)
x = similar(X, nn)
for i in 1:nv
x[i] = TI.init_tps(TI.numtype(eltype(X)), def)
end
# use same parameters if use isa TaylorMap and eltype(x) == eltype(reuse.x)
if reuse isa TaylorMap && eltype(x) == eltype(reuse.x)
x[nv+1:nn] .= view(reuse.x, nv+1:nn)
else # allocate
for i in nv+1:nn
x[i] = TI.init_tps(TI.numtype(eltype(X)), def)
TI.seti!(x[i], 1, i)
end
end
return x
end

function init_q(::Type{Q}, def::AbstractTPSADef) where {Q}
if Q != Nothing
q0 = TI.init_tps(TI.numtype(eltype(Q)), def)
q1 = TI.init_tps(TI.numtype(eltype(Q)), def)
q2 = TI.init_tps(TI.numtype(eltype(Q)), def)
q3 = TI.init_tps(TI.numtype(eltype(Q)), def)
q = Quaternion(q0,q1,q2,q3)
else
q = nothing
end
return q
end

function init_s(::Type{S}, def::AbstractTPSADef) where {S}
if S != Nothing
nv = nvars(def)
s = similar(S, nv, nv)
s .= 0
else
s = nothing
end
return s
end



Expand Down

0 comments on commit 09a139a

Please sign in to comment.