Feedback on a 2d spring simulation

This post does not present a specific question or problem. I have working code—a tiny simulation of a dynamic system involving some springs in 2d space—and I am hoping for general feedback regarding code idiomaticity, performance, and correctness of the physical model.

My main goals at present are to learn the DifferentialEquations.jl ecosystem and to dust off my knowledge of high school physics.

The system I’m simulating is as follows:

  • There are two springs, each with one end anchored at a known point (given by as), and with known spring constants (ks). The spring’s obey Hooke’s law, and their resting lengths (Ls) are known.
  • The free ends of the spring are attached to a single point mass (m).
  • There is gravity (g).
  • We know the initial position of the point mass (u0).

Here is a diagram (graphic design is my passion, etc.):

Here is my code:

using DifferentialEquations
using Plots

# Picture: We have a point mass affixed to a number of springs
# The springs themselves are anchored to various points in space
# We want to visualize how the point mass moves around in time

mutable struct Params
    as::Vector{Vector{Float64}} # Points to which springs are anchored
    ks::Vector{Float64}         # Spring coefficients
    Ls::Vector{Float64}         # Spring length at rest
    m::Float64                  # Mass of the point mass
    g::Float64                  # Gravity in positive y direction
end

# Euclidean distance between two points
function dist(x1::T, y1::T, x2::U, y2::U) where {T<:Number, U<:Number}
    return sqrt((x1 - x2)^2 + (y1 - y2)^2)
end

function f!(du, u, p, t)
    x, dx, y, dy = u

    du[1] = dx
    du[2] = 0
    du[3] = dy
    du[4] = p.g     # Gravity

    for (a, k, L) in zip(p.as, p.ks, p.Ls)
        L_t = dist(x, y, a[1], a[2])          # Current length of the spring
        F = k * (L_t - L)                     # Hooke's law

        # F is the magnitude of the force the spring exerts on the mass
        # It is directed from the mass toward the anchor, so:
        
        # Its x component is F cos θ = - F * (x - a[1]) / L_t
        du[2] -= F * (x - a[1]) / (L_t * p.m)     # = d^2 x

        # Its y component is F sin θ = - F * (y - a[2]) / L_t
        du[4] -= F * (y - a[2]) / (L_t * p.m)     # = d^2 y
    end
end

function solveexampleproblem()
    as = [[0.0, 0.0], [1.0, 1.0]]
    ks = [100.0, 50.0]
    Ls = [0.5, 0.5]
    m = 10.0
    g = -9.81

    p = Params(as, ks, Ls, m, g)

    u0 = [0.3, 0.0, 0.4, 0.0]
    tspan = (0.0, 50.0)

    prob = ODEProblem(f!, u0, tspan, p)

	@info "Solving example problem ..."
    solve(prob, AutoTsit5(Rosenbrock23()))
end

function makeanimation(
    sol::OrdinaryDiffEq.ODECompositeSolution,
    filename::String
)
	@info "Making animation ..."
    animate(
        sol,
        filename,
        idxs=(1, 3),
        c=:crimson,
        label=nothing,
        figsize=(500, 800)
    )
end

let sol = solveexampleproblem()
    makeanimation(sol, "2dHookeAnim.gif")
    @info "Done"
end

And here is the gif it outputs:

2dHookeAnim

A general code review, as well as suggestions for follow-up projects, would be more than welcome. Thank you in advance for any feedback.

At first sight the code is fine IMO. Some type-annotations seem odd ... where {T<:Number, U<:Number}.

What can make it perhaps faster and also more idiomatic in some lines is to use StaticArrays to represent all types of coordinates. For instance, that would allow you to write:

using LinearAlgebra: norm
dist(x,y) = norm(x-y)

#and
du -= F * (x - a) / L_t * p.m

where x, a would be static arrays with both coordinates, and the result is also an array. That change would require some reworking of the code, but the equation would be closer to what they look on paper.

Also, you could simplify a bit the code using things like:

du .= (dx, 0, dy, p.g)

instead of assigning each element in a different line, index by index.

1 Like

Thank you so much! All these tips are helpful.

Originally, I had dist(x1::T, y1::T, x2::T, y2::T) where {T<:Number} but this caused problems (for some solvers) in autodiff where one of the points was in Float64s and the other was in Dual{Float64}s. But yeah, a smarter point type from StaticArrays would obviate the need for this workaround.

It would be more idiomatic to remove the annotation completely, I think.

1 Like

OK, I am working on redoing this with a more legible Point type, but I am running into missing method for oneunit(::Type{Any}) error deep inside the solver algorithm.

Here is a minimal example of the problem I am running into:

using DifferentialEquations

struct Point{T<:Number}
    x::T
    y::T
end

function f!(du::Point, u::Point, _, _)
    du.x = -u.y
    du.y = u.x
end

let
    u0 = Point(1.0, 1.0)
    prob = ODEProblem(f!, u0, (0.0, 1.0))
    sol = solve(prob, Tsit5())
end

Error:

ERROR: MethodError: no method matching oneunit(::Type{Any})
Closest candidates are:
  oneunit(::Type{Union{Missing, T}}) where T at missing.jl:105
  oneunit(::Type{T}) where T at number.jl:358
  oneunit(::T) where T at number.jl:357
  ...
Stacktrace:
  [1] oneunit(#unused#::Type{Any})
    @ Base .\missing.jl:106
  ...

(Here I have used a Point type with x and y fields for simplicity, but in my actual implementation I am using StaticArrays as suggested and having the same issue.)

Am I using the wrong approach by trying to use a struct as my state variable instead of a vector?

It seems like being able to use a custom struct for u and du would enable a very expressive coding style, because your state variable could have legible fields corresponding to physical “objects” in the scene, etc.

x and y are not defined here, maybe you meant u.x and u.y?

Apart from that it feels strange to me that f! has four parameters while your initial point has two. But I don’t know the interface of ODEProblem from the top of my head, so this comment may be irrelevant.

Yes, will edit.

This is required (I think?) by the interface of ODEProblem. The remaining two arguments of f! are for parameters (e.g. the spring coefficients in my OP) and time (in case there is a forcing function). Neither of these are necessary in my minimal example.

Note that Point is immutable, so your f! will not work as you expect. To use it like that you need to define mutable struct. (but that is not related to that specific issue).

Yes, seems that the interface of ODEProblem is fine as it is.

The example runs if you define:

julia> using StaticArrays

julia> mutable struct Point{T} <: FieldVector{2,T}
           x::T
           y::T
       end

Which will give you all the algebra of arrays to your point. The error of missing oneunit(::Any) is pretty bad, though.

Yikes, so how do I initialize a FieldVector? Not in the docstring AFAICT.

julia> a = FieldVector{2}(1.0, 1.0)
ERROR: The constructor for FieldVector{2, Float64}(::Float64, ::Float64) is missing!
...

julia> FieldVector{2, Float64}([2.0, 4.0])
ERROR: The constructor for FieldVector{2, Float64}(::Float64, ::Float64) is missing!
...

julia> FieldVector{2, Float64}(2.0, 4.0)
ERROR: The constructor for FieldVector{2, Float64}(::Float64, ::Float64) is missing!
...

julia> FieldVector{2, Float64}((2.0, 4.0))
ERROR: The constructor for FieldVector{2, Float64}(::Float64, ::Float64) is missing!
...

julia> FieldVector(2.0, 4.0)
ERROR: FieldVector has no static size!

Edit: Nevermind, FieldVector is an abstract type, so you have to define your own concrete type like Point.