ModelingToolkit.jl mechanical example.

limited_oscillator.jl

using ModelingToolkit, OrdinaryDiffEq

@variables t
D = Differential(t)

"The basic unit of the mechanical systems is a body with a position and a force"

@connector function Body(; name)
    sts = @variables x(t) F(t)
    ODESystem(Equation[], t, sts, []; name = name)
end

"The forces on connected bodies sum to zero. Bodies that are connected are in
the same place."
function ModelingToolkit.connect(::Type{Body}, ps...)
    eqs = [0 ~ sum(p -> p.F, ps)]
    for i = 1:(length(ps) - 1)
        push!(eqs, ps[i].x ~ ps[i + 1].x)
    end

    eqs
end

"A body that never moves."
function Fixed(; name)
    @named ground = Body()
    eqs = [D(ground.x) ~ 0]
    ODESystem(eqs, t, [], []; systems = [ground], name = name)
end

"A massive object obeys Newtons second law."
function Mass(; name, m = 1.0)
    @named mass = Body()

    val = [m]
    ps = @parameters m

    sts = @variables v(t)

    eqs = [D(mass.x) ~ v, D(v) ~ mass.F / m]

    ODESystem(eqs, t, sts, ps; systems = [mass], defaults = Dict(zip(ps, val)), name = name)
end

"The force applied by a linear spring is proportional to the distance between
its ends. This spring has zero length."
function Spring(; name, k = 1.0)
    @named left = Body()
    @named right = Body()

    vals = [k]
    ps = @parameters k

    eqs = [left.F ~ -k * (right.x - left.x), 0 ~ right.F + left.F]

    ODESystem(
        eqs,
        t,
        [],
        ps;
        systems = [left, right],
        defaults = Dict(zip(ps, vals)),
        name = name,
    )
end

"BOING!!!"
function simple_example()
    @named g = Fixed()
    @named s = Spring()
    @named m = Mass()

    eqs = [connect(g.ground, s.left)..., connect(s.right, m.mass)...]
    @named oscillator = ODESystem(eqs, t, [], []; systems = [g, s, m])
    simplified_oscillator = structural_simplify(oscillator)
    u0 = Dict(g.ground.x => 0.0, m.mass.x => 0.0, m.v => 1.0)
    prob = ODEProblem(simplified_oscillator, u0, (0.0, 10.0))
    soln = solve(prob, Tsit5())

    @assert(all([abs(soln.u[j][2] - sin(soln.t[j])) < 5e-4 for j = 1:length(soln)]))
    @assert(all([abs(soln.u[j][3] - cos(soln.t[j])) < 5e-4 for j = 1:length(soln)]))

    soln
end

"The distance from the moving body to the fixed body must stay between min and max."
function Limits(; name, min = 0.0, max = 1.0, cor = 0)
    @named fixed = Body()
    @named moving = Body()

    val = [min, max, cor]
    ps = @parameters min max cor

    sts = @variables v(t)

    l = moving.x - fixed.x

    eqs = [
        D(moving.x) ~ v,
        # when against the stop, forces balance, otherwise zero
        fixed.F ~ -moving.F * ((l == min) + (l == max)),
        moving.F ~ -fixed.F * ((l == min) + (l == max)),
    ]

    # Bounce off limits with coefficinet of restitution
    continuous_events = [[0 ~ min - l, 0 ~ max - l] => [v ~ -cor * v]]

    ODESystem(
        eqs,
        t,
        sts,
        ps;
        systems = [fixed, moving],
        defaults = Dict(zip(ps, val)),
        continuous_events = continuous_events,
        name = name,
    )
end

function limit_example()
    @named g = Fixed()
    @named s = Spring(k=10)
    @named m = Mass()
    @named l = Limits(min = -0.5, max = 0.5)
    eqs = [connect(g.ground, s.left, l.fixed)..., connect(s.right, m.mass, l.moving)...]
    @named limited_oscillator = ODESystem(eqs, t, [], []; systems = [g, s, m, l])
    limited_oscillator = structural_simplify(limited_oscillator)
    u0 = Dict(
        g.ground.x => 0.0,
        m.mass.x => 0.0,
        m.v => 1.0,
        l.moving.x => 0.0,
        l.moving.F => 0.0,
    )
    prob = ODEProblem(limited_oscillator, u0, (0.0, 10.0))
    soln = solve(prob, Rodas5())
    limited_oscillator, prob, soln
end