Component-Based Modeling a Spring-Mass System

In this tutorial we will build a simple component-based model of a spring-mass system. A spring-mass system consists of one or more masses connected by springs. Hooke's law gives the force exerted by a spring when it is extended or compressed by a given distance. This specifies a differential-equation system where the acceleration of the masses is specified using the forces acting on them.

Copy-Paste Example

using ModelingToolkit, Plots, DifferentialEquations, LinearAlgebra
using Symbolics: scalarize

@variables t
D = Differential(t)

function Mass(; name, m = 1.0, xy = [0., 0.], u = [0., 0.])
    ps = @parameters m=m
    sts = @variables pos[1:2](t)=xy v[1:2](t)=u
    eqs = scalarize(D.(pos) .~ v)
    ODESystem(eqs, t, [pos..., v...], ps; name)
end

function Spring(; name, k = 1e4, l = 1.)
    ps = @parameters k=k l=l
    @variables x(t), dir[1:2](t)
    ODESystem(Equation[], t, [x, dir...], ps; name)
end

function connect_spring(spring, a, b)
    [
        spring.x ~ norm(scalarize(a .- b))
        scalarize(spring.dir .~ scalarize(a .- b))
    ]
end

spring_force(spring) = -spring.k .* scalarize(spring.dir) .* (spring.x - spring.l)  ./ spring.x

m = 1.0
xy = [1., -1.]
k = 1e4
l = 1.
center = [0., 0.]
g = [0., -9.81]
@named mass = Mass(m=m, xy=xy)
@named spring = Spring(k=k, l=l)

eqs = [
    connect_spring(spring, mass.pos, center)
    scalarize(D.(mass.v) .~ spring_force(spring) / mass.m .+ g)
]

@named _model = ODESystem(eqs, t, [spring.x; spring.dir; mass.pos], [])
@named model = compose(_model, mass, spring)
sys = structural_simplify(model)

prob = ODEProblem(sys, [], (0., 3.))
sol = solve(prob, Rosenbrock23())
plot(sol)

Explanation

Building the components

For each component we use a Julia function that returns an ODESystem. At the top, we define the fundamental properties of a Mass: it has a mass m, a position pos and a velocity v. We also define that the velocity is the rate of change of position with respect to time.

function Mass(; name, m = 1.0, xy = [0., 0.], u = [0., 0.])
    ps = @parameters m=m
    sts = @variables pos[1:2](t)=xy v[1:2](t)=u
    eqs = scalarize(D.(pos) .~ v)
    ODESystem(eqs, t, [pos..., v...], ps; name)
end

Mass (generic function with 1 method)

Note that this is an incompletely specified ODESystem. It cannot be simulated on its own since the equations for the velocity v[1:2](t) are unknown. Notice the addition of a name keyword. This allows us to generate different masses with different names. A Mass can now be constructed as:

Mass(name = :mass1)

\begin{align} \frac{dpos1(t)}{dt} =& \mathrm{v1}\left( t \right) \ \frac{dpos2(t)}{dt} =& \mathrm{v2}\left( t \right) \end{align}

Or using the @named helper macro

@named mass1 = Mass()

\begin{align} \frac{dpos1(t)}{dt} =& \mathrm{v1}\left( t \right) \ \frac{dpos2(t)}{dt} =& \mathrm{v2}\left( t \right) \end{align}

Next we build the spring component. It is characterised by the spring constant k and the length l of the spring when no force is applied to it. The state of a spring is defined by its current length and direction.

function Spring(; name, k = 1e4, l = 1.)
    ps = @parameters k=k l=l
    @variables x(t), dir[1:2](t)
    ODESystem(Equation[], t, [x, dir...], ps; name)
end

Spring (generic function with 1 method)

We now define functions that help construct the equations for a mass-spring system. First, the connect_spring function connects a spring between two positions a and b. Note that a and b can be the pos of a Mass, or just a fixed position such as [0., 0.]. In that sense, the length of the spring x is given by the length of the vector dir joining a and b.

function connect_spring(spring, a, b)
    [
        spring.x ~ norm(scalarize(a .- b))
        scalarize(spring.dir .~ scalarize(a .- b))
    ]
end

connect_spring (generic function with 1 method)

