【UE5】Niagara SimulationStageでNormalMapからHeightMapを復元する(Poisson方程式)

Unreal Engine Advent Calendar 2023 の16日の記事です. Unreal Engine (UE)のカレンダー | Advent Calendar 2023 - Qiita

最近HoudiniでNormalMapからHeightMapを復元するというものを見かけ,
Poisson方程式を解くのであればSimulationStageで実装することができるのではと考えました.

というわけで、@TearsOfJake さんのアドバイスを元に、ポアソン方程式を用いたノーマルマップからハイトマップの復元。昨日の勾配のみの復元に比べて圧倒的に安定している。
左：入力メッシュ
中：入力メッシュから求めたノーマルマップ
右：ノーマルマップから復元した高さ情報 pic.twitter.com/NiNyVvvrbS
— Akira Saito (@a_saito) 2023年11月5日

この記事では「NormalMapからHeightMapを復元する」という問題をSimulationStageで実装する方法を紹介します.
サンプルプロジェクトは以下のGoogleDriveになります (動作確認バージョン UE5.3).
ReconstructHeightMapFromNormal.zip - Google ドライブ

Niagara SimulationStage

Niagara SimulationStage はComputeShader的なGPGPU処理をNiagaraモジュールとして作成できる仕組みです.
Epic Games公式のNiagara Fluids等でも多用されています(本当にSimulationStageでほとんどの処理が作られていますまじかよ).

SimulationStage は2D/3Dのグリッド構造に対するGPGPU処理を比較的簡単に作ることができるように設計されています.
これには数値シミュレーションで必要とされる反復実行なども含まれています.
目的のタスクを適切な形に変形してSimulationStageで実装することで, GPUによる高速な実行が可能になります.

NormalMapからHeightMapを復元する問題

Poisson方程式

ある関数の二階の空間微分と別の関数との関係式をPoisson方程式といいます.
Poisson方程式は離散化によって連立一次方程式へ変形することができ, 連立一次方程式を解く方法は様々に研究されています.
対象の問題をPoisson方程式, ひいては連立一次方程式の形にすることができれば, Solverや解法を利用して解くことができます.

$\Delta \phi =-\dfrac{\rho}{\epsilon_0}$

Poisson方程式への変形

今回は NormalMapとHeightMapの関係をPoisson方程式で表現して離散化, 連立一次方程式として解くことにします.

NormalMapのxy成分はHeightMapの空間微分(勾配)から求められます.
$N_x(x,y)= -\dfrac { \partial H(x,y)}{\partial x}$
$N_y(x,y)= -\dfrac { \partial H(x,y)}{\partial y}$
さらにそれぞれ x, y で偏微分します.

$\dfrac {\partial N_x(x,y)}{\partial x}= -\dfrac { \partial^2 H(x,y)}{\partial x^2}$

$\dfrac {\partial N_y(x,y)}{\partial y}= -\dfrac { \partial^2 H(x,y)}{\partial y^2}$
ここで上の式はPoisson方程式の形になっています.

差分方程式

このPoisson方程式を離散化すると以下のようになります.
(東海大学理学部遠藤研究室さんのページ等がわかりやすいです. -> 電磁場(SP) )

$\frac {N_x(x+h,y) - N_x(x-h,y)}{2h}= -\frac {H(x+h,y)+H(x-h,y) - 2H(x,y)}{h^2}$
$\frac {N_x(x,y+h) - N_x(x,y-h)}{2h}= -\frac {H(x,y+h)+H(x,y-h) - 2H(x,y)}{h^2}$
(左辺は中心差分)

2つの式を足して左辺のNormalMap, 右辺にHeightMapに関わる項をまとめて整理します.
またテクセルでの微分(差分)のため $h = 1$ とします.
$- \frac {N_x(x+1,y) - N_x(x-1,y) + N_x(x,y+1) - N_x(x,y-1)}{2} \\= H(x+1,y)+H(x-1,y) + H(x,y+1)+H(x,y-1) - 4H(x,y)$

上の式の左辺は入力NormalMapとして与えられる既知の値, 右辺は求めたいHeightMapのいくつかの地点の値に係数をかけたものの和です.
よってこれは左辺を定数 $b$ , 未知のHeightMapを $x$ , 近傍HeightMap値への係数を行列 $A$ とした連立一次方程式ということになります.
$b = Ax$

連立一次方程式を解くために更に変形していきます.

Jacobi法

見通しをよくするために左辺は記号で置き換えておきます.
$-N_{d} = H(x+1,y)+H(x-1,y) + H(x,y+1)+H(x,y-1) - 4H(x,y)$
(ここで $N_{d} =\frac {N_x(x+1,y) - N_x(x-1,y) + N_x(x,y+1) - N_x(x,y-1)}{2}$ )

