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

F#还能干点啥

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

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

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

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

在这里插入图片描述

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

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

拆成一阶微分方程组就是：
$\begin{cases} \frac{dx}{dt} = v_x \\ \frac{dv_x}{dt} = 0 \\ \frac{dy}{dt} = v_y \\ \frac{dv_y}{dt} = -g \end{cases}$

其中 $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{cases} x = v_{x0}t + x_0 \\ v_x = v_{x0} \\ y = -\frac{1}{2}gt^2 + v_{y0}t + y_0 \\ v_y = -gt + v_{y0} \end{cases}$

好无聊啊……这么简单的结果我不能接受！我要用更加复杂的办法，并且用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{cases} \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{cases}$

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{cases} \frac{dy}{dt} = v_y \\ \frac{dv_y}{dt} = -g \end{cases}$

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

$\begin{cases} y = -\frac{1}{2}gt^2 + 10 \\ v_y = -gt \end{cases}$

求解代码

求解的代码很简单。

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)