概要

GPGPU（General Purpose GPU）を使用して、大量のParticlesを流動的に動かす実装についてまとめました。

GPGPU（General Purpose GPU）とは、通常グラフィクス処理で使われるGPUを、その他さまざまな目的で利用する考え方です。
特に機械学習分野では、よく目にする言葉です。

GPGPUの使用

このアプリケーションでは、GPGPUの考えに基づいて、通常はColor情報をpixel単位で格納するTextureに対して、位置座標情報を格納して使用しています。

通常のShader

まず、GPGPUを使わない実装を考えてみます。

Three.jsでParticelsを扱うときは、**Points**オブジェクトを使用します。
Vertex Shaderで参照されるpositionは、Pointsにセットしたgeometryの頂点座標（attribute position）が取られます。このpositionに従って、point（particle）が描画されます。
attribute positionは、geometryに対して必須ですが、ビルトインgeometry（PlaneGeometryなど）を使う場合はデフォルトで定義されているので、特に意識する必要がないです。

何らかの処理で、posを変化させることによって、particleに動きをつけることができます。例えば、sin関数を使えば周期性のある動きを表現できます。
また、Vertex Shaderの処理は各フレームで行われますが、positionは前のフレームで更新された座標ではなく、geometryを定義したときの座標になります。

GPGPUを使用したShader

次に、Textureに位置情報を入れて、それを参照する実装を考えてみます。

この方法では、2つの役割を持つRenderTargetを使用します。
1つ目は、座標計算をして、その結果をTextureに出力することを目的とするRenderTarget（Texture RenderTarget）です。このRenderTargetは内部的にレンダリングされるため、表示はされません。
2つ目は、Texture RenderTargetで出力されたTextureを受け取って、表示することを目的とするRenderTarget（Main RenderTarget）です。

もう少し詳しく見てみます。
Texture RenderTargetには、以下のようなレンダリングがされます。

例えば、サイズが10×10のTextureの場合、100pixelの1つ1つに位置情報（x, y, z, w）が格納されます。
Textureは、Fragment ShaderでtexturePositionとして受け取り、何らかの処理によって座標を更新します。
ここで受け取ったTextureは、前のフレームでの位置情報が格納されています。

Main RenderTargetには、以下のようなレンダリングがされています。

ここで使われているShaderは、PointsオブジェクトのShaderMaterialを想定しています。
レンダリングされたTexture RenderTargetをTexture（u_texture）として受け取ります。u_textureから位置情報（positionInfo）を取得して、particleを描画します。
画像の例では、位置情報のwをparticleの生存状態（v_life）として、FragmentShaderに渡しています。FragmentShaderでは、v_lifeをparticleの透明度として使っていて、生存状態が0に近いほど透過するようにしています。

GPGPUを使う利点

ここまでの話だと、「通常のShader」と「GPGPUを使用したShader」で大きな違いは、前のフレームの位置情報を保持できるかどうかでしかありません。
私も色々調べながら実装をしていたときに、「通常のShaderで良くないか？」となりましたが、そのうち大きな利点に気が付きました。

利点：1
実装したアプリケーションでは、262,144（512×512）個の頂点を持つPointsオブジェクトを2つ配置しています。つまり、頂点数は262,144 × 2となっています。
「通常のShader」を使用する場合、262,144 × 2回GPUで座標計算をする必要があります。一方、「GPGPUを使用したShader」を使用する場合、Textureに対して262,144回の計算を行えば、どれだけPointsオブジェクトが増えたとしても、同じTextureから位置情報を取得するならGPUの計算量はかわりません。

利点：2
このアプリケーションでは実装していませんが、前のフレームでの位置情報を使うことでモーションブラーエフェクトを実装することができます。
以下のリポジトリでは、GPGPUを使用したParticlesの素晴らしい実装がなされています。（コードの可読性が高いです！）
モーションブラーも実装されています。

利点：3
これまでの説明では、Textureは位置情報を持った1枚のみを考えてきました。ですが、枚数に関しては複数枚でもいいですし、持たせる情報も4次元のfloat情報なら制約はないはずです。
Three.jsの以下のサンプルでは、位置座標用のTextureの他に、速度の情報をもつTextureを使用しています。

実装

コード全体を載せると、ボリュームが多くなりすぎるので要点だけ解説します。

simulator.ts

「GPGPUを使用したShader」を実装したクラスです。
Three.jsでは、GPGPUの使用を想定した便利なクラス（GPUComputationRenderer）が提供されています。
使用方法は、以下の記事にまとめられています。

セットアップの順番は、以下のようになっています。
1）コンストラクタを生成する
2）Textureを生成する
3）FragmentShaderで参照するTextureを指定する
4）initを実行する