Lastly, we define the spring_force function that takes a spring and returns the force exerted by this spring.

spring_force(spring) = -spring.k .* scalarize(spring.dir) .* (spring.x - spring.l)  ./ spring.x

spring_force (generic function with 1 method)

To create our system, we will first create the components: a mass and a spring. This is done as follows:

m = 1.0
xy = [1., -1.]
k = 1e4
l = 1.
center = [0., 0.]
g = [0., -9.81]
@named mass = Mass(m=m, xy=xy)
@named spring = Spring(k=k, l=l)

\begin{align} \end{align}

We can now create the equations describing this system, by connecting spring to mass and a fixed point.

eqs = [
    connect_spring(spring, mass.pos, center)
    scalarize(D.(mass.v) .~ spring_force(spring) / mass.m .+ g)
]

\begin{align} \mathrm{spring{+}x}\left( t \right) =& \sqrt{\left|\mathrm{mass{+}pos1}\left( t \right)\right|^{2} + \left|\mathrm{mass{+}pos2}\left( t \right)\right|^{2}} \ \mathrm{spring{+}dir1}\left( t \right) =& \mathrm{mass{+}pos1}\left( t \right) \ \mathrm{spring{+}dir2}\left( t \right) =& \mathrm{mass{+}pos2}\left( t \right) \ \frac{dmass{+}v1(t)}{dt} =& \frac{ - spring{+}k \left( - spring{+}l + \mathrm{spring{+}x}\left( t \right) \right) \mathrm{spring{+}dir1}\left( t \right)}{mass{+}m \mathrm{spring{+}x}\left( t \right)} \ \frac{dmass{+}v2(t)}{dt} =& -9.81 + \frac{ - spring{+}k \left( - spring{+}l + \mathrm{spring{+}x}\left( t \right) \right) \mathrm{spring{+}dir2}\left( t \right)}{mass{+}m \mathrm{spring_{+}x}\left( t \right)} \end{align}

Finally, we can build the model using these equations and components.

@named _model = ODESystem(eqs, t, [spring.x; spring.dir; mass.pos], [])
@named model = compose(_model, mass, spring)

\begin{align} \mathrm{spring{+}x}\left( t \right) =& \sqrt{\left|\mathrm{mass{+}pos1}\left( t \right)\right|^{2} + \left|\mathrm{mass{+}pos2}\left( t \right)\right|^{2}} \ \mathrm{spring{+}dir1}\left( t \right) =& \mathrm{mass{+}pos1}\left( t \right) \ \mathrm{spring{+}dir2}\left( t \right) =& \mathrm{mass{+}pos2}\left( t \right) \ \frac{dmass{+}v1(t)}{dt} =& \frac{ - spring{+}k \left( - spring{+}l + \mathrm{spring{+}x}\left( t \right) \right) \mathrm{spring{+}dir1}\left( t \right)}{mass{+}m \mathrm{spring{+}x}\left( t \right)} \ \frac{dmass{+}v2(t)}{dt} =& -9.81 + \frac{ - spring{+}k \left( - spring{+}l + \mathrm{spring{+}x}\left( t \right) \right) \mathrm{spring{+}dir2}\left( t \right)}{mass{+}m \mathrm{spring{+}x}\left( t \right)} \ \frac{dmass{+}pos1(t)}{dt} =& \mathrm{mass{+}v1}\left( t \right) \ \frac{dmass{+}pos2(t)}{dt} =& \mathrm{mass{+}v_2}\left( t \right) \end{align}

We can take a look at the equations in the model using the equations function.

equations(model)

