Skip to content

Commit ea69f90

Browse files
AayushSabharwalAayushSabharwal
committed
Add spring-mass system tutorial
1 parent 8d36632 commit ea69f90

File tree

2 files changed

+230
-0
lines changed

2 files changed

+230
-0
lines changed

‎docs/make.jl

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ makedocs(
1212
"Home" => "index.md",
1313
"Symbolic Modeling Tutorials" => Any[
1414
"tutorials/ode_modeling.md",
15+
"tutorials/spring_mass.md",
1516
"tutorials/acausal_components.md",
1617
"tutorials/higher_order.md",
1718
"tutorials/tearing_parallelism.md",

‎docs/src/tutorials/spring_mass.md

Lines changed: 229 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,229 @@
1+
# Component-Based Modeling a Spring-Mass System
2+
3+
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](https://en.wikipedia.org/wiki/Hooke%27s_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.
4+
5+
## Copy-Paste Example
6+
7+
```julia
8+
using ModelingToolkit, Plots, DifferentialEquations, LinearAlgebra
9+
10+
@variables t
11+
D = Differential(t)
12+
13+
function Mass(; name, m = 1.0, xy = [0., 0.], u = [0., 0.])
14+
ps = @parameters m=m
15+
sts = @variables pos[1:2](t)=xy v[1:2](t)=u
16+
eqs = collect(D.(pos) .~ v)
17+
ODESystem(eqs, t, [pos..., v...], ps; name)
18+
end
19+
20+
function Spring(; name, k = 1e4, l = 1.)
21+
ps = @parameters k=k l=l
22+
@variables x(t), dir[1:2](t)
23+
ODESystem(Equation[], t, [x, dir...], ps; name)
24+
end
25+
26+
function connect_spring(spring, a, b)
27+
[
28+
spring.x ~ norm(collect(a .- b))
29+
collect(spring.dir .~ collect(a .- b))
30+
]
31+
end
32+
33+
spring_force(spring) = -spring.k .* collect(spring.dir) .* (spring.x - spring.l) ./ spring.x
34+
35+
m = 1.0
36+
xy = [1., -1.]
37+
k = 1e4
38+
l = 1.
39+
center = [0., 0.]
40+
g = [0., -9.81]
41+
@named mass = Mass(m=m, xy=xy)
42+
@named spring = Spring(k=k, l=l)
43+
44+
eqs = [
45+
connect_spring(spring, mass.pos, center)
46+
collect(D.(mass.v) .~ spring_force(spring) / mass.m .+ g)
47+
]
48+
49+
@named _model = ODESystem(eqs, t)
50+
@named model = compose(_model, mass, spring)
51+
sys = structural_simplify(model)
52+
53+
prob = ODEProblem(sys, [], (0., 3.))
54+
sol = solve(prob, Rosenbrock23())
55+
plot(sol)
56+
```
57+
58+
## Explanation
59+
### Building the components
60+
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 `vel`. We also define that the velocity is the rate of change of position with respect to time.
61+
62+
```julia
63+
function Mass(; name, m = 1.0, xy = [0., 0.], u = [0., 0.])
64+
ps = @parameters m=m
65+
sts = @variables pos[1:2](t)=xy v[1:2](t)=u
66+
eqs = collect(D.(pos) .~ v)
67+
ODESystem(eqs, t, [pos..., v...], ps; name)
68+
end
69+
```
70+
71+
Note that this is an incompletely specified `ODESystem`. It cannot be simulated on its own since the equations for `pos[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:
72+
73+
```julia
74+
Mass(name = :mass1)
75+
```
76+
77+
Or using the `@named` helper macro
78+
79+
```julia
80+
@named mass1 = Mass()
81+
```
82+
83+
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.
84+
85+
```julia
86+
function Spring(; name, k = 1e4, l = 1.)
87+
ps = @parameters k=k l=l
88+
@variables x(t), dir[1:2](t)
89+
ODESystem(Equation[], t, [x, dir...], ps; name)
90+
end
91+
```
92+
93+
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.]`.
94+
95+
```julia
96+
function connect_spring(spring, a, b)
97+
[
98+
spring.x ~ norm(collect(a .- b))
99+
collect(spring.dir .~ collect(a .- b))
100+
]
101+
end
102+
```
103+
104+
Lastly, we define the `spring_force` function that takes a `spring` and returns the force exerted by this spring.
105+
106+
```julia
107+
spring_force(spring) = -spring.k .* collect(spring.dir) .* (spring.x - spring.l) ./ spring.x
108+
```
109+
110+
To create our system, we will first create the components: a mass and a spring. This is done as follows:
111+
112+
```julia
113+
m = 1.0
114+
xy = [1., -1.]
115+
k = 1e4
116+
l = 1.
117+
center = [0., 0.]
118+
g = [0., -9.81]
119+
@named mass = Mass(m=m, xy=xy)
120+
@named spring = Spring(k=k, l=l)
121+
```
122+
123+
We can now create the equations describing this system, by connecting `spring` to `mass` and a fixed point.
124+
125+
```julia
126+
eqs = [
127+
connect_spring(spring, mass.pos, center)
128+
collect(D.(mass.v) .~ spring_force(spring) / mass.m .+ g)
129+
]
130+
```
131+
132+
Finally, we can build the model using these equations and components.
133+
134+
```julia
135+
@named _model = ODESystem(eqs, t)
136+
@named model = compose(_model, mass, spring)
137+
```
138+
139+
We can take a look at the equations in the model using the `equations` function.
140+
141+
```julia
142+
equations(model)
143+
144+
7-element Vector{Equation}:
145+
Differential(t)(mass₊v[1](t)) ~ -spring₊k*spring₊dir[1](t)*(mass₊m^-1)*(spring₊x(t) - spring₊l)*(spring₊x(t)^-1)
146+
Differential(t)(mass₊v[2](t)) ~ -9.81 - (spring₊k*spring₊dir[2](t)*(mass₊m^-1)*(spring₊x(t) - spring₊l)*(spring₊x(t)^-1))
147+
spring₊x(t) ~ sqrt(abs2(mass₊pos[1](t)) + abs2(mass₊pos[2](t)))
148+
spring₊dir[1](t) ~ mass₊pos[1](t)
149+
spring₊dir[2](t) ~ mass₊pos[2](t)
150+
Differential(t)(mass₊pos[1](t)) ~ mass₊v[1](t)
151+
Differential(t)(mass₊pos[2](t)) ~ mass₊v[2](t)
152+
```
153+
154+
The states of this model are:
155+
156+
```julia
157+
states(model)
158+
159+
7-element Vector{Term{Real, Base.ImmutableDict{DataType, Any}}}:
160+
mass₊v[1](t)
161+
mass₊v[2](t)
162+
spring₊x(t)
163+
mass₊pos[1](t)
164+
mass₊pos[2](t)
165+
spring₊dir[1](t)
166+
spring₊dir[2](t)
167+
```
168+
169+
And the parameters of this model are:
170+
171+
```julia
172+
parameters(model)
173+
174+
6-element Vector{Sym{Real, Base.ImmutableDict{DataType, Any}}}:
175+
spring₊k
176+
mass₊m
177+
spring₊l
178+
mass₊m
179+
spring₊k
180+
spring₊l
181+
```
182+
183+
### Simplifying and solving this system
184+
185+
This system can be solved directly as a DAE using [one of the DAE solvers from DifferentialEquations.jl](https://diffeq.sciml.ai/stable/solvers/dae_solve/). 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.
186+
187+
```julia
188+
sys = structural_simplify(model)
189+
equations(sys)
190+
191+
4-element Vector{Equation}:
192+
Differential(t)(mass₊v[1](t)) ~ -spring₊k*mass₊pos[1](t)*(mass₊m^-1)*(sqrt(abs2(mass₊pos[1](t)) + abs2(mass₊pos[2](t))) - spring₊l)*(sqrt(abs2(mass₊pos[1](t)) + abs2(mass₊pos[2](t)))^-1)
193+
Differential(t)(mass₊v[2](t)) ~ -9.81 - (spring₊k*mass₊pos[2](t)*(mass₊m^-1)*(sqrt(abs2(mass₊pos[1](t)) + abs2(mass₊pos[2](t))) - spring₊l)*(sqrt(abs2(mass₊pos[1](t)) + abs2(mass₊pos[2](t)))^-1))
194+
Differential(t)(mass₊pos[1](t)) ~ mass₊v[1](t)
195+
Differential(t)(mass₊pos[2](t)) ~ mass₊v[2](t)
196+
```
197+
198+
We are left with only 4 equations involving 4 state variables (`mass.pos[1]`, `mass.pos[2]`, `mass.vel[1]`, `mass.vel[2]`). We can solve the system by converting it to an `ODEProblem` in mass matrix form and solving with an [`ODEProblem` mass matrix solver](https://diffeq.sciml.ai/stable/solvers/dae_solve/#OrdinaryDiffEq.jl-(Mass-Matrix)). This is done as follows:
199+
200+
```julia
201+
prob = ODEProblem(sys, [], (0., 3.))
202+
sol = solve(prob, Rosenbrock23())
203+
plot(sol)
204+
```
205+
206+
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.
207+
208+
```julia
209+
observed(sys)
210+
211+
3-element Vector{Equation}:
212+
spring₊dir[2](t) ~ mass₊pos[2](t)
213+
spring₊dir[1](t) ~ mass₊pos[1](t)
214+
spring₊x(t) ~ sqrt(abs2(mass₊pos[1](t)) + abs2(mass₊pos[2](t)))
215+
```
216+
217+
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 signficantly better!
218+
219+
We can access these variables using the solution object. For example, let's retrieve the length of the spring over time:
220+
221+
```julia
222+
sol[spring.x]
223+
```
224+
225+
We can also plot its timeseries:
226+
227+
```julia
228+
plot(sol, vars = [spring.x])
229+
```

0 commit comments

Comments
 (0)