❄️

RustとBevyで雪の結晶のシミュレーションをする

necocen

2024/12/18に公開

静けさに抱かれながら　また今日も待っている
ゆるやかに降る　水じゃなくてもっと寂しい粒 ^[1]

イントロダクション

こんにちは。冬ですね。
サービス開発パートナー事業部　バックエンドエンジニアの@necocenです。

冬といえば雪、雪といえばきれいな結晶です。
今回は、雪の結晶（雪片）の成長シミュレーションモデルをRustで実装して、Rayonで並列化を行い、さらにゲームエンジンBevyを使って可視化し、Webブラウザで表示できるようにした話をします。特に仕事とは関係がありません^[2]。

この記事でわかること

Rustでちょっとした数値シミュレーションを書く一つの方法
BevyとRayonの組み合わせをWebブラウザで動作させるための工夫
- Bevy単独なら簡単ですがここはちょっと工夫がいりますという話
雪片の成長の（ある程度物理的にリアリティのある）モデル化

もし上記のどれにも興味がなければ、記事の最後らへんに生成結果を何枚か載せているので、それを見てクリスマス気分を楽しんでいってください。

実装したもの

コードは以下のGitHubリポジトリにあります。
この記事で説明しなかった要素（ファイルへの書き出しなど）もありますが、大筋では同じものです。

https://github.com/necocen/snowflake

ブラウザで動くデモサイトもあります。

https://yuki.necocen.info/

論文について

今回実装した論文は、
Gravner, J., Griffeath, D. (2008). Modeling snow crystal growth II: A mesoscopic lattice map with plausible dynamics. Physica D: Nonlinear Phenomena, 237(3), 385-404.
です。この論文では、-20℃〜-10℃における雪片の形成過程を三角格子の上でモデル化し、かなりそれっぽい再現に成功しています。

論文から引用した、さまざまなシミュレーション結果を示す図

モデルでは、核となる一点を初期状態として周囲に氷が形成され、周囲の水蒸気が拡散・凝縮することで、六角形状の雪片パターンが形成されていきます。格子上に「氷」「水蒸気」「疑似液体層」の質量を表現し、それらが拡散→凍結→付着→融解のステップを繰り返しながら成長していくことで、複雑に枝分かれした雪片の形状が生まれていきます。
詳しいアルゴリズムについては記事末尾の補遺でもうすこし説明したので、興味のある方は読んでください。

実装の概略

実際の実装では、三角格子を斜交座標系で表現し、端では周期境界条件（あっちの端から飛び出したらこっちの端から出てくる）をとりました。

計算を効率的に行えるとうれしいので、ndarrayというクレートを使うことにしました。これはPythonでよく使われるNumPyのRust版みたいなもので、筆者はNumPyがすこしならわかることもあって採用しました^[3]^[4]。

まず状態を定義します。Array2<T>はT型の2次元配列を表現するndarrayの構造体です。

gravner-griffeath.rs

struct State {
    /// セルが雪片に含まれていればtrue、いなければfalse
    a: Array2<bool>,
    /// 氷の表面にある疑似液体層（quasi-liquid layer）の質量
    b: Array2<f32>,
    /// 氷の質量
    c: Array2<f32>,
    /// 水蒸気の質量
    d: Array2<f32>,
}

初期状態では、核となる中心の一点が凍結しています。引数nが格子全体のサイズ、rhoが初期状態で周囲に漂っている水蒸気の質量（ $\rho$ ）です。

gravner-griffeath.rs

impl State {
    fn new(n: usize, rho: f32) -> Self {
        let center = [n / 2, n / 2];
        let mut a = Array2::<bool>::default((n, n));
        a[center] = true;
        let b = Array2::<f32>::zeros((n, n));
        let mut c = Array2::<f32>::zeros((n, n));
        c[center] = 1.0;
        let mut d = Array2::<f32>::ones((n, n)) * rho;
        d[center] = 0.0;
        Self { a, b, c, d }
    }
}

