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)

Or using the @named helper macro

@named mass1 = Mass()

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)

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)
]

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)

And the parameters of this model are:

parameters(model)

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)

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)

This is done as follows:

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

What if we want the timeseries of a different variable? That information is not lost! Instead, structural_simplify simply changes state variables into observed variables.

observed(sys)

These are explicit algebraic equations which can be used to reconstruct the required variables on the fly. This leads to dramatic computational savings since implicitly solving an ODE scales as O(n^3), so fewer states are significantly better!

We can access these variables using the solution object. For example, let's retrieve the x-position of the mass over time:

sol[mass.pos[1]]

We can also plot the path of the mass:

plot(sol, vars = (mass.pos[1], mass.pos[2]))