🐥
ブラウザ上でヌルヌル動作する流体シミュレーションを Rust + wasm-bindgen-rayon で実装する

2024/12/04に公開
Rust
この記事は Rust Advent Calender 2024 の 4 日目の記事です．
Rust と wasm-bindgen-rayon を使って，ブラウザ上で動作する流体シミュレーションを実装しました．マウスで流体をグリグリいじれるようになっているので，遊んでみてください．
Demo（キャンバスの描画に数秒かかるかもしれません．動作には SharedArrayBuffer に対応しているブラウザが必要です）:

https://fluid-simulation-test.netlify.app/
Repository :

https://github.com/matsuoka-601/wasm-fluid-simulation
自分の環境では，12 スレッドを使って 10,000 個の粒子を 60FPS で処理できているようです．粒子の個数を変えられるようになっているので，パラメータをいじりながらパフォーマンスの変化を見てみるのも面白いと思います．

 作った動機この流体シミュレーションを作ったきっかけは，Sebastian Lague の以下の動画を見たことです．
https://www.youtube.com/watch?v=rSKMYc1CQHE&t=4s
この動画では，粒子法に基づく流体シミュレーションを 1 から実装していき，最終的には素晴らしいクオリティーのシミュレーションを実現しています．
ただ，上記のシミュレーションは Unity 上で動くように作られているため，手元で手軽に実行するにはややハードルが高いです．ブラウザ上で手軽に動かせるシミュレーションがあれば，もっと気軽に試せるはずです．そこで，似たような流体シミュレーションを Rust で実装して WASM にコンパイルし，ブラウザ上で動かせるようにしたのが今回の成果物となります．
シミュレーションを実装するにあたって，粒子の数がある程度多くなってもヌルヌル動くようにしたいです．そこで，rayon クレートを用いてシミュレーションを並列化することにしました．

 流体シミュレーションの概要この記事では，流体シミュレーションの物理的側面に深く立ち入ることはしませんが，今回実装したソルバがどういう計算をしているのかざっくり把握できるように，シミュレーションの概要を簡単に説明します．
今回作った流体シミュレーションは，Smoothed Particle Hydrodynamics（SPH, SPH 法）という手法に基づくものです．SPH 法では，流体を有限個の粒子の集団とみなして，それらの粒子の動きを通じて流体の物理的挙動をシミュレーションします．また，それらの各粒子は質量，圧力，密度といった物理量を持ちます．シミュレーションの各タイムステップでは，各粒子について，近傍粒子の物理量を適当に重みづけして足し合わせていくことで，次のタイムステップのための物理量を計算します．
SPH 法は，以下の 3 ステップからなります．
各粒子の圧力・密度の計算
各粒子の力の計算
各粒子の速度の更新
各ステップについて，対応する Rust のコードも交えながら簡単に説明していきます．なお，ソルバの処理内容を把握するための最低限の説明しかしないので，SPH 法自体についてもっと詳しく知りたい場合には末尾の 参考文献 のところを参照してください．

 密度・圧力の計算SPH 法では，各粒子の密度を，その粒子から一定距離 R 以内にある粒子の質量を適当に重みづけして足し合わせることで求めます．そして，その重みづけは，「遠いほど影響が弱くなる」ように行われます．具体的には，|\bm{x}| が大きくなるにつれ減衰する関数 W(\bm{x}) を用いて，粒子 i に対する粒子 j の「影響度合い」を W(\bm{r}_i-\bm{r}_j) と表します（\bm{r}_i,\bm{r}_j はそれぞれ粒子 i,j の位置ベクトル）．これに粒子 j の質量 m_j を掛けたものが，粒子 i の密度における粒子 j の寄与になります．
まとめると，粒子 i の密度 \rho_i は，以下のように計算されます．

    \rho_i = \sum_{j\in N(i)}m_j\times W(\bm{r}_i-\bm{r}_j)
ここで，N(i) は粒子 i から距離 R 以内にある粒子の集合であり，m_j は粒子 j の質量です．なお，W(\bm{x}) はカーネル関数と呼ばれ，R はカーネル半径と呼ばれます．具体的にどういうカーネル関数を用いるかについては，Sebastian の動画や参考文献の論文を参照してください．
SPH 法で粒子にかかる圧力を計算する方法はいくつか提案されていますが，今回は粒子 i の圧力 p_i を，密度 \rho_i を用いて以下のように定義することにします．

    p_i=k(\rho_i-\rho^\prime)
