Skip to content

Commit

Permalink
Merge pull request #17 from baggepinnen/args
Browse files Browse the repository at this point in the history 
pass along additional args to f
  • Loading branch information
@baggepinnen
baggepinnen authored Dec 17, 2024
2 parents 2e6ace4 + c6916d3 commit d0ab694
Show file tree
Hide file tree
Showing 2 changed files with 54 additions and 50 deletions.
2 changes: 1 addition & 1 deletion Project.toml
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
name = "SeeToDee"
uuid = "1c904df7-48cd-41e7-921b-d889ed9a470c"
authors = ["Fredrik Bagge Carlson"]
version = "1.4.0"
version = "1.5.0"

[deps]
FastGaussQuadrature = "442a2c76-b920-505d-bb47-c5924d526838"
Expand Down
102 changes: 53 additions & 49 deletions src/SeeToDee.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,11 @@ Works for both continuous and discrete-time dynamics.
"""
function linearize(f, x, u, args...)
A = ForwardDiff.jacobian(x->f(x, u, args...), x)
B = ForwardDiff.jacobian(u->f(convert(typeof(u), x), u, args...), u)
if u isa SVector
B = ForwardDiff.jacobian(u->f(x, u, args...), u)
else
B = ForwardDiff.jacobian(u->f(convert(typeof(u), x), u, args...), u)
end
A, B
end

Expand All @@ -40,55 +44,55 @@ struct Rk4{F,TS}
end


function (integ::Rk4{F})(x, u, p, t; Ts=integ.Ts, supersample=integ.supersample) where F
function (integ::Rk4{F})(x, u, p, t, args...; Ts=integ.Ts, supersample=integ.supersample) where F
f = integ.f
f1 = f(x, u, p, t)
_inner_rk4(integ, f1, x, u, p, t; Ts, supersample) # Dispatch depending on return type of dynamics
f1 = f(x, u, p, t, args...)
_inner_rk4(integ, f1, x, u, p, t, args...; Ts, supersample) # Dispatch depending on return type of dynamics
end


function _inner_rk4(integ::Rk4{F}, f1::SArray, x, u, p, t; Ts=integ.Ts, supersample=integ.supersample) where F
function _inner_rk4(integ::Rk4{F}, f1::SArray, x, u, p, t, args...; Ts=integ.Ts, supersample=integ.supersample) where F
Ts2 = Ts / supersample
f = integ.f
f2 = f(x + Ts2 / 2 * f1, u, p, t + Ts2 / 2)
f3 = f(x + Ts2 / 2 * f2, u, p, t + Ts2 / 2)
f4 = f(x + Ts2 * f3, u, p, t + Ts2)
f2 = f(x + Ts2 / 2 * f1, u, p, t + Ts2 / 2, args...)
f3 = f(x + Ts2 / 2 * f2, u, p, t + Ts2 / 2, args...)
f4 = f(x + Ts2 * f3, u, p, t + Ts2, args...)
add = Ts2 / 6 .* (f1 .+ 2 .* f2 .+ 2 .* f3 .+ f4)
# This gymnastics with changing the name to y is to ensure type stability when x + add is not the same type as x. The compiler is smart enough to figure out the type of y
y = x + add
for i in 2:supersample
t += Ts2
f1 = f(y, u, p, t)
f2 = f(y + Ts2 / 2 * f1, u, p, t + Ts2 / 2)
f3 = f(y + Ts2 / 2 * f2, u, p, t + Ts2 / 2)
f4 = f(y + Ts2 * f3, u, p, t + Ts2)
f1 = f(y, u, p, t, args...)
f2 = f(y + Ts2 / 2 * f1, u, p, t + Ts2 / 2, args...)
f3 = f(y + Ts2 / 2 * f2, u, p, t + Ts2 / 2, args...)
f4 = f(y + Ts2 * f3, u, p, t + Ts2, args...)
add = Ts2 / 6 .* (f1 .+ 2 .* f2 .+ 2 .* f3 .+ f4)
y += add
end
return y
end

function _inner_rk4(integ::Rk4{F}, f1, x, u, p, t; Ts=integ.Ts, supersample=integ.supersample) where F
function _inner_rk4(integ::Rk4{F}, f1, x, u, p, t, args...; Ts=integ.Ts, supersample=integ.supersample) where F
Ts2 = Ts / supersample
f = integ.f
xi = x .+ (Ts2 / 2) .* f1
f2 = f(xi, u, p, t + Ts2 / 2)
f2 = f(xi, u, p, t + Ts2 / 2, args...)
xi .= x .+ (Ts2 / 2) .* f2
f3 = f(xi, u, p, t + Ts2 / 2)
f3 = f(xi, u, p, t + Ts2 / 2, args...)
xi .= x .+ Ts2 .* f3
f4 = f(xi, u, p, t + Ts2)
f4 = f(xi, u, p, t + Ts2, args...)
xi .= (Ts2 / 6) .* (f1 .+ 2 .* f2 .+ 2 .* f3 .+ f4)
# This gymnastics with changing the name to y is to ensure type stability when x + add is not the same type as x. The compiler is smart enough to figure out the type of y
y = _mutable(x + xi) # If x is non-static but xi isn't, we get a static array and adding to y won't work
for i in 2:supersample
t += Ts2
f1 = f(y, u, p, t)
f1 = f(y, u, p, t, args...)
xi = y .+ (Ts2 / 2) .* f1
f2 = f(xi, u, p, t + Ts2 / 2)
f2 = f(xi, u, p, t + Ts2 / 2, args...)
xi .= y .+ (Ts2 / 2) .* f2
f3 = f(xi, u, p, t + Ts2 / 2)
f3 = f(xi, u, p, t + Ts2 / 2, args...)
xi .= y .+ Ts2 .* f3
f4 = f(xi, u, p, t + Ts2)
f4 = f(xi, u, p, t + Ts2, args...)
xi .= (Ts2 / 6) .* (f1 .+ 2 .* f2 .+ 2 .* f3 .+ f4)
y .+= xi
end
Expand Down Expand Up @@ -120,25 +124,25 @@ struct Rk3{F,TS}
end


function (integ::Rk3{F})(x, u, p, t; Ts=integ.Ts, supersample=integ.supersample) where F
function (integ::Rk3{F})(x, u, p, t, args...; Ts=integ.Ts, supersample=integ.supersample) where F
f = integ.f
f1 = f(x, u, p, t)
_inner_rk3(integ, f1, x, u, p, t; Ts, supersample) # Dispatch depending on return type of dynamics
f1 = f(x, u, p, t, args...)
_inner_rk3(integ, f1, x, u, p, t, args...; Ts, supersample) # Dispatch depending on return type of dynamics
end

function _inner_rk3(integ::Rk3{F}, f1, x, u, p, t; Ts=integ.Ts, supersample=integ.supersample) where F
function _inner_rk3(integ::Rk3{F}, f1, x, u, p, t, args...; Ts=integ.Ts, supersample=integ.supersample) where F
Ts2 = Ts / supersample
f = integ.f
f2 = f(x + Ts2 / 2 * f1, u, p, t + Ts2 / 2)
f3 = f(x - Ts2 * f1 .+ 2 * Ts2 * f2, u, p, t + Ts2)
f2 = f(x + Ts2 / 2 * f1, u, p, t + Ts2 / 2, args...)
f3 = f(x - Ts2 * f1 .+ 2 * Ts2 * f2, u, p, t + Ts2, args...)
add = Ts2 / 6 .* (f1 .+ 4 .* f2 .+ f3)
# This gymnastics with changing the name to y is to ensure type stability when x + add is not the same type as x. The compiler is smart enough to figure out the type of y
y = x + add
for i in 2:supersample
t += Ts2
f1 = f(y, u, p, t)
f2 = f(y + Ts2 / 2 * f1, u, p, t + Ts2 / 2)
f3 = f(y - Ts2 * f1 .+ 2 * Ts2 * f2, u, p, t + Ts2)
f1 = f(y, u, p, t, args...)
f2 = f(y + Ts2 / 2 * f1, u, p, t + Ts2 / 2, args...)
f3 = f(y - Ts2 * f1 .+ 2 * Ts2 * f2, u, p, t + Ts2, args...)
add = Ts2 / 6 .* (f1 .+ 4 .* f2 .+ f3)
y += add
end
Expand Down Expand Up @@ -169,22 +173,22 @@ struct ForwardEuler{F,TS}
end


function (integ::ForwardEuler{F})(x, u, p, t; Ts=integ.Ts, supersample=integ.supersample) where F
function (integ::ForwardEuler{F})(x, u, p, t, args...; Ts=integ.Ts, supersample=integ.supersample) where F
f = integ.f
f1 = f(x, u, p, t)
_inner_forwardeuler(integ, f1, x, u, p, t; Ts, supersample) # Dispatch depending on return type of dynamics
f1 = f(x, u, p, t, args...)
_inner_forwardeuler(integ, f1, x, u, p, t, args...; Ts, supersample) # Dispatch depending on return type of dynamics
end


function _inner_forwardeuler(integ::ForwardEuler{F}, f1, x, u, p, t; Ts=integ.Ts, supersample=integ.supersample) where F
function _inner_forwardeuler(integ::ForwardEuler{F}, f1, x, u, p, t, args...; Ts=integ.Ts, supersample=integ.supersample) where F
Ts2 = Ts / supersample
f = integ.f
add = Ts2 .* f1
# This gymnastics with changing the name to y is to ensure type stability when x + add is not the same type as x. The compiler is smart enough to figure out the type of y
y = x + add
for i in 2:supersample
t += Ts2
f2 = f(y, u, p, t)
f2 = f(y, u, p, t, args...)
add = Ts2 .* f2
y += add
end
Expand Down Expand Up @@ -214,22 +218,22 @@ struct Heun{F,TS}
end
end

function (integ::Heun{F})(x, u, p, t; Ts=integ.Ts, supersample=integ.supersample) where F
function (integ::Heun{F})(x, u, p, t, args...; Ts=integ.Ts, supersample=integ.supersample) where F
f = integ.f
f1 = f(x, u, p, t)
_inner_heun(integ, f1, x, u, p, t; Ts, supersample) # Dispatch depending on return type of dynamics
f1 = f(x, u, p, t, args...)
_inner_heun(integ, f1, x, u, p, t, args...; Ts, supersample) # Dispatch depending on return type of dynamics
end

function _inner_heun(integ::Heun{F}, f1, x, u, p, t; Ts=integ.Ts, supersample=integ.supersample) where F
function _inner_heun(integ::Heun{F}, f1, x, u, p, t, args...; Ts=integ.Ts, supersample=integ.supersample) where F
Ts2 = Ts / supersample
f = integ.f
f2 = f(x + Ts2 * f1, u, p, t + Ts2)
f2 = f(x + Ts2 * f1, u, p, t + Ts2, args...)
add = (Ts2 / 2) .* (f1 .+ f2)
y = x + add
for i in 2:supersample
t += Ts2
f1 = f(y, u, p, t)
f2 = f(y .+ Ts2 .* f1, u, p, t + Ts2)
f1 = f(y, u, p, t, args...)
f2 = f(y .+ Ts2 .* f1, u, p, t + Ts2, args...)
add = (Ts2 / 2) .* (f1 .+ f2)
y += add
end
Expand Down Expand Up @@ -347,7 +351,7 @@ function SimpleColloc(dyn, Ts::T0, x_inds::AbstractVector{Int}, a_inds, nu::Int;
SimpleColloc(dyn, Ts, nx, x_inds, a_inds, nu, T.(D), T.(τ*Ts), abstol, cache, problem, solver, residual)
end

function coldyn(xv::AbstractArray{T}, (integ, x0, u, p, t)) where T
function coldyn(xv::AbstractArray{T}, (integ, x0, u, p, t, args...)) where T
(; dyn, x_inds, a_inds, D, τ, residual) = integ
nx, na = length(x_inds), length(a_inds)
cv, x_cache, ẋ, x = get_cache!(integ, xv)
Expand All @@ -367,13 +371,13 @@ function coldyn(xv::AbstractArray{T}, (integ, x0, u, p, t)) where T
if residual
allinds = 1:(nx+na)
@views for k in 1:n_c
res = dyn(ẋ[:,k], x[:,k], u, p, t+τ[k])
res = dyn(ẋ[:,k], x[:,k], u, p, t+τ[k], args...)
cv[allinds] .= res
allinds = allinds .+ (nx+na)
end
else
@views for k in 1:n_c
temp_dyn = dyn(x[:,k], u, p, t+τ[k])
temp_dyn = dyn(x[:,k], u, p, t+τ[k], args...)
cv[inds_c] .-= temp_dyn[inds_x]
inds_c = inds_c .+ (nx+na)

Expand All @@ -384,7 +388,7 @@ function coldyn(xv::AbstractArray{T}, (integ, x0, u, p, t)) where T
cv
end

function (integ::SimpleColloc)(x0::T, u, p, t; abstol=integ.abstol)::T where T
function (integ::SimpleColloc)(x0::T, u, p, t, args...; abstol=integ.abstol)::T where T
nx, na = length(integ.x_inds), length(integ.a_inds)
n_c = length(integ.τ)
_, _, _, u00 = get_cache!(integ, x0)
Expand All @@ -393,7 +397,7 @@ function (integ::SimpleColloc)(x0::T, u, p, t; abstol=integ.abstol)::T where T
for i = 1:nx0, j = 1:n_c
u0[i + (j-1)*nx0] = x0[i]
end
problem = SciMLBase.remake(integ.nlproblem, u0=u0,p=(integ, x0, u, p, t))
problem = SciMLBase.remake(integ.nlproblem, u0=u0,p=(integ, x0, u, p, t, args...))
solution = solve(problem, integ.solver; abstol)
if !SciMLBase.successful_retcode(solution)
@warn "Nonlinear solve failed to converge" solution.retcode maxlog=10
Expand All @@ -411,18 +415,18 @@ Given the differential state variables in `x0`, initialize the algebraic variabl
- `integ`: An intergrator like [`SeeToDee.SimpleColloc`](@ref)
- `x0`: Initial state descriptor (differential and algebraic variables, where the algebraic variables comes last)
"""
function initialize(integ, x0, u, p, t=0.0; solver = integ.solver, abstol = integ.abstol)
function initialize(integ, x0, u, p, t=0.0, args...; solver = integ.solver, abstol = integ.abstol)
(; dyn, nx, nu) = integ
# u = zeros(nu)
x0 = copy(x0)
diffinds = integ.x_inds
alginds = integ.a_inds
res0 = dyn(x0, u, p, t)
res0 = dyn(x0, u, p, t, args...)
norm(res0[alginds]) < abstol && return x0
res = function (z, _)
x0T = eltype(z).(x0)
x0T[alginds] .= z
dyn(x0T, u, p, t)[alginds]
dyn(x0T, u, p, t, args...)[alginds]
end
problem = NonlinearProblem(res, x0[alginds], p)
solution = solve(problem, solver; abstol)
Expand Down

2 comments on commit d0ab694

@baggepinnen
Copy link
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@JuliaRegistrator register

@JuliaRegistrator
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Registration pull request created: JuliaRegistries/General/121529

Tip: Release Notes

Did you know you can add release notes too? Just add markdown formatted text underneath the comment after the text
"Release notes:" and it will be added to the registry PR, and if TagBot is installed it will also be added to the
release that TagBot creates. i.e.

@JuliaRegistrator register

Release notes:

## Breaking changes

- blah

To add them here just re-invoke and the PR will be updated.

Tagging

After the above pull request is merged, it is recommended that a tag is created on this repository for the registered package version.

This will be done automatically if the Julia TagBot GitHub Action is installed, or can be done manually through the github interface, or via:

git tag -a v1.5.0 -m "<description of version>" d0ab69487d346931f00ee93275ad9638dd90aec5
git push origin v1.5.0

Please sign in to comment.