更新処理の詳細は割愛しますが、上で述べたような拡散→凍結→付着→融解のステップごとに「計算に必要なArray2をZipでまとめてfor_eachでセルごとに計算する」という形の処理をしています。たとえば凍結ステップはこんな風になります^[5]。

gravner-griffeath.rs

// (ii) Freezing
let mut b_new = self.b.clone();
let mut c_new = self.c.clone();
Zip::from(&self.a)
    .and(&mut b_new)
    .and(&mut c_new)
    .and(&mut d_new)
    .and(&neighbors)
    .for_each(|&a, b, c, d, &neighbors| {
        // 雪片に隣接するセルの
        if !a && neighbors > 0 {
            // 水蒸気が凍結して疑似液体層と氷になる
            *b += (1.0 - kappa) * *d;
            *c += kappa * *d;
            *d = 0.0;
        }
    });

Rayonによる並列化

Rustで並列計算をしたい場合によく使われるのはRayonクレートです。そしてndarrayクレートはRayonと連携して計算をスレッドで並列化することができるので、やってみます。

まずCargo.tomlでndarrayを追加している部分でrayonフィーチャを有効にします。

Cargo.toml

ndarray = { version = "0.16.1", features = ["rayon"] }

次にndarrayの計算で.for_eachを呼んでいる部分を.par_for_eachにします。

gravner-griffeath.rs

 Zip::from(&self.a)
     .and(&mut b_new)
     .and(&mut c_new)
     .and(&mut d_new)
     .and(&neighbors)