ここで，k は剛性を表す係数であり，\rho^\prime は基準密度です．粒子たちは常に基準密度 \rho^\prime を目指すようにふるまうことになります．
以上の密度・圧力の計算を Rust のコードに落とし込むと次のようになります．なお，今回は粒子によらず質量が一定であるとし，その質量はコード中でMASSと表します．
particles.iter_mut().enumerate().for_each(|(i, particle)|{
    let pi = &particles_copy[i];
    particle.density = 0.0;
    particles_copy.iter().for_each(|pj|{
        let r = (pi.position - pj.position).length();
        if r < KERNEL_RADIUS {
            particle.density += MASS * DensityKernel(pi.position - pj.position);
        }
    });
    particle.pressure = k * (particle.density - TARGET_DENSITY);
});
DensityKernelがカーネル関数 W(\bm{x}) に相当し，TARGET_DENSITYが基準密度 \rho^\prime に相当します．また，for_eachの中でparticlesを借用することはできないので，particlesのコピーを借用するようにしています．

 力の計算粒子の運動方程式は，以下のナビエ・ストークス方程式により表されます．

    \rho\frac{D\bm{v}}{Dt}=-\nabla p+\mu\nabla^2\bm{v}+\rho\bm{g}
右辺は，それぞれ圧力項，粘性項，外力項であり，\mu は粘性の大きさを表す係数です．また，\bm{g} は重力加速度を表すベクトルです．
圧力項，粘性項の離散化には様々な定式化がありますが，今回用いたのは以下のような式です．i 番目の粒子の圧力項を \bm{f}_i^\mathrm{pressure}，粘性項を \bm{f}_i^\mathrm{viscosity} とします．

    \begin{aligned}
    \bm{f}_i^\mathrm{pressure}&=-\nabla p(\bm{r}_i)=-\sum_{j\in N(i)\setminus\{i\}}m_j\frac{p_i+p_j}{2\rho_j}\nabla W(\bm{r}_i-\bm{r}_j)\\
    \bm{f}_i^\mathrm{viscosity}&=\mu\nabla^2\bm{v}(\bm{r}_i)=\mu\sum_{j\in N(i)\setminus\{i\}}m_j\frac{\bm{v}_j-\bm{v}_i}{\rho_j}\nabla^2 W(\bm{r}_i-\bm{r}_j)
    \end{aligned}
つまり，各粒子の力の計算は以下のようなコードで表されます．圧力項がfpress，粘性項がfviscです．
particles.iter_mut().enumerate().for_each(|(i, particle)|{
    let pi = &particles_copy[i];
    let mut fpress = Vec2::new(0.0, 0.0);
    let mut fvisc = Vec2::new(0.0, 0.0);

    particles_copy.iter().enumerate().for_each(|(j, pj)|{
        if i != j { 
            let r = (pi.position - pj.position).length();
            if r < KERNEL_RADIUS { 
                let shared_pressure = (pi.pressure + pj.pressure) / 2;
                fpress += -MASS * shared_pressure * PressureKernelGradient(pi.position - pj.position) / pj.density;
                fvisc += MASS * (pj.velocity - pi.velocity) * ViscosityKernelLaplacian(pi.position - pj.position) / pj.density
            }
        }
    });

    fvisc *= VISCOSITY;
    let fgrv = pi.density * GRAVITY;
    particle.force = fpress + fvisc + fgrv;
});
式ではカーネル関数をすべて W(\bm{x}) と表しましたが，実際には，密度・圧力・粘性の計算で異なるカーネル関数を使い分けることが多いようです．そのため，コード中でもDensityKernel, PressureKernelGradient, ViscosityKernelLaplacianのようにカーネル関数の使い分けをしています．

 速度・位置の更新力が求まったので，それを密度 \rho で割ると加速度 \frac{D\bm{v}}{Dt} が求まります．今回は，加速度から位置や速度を計算するのには単純なオイラー法を用いています．つまり，時刻 t における加速度，速度，位置をそれぞれ \bm{a}_t, \bm{v}_t, \bm{x}_t とすると，

    \begin{aligned}
        \bm{v}_{t} &= \bm{v}_{t-1} + \Delta t\times\bm{a}_{t}\\
        \bm{x}_{t} &= \bm{x}_{t-1} + \Delta t\times\bm{v}_{t}
    \end{aligned}
