Calculating a cannon shell trajectory, the fun way

In their very nice book, Computational Physics 2nd Ed., Nicholas Giordano and Hisao Nakanishi take us through many physics problems, solving them using numerical methods; they use the venerable FORTRAN language and it works, but I am convinced that using a functional language works as well while offering a few advantages, so to try and prove my hypothesis I started to slowly translate the code from FORTRAN to F# (when learning I definitively prefer the scenic route).

In this post we will see how to calculate the trajectory of a cannon shell with air resistance, taken from Chapter 2 of the book. The air resistance factor is important because due to it we don't get a simple parabola with an easy formula, but something like this:
A cannon shell trajectory with air resistance Do you notice how the curve is slightly bent to the left due to air resistance? It turns out that for getting it right we have to solve a differential equation.

But first, let's take care of a few details. The book uses mostly SI units and degrees for angles, but people in the United States is used to imperial units, and most math libraries use radians for angles, so you can see that unit confusion can easily get us in trouble; F# units to the rescue! Let's define the units we plan to use, including a conversion function from degrees to radians:

open Microsoft.FSharp.Data.UnitSystems.SI.UnitSymbols

type rad

type deg =
    static member toRad x =
        (System.Math.PI / 180.0<deg/rad>) * x

Furthermore, we will be using vectors, which will represent the position and the velocity of the cannon ball at any given time, so it comes handy to define a class that represent a two-dimensional vector:

type Vector2<[<Measure>] 'u>(x : float<'u>, y : float<'u>) =
    member this.x = x
    member this.y = y
    member this.norm = x * x + y * y |> sqrt

    static member (*) (v : Vector2<_>, a) =
       Vector2(v.x * a, v.y * a)
    static member (+) (v : Vector2<_>, w : Vector2<_>) =
        Vector2(v.x + w.x, v.y + w.y)
    static member (-) (v : Vector2<_>, w : Vector2<_>) =
        Vector2(v.x - w.x, v.y - w.y)
    static member fromPolar (r : float<'v>) (a : float<rad>) =
        Vector2<_>(r * cos (float a), r * sin (float a))

Note that every vector will have a physical unit, that is what the <[<Measure>] 'u> is for, 'u will be replaced by something specific when we create an instance, for example a position vector would be in meters (Vector2<m>), and a velocity vector would be in meters per second (Vector2<m/s>), we have also defined a few basic vector operations: multiplication by a scalar, sum, and difference, finally we have a function to convert from polar to Cartesian coordinates. Do you start to see how expressive (and safer) you can be with a few lines of F#?

Now for the central problem, Giordano & Nakanishi explain with detail the math behind the formulas, so I will just resume the results:

  1. We will calculate the position and velocity of the shell at successive small intervals of time Δt (say 1/10 of a second)
  2. Given the current position, we can approximate the next position multiplying the current velocity by the time interval, and adding this to the current position, something like this: pi+1 = pi + v x Δt. Please note that p is a position vector, v is a velocity vector, and Δt is a scalar (a seconds scalar, for example 0.1<s>)
  3. Given the current vertical speed, we can approximate the next speed by subtracting the gravity and the air resistance during the time interval, something like this: vy,i+1 = vy,i - g x Δt - vy x Δt x Am x |v| (Am is the air resistance)
  4. Finally, given the current horizontal speed, we can approximate the next speed by subtracting the air resistance during the time interval, something like this: vx,i+1 = vx,i - vx x Δt x Am x |v|

As I said before, you can check the detailed justification for all the formulas in Chapter 2 of the book. Now, let's put on our programmer's hat and consider the situation: you have to generate the first pair of values (position and velocity), and from this generate the second pair, and from this the third pair, and so on. You immediately think of a loop, but F# proposes you and intriguing alternative: Seq.unfold. As explained in the documentation, you must pass a function as a parameter to Seq.unfold, that function must generate some next value from the current state; Seq.unfold repeatedly calls the function thus generating an ever growing sequence, when there is no next value (in our case, when the shell hits the ground) your logic should return None which signals Seq.unfold to stop. So Seq.unfold comes handy to generate a sequence of times, positions, and velocities, using the formulas described above. The whole code looks like this:

let calculate (v_0, angle, p_0, a_m, g, dt) =
    let dg = Vector2(0.0<_>, g * dt)
        (fun (t, p : Vector2<_>, v) -> 
            if p.y <= 0.0<m> && t > 0.0<s> then None 
                let dv = v * dt 
                    (t, p, v),
                    (t + dt, p + dv, v - dg - dv * (a_m * v.norm))) )
        (0.0<_>, p_0, Vector2<_>.fromPolar v_0 angle)

The generated sequence has this structure: { (t0, p0, v0), (t1, p1, v1), ...} with the position and velocity of the shell at every point in time, ready to be plotted. Note that the code also knows the dimension of every element (s, m, m/s, m/s2, etc.) and the F# compiler makes you respect that. Not bad for a few lines of functional code :)

Oh, and as we don't use any mutable variables, the code is ready to be parallelized, so we could use 64 cores to calculate the trajectory of ten thousand shells, but that's a subject for another post.