F#奇妙游（17）：F#与真空一维平面大地抛石飞行力学

F#还能干点啥

距离上一次更新已经过去了很久（40分钟之久！），这段时间我在学习F#，并且在工作（划掉，躺肥并没有工作要做）中使用F#。

那干点啥呢？还是老本行吧，搞点飞行力学。

有一个球（质点），在一维平面大地、真空两个假设条件下，以一定的初始速度和初始角度抛出，求其运动轨迹。什么？这也是飞行力学？当然是啊，飞行力学不就是研究物体在空间中的运动吗？我假设真空怎么了？我假设一维平面大地怎么了？我假设球是质点怎么了？还不是飞行力学啊！

扔石头就是最早的飞行力学应用！……好吧，我承认是我说的。但是，不能因为是我说的就说不对啊！

问题超级简单，球的位置可以用一个二维向量表示，速度也可以用一个二维向量表示，那么问题就是求解一个二阶微分方程组：

$$\frac{d^2\vec{r}}{d{t}^2}=\vec{a}$$

拆成一阶微分方程组就是：
$$\{\begin{array}{l}\frac{dx}{dt}=v_x\\ \frac{d{v}_x}{dt}=0\\ \frac{dy}{dt}=v_y\\ \frac{d{v}_y}{dt}=-g\end{array}$$

其中$g$是重力加速度，$g=9.8m/s^2$。在时刻$t=0$有，球的位置和速度分别为：${\vec{r}}_0=[x_0,y_0{]}^T$和${\vec{v}}_0=[v_{x0},v_{y0}{]}^T$。问题的解析解是：

$$\{\begin{array}{l}x=v_{x0}t+x_0\\ v_x=v_{x0}\\ y=-\frac{1}{2}g{t}^2+v_{y0}t+y_0\\ v_y=-gt+v_{y0}\end{array}$$

好无聊啊……这么简单的结果我不能接受！我要用更加复杂的办法，并且用F#来实现！

飞行力学的数学模型

这一类问题，在数学上都可以归结为Ordinary Differential Equation，ODE，常微分方程。

$$\frac{d\vec{y}}{dt}=\vec{f}(t,\vec{y})$$

或者写为：

$$\dot{\vec{y}}=\vec{f}(t,\vec{y})$$

求解ODE的解，就是在$\vec{y}(0)$已知的情况下，求解$\vec{y}(t),t\in [0,t_f]$。

常微分方程的求解方法有很多，比如解析解、数值解、级数解等等。解析解就是上面那个解，数值解就是用数值方法求解，级数解就是用级数展开求解。

这里我使用四阶龙格库塔法来求解这个问题。

龙格库塔法

龙格库塔法是一种数值解常微分方程的方法，它的思想是，用一系列的中间变量来逼近微分方程的解。

$$\{\begin{array}{l}{\vec{k}}_1=\vec{f}(t_n,{\vec{y}}_n)\\ {\vec{k}}_2=\vec{f}(t_n+\frac{h}{2},{\vec{y}}_n+\frac{h}{2}{\vec{k}}_1)\\ {\vec{k}}_3=\vec{f}(t_n+\frac{h}{2},{\vec{y}}_n+\frac{h}{2}{\vec{k}}_2)\\ {\vec{k}}_4=\vec{f}(t_n+h,{\vec{y}}_n+h{\vec{k}}_3)\\ {\vec{y}}_{n+1}={\vec{y}}_n+\frac{h}{6}({\vec{k}}_1+2{\vec{k}}_2+2{\vec{k}}_3+{\vec{k}}_4)\end{array}$$

F#实现

运算符重载

可以看到上面的算法里面最突出的特征就是有很多向量的计算。那么首先我们先定义浮点数向量之间的算术计算。F#有一套很优雅的算子符号重载机制，我们可以用这个机制来实现向量的算术计算。

以加法为例，需要处理浮点数与向量的加法，向量与浮点数的加法，以及向量与向量的加法。

module arithmetic =
    let inline (.+) (a: obj) (b: obj) =
        match a, b with
        | :? float as op1, (:? array<float> as op2) -> Array.map (fun x -> op1 + x) op2
        | :? array<float> as op1, (:? float as op2) -> Array.map (fun x -> x + op2) op1
        | :? array<float> as op1, (:? array<float> as op2) -> Array.map2 (+) op1 op2
        | _ -> failwith "Type mismatch"

这个代码里使用了F#的类型判断match语法。唯一需要注意的就是第二个操作数的类型判断前后的括号，如果这里没有这个括号，op2就会映射到元组上，而不是float或者array上。

数组的计算，F#已经通过Array.map和Array.map2提供了支持，我们只需要在这个基础上进行封装即可。

通过这样的方法，就可以很容易把其他运算定义出来。

module arithmetic =
    let inline (.+) (a: obj) (b: obj) =
        match a, b with
        | :? float as op1, (:? array<float> as op2) -> Array.map (fun x -> op1 + x) op2
        | :? array<float> as op1, (:? float as op2) -> Array.map (fun x -> x + op2) op1
        | :? array<float> as op1, (:? array<float> as op2) -> Array.map2 (+) op1 op2
        | _ -> failwith "Type mismatch"

    let inline (.-) (a: obj) (b: obj) =
        match a, b with
        | :? float as op1, (:? array<float> as op2) -> Array.map (fun x -> op1 - x) op2
        | :? array<float> as op1, (:? float as op2) -> Array.map (fun x -> x - op2) op1
        | :? array<float> as op1, (:? array<float> as op2) -> Array.map2 (-) op1 op2
        | _ -> failwith "Type mismatch"

    let inline (.*) (a: obj) (b: obj) =
        match a, b with
        | :? float as op1, (:? (array<float>) as op2) -> Array.map (fun x -> op1 * x) op2
        | :? array<float> as op1, (:? float as op2) -> Array.map (fun x -> x * op2) op1
        | :? array<float> as op1, (:? array<float> as op2) -> Array.map2 (*) op1 op2
        | _ -> failwith "Type mismatch"

    let inline (./) (a: obj) (b: obj) =
        match a, b with
        | :? float as op1, (:? (array<float>) as op2) -> Array.map (fun x -> op1 / x) op2
        | :? array<float> as op1, (:? float as op2) -> Array.map (fun x -> x / op2) op1
        | :? (array<float>) as op1, (:? (array<float>) as op2) -> Array.map2 (/) op1 op2
        | _ -> failwith "Type mismatch"

    let inline (.^) (a: obj) (b: obj) =
        match a, b with
        | :? float as op1, (:? (array<float>) as op2) -> Array.map (fun x -> op1 ** x) op2
        | :? array<float> as op1, (:? float as op2) -> Array.map (fun x -> x ** op2) op1
        | :? array<float> as op1, (:? array<float> as op2) -> Array.map2 (fun x y -> x ** y) op1 op2
        | _ -> failwith "Type mismatch"

    let inline dotSin (a: array<float>) = Array.map sin

    let inline dotCos (a: array<float>) = Array.map cos

    let inline dotTan (a: array<float>) = Array.map tan

    let inline dotAsin (a: array<float>) = Array.map asin

    let inline dotAcos (a: array<float>) = Array.map acos

    let inline dotAtan (a: array<float>) = Array.map atan

    let inline dotSinh (a: array<float>) = Array.map sinh

    let inline dotCosh (a: array<float>) = Array.map cosh

    let inline dotTanh (a: array<float>) = Array.map tanh

    let inline dotLn (a: array<float>) = Array.map log

    let inline dotLog10 (a: array<float>) = Array.map log10

    let inline dotExp (a: array<float>) = Array.map exp

    let inline dotSqrt (a: array<float>) = Array.map sqrt

    let inline dotAbs (a: array<float>) = Array.map abs

    let inline dotSign (a: array<float>) = Array.map sign

    let inline dotCeil (a: array<float>) = Array.map ceil

    let inline dotFloor (a: array<float>) = Array.map floor

    let inline dotRound (a: array<float>) = Array.map round

    let inline dotTruncate (a: array<float>) = Array.map truncate

龙格库塔法

利用上面的算子符号重载，我们可以很容易地实现龙格库塔法的单步推进。

    open arithmetic


    let rk4Step
        (t0: float)
        (dt: float)
        (x0: array<float>)
        (para: obj)
        (f: obj -> float -> array<float> -> array<float>)
        =
        let k1 = f para t0 x0
        let k2 = f para (t0 + dt / 2.0) (x0 .+ dt / 2.0 .* k1)
        let k3 = f para (t0 + dt / 2.0) (x0 .+ dt / 2.0 .* k2)
        let k4 = f para (t0 + dt) (x0 .+ dt .* k3)

        x0
        .+ dt / 6.0 .* (k1 .+ 2.0 .* k2 .+ 2.0 .* k3 .+ k4)

可以看到，这就是数学公式的直接翻译，非常简单。

单步推进的基础上，就可以实现前面提到的ODE边值问题的解法。

    let rk4
        (t0: float)
        (dt: float)
        (x0: array<float>)
        (para: obj)
        (f: obj -> float -> array<float> -> array<float>)
        (t: float)
        =
        let n = int (t / dt)
        let rec loop (x: array<float>) (t: float) (n: int) =
            if n = 0 then
                x
            else
                let x1 = rk4Step t dt x para f
                loop x1 (t + dt) (n - 1)

        loop x0 t0 n

尾迭代虽然很酷炫，但是对于我们只能理解扔石头飞行力学的头脑而言，还是太难了……那么我们就是用F#不太纯一面来实现。

    let rk4
        (tspan: float * float)
        (dt: float)
        (x0: array<float>)
        (para: obj)
        (f: obj -> float -> array<float> -> array<float>)
        =
        let t0, tf = tspan
        let t = [| t0..dt..tf |]

        let t =
            if t[t.Length - 1] < tf then
                Array.append t [| tf |]
            else
                t

        let n = t.Length

        let x = Array.zeroCreate n

        x[0] <- x0

        for i = 0 to n - 2 do
            x[i + 1] <- rk4Step t[i] (t[i + 1] - t[i]) x.[i] para f

        t, x

这个代码就是把前面的rk4Step循环调用，但是对于时间不长和时间计算点进行了一点点处理，确保积分达到$t_f$，实现了对ODE的求解。

问题的解

定高释放石头飞行力学

对于上面的飞行力学问题，$\vec{y}=[x,v_x,y,v_y{]}^T$，$\vec{f}(t,\vec{y})=[v_x,0,v_y,-g{]}^T$。我们首先来做一个更加简化的问题，假设没有初始速度，就在$x_0,y_0=10$的地方释放这个石头（这也是飞行力学！）。问题就变为：

$$\{\begin{array}{l}\frac{dy}{dt}=v_y\\ \frac{d{v}_y}{dt}=-g\end{array}$$

而$y_0=10,v_{y0}=0$，那么解析解就是：

$$\{\begin{array}{l}y=-\frac{1}{2}g{t}^2+10\\ v_y=-gt\end{array}$$

求解代码

求解的代码很简单。

let func (para: obj) (t: float) (y: float []) =
    let g = para :?> float
    [| y[1]; g |]

let y0 = [| 10.0; 0.0 |]
let t0 = 0.0
let dt = 0.0001

let t1 = 1.0000001

首先定义了被积函数，然后定义了初始条件，最后定义了积分区间。然后是调用rk4函数进行求解。注意，这里把重力加速度作为参数传递给了被积函数。

let t, yy = rk4 (t0, t1) dt y0 -9.8 func

最后是数据处理，绘制图形，绘图用的ScottPlot，前面已经介绍过。

                     
let n = yy.Length - 1


let get_y (y: float array array) i = (array2D [for j in 0 .. y.Length-1 -> y.[j]]).[*, i]

let yTherory t =
    0.5 * -9.8 .* t .* t .+ 10.0

let vTherory t =
    -9.8 .* t
    

let plt = Plot(600, 400)
plt.AddScatter(t, yTherory t, label = "Theoretical result")
plt.AddScatter(t, get_y yy 0, label = "Numerical result")
plt.XAxis.Label("t (s)")
plt.YAxis.Label("y (m)")
plt.Legend()
plt.Title("Position")
plt.SaveFig("fall_y.png", scale = 2.0)

plt.Clear()
plt.AddScatter(t, vTherory t, label = "Theoretical result")
plt.AddScatter(t, get_y yy 1, label = "Numerical result")
plt.XAxis.Label("t (s)")
plt.YAxis.Label("v (m/s)")
plt.Title("Velocity")
plt.SaveFig("fall_v.png", scale = 2.0)

结果对比

我们来看看数值解和解析解的对比：

一致性还是很强的，当然这里步长取得非常小，所以误差也很小。当我们把步长放到0.2的时候，还是能得到很好的结果。

总结

1. 完美解决了真空一维平面大地的定高释放石头飞行力学问题；
2. 学习了F#的运算符重载机制；
3. 学习了F#复杂的类型测试match语法，注意括号的使用；
4. F#进行数值计算果然还算是好用。