という式に従って速度や位置を更新していきます．

 近傍探索先ほど示したコードでは，各粒子の近傍粒子をナイーブに全探索していましたが，これでは粒子数が多くなったときにパフォーマンスが急激に悪化してしまいます．それを避けるために非常に有効なのが，平面を格子状に分割し，探索空間を限定することです．以下のように，平面全体を，カーネル半径 R を一辺とする正方形のマスで分割してみましょう．
すると，各粒子について，その粒子から距離 R 以内にある粒子は，必ず周辺の 3x3 = 9 マスの中に収まっています（図で赤く塗った示したマスです）．ゆえに，先ほどは近傍粒子の候補としてすべての粒子を探索したのに対し，平面を格子状に分割すると周辺 9 マスの中にある粒子のみを探索すればよいことになります．この分割による高速化の効果は絶大で，ナイーブな全探索と比べると数十倍も実行時間が違ってくることもあります．
今回は，各マスについて，そのマスに含まれる粒子のインデックスをVec<u32>として持つように実装しています．格子全体はVec<Vec<u32>>のような 2 次元ベクタとして保持します．

 並列化今回は，シミュレーションの並列化のために rayon クレートを使いました．rayon は，Rust における主要な並列処理ライブラリの一つであり，イテレータを少し書き換えるだけで簡単にマルチスレッドによる並列化ができます．
SPH 法は，各粒子の物理量の計算を独立に行うため，割と素直に並列化できます．たとえば，密度や力を計算するコードであれば，以下のようにiter_mut()をpar_iter_mut()に変えるだけで並列化ができます．
particles.par_iter_mut().enumerate().for_each(|(i, particle)|{
    ...
});

 wasm-bindgen と wasm-bindgen-rayon についてここでは，Rust と JavaScript の相互運用に用いた wasm-bindgen と、WASM でマルチスレッドによる並列処理をするために用いた wasm-bindgen-rayon について説明します．

 wasm-bindgenRust コードをコンパイルしてできた WASM モジュールをブラウザ上で実際に動かすためには，WASM モジュールと JavaScript との間でのやりとりが必要になります．その橋渡しをよしなにやってくれるのが，wasm-bindgen です．wasm-bindgen を使うと，Rust で記述した関数を JavaScript 側から呼び出したり，Rust で定義した構造体を JavaScript 側で使うことができるようになります．たとえば，今回の流体シミュレーションでは，以下のように#[wasm_bindgen]をつけてstart関数を宣言しています．
#[wasm_bindgen]
pub fn start() -> Result<(), JsValue> {
    // シミュレーションを開始するコード
}
これにより，JavaScript 側から上記のstart関数を呼び出せるよう，WASM モジュールと JavaScript との橋渡しをするコードを wasm-bindgen が生成してくれます．JavaScript からは，そのコードを import するだけでstart関数を使えます．
なお，wasm-bindgen は，実際にはwasm-packというツールを通して間接的に呼ばれることが標準的です（今回のケースでもそうです）．ここでは詳しい説明を省きますが，wasm-bindgen や wasm-pack など，Rust ⇒ WASM の変換に関するツールの関係性については MDN のチュートリアルが詳しいです．

 wasm-bindgen-rayonWASM でマルチスレッドによる処理をする場合，現状では WebAssembly threads を使うことになります．Rust で書いた並列処理のコードから WebAssembly threads を使う方法はいくつかあるのですが，wasm-bindgen-rayon を使うのが現時点で最もポピュラーな選択肢だと思います。wasm-bindgen-rayon を使うと，rayon を使って書いた並列処理のコードをそのままブラウザ上でも使えるので便利です．