modules/simulator.ts

constructor(gl: THREE.WebGLRenderer, private _width: number, private _height: number) {
	this._gpuCompute = new GPUComputationRenderer(this._width, this._height, gl)
	this._setTexturePosition()
	this._setVariableDependencies()
	this._gpuCompute.init()
	const gui = GUIController.instance
	gui.initParticlePositionParams()
}

ここでは、uniformに対するコントローラーも追加しています。

compute関数を呼ぶことで、Textureを更新します。

modules/simulator.ts

compute = () => {
	this._gpuCompute.compute()
	updateParticlesPositionUniforms(this._positionMaterial)
}

updateParticlesPositionUniformsでは、FragmentShaderで使用されるuniformを更新します。

Textureは、texturePositionから取得します。

modules/simulator.ts

get texturePosition() {
	const variable = this._variables.find(v => v.name === 'texturePosition')!
	const target = this._gpuCompute.getCurrentRenderTarget(variable) as THREE.WebGLRenderTarget
	return target.texture
}

positionShader.ts

simulator.tsで位置情報のTextureに使用しているFragmentShaderは、以下のようになっています。

glsl/positionShader.ts

import { curl } from './curl';

export const positionFragmentShader = `
uniform sampler2D u_defaultTexture;
uniform float u_time;
uniform float u_frequency; // 一体性
uniform	float u_amplitude; // 流動速度
uniform float u_divergence; // 発散性

const float dieSpeed = 0.99;

${curl}

void main()	{
	vec2 uv = gl_FragCoord.xy / resolution.xy;
	vec4 tmpPos = texture2D(texturePosition, uv);
	vec3 position = tmpPos.xyz;
	float life = tmpPos.w;

	if (life < 0.1) {
		vec4 defPos = texture2D(u_defaultTexture, uv);
		position = defPos.xyz;
		life = defPos.w;
	}

	float seed = u_time * 10.0; // 変化
	vec3 pos = position * u_frequency;

	vec3 target = position + u_amplitude * curl(seed, pos.x, pos.y, pos.z);
	target *= u_divergence;
	
	gl_FragColor = vec4(target, life * dieSpeed);
}
`

uvを取得するためのresolutionは、GPUComputationRenderer側で自動的に割り当ててくれます。
位置情報を格納しているTexture（texturePosition）は、simulator.tsのthis._setVariableDependencies()で定義したTextureが自動で割り当てられます。今回はtexturePositionだけですが、複数のTextureを参照することもできます。
位置情報は、Curl Noiseを使用して更新しています。引数には、seed値と現在の座標を渡します。
seedは、時間経過（u_time）で変更されます。これによってノイズの形状が徐々に変化していきます。
u_frequencyは、Particlesの一体性を表します。数値が低いほどParticlesが一体として動き、高いほどバラバラに動きます。
u_amplitudeは、Particlesの移動速度を表します。
u_divergenceは、Particlesの発散性を表します。数値が高いほど発散します。

他に、texturePositionのw成分を、Particlesの生存状態として扱っています。時間経過で徐々に0に近付き、0.1を下回るとParticleが初期位置から再生成されます。

Particles.tsx

このコンポーネントでは、Pointsオブジェクトを生成しています。
以下のことを行っています。

simulator.tsで定義したSimulatorクラスのインスタンスの生成。サイズは512×512。
Pointsオブジェクトに割り当てるBufferGeometryの頂点（attribute position）の生成
ShaderMaterialの生成。このアプリケーションではPointsオブジェクトを2つ生成していて、それぞれ色を分けたかったので、Materialは2つ生成しています。
geometryの生成。頂点数とその位置情報が2つのPointsで変わらないので、共有しています。
Textureの更新と、uniformの更新。

特に、Pointsの位置情報はTextureが持っているので、BufferGeometryの頂点はその役割を持っていません。
ただし、BufferGeometryに対してattribute positionの定義は必須なので、この情報をTextureから位置情報を取るためのuvとして定義しています。

componets/Particles.tsx

const vertices = useMemo(() => {
	const vertices: number[] = []
	for (let ix = 0; ix < width; ix++) {
		for (let iy = 0; iy < height; iy++) {
			// uv to access texture
			vertices.push(ix / width, iy / height, 0)
		}
	}
	return Float32Array.from(vertices)
}, [])

uvは0 ~ 1の範囲で定義する必要があるため、基準化しています。z座標の情報は使わないので、0としています。

フレームループでは、SimulatorのTextureを更新して、それをuniformに割り当てています。

componets/Particles.tsx