-    .for_each(|&a, b, c, d, &neighbors| {
+    .par_for_each(|&a, b, c, d, &neighbors| {
         if !a && neighbors > 0 {
             *b += (1.0 - kappa) * *d;
             *c += kappa * *d;
             *d = 0.0;
         }
     });

これで並列化されました。手許のマシン（M1 MBP、10コア）で速度を比較してみます。

比較に使ったコード

1000×1000格子で1000ステップまで実行させるコードを書きました。

let mut state = State::new(1000, 0.5);
let config = SimulationConfig::default();
for _ in 0..1000 {
    state.update(config);
}

❯ hyperfine --warmup 3 target/release/snowflake-rs target/release/snowflake-rs-parallel
Benchmark 1: target/release/snowflake-rs
  Time (mean ± σ):     11.865 s ±  0.031 s    [User: 10.763 s, System: 1.083 s]
  Range (min … max):   11.810 s … 11.901 s    10 runs

Benchmark 2: target/release/snowflake-rs-parallel
  Time (mean ± σ):      3.774 s ±  0.044 s    [User: 17.926 s, System: 5.778 s]
  Range (min … max):    3.726 s …  3.850 s    10 runs

Summary
  target/release/snowflake-rs-parallel ran
    3.14 ± 0.04 times faster than target/release/snowflake-rs

数ヶ所書き替えるだけで3倍程度の高速化になりました。もっと真面目にチューニングすればもっと速くなるかもしれません。

Bevyによる可視化

計算だけして表示がないとつまらないので、簡易的な可視化を行います。ここではゲームエンジンBevyを利用しました。

BevyはECS（Entity-Component-System）アーキテクチャを採用しており、エンティティにコンポーネントを付与し、システム（関数）でそれらを更新していく流れでゲームを構築します。今回は複雑なゲームロジックは不要なので、「シミュレーション状態を読み取って、対応する描画メッシュやマテリアルを更新する」だけの使いかたをしています。

Bevyでは相互に関連するシステムやコンポーネントをプラグインとして分離することができるので、今回はシミュレーション部分と可視化部分をそれぞれプラグインとして切り出しました。

main.rs

fn main() {
    App::new()
        // シミュレーション結果やステップ数などアルゴリズムに依存しない状態を保持
        .init_resource::<Field>()
        // 一時停止ボタン操作などを通知するイベントを宣言
        .add_event::<ControlEvent>()
        // デフォルトプラグイン群とUIプラグイン
        .add_plugins((DefaultPlugins, EguiPlugin))
        // シミュレーションを実行するプラグイン
        .add_plugins(gravner_griffeath::GravnerGriffeathSimulatorPlugin)
        // シミュレーション結果を可視化するプラグイン
        .add_plugins(visualization::VisualizationPlugin)
        // アプリケーションの初期化処理
        .add_systems(Startup, (start_simulation, set_window_title))
        // 共通UI（ステップ数表示や一時停止ボタンなど）の更新
        .add_systems(Update, configure_ui)
        .run();
}

こうすることでシミュレーション部分はシミュレーションアルゴリズムとその設定UIに集中できるし、別のアルゴリズムへの差し替え^[6]も容易になります。

プラグインの定義もシンプルで、Pluginトレイトを実装して、build関数の中でプラグインが使用するリソースや実行するシステムを宣言することでプラグインになります。シミュレーションプラグインの定義は以下のようになります。

gravner-griffeath.rs

pub struct GravnerGriffeathSimulatorPlugin;

impl Plugin for GravnerGriffeathSimulatorPlugin {
    fn build(&self, app: &mut App) {
        app.init_resource::<SimulationConfig>();
        app.init_resource::<State>();
        app.add_systems(Main, update_simulation);
        app.add_systems(Update, (event_listener, configure_ui));
    }
}

update_simulationの中でシミュレーションの状態更新を行います。毎フレームごとに描画まですると効率が悪いので、1フレームごとに5回更新されるようにしておきます^[7]。

実行中に各種パラメータを変更できるUIがあるとうれしいので、bevy_eguiでスライダーのUIをつけています。configure_uiではその定義をしています。

シミュレーションパラメータを変更するスライダー群のUI

描画部分も同様にプラグイン化します。ゴチャゴチャやっているので詳細の説明は割愛しますが、

セルの値（＝氷の質量）を0〜1に正規化（雪片内の最大値を1とする（雑な）計算）してセルのアルファ値とする
- 質量が大きいセルほど不透明になるようにしたい
六角形メッシュの座標と透明度を変えながら各セルを描画する
- GPUインスタンシングで効率化するためにシェーダを書いたりレンダリングパイプラインをカスタムしたり

という風になっています。

実行中の表示内容は以下のような感じです。ゆっくりですが、新たな枝分かれが発生して成長し始める様子がわかると思います。

Bevyでシミュレータが動作している様子

Webブラウザでの実行

Bevyで可視化することの嬉しさの一つに、Webブラウザでの実行に対応しているということがあります。せっかくなので試してみたいですよね。
しかし、今回はRayonを使っているので、ブラウザで実行するためにはすこし工夫が要ります。また、Rustはnightly版を使う必要があります^[8]。

まず、WASMを生成するための設定と、wasm-bindgenとwasm-bindgen-rayonをCargo.tomlに追加します^[9]。

Cargo.toml

[lib]
crate-type = ["cdylib", "rlib"]

[target.'cfg(target_arch = "wasm32")'.dependencies]
wasm-bindgen = "0.2.99"
wasm-bindgen-rayon = { version = "1.2.2", features = ["no-bundler"] }

次に、これまでmain.rsに書いていた諸々をlib.rsに移します。fn main()に書いていたものをpub fn run()に移し、fn main()からはrun()を呼ぶようにします^[10]。run()はJavaScriptから呼び出せるようにwasm_bindgen属性をつけます。

lib.rs

#[cfg_attr(target_arch = "wasm32", wasm_bindgen)]
pub fn run() {
    App::new()
        .init_resource::<Field>()
        .add_event::<ControlEvent>()
        .add_plugins((DefaultPlugins, EguiPlugin))
        .add_plugins(gravner_griffeath::GravnerGriffeathSimulatorPlugin)
        .add_plugins(visualization::VisualizationPlugin)
        .add_systems(Startup, (start_simulation, set_window_title))
        .add_systems(Update, configure_ui)
        .run();
}

wasm-bindgen-rayonのドキュメントに従ってinit_thread_pool関数をre-exportします。

lib.rs

#[cfg(target_arch = "wasm32")]
pub use wasm_bindgen_rayon::init_thread_pool;

各種ビルドオプションを.cargo/config.tomlに追記します。

.cargo/config.toml

[target.wasm32-unknown-unknown]
rustflags = ["--cfg=web_sys_unstable_apis", "-C", "target-feature=+atomics,+bulk-memory,+mutable-globals"]

[unstable]
build-std = ["panic_abort", "std"]

ここまで準備できたら、wasm-packを使ってビルドします。上述のようにRustはnightly版を使う必要があることに注意してください。

rustup run nightly wasm-pack build --target web --release

これで./pkg配下に諸々が出力されるので、これを呼び出すHTMLファイルをクレートのルートに作成します。

index.html

<!DOCTYPE html>
<html lang="ja">
    <head>
        <meta charset="UTF-8">
        <meta http-equiv="X-UA-Compatible" content="IE=edge">
        <meta name="viewport" content="width=device-width, initial-scale=1.0">
        <title>Snowflake Simulator</title>
    </head>
    <body>
        <script type="module">
            import init, { initThreadPool, run } from "./pkg/snowflake_rs.js";
            try {
                await init();
                console.log("WASM initialized");
                // スレッドプールを初期化する
                await initThreadPool(navigator.hardwareConcurrency);
                console.log("Web Workers initialized");
                // 起動
                await run();
            } catch (error) {
                if (!error.message.startsWith("Using exceptions for control flow, don't mind me. This isn't actually an error!")) {
                    throw error;
                }
            }
        </script>
    </body>
</html>

あとはこのHTMLを配信すればブラウザで動きます。

ただし、スレッド間のデータ共有でSharedArrayBufferを使うため、HTTPヘッダーにCross-Origin-Opener-Policy='same-origin'とCross-Origin-Embedder-Policy='require-corp'の設定が必要になります^[11]。

筆者が手許で確認する時はdevserverの起動オプションで設定しました。

devserver --header Cross-Origin-Opener-Policy='same-origin' --header Cross-Origin-Embedder-Policy='require-corp' --noreload

あとはhttp://localhost:8080にアクセスすればブラウザ版が動作するのを確認できます^[12]。
Firefoxで開いたWebページ上でシミュレータが動作している様子

デモサイトはCloudflare Pagesで公開しています^[13]。WASMファイルのサイズが、公開できるファイルの最大サイズ（25MiB）にギリギリ収まったので公開できました。

ギャラリー

記事のおわりに、いろいろなパラメータで生成した雪片の画像を何枚か貼ります^[14]。

実際の雪片は、成長途中で空のさまざまな場所を移動し、温度や水蒸気量などの変化を経験しているものと考えられます。
このシミュレータも、実行中にパラメータを変化させることができます。パラメータの変化に合わせて成長の様子が変わるのがわかると思います^[15]^[16]。
実行中にパラメータを変化させる様子

きれいですね。

まとめ

この記事では、

雪片の成長シミュレーションモデルを提案する論文からアルゴリズムをRustで実装し、
Rayonによって並列化し、
Bevyによって可視化を行い、
Bevy+Rayonの構成をWebブラウザで動作させるための工夫について説明し、
いくつかの出力例を示すことでクリスマスの雰囲気を盛り上げました。

論文で提案されている具体的なアルゴリズムについては、補遺で簡単に説明しています。

それでは、よいお年をお迎えください。

謝辞

暇つぶしにちょうどいい題材を提供してくれ、ギャラリーの作成を手伝ってくれた妻に感謝します。

補遺：雪片成長モデル

論文で提案されているアルゴリズムの概要をとっても駆け足で説明します。詳しく正確に知りたい場合は論文を読んでください。

まず以下のような初期化を行います。

三角格子 $\mathbb{T}$ 上の変数 $a$ , $b$ , $c$ , $d$ を用意する^[17]
- $a(x)$ : セル（格子点） $x\in\mathbb{T}$ が凍結している（雪片の一部になっている）かどうか（真偽値）
- $b(x)$ : セル $x$ に含まれる疑似液体層（quasi-liquid layer）の質量（実数）
- $c(x)$ : セル $x$ に含まれる氷の質量（実数）
- $d(x)$ : セル $x$ に含まれる水蒸気の質量（実数）
格子の中心のセル（ $x=\mathbf{0}$ ）だけを凍結状態（ $a(\mathbf{0})=\mathrm{True}, b(\mathbf{0})=0, c(\mathbf{0})=1, d(\mathbf{0})=0$ ）とし、残りの全セル（ $x\neq\mathbf{0}$ ）を質量 $\rho$ の水蒸気で満たす（ $a(x)=\mathrm{False}, b(x)=c(x)=0, d(x)=\rho$ ）
- これが雪片の核になる

あとは以下を心ゆくまで順番に繰り返します。

拡散：凍っていないすべてのセルで、隣接セル（6つ）との間で水蒸気を交換させる
凍結：凍結セルに隣接する非凍結セル（境界セル）で、水蒸気を疑似液体層と氷に分配する

\begin{aligned} b(x) &\leftarrow b(x)\ +&(1-\kappa)d(x) \\ c(x) &\leftarrow c(x)\ +&\kappa d(x) \\ d(x) &\leftarrow 0 \end{aligned}

付着： $b(x)$ が閾値を超えた境界セルは凍結したものとして雪片に含め（ $a(x)\leftarrow\mathrm{True}$ ）、その場合は疑似液体層を氷にする（ $c(x) \leftarrow b(x) + c(x),\,b(x) \leftarrow 0$ ）
融解：境界セルの疑似液体層と氷の一部を水蒸気に戻す

\begin{aligned} b(x)&\leftarrow(1-\mu)b(t)\\ c(x)&\leftarrow(1-\gamma)c(t)\\ d(x)&\leftarrow d(x)+\mu b(x)+\gamma c(x) \end{aligned}

付着ステップで、疑似液体層が凍るかどうかを判定する閾値は、セルの周囲に凍結セルがいくつあるかによって変えます。

隣接する凍結セルが2個以下の場合： $b(x)\ge \beta$ のとき凍結
3個の場合： $b(x)\ge 1$ または $\left(\sum_{y:\,\mathrm{neighbors\ of}\ x} d(y) \lt\theta\ \mathrm{and}\ b(x) \ge \alpha \right)$
4個以上の場合：かならず凍結

この部分のパラメータ（ $\alpha, \beta, \theta$ ）は、境界の形状（凹であるか凸であるか）ごとに雪片の成長しやすさに差があることを表現しています^[18]。これが雪片の複雑な枝分かれパターンに影響を及ぼしています。

脚注

『雪、無音、窓辺にて。』長門有希（CV.茅原実里）（作詞：畑亜貴） ↩︎
この記事の内容は仕事とは関係がありませんが、バックエンドチームとしてはなんとかバックエンドをRustで書けないか画策しています。それかScala。 ↩︎
最初のプロトタイプをPythonで作っていた名残でもあります。 ↩︎
とはいったものの、言うほどNumPyっぽくはない気もしています。 ↩︎
基本はセル単位の計算なのでこの例ではきれいに書けているのですが、他のステップでは近傍セルも見に行く必要があるので添字によるアクセスがけっこう発生しています。いろいろ工夫すればもうすこし気持ちよく書けるかもしれませんが、試していません。 ↩︎
実は一番最初は別の論文を実装していました（参考: https://doi.org/10.1016/j.chaos.2004.06.071 ）。 ↩︎
シミュレーションは独立したスレッドで動作して、描画のタイミングで最新の結果だけを表示するようにするのが一番効率がよいと思われますが、あとでWASMに対応するときに困難が生じるのでその方針はとりません。正確には、バックグラウンドスレッドを立ち上げること自体はwasm_threadを使えば可能になりますが、描画スレッドとの間で結果やパラメータをやりとりしようとするとmutexが必要になり、メインスレッドをブロックすることになるのでブラウザでは動かせないようです（参考: https://rustwasm.github.io/docs/wasm-bindgen/examples/raytrace.html ）。 ↩︎
必要に応じてnightly版のインストール（rustup install nightly）と、WASMターゲットの追加（rustup target add wasm32-unknown-unknown）が必要です。 ↩︎
今回はJavaScript側でバンドラーを使わないのでwasm-bindgen-rayonでno-bundlerフィーチャを設定しています。バンドラーを使う場合はこの設定は不要になりますが、その場合（？）workerの起動でコケる問題があるようです（GitHub Issue: https://github.com/RReverser/wasm-bindgen-rayon/issues/16 ）。Issueにある解決方法の一つはwasm-bindgenを0.2.93まで戻すことですが、そうするとbevy_eguiとbevyのバージョンも現在の最新から戻す必要があります。もう一つはworkerHelpers.worker.jsを手動でコピーすることです。 ↩︎
なぜこうするかというと、run()を明示的に呼べる形にすることで、アプリが立ち上がるより前にinitThreadPool()を呼んでRayonのスレッドプールを初期化したいからです。Rayonを使わないのであればmain.rsに書いたままで問題ありません（その場合crate-typeの設定もwasm_bindgen属性も不要ですが、バイナリクレートになるのでwasm-packではビルドできないと思います）。 ↩︎
参考: https://web.dev/articles/coop-coep?hl=ja ↩︎
現在のWASMはアドレス空間が32bitであるため、2GB以上のメモリを確保することはできません。そのため、（ビルドする環境や依存クレートのバージョンによっては）OOMエラーになる場合があります（"unreachable"という例外を投げるので原因に気づくのにだいぶ手間取りましたが……）。そういう場合はシミュレーションする格子のサイズを小さくすれば動作します。 ↩︎
ハマった点があって、Cloudflare Pagesは存在しないファイルに対して404を返さず適当にリダイレクトするので、Bevyがシェーダーのメタデータファイル（存在しない）を読みに行ってうまく読み込めないというエラーを出す問題がありました。仕方がないのでアセット読み込みの設定でメタデータのチェックを無効にして回避しました。 ↩︎
アルゴリズムの性質上、これらの雪片は正確に6回対称性を持つはずですが、部分的に壊れている例もあります。これが実装上のミスなのか計算上の誤差なのかは調べていません。 ↩︎
ただし、空気中の水蒸気量 $\rho$ だけは初期値にのみ反映されます。 ↩︎
$\beta$ を小さくすると枝分かれが起きやすく、大きくするとのっぺりと育つ傾向があります。これは雪片表面の尖っている部分に対して付着が起こるための閾値が $\beta$ によって決まっているためです。 ↩︎
このモデルの嬉しい特徴として、総質量 $\sum_{x\in\mathbb{T}} (b(x)+c(x)+d(x))$ が（保存則を破るノイズ項を入れない限り）保存するようになっています。実際には浮動小数点演算の誤差分だけ（？）変動していきますが、デバッグの参考にできます。 ↩︎
特に $\alpha$ と $\theta$ は、実験的には知られているが理論的には説明できていない効果をモデルに取り込むためのものとされています。 ↩︎

株式会社ガラパゴス（有志）Publication

株式会社ガラパゴス（有志）株式会社ガラパゴスのエンジニアを中心に情報発信をしている Publication です。各記事の内容は個人の意見であり、企業を代表するものではございません。

イントロダクション

この記事でわかること

実装したもの

論文について

実装の概略

Rayonによる並列化

Bevyによる可視化

Webブラウザでの実行

ギャラリー

まとめ

謝辞

補遺： 雪片成長モデル

Discussion

補遺：雪片成長モデル