基本的には，公式レポジトリ の README に沿ってセッティングしていけば特に問題なく使えます．ただ，執筆時点の wasm-bindgen の最新版である0.2.95に問題があり，workerHelpers.worker.jsが正しく生成されないという問題があるようです（関連する issue）．0.2.93の wasm-bindgen だとこの問題は起きないので，問題が修正されるまでは0.2.93の wasm-bindgen を使っておくのが安全そうです．
また，セキュリティ上の理由により，WebAssembly threads をブラウザ上で使うためには COEP ヘッダーや COOP ヘッダーの設定を適切に行う ことが必要になります．

 実際に動かしてみる & ベンチマークシングルスレッドの場合と 12 スレッドを使った場合で，10,000 粒子をシミュレーションするパフォーマンスの比較をしてみました．なお，以下のシミュレーションは AMD Ryzen 7 5825U 上で動作させています．
シングルスレッド

12 スレッド

60FPS でシミュレーションを動作させたい場合，1 フレームの処理にかかる時間が 1000/60 = 16.6ms を下回るようにすることが目標になります．この時間を計測したところ，シングルスレッドの場合は約 90ms 程度だったのに対し，12 スレッドを使った場合には約 14ms 程度でした．上記の画像を見ると，その差がシミュレーションのスムーズさにもたらす影響がよくわかると思います．実装した肌感覚として，SPH 法による流体シミュレーションでは，並列化なしでこのパフォーマンスを出すのはかなり厳しそうです．

 まとめ & 所感以下がまとめです．
ブラウザ上で動作する流体シミュレーションを Rust で実装した．
wasm-bindgen-rayon 経由で WebAssembly threads を用いることにより，マルチスレッドによる並列化を行った．
シミュレーションには SPH 法を用いた．
12 スレッドを用いることで，10,000 個の粒子を 60FPS で処理できている．
今回作ったシミュレーションのたたき台は，Rust ではなく JavaScript を使って実装したのですが，型に関するエラーを見逃しやすく，バグの原因を特定するのに結構な時間を費やしました．また，演算子オーバーロードがないため，ベクトルの計算にいちいちメソッドを呼ばなければならないのも，地味にきつかったです．一方，思い切って Rust に切り替えた後は，ガチガチの型チェックのもと，ソルバのロジック部分に集中することができて快適でした．また，Rust を WASM に変換し，ブラウザ上で動かすためのエコシステム（wasm-pack や wasm-bindgen まわり）がかなり整備されてきているのも良かったです．
今後は，今回作ったシミュレーションを三次元に拡張しようと目論んでいます．その際，計算量が二次元の場合よりも増大することは避けられないので，WebGPU という API を用いて GPGPU の並列処理能力を借りようと思っています．実は，すでにたたき台はできており，近傍粒子を全探索する場合でも 20,000 粒子をそこそこの速度で処理できているようです．GPGPU すごい．


WebGPU を用いた 20,000 粒子のシミュレーション
これから実装する，空間の格子分割による高速化により，どれくらいの規模の粒子を扱えるようになるのかが楽しみなところです．

 参考文献
Coding Adventure: Simulating Fluids
今回のシミュレーションを作るモチベーションとなった，Sebastian Lague の動画です．シミュレーションの各ステップについて，実装に落とし込める粒度で詳しく説明されています．


Particle-Based Fluid Simulation for Interactive Applications, Müller et al. 2003
Sebastian Lague の動画では，この論文に基づいて SPH 法の実装を行っています．非常に有名な論文です．


Particle-based Viscoelastic Fluid Simulation, Clavet et al. 2005
記事では説明しませんでしたが，今回の流体シミュレーションでは，この論文で提案された near-density と near-pressure というテクニックを実装しています（Sebastian Lague の動画内でも言及があります）．このテクニックは表面張力をシミュレーションするのに効果的で，実装すると液体っぽさがグッと増します．

Lucas V. Schuermann という方の，SPH 法に関する記事や実装を大いに参考にしました．

Implementing SPH in 2D （SPH 法に基づくシンプルなソルバを実装する記事）

mueller-sph （Müller の論文に基づく流体シミュレーションの簡潔な実装）


Rust+WebAssemblyでMDNのWebGLチュートリアルを写経する
今回のシミュレーションでは，Rust から WebGL の API を呼んで描画を行っています．その実装にあたって，この記事を参考にしました．
作った動機

流体シミュレーションの概要

密度・圧力の計算

力の計算

速度・位置の更新

近傍探索

並列化

wasm-bindgen と wasm-bindgen-rayon について

wasm-bindgen

wasm-bindgen-rayon

実際に動かしてみる & ベンチマーク

まとめ & 所感

参考文献

Discussion