\begin{align} \mathrm{spring{+}x}\left( t \right) =& \sqrt{\left|\mathrm{mass{+}pos1}\left( t \right)\right|^{2} + \left|\mathrm{mass{+}pos2}\left( t \right)\right|^{2}} \ \mathrm{spring{+}dir1}\left( t \right) =& \mathrm{mass{+}pos1}\left( t \right) \ \mathrm{spring{+}dir2}\left( t \right) =& \mathrm{mass{+}pos2}\left( t \right) \ \frac{dmass{+}v1(t)}{dt} =& \frac{ - spring{+}k \left( - spring{+}l + \mathrm{spring{+}x}\left( t \right) \right) \mathrm{spring{+}dir1}\left( t \right)}{mass{+}m \mathrm{spring{+}x}\left( t \right)} \ \frac{dmass{+}v2(t)}{dt} =& -9.81 + \frac{ - spring{+}k \left( - spring{+}l + \mathrm{spring{+}x}\left( t \right) \right) \mathrm{spring{+}dir2}\left( t \right)}{mass{+}m \mathrm{spring{+}x}\left( t \right)} \ \frac{dmass{+}pos1(t)}{dt} =& \mathrm{mass{+}v1}\left( t \right) \ \frac{dmass{+}pos2(t)}{dt} =& \mathrm{mass{+}v_2}\left( t \right) \end{align}

The states of this model are:

states(model)

7-element Vector{Term{Real, Base.ImmutableDict{DataType, Any}}}:
 spring₊x(t)
 spring₊dir[1](t)
 spring₊dir[2](t)
 mass₊pos[1](t)
 mass₊pos[2](t)
 mass₊v[1](t)
 mass₊v[2](t)

And the parameters of this model are:

parameters(model)

3-element Vector{Any}:
 mass₊m
 spring₊k
 spring₊l

Simplifying and solving this system

This system can be solved directly as a DAE using one of the DAE solvers from DifferentialEquations.jl. However, we can symbolically simplify the system first beforehand. Running structural_simplify eliminates unnecessary variables from the model to give the leanest numerical representation of the system.

sys = structural_simplify(model)
equations(sys)

\begin{align} \frac{dmass{+}v1(t)}{dt} =& \frac{ - spring{+}k \left( - spring{+}l + \mathrm{spring{+}x}\left( t \right) \right) \mathrm{spring{+}dir1}\left( t \right)}{mass{+}m \mathrm{spring{+}x}\left( t \right)} \ \frac{dmass{+}v2(t)}{dt} =& -9.81 + \frac{ - spring{+}k \left( - spring{+}l + \mathrm{spring{+}x}\left( t \right) \right) \mathrm{spring{+}dir2}\left( t \right)}{mass{+}m \mathrm{spring{+}x}\left( t \right)} \ \frac{dmass{+}pos1(t)}{dt} =& \mathrm{mass{+}v1}\left( t \right) \ \frac{dmass{+}pos2(t)}{dt} =& \mathrm{mass{+}v2}\left( t \right) \end{align}

We are left with only 4 equations involving 4 state variables (mass.pos[1], mass.pos[2], mass.v[1], mass.v[2]). We can solve the system by converting it to an ODEProblem. Some observed variables are not expanded by default. To view the complete equations, one can do

full_equations(sys)

\begin{align} \frac{dmass{+}v1(t)}{dt} =& \frac{ - spring{+}k \left( - spring{+}l + \sqrt{\left|\mathrm{mass{+}pos1}\left( t \right)\right|^{2} + \left|\mathrm{mass{+}pos2}\left( t \right)\right|^{2}} \right) \mathrm{mass{+}pos1}\left( t \right)}{mass{+}m \sqrt{\left|\mathrm{mass{+}pos1}\left( t \right)\right|^{2} + \left|\mathrm{mass{+}pos2}\left( t \right)\right|^{2}}} \ \frac{dmass{+}v2(t)}{dt} =& -9.81 + \frac{ - spring{+}k \left( - spring{+}l + \sqrt{\left|\mathrm{mass{+}pos1}\left( t \right)\right|^{2} + \left|\mathrm{mass{+}pos2}\left( t \right)\right|^{2}} \right) \mathrm{mass{+}pos2}\left( t \right)}{mass{+}m \sqrt{\left|\mathrm{mass{+}pos1}\left( t \right)\right|^{2} + \left|\mathrm{mass{+}pos2}\left( t \right)\right|^{2}}} \ \frac{dmass{+}pos1(t)}{dt} =& \mathrm{mass{+}v1}\left( t \right) \ \frac{dmass{+}pos2(t)}{dt} =& \mathrm{mass{+}v2}\left( t \right) \end{align}

This is done as follows:

prob = ODEProblem(sys, [], (0., 3.))
sol = solve(prob, Rosenbrock23())
plot(sol)