$H(x,y)$ に関する形に変形します.

$H(x,y) = \frac{(H(x+1,y)+H(x-1,y)+H(x,y+1)+H(x,y-1)) + N_{d} }{4}$

この式を反復的に計算することで $H(x,y)$ を解に近づけていく方法をJacobi法といいます.
プログラム的に考えるとHeightMapの各テクセルについて上下左右の近傍値とNormalMapの値を使って更新しているといえます.
現在のHeightMapから次のHeightMapの各テクセルを並列に計算できるためGPGPUとも相性が良いです.
次はJacobi法によるHeightMapの計算をNiagara SimulationStageで実装します.

Niagara SimulationStageでの実装

先に必要となるデータや処理を整理しておきます.

外部データ一覧

外部からNiagaraに設定するデータになります.

入力NormalMap Texture
- HeightMap復元の元になる入力データ
HeightMap出力用RenderTargetTexture
- 結果のHeightMap出力先として Floatフォーマットの適当なサイズで用意しておきます.
プレビュー用マテリアル
- 計算結果のプレビューをするためのマテリアル M_PreviewHeight
- InTex という名前でテクスチャパラメータを引き取って表示するだけのものです.

処理ステージ一覧

これらのステージをSimulationStageとして実装します.

初期化
- 作業用のGrid2DCollectionのクリアなど, メインの処理の準備をします.
- 実行タイミングをリセット時の一度のみに設定することで, フレームを跨いでJacobi反復を継続させます.
Jacobi反復処理
- HeightMapの反復計算の1回分を計算します.
- SimulationStageの反復数指定によって1フレームに複数回実行可能です.
出力
- HeightMapを出力先RenderTargetに書き出します.

ここからNiagaraSystemのEmitterを実装していきます.

EmptyEmitterを追加

雛形のEmitterを追加します.
最低限の設定が含まれたテンプレートの Empty Emitterを使います.
(コンテキストメニューの「空のエミッタを追加」は本当に無なので面倒です. )

Emitterの設定

SimulationStageを使うためにGPUComputeに変更.

プレビューを簡単にするためにSpriteRendererのSourceModeをEmitterに変更.

Spriteを見やすい大きさにするためにEmitterSpawnステージにSpriteSizeの設定を追加.

ユーザーパラメータの作成と既定値

外部から入力NormalMapと出力先RenderTargetを受け取るためのユーザーパラメータです.
Systemユーザーバラメータの + から新規作成>オブジェクトで
Texture と Texture Render Target を一つずつ作成します.
それぞれ InputNormalTex, OutputHeight と名前をつけておきます.

作業中のプレビューをしやすくするためにユーザーパラメータに既定値を設定しておきます.
InputNormalTex には適当なNormalMapアセット.
OutputHeight には準備しておいた出力用RenderTargetTexture.

ユーザーパラメータとDataInterfaceのバインド

内部でデータを参照できるように入力NormalMap等をバインドします.

EmitterSpawnステージの + から「新規または既存のパラメータを直接設定」でSetParameter項目を追加.

SetParameterの + から Texture Sample を選択.

名前を DI_InputNormalTexに変更.
Texture User Parameter にユーザーパラメータの InputNormalTex を設定.

同様に SetParameter を追加して Render Target 2D に設定.
名前を DI_OutputRenderTarget に変更.
ユーザーパラメータを利用するために Inherit User Parameter Settings をON.
Render Target User Parameter に OutputHeight を設定.

プレビューマテリアルの設定

EmitterのSpriteRendererにプレビュー用マテリアル M_PreviewHeight を設定します.

マテリアルのテクスチャパラメータ InTex にプレビュー対象のテクスチャをバインドします.
SpriteRendererのバインディング>Material Parameters>Attribute Bindings に要素を追加.
Material Parameter Name からプレビュー用マテリアルのテクスチャパラメータ InTex が選択可能です.
Niagara Variable に出力先RenderTarget2Dである DI_OutputRenderTarget をバインドします.
これで計算結果が書き込まれたRenderTargetの内容をスプライトに表示できます.

Grid2DCollectionの準備

Grid2DCollectionは汎用2Dグリッドデータインターフェイスです.
指定解像度グリッドの集合(Collection)として様々なデータの格納ができます.
今回は反復更新中のHeightMap値を格納するために使用します.

EmitterSpawnステージの + からSetParameter項目を追加し,
SetParameterの + からGrid2DCollection を選択する.
名前を WorkGrid2D に変更.
パラメータを展開して Num Attributes を 1 以上に設定.
(本来必要ないと思われますが, ここで1以上の値を設定していないとGrid2Dが値を保持してくれない場合があります)

次にグリッド解像度を設定します.
今回は入力NormalMapと同じ解像度で計算をしたいのでそのための設定もします.
標準では解像度設定モジュールが無いのでスクラッチパッドモジュールで処理します.
EmitterSpawnステージの + から「新しいスクラッチパッドモジュール」.

