yum/3ner

A toon shader for Unity's BIRP.

git clone https://git.yummers.dev/yum/3ner

yumImplement mixed-radix computationdf08b1f

master
branchesmasterall refs →
13.9 KiB356 linesraw
1using System.IO;
2using System.Collections.Generic;
3using System.Reflection;
4using UnityEditor;
5using UnityEngine;
6
7
8namespace FourierWater.Editor
9{
10    public class FourierWater : EditorWindow
11    {
12        [Tooltip("Size in meters to simulate")]
13        public float sizeM = 256f;
14
15        [SerializeField]
16        private int bigN = 1024;
17
18        [Tooltip("Gravitational acceleration (m/s^2)")]
19        private float g = 9.81f;
20        [Tooltip("Surface tension (N/m)")]
21        private float surfaceTension = 0.076f;
22        [Tooltip("Density (kg / m^3)")]
23        private float density = 1000f;
24        [Tooltip("Wind speed (m/s)")]
25        private float U = 30f;
26        [Tooltip("Fetch (m)")]
27        private float F = 100f * 1000f;
28        [Tooltip("Depth (m)")]
29        private float h = 60f;
30        [Tooltip("Peak enhancement (unitless)")]
31        private float gamma = 4f;
32
33        [SerializeField]
34        private MeshFilter waterMesh;
35
36        // Viscous deep water wave relation. See:
37        //   https://yummers.dev/wave-dispersion-relations-with-derivatives.html
38        private float Omega(float k)
39        {
40            float k3 = k * k * k;
41            return Mathf.Sqrt(g * k + surfaceTension * k3 / density);
42        }
43
44        private float DOmegaDK(float k)
45        {
46            float k2 = k * k;
47            float k3 = k2 * k;
48            return (g + 3 * surfaceTension * k2 / density) / (2 * Mathf.Sqrt(g * k + surfaceTension * k3 / density));
49        }
50
51        // JONSWAP spectrum
52        private float SJonswap(float omega) {
53          // Correction from equation 8 of Barfuss et. al. "The Impact of
54          // Offshore Wind Farms on Sea State Demonstrated by Airborne LiDAR
55          // Measurements."
56          float omega_p = 22.0f * Mathf.Pow((g * g / (U * F)), 0.33333f);
57          float sigma = (omega <= omega_p) ? 0.07f : 0.09f;
58          float alpha = 0.076f * Mathf.Pow(U * U / (F * g), 0.22f);
59          float r = Mathf.Exp(-(omega - omega_p) * (omega - omega_p) / (2 * sigma * sigma * omega_p * omega_p));
60          float omega2 = omega * omega;
61          float omega4 = omega2 * omega2;
62          float omega5 = omega4 * omega;
63          float meow4 = (omega_p / omega);
64          meow4 = meow4 * meow4;
65          meow4 = meow4 * meow4;
66          float peak_enhancement = Mathf.Pow(gamma, r);
67          return (alpha * g * g / omega5) * Mathf.Exp(-1.25f * meow4) * peak_enhancement;
68        }
69
70        // TMA spectrum
71        private float STma(float omega) {
72          float omega_h = omega * Mathf.Sqrt(h / g);
73          if (omega_h <= 1.0f ) {
74            return SJonswap(omega) * 0.5f * omega_h * omega_h;
75          } else if (omega_h < 2.0f) {
76            return SJonswap(omega) * (1.0f - 0.5f * (2.0f - omega_h) * (2.0f - omega_h));
77          } else {
78            return SJonswap(omega);
79          }
80        }
81
82        // Flat directional spreading function
83        private float DFlat()
84        {
85          return 1.0f / (3.14159265f * 2.0f);
86        }
87
88        // https://stackoverflow.com/a/218600
89        private float RandomGaussian(System.Random r)
90        {
91          float u1 = (float) (1.0f - r.NextDouble());
92          float u2 = (float) (1.0f - r.NextDouble());
93          float randStdNormal = Mathf.Sqrt(-2.0f * Mathf.Log(u1)) * Mathf.Sin(2.0f * Mathf.PI * u2);
94          // Want mean=0 stdev=1 so just return randStdNormal
95          return randStdNormal;
96        }
97
98        private Texture2D MakePowerSpectrum(string name, int sizePx, float sizeM)
99        {
100            Texture2D tex = new Texture2D(sizePx, sizePx, TextureFormat.RGBAHalf,
101                mipChain: false, linear: true)
102            {
103                name = name,
104                filterMode = FilterMode.Point,
105                wrapMode = TextureWrapMode.Clamp,
106            };
107
108            System.Random rand = new System.Random();
109            float tau = Mathf.PI * 2.0f;
110            for (int y = 0; y < sizePx; y++)
111            {
112                // y-axis wave number
113                float kY = (tau / sizeM) * (y < sizePx / 2 ? y : y - sizePx);
114                for (int x = 0; x < sizePx; x++)
115                {
116                    // x-axis wave number
117                    float kX = (tau / sizeM) * (x < sizePx / 2 ? x : x - sizePx);
118                    float k = Mathf.Sqrt(kX * kX + kY * kY);
119
120                    // Per Horvath 2015 eq. 61:
121                    //   S(k_x, k_y) = S(omega, theta) * DOmegaDK(k) / k.
122                    float omega = Omega(k);
123                    float s = DFlat() * STma(omega) * DOmegaDK(k) / k;
124
125                    float dk = tau / sizeM;
126                    float nu = RandomGaussian(rand);
127                    // Horvath eq. 47
128                    float amp = nu * Mathf.Sqrt(2.0f * s * dk * dk);
129                    float phase = (float) rand.NextDouble() * tau;
130
131                    Color res = new Color(amp * Mathf.Cos(phase), amp * Mathf.Sin(phase), 0, 1);
132                    if (k == 0) {
133                      res = new Color(0, 0, 0, 1);
134                    }
135                    tex.SetPixel(x, y, res);
136                }
137            }
138
139            tex.Apply(updateMipmaps: false);
140            return tex;
141        }
142
143        readonly string pipelineName = "Fourier Water";
144
145        [MenuItem("Tools/yum_food/Fourier Water")]
146        public static void ShowWindow()
147        {
148            GetWindow<FourierWater>("Fourier Water");
149        }
150
151        private string GetGeneratedAssetsDir()
152        {
153            MonoScript script = MonoScript.FromScriptableObject(this);
154            string top = Path.GetDirectoryName(Path.GetDirectoryName(
155                AssetDatabase.GetAssetPath(script)));
156            string genDir = Path.Combine(top, "Generated").Replace('\\', '/');
157            if (AssetDatabase.IsValidFolder(genDir))
158            {
159                AssetDatabase.DeleteAsset(genDir);
160            }
161            Directory.CreateDirectory(genDir);
162            AssetDatabase.Refresh();
163            return genDir;
164        }
165
166        private CustomRenderTexture MakeCRT(string name, int size, Material mat)
167        {
168            CustomRenderTexture tex = new CustomRenderTexture(
169                size, size,
170                RenderTextureFormat.ARGBHalf,
171                RenderTextureReadWrite.Linear)
172            {
173                name = name,
174                material = mat,
175                shaderPass = 0,
176                filterMode = FilterMode.Bilinear,
177                wrapMode = TextureWrapMode.Repeat,
178                useMipMap = false,
179                autoGenerateMips = false,
180                doubleBuffered = false,
181                initializationMode = CustomRenderTextureUpdateMode.OnDemand,
182                initializationSource = CustomRenderTextureInitializationSource.TextureAndColor,
183                initializationColor = Color.black,
184                updateMode = CustomRenderTextureUpdateMode.OnDemand,
185                updatePeriod = 0f
186            };
187
188            tex.Create();
189            return tex;
190        }
191
192        private static int[] PlanRadices(int size)
193        {
194            if (size < 1 || (size & (size - 1)) != 0)
195            {
196                throw new System.ArgumentException("FFT size must be a power of two.");
197            }
198            var stages = new List<int>();
199            while (size > 1) {
200                int radix = size % 8 == 0 ? 8 : size % 4 == 0 ? 4 : 2;
201                stages.Add(radix);
202                size /= radix;
203            }
204            return stages.ToArray();
205        }
206
207        private void OnGUI()
208        {
209            bigN = EditorGUILayout.IntField("FFT resolution", bigN);
210            sizeM = EditorGUILayout.FloatField("Patch size (meters)", sizeM);
211            U = EditorGUILayout.FloatField("Wind speed (m/s)", U);
212            F = EditorGUILayout.FloatField("Fetch (meters)", F);
213            h = EditorGUILayout.FloatField("Depth (meters)", h);
214            gamma = EditorGUILayout.FloatField("Peak enhancement", gamma);
215
216            waterMesh = (MeshFilter) EditorGUILayout.ObjectField(
217                "Water mesh", waterMesh, typeof(MeshFilter), true);
218
219            if (!GUILayout.Button("Generate"))
220            {
221                return;
222            }
223
224            if (bigN < 1 || bigN > 1024 || (bigN & (bigN - 1)) != 0)
225            {
226                Debug.LogError("FFT resolution must be a power of two from 1 to 1024.");
227                return;
228            }
229
230            if (waterMesh == null || waterMesh.sharedMesh == null)
231            {
232                Debug.LogError("Assign a water mesh before generating.");
233                return;
234            }
235            Renderer waterRenderer = waterMesh.GetComponent<Renderer>();
236            if (waterRenderer == null || waterRenderer.sharedMaterial == null)
237            {
238                Debug.LogError("The water mesh needs a renderer and material.");
239                return;
240            }
241            Material waterMaterial = waterRenderer.sharedMaterial;
242            float meshSize = waterMesh.sharedMesh.bounds.size.x;
243            if (sizeM <= 0 || meshSize <= 0)
244            {
245                Debug.LogError("Patch size and mesh size must be positive.");
246                return;
247            }
248
249            // Delete old shit
250            GameObject oldPipeline = GameObject.Find(pipelineName);
251            if (oldPipeline != null)
252            {
253                Undo.DestroyObjectImmediate(oldPipeline);
254            }
255            string genDir = GetGeneratedAssetsDir();
256
257            GameObject parent = new GameObject(pipelineName);
258
259            Shader shader = Shader.Find("yum_food/fourier_water_compute");
260            Shader displayShader = Shader.Find("yum_food/fourier_water_display");
261
262            Texture spectrum = MakePowerSpectrum("spectrum", bigN, sizeM);
263            string spectrumPath = Path.Combine(genDir, "spectrum.asset");
264            AssetDatabase.CreateAsset(spectrum, spectrumPath);
265
266            string[] outputProperties = {
267                "_Fourier_Water_Height",
268                "_Fourier_Water_Normals",
269                "_Fourier_Water_Displacement",
270                "_Fourier_Water_Displacement_Derivatives"
271            };
272            string[] inputKeywords = {
273                "_TIME_EVOLVE", "_DERIVATIVE", "_DISPLACEMENT", "_DISPLACEMENT_DERIVATIVE"
274            };
275            int n_pipelines = outputProperties.Length;
276            var outputs = new CustomRenderTexture[n_pipelines];
277            int[] radices = PlanRadices(bigN);
278            for (int which_pipeline = 0; which_pipeline < n_pipelines; which_pipeline++) {
279              Texture input = spectrum;
280              int n_stages = radices.Length;
281              for (int axis = 0; axis < 2; axis++)
282              {
283                string axisStr = axis == 0 ? "x" : "y";
284                int stageN = 1;
285                int stride = bigN;
286                for (int i = (axis == 0 ? -1 : 0); i < n_stages; i++)
287                {
288                  if (i > -1) {
289                    stride /= radices[i];
290                    stageN *= radices[i];
291                  }
292
293                  string stageName = $"pipeline_{which_pipeline}_axis_{axisStr}_stage_{i}";
294                  Material mat = new Material(shader);
295                  mat.EnableKeyword($"_FFT_SIZE_{bigN}");
296
297                  GameObject quad = GameObject.CreatePrimitive(PrimitiveType.Quad);
298                  quad.transform.SetParent(parent.transform);
299                  quad.transform.localPosition = new Vector3(i, 2 * (n_pipelines - which_pipeline) - axis, 0);
300
301                  if (i == -1) {
302                    mat.EnableKeyword(inputKeywords[which_pipeline]);
303                  } else if (radices[i] == 8) {
304                    mat.EnableKeyword("_RADIX_8");
305                  } else if (radices[i] == 4) {
306                    mat.EnableKeyword("_RADIX_4");
307                  }
308                  mat.SetTexture("_Signal", input);
309                  mat.SetInt("_Stage", i);
310                  mat.SetInt("_Big_N", stageN);
311                  mat.SetFloat("_Size_M", sizeM);
312                  mat.SetInt("_Stride", stride);
313                  mat.SetInt("_Direction", axis);
314
315                  string matPath = Path.Combine(genDir,
316                      $"{stageName}.mat").Replace('\\', '/');
317                  AssetDatabase.CreateAsset(mat, matPath);
318
319                  CustomRenderTexture crt = MakeCRT(stageName, bigN, mat);
320                  string crtPath = Path.Combine(genDir,
321                      $"{stageName}.asset").Replace('\\', '/');
322                  AssetDatabase.CreateAsset(crt, crtPath);
323
324                  Material displayMat = new Material(displayShader);
325                  displayMat.SetTexture("_MainTex", crt);
326                  displayMat.SetInt("_Stage", i);
327                  displayMat.SetInt("_Num_Stages", Mathf.Max(1, n_stages));
328                  displayMat.SetInt("_Direction", axis);
329                  string displayMatPath = Path.Combine(genDir,
330                      $"{stageName}_display.mat").Replace('\\', '/');
331                  AssetDatabase.CreateAsset(displayMat, displayMatPath);
332                  quad.GetComponent<Renderer>().sharedMaterial = displayMat;
333
334                  var updater = quad.AddComponent<ShaderFes.CustomRenderTextureUpdater>();
335                  var serializedUpdater = new SerializedObject(updater);
336                  serializedUpdater.FindProperty("_targetTexture").objectReferenceValue = crt;
337                  serializedUpdater.FindProperty("_updateCount").intValue = 1;
338                  serializedUpdater.FindProperty("_updatePeriod").floatValue = 0;
339                  serializedUpdater.ApplyModifiedPropertiesWithoutUndo();
340
341                  input = crt;
342                }
343              }
344              outputs[which_pipeline] = (CustomRenderTexture)input;
345            }
346
347            for (int pipeline = 0; pipeline < n_pipelines; pipeline++) {
348                waterMaterial.SetTexture(outputProperties[pipeline], outputs[pipeline]);
349            }
350            waterMaterial.SetFloat("_Fourier_Water_Patch_Size", sizeM);
351            waterMaterial.SetFloat("_Fourier_Water_Mesh_Size", meshSize);
352            EditorUtility.SetDirty(waterMaterial);
353            AssetDatabase.SaveAssets();
354        }
355    }
356}