using System; using FloodSWE.Preprocess; using FloodSWE.Streaming; using FloodSWE.TileGraph; using Unity.Collections; using UnityEngine; using UnityEngine.Rendering; public sealed class SweTileSimulator : MonoBehaviour { [Header("Resources")] public ComputeShader ghostExchangeShader; public ComputeShader fluxShader; [Header("Grid")] public int gridRes = 256; public float tileSizeMeters = 1000f; [Header("Initial Conditions")] public bool damBreak = true; public float initialDepthLeft = 1.0f; public float initialDepthRight = 0.0f; public Vector2 initialVelocity = Vector2.zero; [Header("Rain")] public float rainRateMmPerHr = 0.0f; [Header("Simulation")] public float cfl = 0.5f; public float tickSeconds = 0.5f; public bool logCfl = true; public bool useFullPrecisionTextures = true; [Header("Static Data")] public Texture2D terrainHeight; public Texture2D porosity; public float flatTerrainHeightMeters = 0.0f; public float defaultPorosity = 1.0f; public bool resampleTerrainToGrid = true; public bool resamplePorosityToGrid = true; [Header("Streaming (Debug)")] public bool emitHeightmapPackets = false; public int packetEveryNTicks = 1; public TileId packetTileId = new TileId(1, 0, 0); public HeightmapPacket lastHeightmapPacket; public event Action HeightmapPacketReady; [Header("Debug Outputs")] public RenderTexture debugWater; public RenderTexture debugVelocity; [Header("Debug - Conservation")] public bool logConservation = false; public bool debugClampStats = true; public bool applyMassCorrection = false; [Range(0.5f, 2.0f)] public float massCorrectionMinScale = 0.8f; [Range(0.5f, 2.0f)] public float massCorrectionMaxScale = 1.2f; private const float Gravity = 9.81f; private RenderTexture waterA; private RenderTexture waterB; private RenderTexture velA; private RenderTexture velB; private RenderTexture waterGhost; private RenderTexture velGhost; private Texture2D resolvedTerrain; private Texture2D resolvedPorosity; private RenderTexture clampMask; private RenderTexture dummyRw; private int ghostKernel; private int fluxKernel; private int initKernel; private float dx; private float accumulatedTime; private float lastMaxDepth; private float lastMaxSpeed; private float lastTotalDepth; private float lastTotalVolume; private int lastClampedCells; private float lastClampedRatio; private int lastNanCells; private float lastDtMax; private float lastDt; private int lastSubsteps; private bool useA = true; private bool isInitialized; private int tickIndex; private float pendingMassScale = 1.0f; public float LastMaxDepth { get { return lastMaxDepth; } } public float LastMaxSpeed { get { return lastMaxSpeed; } } public float LastTotalDepth { get { return lastTotalDepth; } } public float LastTotalVolume { get { return lastTotalVolume; } } public int LastClampedCells { get { return lastClampedCells; } } public float LastClampedRatio { get { return lastClampedRatio; } } public int LastNanCells { get { return lastNanCells; } } public float LastDtMax { get { return lastDtMax; } } public float LastDt { get { return lastDt; } } public int LastSubsteps { get { return lastSubsteps; } } public float CellSizeMeters { get { return dx; } } private void OnEnable() { Initialize(); } private void OnDisable() { Release(); } private void Update() { if (!isInitialized) { return; } accumulatedTime += Time.deltaTime; if (accumulatedTime < tickSeconds) { return; } accumulatedTime -= tickSeconds; float maxDepth = Mathf.Max(lastMaxDepth, 1e-4f); float maxSpeed = Mathf.Max(lastMaxSpeed, 0.0f); float dtMax = cfl * dx / (maxSpeed + Mathf.Sqrt(Gravity * maxDepth)); if (float.IsNaN(dtMax) || dtMax <= 0.0f) { dtMax = tickSeconds; } int substeps = Mathf.Max(1, Mathf.CeilToInt(tickSeconds / dtMax)); float dt = tickSeconds / substeps; lastDtMax = dtMax; lastSubsteps = substeps; lastDt = dt; if (logCfl) { Debug.Log($"SWE CFL: maxDepth={maxDepth:F3}m maxSpeed={maxSpeed:F3}m/s dtMax={dtMax:F4}s substeps={substeps} dt={dt:F4}s"); } float rainRate = rainRateMmPerHr / 1000.0f / 3600.0f; for (int i = 0; i < substeps; i++) { DispatchGhost(); DispatchFlux(dt, rainRate); SwapBuffers(); } if (applyMassCorrection && Mathf.Abs(pendingMassScale - 1.0f) > 0.001f) { DispatchScale(pendingMassScale); pendingMassScale = 1.0f; } debugWater = CurrentWater; debugVelocity = CurrentVelocity; tickIndex++; TryEmitHeightmapPacket(); RequestReadback(); } private RenderTexture CurrentWater { get { return useA ? waterA : waterB; } } private RenderTexture CurrentVelocity { get { return useA ? velA : velB; } } private RenderTexture NextWater { get { return useA ? waterB : waterA; } } private RenderTexture NextVelocity { get { return useA ? velB : velA; } } private void Initialize() { if (ghostExchangeShader == null || fluxShader == null) { Debug.LogError("SweTileSimulator: missing compute shaders."); return; } if (useFullPrecisionTextures && SystemInfo.graphicsDeviceType == UnityEngine.Rendering.GraphicsDeviceType.OpenGLES3) { Debug.LogWarning("SweTileSimulator: full precision UAVs are unsupported on GLES3; falling back to half precision."); useFullPrecisionTextures = false; } gridRes = Mathf.Max(4, gridRes); tileSizeMeters = Mathf.Max(1.0f, tileSizeMeters); dx = tileSizeMeters / gridRes; ghostKernel = ghostExchangeShader.FindKernel("GhostExchange"); fluxKernel = fluxShader.FindKernel("FluxUpdate"); initKernel = fluxShader.FindKernel("InitDamBreak"); RenderTextureFormat waterFormat = useFullPrecisionTextures ? RenderTextureFormat.RFloat : RenderTextureFormat.RHalf; RenderTextureFormat velFormat = useFullPrecisionTextures ? RenderTextureFormat.RGFloat : RenderTextureFormat.RGHalf; waterA = CreateRenderTexture(gridRes, gridRes, waterFormat, "SWE_WaterA"); waterB = CreateRenderTexture(gridRes, gridRes, waterFormat, "SWE_WaterB"); velA = CreateRenderTexture(gridRes, gridRes, velFormat, "SWE_VelA"); velB = CreateRenderTexture(gridRes, gridRes, velFormat, "SWE_VelB"); waterGhost = CreateRenderTexture(gridRes + 2, gridRes + 2, waterFormat, "SWE_WaterGhost"); velGhost = CreateRenderTexture(gridRes + 2, gridRes + 2, velFormat, "SWE_VelGhost"); dummyRw = CreateRenderTexture(1, 1, waterFormat, "SWE_DummyRW"); if (debugClampStats) { clampMask = CreateRenderTexture(gridRes, gridRes, waterFormat, "SWE_ClampMask"); } ResolveStaticTextures(); ClearRenderTexture(waterB); ClearRenderTexture(velB); ClearRenderTexture(waterGhost); ClearRenderTexture(velGhost); fluxShader.SetInt("_GridRes", gridRes); fluxShader.SetFloat("_InitDepthLeft", initialDepthLeft); fluxShader.SetFloat("_InitDepthRight", initialDepthRight); fluxShader.SetVector("_InitVelocity", initialVelocity); fluxShader.SetInt("_InitDamBreak", damBreak ? 1 : 0); fluxShader.SetTexture(initKernel, "_WaterOut", waterA); fluxShader.SetTexture(initKernel, "_VelOut", velA); DispatchKernel(fluxShader, initKernel, gridRes, gridRes); useA = true; debugWater = CurrentWater; debugVelocity = CurrentVelocity; lastMaxDepth = Mathf.Max(initialDepthLeft, initialDepthRight); lastMaxSpeed = Mathf.Sqrt(Gravity * Mathf.Max(lastMaxDepth, 0.0f)); isInitialized = true; } private void Release() { isInitialized = false; ReleaseRenderTexture(waterA); ReleaseRenderTexture(waterB); ReleaseRenderTexture(velA); ReleaseRenderTexture(velB); ReleaseRenderTexture(waterGhost); ReleaseRenderTexture(velGhost); ReleaseRenderTexture(clampMask); ReleaseRenderTexture(dummyRw); waterA = null; waterB = null; velA = null; velB = null; waterGhost = null; velGhost = null; clampMask = null; dummyRw = null; if (resolvedTerrain != null && resolvedTerrain != terrainHeight) { Destroy(resolvedTerrain); } if (resolvedPorosity != null && resolvedPorosity != porosity) { Destroy(resolvedPorosity); } resolvedTerrain = null; resolvedPorosity = null; } private void ResolveStaticTextures() { resolvedTerrain = ResolveTerrainTexture(); resolvedPorosity = ResolvePorosityTexture(); } private Texture2D ResolveTerrainTexture() { if (terrainHeight == null) { return CreateFlatTexture(gridRes, flatTerrainHeightMeters); } if (!resampleTerrainToGrid) { if (terrainHeight.width == gridRes && terrainHeight.height == gridRes) { return terrainHeight; } Debug.LogWarning("SweTileSimulator: terrainHeight size mismatch and resampling disabled. Using flat terrain."); return CreateFlatTexture(gridRes, flatTerrainHeightMeters); } if (!terrainHeight.isReadable) { Debug.LogWarning("SweTileSimulator: terrainHeight texture must be readable. Using flat terrain."); return CreateFlatTexture(gridRes, flatTerrainHeightMeters); } if (terrainHeight.width == gridRes && terrainHeight.height == gridRes) { return terrainHeight; } float minHeight; float maxHeight; float[] resampled = DemResampler.ResampleHeights( terrainHeight, gridRes, flatTerrainHeightMeters, out minHeight, out maxHeight, true); return CreateTextureFromFloats(resampled, gridRes, gridRes, "SWE_TerrainResampled"); } private Texture2D ResolvePorosityTexture() { if (porosity == null) { return CreateFlatTexture(gridRes, Mathf.Clamp01(defaultPorosity)); } if (!resamplePorosityToGrid) { if (porosity.width == gridRes && porosity.height == gridRes) { return porosity; } Debug.LogWarning("SweTileSimulator: porosity size mismatch and resampling disabled. Using default porosity."); return CreateFlatTexture(gridRes, Mathf.Clamp01(defaultPorosity)); } if (!porosity.isReadable) { Debug.LogWarning("SweTileSimulator: porosity texture must be readable. Using default porosity."); return CreateFlatTexture(gridRes, Mathf.Clamp01(defaultPorosity)); } if (porosity.width == gridRes && porosity.height == gridRes) { return porosity; } float minValue; float maxValue; float[] resampled = DemResampler.ResampleHeights( porosity, gridRes, Mathf.Clamp01(defaultPorosity), out minValue, out maxValue, false); Clamp01(resampled); return CreateTextureFromFloats(resampled, gridRes, gridRes, "SWE_PorosityResampled"); } private void TryEmitHeightmapPacket() { if (!emitHeightmapPackets) { return; } int interval = Mathf.Max(1, packetEveryNTicks); if (tickIndex % interval != 0) { return; } if (HeightmapExtractor.TryExtractFromRenderTexture(debugWater, tickIndex, packetTileId, out HeightmapPacket packet)) { lastHeightmapPacket = packet; HeightmapPacketReady?.Invoke(packet); } } private void DispatchGhost() { ghostExchangeShader.SetInt("_GridRes", gridRes); ghostExchangeShader.SetInt("_HasNorth", 0); ghostExchangeShader.SetInt("_HasSouth", 0); ghostExchangeShader.SetInt("_HasEast", 0); ghostExchangeShader.SetInt("_HasWest", 0); ghostExchangeShader.SetTexture(ghostKernel, "_WaterIn", CurrentWater); ghostExchangeShader.SetTexture(ghostKernel, "_VelIn", CurrentVelocity); ghostExchangeShader.SetTexture(ghostKernel, "_WaterNorth", CurrentWater); ghostExchangeShader.SetTexture(ghostKernel, "_VelNorth", CurrentVelocity); ghostExchangeShader.SetTexture(ghostKernel, "_WaterSouth", CurrentWater); ghostExchangeShader.SetTexture(ghostKernel, "_VelSouth", CurrentVelocity); ghostExchangeShader.SetTexture(ghostKernel, "_WaterEast", CurrentWater); ghostExchangeShader.SetTexture(ghostKernel, "_VelEast", CurrentVelocity); ghostExchangeShader.SetTexture(ghostKernel, "_WaterWest", CurrentWater); ghostExchangeShader.SetTexture(ghostKernel, "_VelWest", CurrentVelocity); ghostExchangeShader.SetTexture(ghostKernel, "_WaterOut", waterGhost); ghostExchangeShader.SetTexture(ghostKernel, "_VelOut", velGhost); DispatchKernel(ghostExchangeShader, ghostKernel, gridRes + 2, gridRes + 2); } private void DispatchFlux(float dt, float rainRate) { fluxShader.SetInt("_GridRes", gridRes); fluxShader.SetFloat("_Dx", dx); fluxShader.SetFloat("_Dt", dt); fluxShader.SetFloat("_Gravity", Gravity); fluxShader.SetFloat("_RainRate", Mathf.Max(0.0f, rainRate)); fluxShader.SetInt("_UsePorosity", resolvedPorosity != null ? 1 : 0); fluxShader.SetTexture(fluxKernel, "_WaterIn", waterGhost); fluxShader.SetTexture(fluxKernel, "_VelIn", velGhost); fluxShader.SetTexture(fluxKernel, "_WaterOut", NextWater); fluxShader.SetTexture(fluxKernel, "_VelOut", NextVelocity); fluxShader.SetTexture(fluxKernel, "_TerrainHeight", resolvedTerrain); fluxShader.SetTexture(fluxKernel, "_Porosity", resolvedPorosity); fluxShader.SetInt("_DebugClamp", debugClampStats ? 1 : 0); fluxShader.SetTexture(fluxKernel, "_ClampMask", debugClampStats && clampMask != null ? clampMask : dummyRw); DispatchKernel(fluxShader, fluxKernel, gridRes, gridRes); } private void DispatchScale(float scale) { int scaleKernel = fluxShader.FindKernel("ScaleWater"); fluxShader.SetInt("_GridRes", gridRes); fluxShader.SetFloat("_Scale", Mathf.Clamp(scale, 0.0f, 10.0f)); fluxShader.SetTexture(scaleKernel, "_WaterOut", CurrentWater); DispatchKernel(fluxShader, scaleKernel, gridRes, gridRes); } private void SwapBuffers() { useA = !useA; } private void RequestReadback() { if (!SystemInfo.supportsAsyncGPUReadback) { return; } RenderTexture water = CurrentWater; RenderTexture vel = CurrentVelocity; if (useFullPrecisionTextures) { AsyncGPUReadback.Request(water, 0, request => { if (!request.hasError) { DepthStats stats = ComputeDepthStats(request.GetData()); if (!float.IsNaN(stats.Max) && !float.IsInfinity(stats.Max)) { lastMaxDepth = stats.Max; } if (!float.IsNaN(stats.Sum) && !float.IsInfinity(stats.Sum)) { float previousTotal = lastTotalDepth > 0.0f ? lastTotalDepth : stats.Sum; lastTotalDepth = stats.Sum; lastTotalVolume = stats.Sum * (dx * dx); if (applyMassCorrection && stats.Sum > 1e-6f) { float scale = previousTotal / stats.Sum; pendingMassScale = Mathf.Clamp(scale, massCorrectionMinScale, massCorrectionMaxScale); } if (logConservation) { Debug.Log($"SWE mass: totalDepth={lastTotalDepth:F3}m volume={lastTotalVolume:F3}m^3 scaleNext={pendingMassScale:F3}"); } } } }); AsyncGPUReadback.Request(vel, 0, request => { if (!request.hasError) { float speed = ComputeMaxSpeed(request.GetData()); if (!float.IsNaN(speed) && !float.IsInfinity(speed)) { lastMaxSpeed = speed; } } }); } else { AsyncGPUReadback.Request(water, 0, request => { if (!request.hasError) { DepthStats stats = ComputeDepthStats(request.GetData()); if (!float.IsNaN(stats.Max) && !float.IsInfinity(stats.Max)) { lastMaxDepth = stats.Max; } if (!float.IsNaN(stats.Sum) && !float.IsInfinity(stats.Sum)) { float previousTotal = lastTotalDepth > 0.0f ? lastTotalDepth : stats.Sum; lastTotalDepth = stats.Sum; lastTotalVolume = stats.Sum * (dx * dx); if (applyMassCorrection && stats.Sum > 1e-6f) { float scale = previousTotal / stats.Sum; pendingMassScale = Mathf.Clamp(scale, massCorrectionMinScale, massCorrectionMaxScale); } if (logConservation) { Debug.Log($"SWE mass: totalDepth={lastTotalDepth:F3}m volume={lastTotalVolume:F3}m^3 scaleNext={pendingMassScale:F3}"); } } } }); AsyncGPUReadback.Request(vel, 0, request => { if (!request.hasError) { float speed = ComputeMaxSpeed(request.GetData()); if (!float.IsNaN(speed) && !float.IsInfinity(speed)) { lastMaxSpeed = speed; } } }); } if (debugClampStats && clampMask != null) { AsyncGPUReadback.Request(clampMask, 0, request => { if (!request.hasError) { ClampStats stats = useFullPrecisionTextures ? CountClampStats(request.GetData()) : CountClampStats(request.GetData()); lastClampedCells = stats.Clamped; lastNanCells = stats.NaNs; lastClampedRatio = stats.Clamped / (float)(gridRes * gridRes); } }); } } private static DepthStats ComputeDepthStats(NativeArray data) { float max = 0.0f; double sum = 0.0; for (int i = 0; i < data.Length; i++) { float h = HalfToFloat(data[i]); if (h > max) { max = h; } if (h > 0.0f) { sum += h; } } return new DepthStats(max, (float)sum); } private static DepthStats ComputeDepthStats(NativeArray data) { float max = 0.0f; double sum = 0.0; for (int i = 0; i < data.Length; i++) { float h = data[i]; if (h > max) { max = h; } if (h > 0.0f) { sum += h; } } return new DepthStats(max, (float)sum); } private static float ComputeMaxSpeed(NativeArray data) { float max = 0.0f; int count = data.Length / 2; for (int i = 0; i < count; i++) { float u = HalfToFloat(data[i * 2]); float v = HalfToFloat(data[i * 2 + 1]); float speed = Mathf.Sqrt(u * u + v * v); if (speed > max) { max = speed; } } return max; } private static float ComputeMaxSpeed(NativeArray data) { float max = 0.0f; int count = data.Length / 2; for (int i = 0; i < count; i++) { float u = data[i * 2]; float v = data[i * 2 + 1]; float speed = Mathf.Sqrt(u * u + v * v); if (speed > max) { max = speed; } } return max; } private static ClampStats CountClampStats(NativeArray data) { int clamped = 0; int nans = 0; for (int i = 0; i < data.Length; i++) { float v = HalfToFloat(data[i]); if (v > 1.5f) { nans++; } else if (v > 0.5f) { clamped++; } } return new ClampStats(clamped, nans); } private static ClampStats CountClampStats(NativeArray data) { int clamped = 0; int nans = 0; for (int i = 0; i < data.Length; i++) { float v = data[i]; if (v > 1.5f) { nans++; } else if (v > 0.5f) { clamped++; } } return new ClampStats(clamped, nans); } private readonly struct ClampStats { public readonly int Clamped; public readonly int NaNs; public ClampStats(int clamped, int nans) { Clamped = clamped; NaNs = nans; } } private static float HalfToFloat(ushort half) { uint sign = (uint)(half >> 15) & 0x00000001u; int exp = (half >> 10) & 0x0000001F; int mant = half & 0x000003FF; if (exp == 0) { if (mant == 0) { return sign == 1 ? -0.0f : 0.0f; } exp = 1; while ((mant & 0x00000400) == 0) { mant <<= 1; exp--; } mant &= 0x000003FF; } else if (exp == 31) { uint infNaN = (sign << 31) | 0x7F800000u | ((uint)mant << 13); return BitConverter.ToSingle(BitConverter.GetBytes(infNaN), 0); } uint fexp = (uint)(exp + (127 - 15)); uint fmant = (uint)(mant << 13); uint bits = (sign << 31) | (fexp << 23) | fmant; return BitConverter.ToSingle(BitConverter.GetBytes(bits), 0); } private readonly struct DepthStats { public readonly float Max; public readonly float Sum; public DepthStats(float max, float sum) { Max = max; Sum = sum; } } private static Texture2D CreateTextureFromFloats(float[] data, int width, int height, string name) { Texture2D texture = new Texture2D(width, height, TextureFormat.RFloat, false, true) { name = name, wrapMode = TextureWrapMode.Clamp, filterMode = FilterMode.Point }; texture.SetPixelData(data, 0); texture.Apply(false, true); return texture; } private static Texture2D CreateFlatTexture(int resolution, float value) { int length = resolution * resolution; float[] data = new float[length]; for (int i = 0; i < length; i++) { data[i] = value; } return CreateTextureFromFloats(data, resolution, resolution, "SWE_Flat"); } private static void Clamp01(float[] data) { for (int i = 0; i < data.Length; i++) { data[i] = Mathf.Clamp01(data[i]); } } private static RenderTexture CreateRenderTexture(int width, int height, RenderTextureFormat format, string name) { RenderTexture rt = new RenderTexture(width, height, 0, format) { name = name, enableRandomWrite = true, filterMode = FilterMode.Point, wrapMode = TextureWrapMode.Clamp }; rt.Create(); return rt; } private static void ClearRenderTexture(RenderTexture rt) { if (rt == null) { return; } RenderTexture prev = RenderTexture.active; RenderTexture.active = rt; GL.Clear(false, true, Color.clear); RenderTexture.active = prev; } private static void ReleaseRenderTexture(RenderTexture rt) { if (rt != null) { rt.Release(); } } private static void DispatchKernel(ComputeShader shader, int kernel, int width, int height) { uint x; uint y; uint z; shader.GetKernelThreadGroupSizes(kernel, out x, out y, out z); int groupsX = Mathf.CeilToInt(width / (float)x); int groupsY = Mathf.CeilToInt(height / (float)y); shader.Dispatch(kernel, groupsX, groupsY, 1); } }