空のモジュールが生成されるので, モジュール名は SetResolutionWorkGrid にしておきましょう.

SetResolutionWorkGridの入力Mapに
解像度設定対象のGrid2DCollection
サイズ情報を取得するためのTextureSample
の枠を追加します.

TextureSampleからGetTextureDimensionsでサイズ情報を取得し,
Grid2DCollectionのSetNumCellsで解像度を設定するように組んで適用ボタンを押します.

Emitterのモジュール詳細パネルに戻ると,
入力Mapに追加したGrid2DCollectionとTextureSampleの項目が現れます.
Grid2DCollectionには Link Inputs/エミッタ/WorkGrid2D
TextureSampleには Link Inputs/エミッタ/DI_InputNormalTex
を設定すると, ようやく入力NormalMapと同じ解像度のGridのセットアップが完了です.

初期化ステージ

Emitterの +ステージから GenericSimulationStage を追加します.

詳細パネルでこのステージの設定をします. 名前はInitStage等にしておきます.

Simulation Stage Name は InitStage
- わかりやすい名前ならなんでも良い
Iteration Source は Data Interface
- Grid2DなどのDataInterfaceの要素を実行単位にするため.
Execute Behavior は On Simulation Reset
- このステージはEmitterがSpawnしたタイミングなどの初回一度だけ実行するため.
Data Interface は WorkGrid2D
- このGrid2Dの要素毎に処理を実行するため
  - ComputeShaderとして考えると, このGrid2Dの解像度分のThreadが起動する.

Stageの + からこのStageの処理を記述するスクラッチパッドモジュールを追加します.

入力Mapに Grid2DCollection と TextureSample を追加します.
それぞれ作業用GridとNormalMapを取り込むために使います.
出力Mapに Vector2D と float を追加します.
それぞれ名前を NormalXY と Height に変更します.
さらに右クリックから名前空間を変更>STACKCONTEXT に設定します.
STACKCONTEXT はStageやフレームを跨いで読み書きできる名前空間(カテゴリ)です.
使用する場所でもっとも適切な名前空間になってくれるようです.
まとめると以下のようになります.

入力Map
- Grid2DCollection
  - Grid2D上での作業のため
- TextureSample
  - NormalMap取り込みのため
