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(t)[1:2]=xy v(t)[1:2]=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(t)[1:2]
    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(t)[1:2]=xy v(t)[1:2]=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{\mathrm{d}}{\mathrm{d}t} pos(t)_1 =& v(t)_1 \\ \frac{\mathrm{d}}{\mathrm{d}t} pos(t)_2 =& v(t)_2 \end{align} \]

Or using the @named helper macro

@named mass1 = Mass()

\[ \begin{align} \frac{\mathrm{d}}{\mathrm{d}t} pos(t)_1 =& v(t)_1 \\ \frac{\mathrm{d}}{\mathrm{d}t} pos(t)_2 =& v(t)_2 \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(t)[1:2]
    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} spring_{+}x\left( t \right) =& \sqrt{\left|mass_{+}pos(t)_1\right|^{2} + \left|mass_{+}pos(t)_2\right|^{2}} \\ spring_{+}dir(t)_1 =& mass_{+}pos(t)_1 \\ spring_{+}dir(t)_2 =& mass_{+}pos(t)_2 \\ \frac{\mathrm{d}}{\mathrm{d}t} mass_{+}v(t)_1 =& \frac{ - spring_{+}k \left( - spring_{+}l + spring_{+}x\left( t \right) \right) spring_{+}dir(t)_1}{mass_{+}m spring_{+}x\left( t \right)} \\ \frac{\mathrm{d}}{\mathrm{d}t} mass_{+}v(t)_2 =& -9.81 + \frac{ - spring_{+}k \left( - spring_{+}l + spring_{+}x\left( t \right) \right) spring_{+}dir(t)_2}{mass_{+}m 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)

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

equations(model)

The states of this model are:

states(model)

7-element Vector{Term{Real}}:
 spring₊x(t)
 (spring₊dir(t))[1]
 (spring₊dir(t))[2]
 (mass₊pos(t))[1]
 (mass₊pos(t))[2]
 (mass₊v(t))[1]
 (mass₊v(t))[2]

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{\mathrm{d}}{\mathrm{d}t} mass_{+}pos(t)_1 =& mass_{+}v(t)_1 \\ \frac{\mathrm{d}}{\mathrm{d}t} mass_{+}v(t)_1 =& mass_{+}pos(t)[1]ˍtt \\ \frac{\mathrm{d}}{\mathrm{d}t} mass_{+}pos(t)_2 =& mass_{+}v(t)_2 \\ \frac{\mathrm{d}}{\mathrm{d}t} mass_{+}v(t)_2 =& mass_{+}pos(t)[2]ˍtt \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{\mathrm{d}}{\mathrm{d}t} mass_{+}pos(t)_1 =& mass_{+}v(t)_1 \\ \frac{\mathrm{d}}{\mathrm{d}t} mass_{+}v(t)_1 =& - \frac{spring_{+}k \left( - spring_{+}l + \sqrt{\left|mass_{+}pos(t)_1\right|^{2} + \left|mass_{+}pos(t)_2\right|^{2}} \right) mass_{+}pos(t)_1}{mass_{+}m \sqrt{\left|mass_{+}pos(t)_1\right|^{2} + \left|mass_{+}pos(t)_2\right|^{2}}} \\ \frac{\mathrm{d}}{\mathrm{d}t} mass_{+}pos(t)_2 =& mass_{+}v(t)_2 \\ \frac{\mathrm{d}}{\mathrm{d}t} mass_{+}v(t)_2 =& -9.81 + \frac{ - spring_{+}k \left( - spring_{+}l + \sqrt{\left|mass_{+}pos(t)_1\right|^{2} + \left|mass_{+}pos(t)_2\right|^{2}} \right) mass_{+}pos(t)_2}{mass_{+}m \sqrt{\left|mass_{+}pos(t)_1\right|^{2} + \left|mass_{+}pos(t)_2\right|^{2}}} \end{align} \]

This is done as follows:

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