useFrame(() => {
	simulator.compute()
	updateParticlesUniforms(matShader, simulator)
	updateParticlesUniforms(matShader2, simulator)
})

particlesShader.ts

PointsオブジェクトのShaderMaterialに割り当てるShaderを定義しています。
特に、particlesVertexShader では、attribute positionをuvとして使用して位置情報（positionInfo）取得しています。

glsl/particlesShader.ts

export const particlesVertexShader = `
uniform sampler2D u_texture;
varying float v_life;

void main() {
	vec4 positionInfo = texture2D(u_texture, position.xy);
	vec4 mvPosition = modelViewMatrix * vec4(positionInfo.xyz, 1.0);

	v_life = positionInfo.w;

	gl_PointSize = positionInfo.w / length(mvPosition.xyz) * 10.0;
	gl_Position = projectionMatrix * mvPosition;
}
`

gui.ts

コントローラーには、**lil-gui**を使用してます。
Reactでコントローラーを使う場合、levaが便利で普段はこちらを使用していますが、今回は以下の理由でlil-guiを使用しました。

levaで値を変更すると、コンポーネントが再生成されるため、アニメーションが一時的に止まる。
コンポーネント（tsx）ではない、Simulatorクラス内の値に対してもコントローラーを使いたい。

ただし、lil-guiはReact用に作られているわけではないので、使うには少し工夫が必要です。

先に結論を書きます。Reactでlil-guiを使用するときは、以下のようにするとうまく実装できます。
・Singletonパターンのクラスでカプセル化する
・項目の重複を防止する処理を入れる
・コントロールに設定するオブジェクトは、コンポーネント外で管理する

lil-guiを使用するためにはインスタンスを生成して、それにコントロールを追加していきます。
注意点として、インスタンスを生成するごとにコントローラーがいくつも追加されるため、アプリケーション全体でひとつのインスタンスを共有する必要があります。

Reactは仕様上、最初にコンポーネントが何回か呼ばれるため、コンポーネント内で普通にインスタンスを生成すると、コントローラーが複数個重なって表示されてしまいます。
そのため、Singletonパターンのクラスを作成して、その中でlil-guiのインスタンスを持つようにするのがベターだと思います。

クラスにおけるSingletonパターンとは、そのクラスのインスタンスをひとつしか保持できないデザインパターンです。

modules/gui.ts

export class GUIController {
	private static _instance: GUIController | null
	private _gui

	private constructor() {
		this._gui = new GUI()
	}

	static get instance() {
		if (!this._instance) {
			this._instance = new GUIController()
		}
		return this._instance
	}
	// ・・・
}

コンストラクタにprivate子をつけることで、外部からインスタンスを生成できないようにします。
外部からインスタンスにアクセスするときは、instanceアクセッサを呼び出します。instanceアクセッサ内では、このクラスのインスタンスがなければ生成します。

外部からは、以下のようにアクセスします。

componets/Particles.tsx

const gui = GUIController.instance
gui.initParticleColors()

これでコントローラー自体が複数作成されることはなくなりました。

ただ、この状態でコンポーネント内でコントロールを追加（.add, .addcolorなど）すると、上述したようにReactのコンポーネントが何回か呼ばれるため、同じ項目がいくつも生成されます。
そのため、コントロールのフォルダ名（title）と項目名（name）を使用して、同じ項目が追加されるのを防いでいます。

modules/gui.ts

private _folder = (title: string) => {
	let folder = this._gui.folders.find(f => f._title === title)
	if (!folder) folder = this._gui.addFolder(title)
	return folder
}

private _uncontainedName = (folder: GUI, name: string) => {
	return !folder.controllers.find(c => c._name === name)
}

コンポーネント内で定義したオブジェクト（Shaderのuniformなど）を直接コントロールに設定すると、コンポーネントが再生成されたときにオブジェクトも再生成される場合（useMemoなどを使用していない場合）、設定したオブジェクトとコンポーネント内のオブジェクトが異なるためコントロールが効かなくなります。

このため、オブジェクトは別ファイル（modules/store.ts）で管理して、これをコントロールに指定します。そして、フレームループでコンポーネント内のオブジェクト（Shaderのuniform）を更新するようにします。

リポジトリ

参考

GPGPU

Particles

まとめ

参考にさせていただいた The Spirit が、コメントがほぼなくJavaScriptで記述されているのにも関わらず、とても可読性のよいコードになっています。
クラス・関数・変数の命名がしっかりしていて、ファイル設計もとてもわかりやすく、感化されてしまいました。

今回のような実験的なアプリケーションの場合、普段はあんまりファイル分けをしないのですが、このリポジトリではあえて細かく分けてみました。