出力Map
- NormalXY (STACKCONTEXT Vector2D
  - NormalMapの値をGrid2D上に保存して読み取るため.
- Height ( STACKCONTEXT float
  - Grid2D上でのHeight更新作業用.

処理内容としては

NormalMapからサンプリングしたベクトルをGrid2DにNormalXYという名前で登録
- 後段のStageではTextureにアクセスせずGrid2Dで完結させるため
Grid2DにHeightという名前でfloat値を登録.
- 初期値は0. このHeightが反復処理で更新されていく.

適用ボタンを押したらEmitterのモジュール詳細から

Grid2DCollection にWorkGrid2D を設定
TextureSample にDI_InputNormalTex を設定

Jacobi反復処理ステージ

GenericSimulationStageを新規追加して以下のように設定します.

Simulation Stage Name は JacobiIteration
Iteration Source は Data Interface
Num Iterations は 1
- ここが1フレームで実行される反復数になります. 1024とか入れるとUEがハングするので注意.
  - 初期化ステージのInitStageを On Simulation Reset にしたことで, フレームを跨いで大量の反復を継続するようにしています.
Execute Behavior は Always
- このステージは毎フレーム実行するため.
Data Interface は WorkGrid2D

スクラッチパッドを追加します.

入力Map
- Grid2DCollection
  - Grid2D上での作業のため
出力Map
- Height ( STACKCONTEXT float
  - Grid2D上のHeightを更新するため.

スクラッチパッドにはJacobi法によるHeight値更新計算をCustomHLSLで記述します.
CustomHLSLノードの入出力は以下のように設定します.

入力
- Grid2DCollection
- Cell座標X(int)
- Cell座標Y(int)
出力
- 更新後のHeight値(float)

CustomHLSLの中身は先のJacobi法をそのままコードに落とし込み, NormalXYやHeightと名付けしたGrid上の値を読み取って計算しています.
面白い点としては GetPreviousVector2DValue() などの関数で, 前回反復時のGridの値 を取得できることでしょうか.
GPGPUで並列計算を書いたことがある方にはお馴染みですが, 並列計算における近傍値読み取りと書き込みの衝突解決のためにダブルバッファ化して読み取りは前回バッファ, 書き込みは今回バッファとするということを仕組みとしてサポートしてくれているようです.
また, 最後のほうでHeightのmax(0, )をとって正の値に制約していますが. こちらはProjectionという値の取りうる値にクランプする手法になります.
(Height値は正の値であるべきなので)

// Get Grid Size.
int NumCellX; int NumCellY;
WorkGrid.GetNumCells(NumCellX, NumCellY);

float cur_height;
WorkGrid.GetPreviousFloatValue<Attribute="Height">(CellPosX, CellPosY, cur_height);

// Neighbor Normal.
float2 nnx0 = float2(0,0);
float2 nnx1 = float2(0,0);
float2 nny0 = float2(0,0);
float2 nny1 = float2(0,0);

// Neighbor Current Height.
float nhx0 = 0;
float nhx1 = 0;
float nhy0 = 0;
float nhy1 = 0;

// Get Neighbor Normal and Height.
if(0 < CellPosX)
{
    WorkGrid.GetPreviousVector2DValue<Attribute="NormalXY">(CellPosX - 1, CellPosY, nnx1);
    WorkGrid.GetPreviousFloatValue<Attribute="Height">(CellPosX - 1, CellPosY, nhx1);
}
if(NumCellX-1 > CellPosX)
{
    WorkGrid.GetPreviousVector2DValue<Attribute="NormalXY">(CellPosX + 1, CellPosY, nnx0);
    WorkGrid.GetPreviousFloatValue<Attribute="Height">(CellPosX + 1, CellPosY, nhx0);
}
if(0 < CellPosY)
{
    WorkGrid.GetPreviousVector2DValue<Attribute="NormalXY">(CellPosX, CellPosY - 1, nny1);
    WorkGrid.GetPreviousFloatValue<Attribute="Height">(CellPosX, CellPosY - 1, nhy1);
}
if(NumCellY-1 > CellPosY)
{
    WorkGrid.GetPreviousVector2DValue<Attribute="NormalXY">(CellPosX, CellPosY + 1, nny0);
    WorkGrid.GetPreviousFloatValue<Attribute="Height">(CellPosX, CellPosY + 1, nhy0);
}

// Calc Jacobi Iteration Step for Height.
const float dn = ((nnx0.x - nnx1.x) + (nny0.y - nny1.y)) / 2.0;
float next_height = (nhx0 + nhx1 + nhy0 + nhy1 + dn) / 4.0;
// Projected. Constrain to a number greater than 0.
OutHeight = max(0.0, next_height);

スクラッチパッドを適用したらInitStageと同様にスクラッチパッド入力のGrid2DCollectionにWorkGrid2Dを設定します.

出力ステージ

GenericSimulationStageを新規追加して以下のように設定します.

Simulation Stage Name は JacobiIteration
Iteration Source は Data Interface
Num Iterations は 1
Execute Behavior は Always
- 毎フレーム結果を外部RenderTargetに出力します.
Data Interface は DI_OutputRenderTarget
- このステージは結果のHeightMap出力を目的のため, 出力RenderTargetである DI_OutputRenderTarget を指定.

スクラッチパッドは特に複雑なことはなくGrid2DからHeightを取得してRenderTargetへ出力します.

入力Map
- Grid2DCollection
  - 計算結果のHeightの取り出しのため
- RenderTarget2D
  - 出力先RenderTargetのData Interfaceとして

注意点としては, 今回のStageはRenderTarget2DをDataInterfaceとしているため, UV座標やテクセル座標を取得するためにGrid2DではなくRenderTarget2Dを使う点です.
(Exec to Unit や Exec to Index)

忘れないようにスクラッチパッド入力の設定をします.
Grid2DCollectionにWorkGrid2D, RenderTarget2DにDI_OutputRenderTarget.

結果

完成するとプレビューや出力先RenderTarget2Dで計算結果のHeightMapが表示されます.
Jacobi反復ステージの反復数を 1 にしているため, フレームごとに1回の反復で徐々に計算が収束していきます.
もっと早く収束させたい場合はJacobi反復ステージの Num Iterations を 16 などに増やしてみてください.
結果のHeightMapをテクスチャとして保存したい場合は
「RenderTarget2Dからスタティックテクスチャを作成」
をしてください.

アドベントカレンダー用 pic.twitter.com/rE2PGLBdMG
— なが (@nagakagachi) 2023年12月16日

まとめ

NormalMapからHeightMapを復元するために
Poisson方程式の反復解法(Jacobi法)を
Niagara SimulationStageで実装する

ということをしました.

SimulationStageはNiagaraモジュールとしてComputeShader的なGPGPU計算が実装できて面白いですね.
ランタイムVFX用途ではなく今回のようなツール的な使い方も模索のし甲斐があると思いました.
ただ一部挙動が怪しかったりUIが親切でなかったりドキュメントがなかったりする点でハードルが高いです.
今後のアップデートでどう変わっていくのか注目したいですね.

~~~

次の17日記事は @razupi さんとのことです!