Assets/scripts/slice.cs - main - molo4ka/wind-simulation

Assets/scripts/slice.cs
458 строк13 KB
molo4kaPlus
Initial check-in
18 май 2026, 12:33
18 май 2026, 12:333216236
using UnityEngine;

public class ParticlePressureSlice : MonoBehaviour
{
    public enum SliceAxis
    {
        X,
        Y,
        Z
    }

    [Header("Source")]
    public ParticleGenerator generator;
    public ComputeShader sliceShader;

    [Header("Slice")]
    public SliceAxis sliceAxis = SliceAxis.X;
    public float slicePosition = 0f;
    public float sliceThickness = 0.12f;
    public int sliceResolution = 64;

    [Header("Display")]
    [Range(0f, 1f)]
    public float opacity = 0.9f;

    public bool updateEveryFrame = true;
    public bool showUI = true;

    RenderTexture pressureTexture;
    RenderTexture blurTexture;
    RenderTexture maxTexture;

    ComputeBuffer maxBuffer;
    Material previewMaterial;

    int clearKernel;
    int buildKernel;
    int blurHorizontalKernel;
    int blurVerticalKernel;
    int reduceKernel;
    int finalReduceKernel;

    float lastSlicePosition;
    float lastSliceThickness;
    int lastResolution;
    SliceAxis lastAxis;

    void OnEnable()
    {
        if (generator == null)
            generator = GetComponent<ParticleGenerator>();

        Initialize();
        MarkDirty();
    }

    void OnDisable()
    {
        ReleaseResources();
    }

    void Update()
    {
        if (generator == null)
            return;

        if (!Application.isPlaying && !updateEveryFrame)
            return;

        if (updateEveryFrame || IsDirty())
            BuildSliceGPU();
    }

    bool IsDirty()
    {
        return
            lastSlicePosition != slicePosition ||
            lastSliceThickness != sliceThickness ||
            lastResolution != sliceResolution ||
            lastAxis != sliceAxis;
    }

    void MarkDirty()
    {
        lastSlicePosition = slicePosition;
        lastSliceThickness = sliceThickness;
        lastResolution = sliceResolution;
        lastAxis = sliceAxis;
    }

    void Initialize()
    {
        if (sliceShader == null || generator == null)
            return;

        clearKernel = sliceShader.FindKernel("ClearPressure");
        buildKernel = sliceShader.FindKernel("BuildPressure");
        blurHorizontalKernel = sliceShader.FindKernel("BlurHorizontal");
        blurVerticalKernel = sliceShader.FindKernel("BlurVertical");
        reduceKernel = sliceShader.FindKernel("ReduceMax");
        finalReduceKernel = sliceShader.FindKernel("FinalReduce");

        CreateResources();

        if (previewMaterial == null)
        {
            Shader s = Shader.Find("Hidden/ParticlePressurePreview");
            if (s != null)
                previewMaterial = new Material(s);
        }
    }

    void CreateResources()
    {
        CreateRT(ref pressureTexture);
        CreateRT(ref blurTexture);

        if (maxTexture == null)
        {
            maxTexture = new RenderTexture(1, 1, 0, RenderTextureFormat.RFloat);
            maxTexture.enableRandomWrite = true;
            maxTexture.filterMode = FilterMode.Point;
            maxTexture.wrapMode = TextureWrapMode.Clamp;
            maxTexture.Create();
        }

        int groups = Mathf.CeilToInt(sliceResolution / 8f);
        int totalGroups = groups * groups;

        if (maxBuffer == null || maxBuffer.count != totalGroups)
        {
            if (maxBuffer != null)
                maxBuffer.Release();

            maxBuffer = new ComputeBuffer(totalGroups, sizeof(float));
        }
    }

    void CreateRT(ref RenderTexture rt)
    {
        if (rt != null && rt.width == sliceResolution)
            return;

        if (rt != null)
            rt.Release();

        rt = new RenderTexture(
            sliceResolution,
            sliceResolution,
            0,
            RenderTextureFormat.RFloat
        );

        rt.enableRandomWrite = true;
        rt.wrapMode = TextureWrapMode.Clamp;
        rt.filterMode = FilterMode.Bilinear;
        rt.Create();
    }

    void ReleaseResources()
    {
        if (pressureTexture != null) pressureTexture.Release();
        if (blurTexture != null) blurTexture.Release();
        if (maxTexture != null) maxTexture.Release();
        if (maxBuffer != null) maxBuffer.Release();

        pressureTexture = null;
        blurTexture = null;
        maxTexture = null;
        maxBuffer = null;

        if (previewMaterial != null)
        {
            DestroyImmediate(previewMaterial);
            previewMaterial = null;
        }
    }

    public void BuildSliceGPU()
    {
        if (sliceShader == null || generator == null)
            return;

        if (generator.ParticleBuffer == null)
            return;

        CreateResources();

        int groups = Mathf.CeilToInt(sliceResolution / 8f);

        sliceShader.SetInt("_Resolution", sliceResolution);
        sliceShader.SetInt("_NumParticles", generator.particleCount);
        sliceShader.SetInt("_GroupCount", groups);

        sliceShader.SetVector("_SimulationCenter", generator.SimulationCenter);
        sliceShader.SetVector("_SimulationSize", generator.SimulationSize);

        sliceShader.SetInt("_SliceAxis", (int)sliceAxis);
        sliceShader.SetFloat("_SlicePosition", slicePosition);
        sliceShader.SetFloat("_SliceThickness", sliceThickness);

        sliceShader.SetBuffer(buildKernel, "_Particles", generator.ParticleBuffer);

        // spatial grid
        var gridHeads = typeof(ParticleGenerator)
            .GetField("gridHeads", System.Reflection.BindingFlags.NonPublic | System.Reflection.BindingFlags.Instance)
            .GetValue(generator) as ComputeBuffer;

        var gridNext = typeof(ParticleGenerator)
            .GetField("gridNext", System.Reflection.BindingFlags.NonPublic | System.Reflection.BindingFlags.Instance)
            .GetValue(generator) as ComputeBuffer;

        var gridResolution = (Vector3Int)typeof(ParticleGenerator)
            .GetField("gridResolution", System.Reflection.BindingFlags.NonPublic | System.Reflection.BindingFlags.Instance)
            .GetValue(generator);

        float cellSize = (float)typeof(ParticleGenerator)
            .GetField("cellSize", System.Reflection.BindingFlags.NonPublic | System.Reflection.BindingFlags.Instance)
            .GetValue(generator);

        sliceShader.SetBuffer(buildKernel, "_GridHeads", gridHeads);
        sliceShader.SetBuffer(buildKernel, "_GridNext", gridNext);
        sliceShader.SetInts("_GridResolution", gridResolution.x, gridResolution.y, gridResolution.z);
        sliceShader.SetFloat("_CellSize", cellSize);

        sliceShader.SetTexture(clearKernel, "_PressureTexture", pressureTexture);
        sliceShader.SetTexture(buildKernel, "_PressureTexture", pressureTexture);

        sliceShader.Dispatch(clearKernel, groups, groups, 1);
        sliceShader.Dispatch(buildKernel, groups, groups, 1);

        sliceShader.SetTexture(blurHorizontalKernel, "_PressureTexture", pressureTexture);
        sliceShader.SetTexture(blurHorizontalKernel, "_BlurTexture", blurTexture);

        sliceShader.SetTexture(blurVerticalKernel, "_PressureTexture", pressureTexture);
        sliceShader.SetTexture(blurVerticalKernel, "_BlurTexture", blurTexture);

        sliceShader.Dispatch(blurHorizontalKernel, groups, groups, 1);
        sliceShader.Dispatch(blurVerticalKernel, groups, groups, 1);

        sliceShader.SetTexture(reduceKernel, "_PressureTexture", pressureTexture);
        sliceShader.SetBuffer(reduceKernel, "_MaxBuffer", maxBuffer);

        sliceShader.Dispatch(reduceKernel, groups, groups, 1);

        sliceShader.SetBuffer(finalReduceKernel, "_MaxBuffer", maxBuffer);
        sliceShader.SetTexture(finalReduceKernel, "_MaxTexture", maxTexture);

        sliceShader.Dispatch(finalReduceKernel, 1, 1, 1);

        MarkDirty();
    }

    void OnRenderObject()
    {
        DrawSlice();
    }

    void DrawSlice()
    {
        if (previewMaterial == null || pressureTexture == null || generator == null)
            return;

        Vector3 center = generator.SimulationCenter;
        Vector3 size = generator.SimulationSize;

        Vector3 planeCenter = center;
        Vector3 right = Vector3.right;
        Vector3 up = Vector3.up;

        float width = 1f;
        float height = 1f;

        switch (sliceAxis)
        {
            case SliceAxis.X:
                planeCenter.x = slicePosition;
                right = Vector3.forward;
                up = Vector3.up;
                width = size.z;
                height = size.y;
                break;

            case SliceAxis.Y:
                planeCenter.y = slicePosition;
                right = Vector3.right;
                up = Vector3.forward;
                width = size.x;
                height = size.z;
                break;

            case SliceAxis.Z:
                planeCenter.z = slicePosition;
                right = Vector3.right;
                up = Vector3.up;
                width = size.x;
                height = size.y;
                break;
        }

        Vector3 r = right * width * 0.5f;
        Vector3 u = up * height * 0.5f;

        Vector3 p0 = planeCenter - r - u;
        Vector3 p1 = planeCenter + r - u;
        Vector3 p2 = planeCenter + r + u;
        Vector3 p3 = planeCenter - r + u;

        previewMaterial.SetTexture("_PressureTexture", pressureTexture);
        previewMaterial.SetTexture("_MaxTexture", maxTexture);
        previewMaterial.SetFloat("_Opacity", opacity);

        previewMaterial.SetPass(0);

        GL.PushMatrix();

        GL.Begin(GL.QUADS);

        GL.TexCoord2(0, 0); GL.Vertex(p0);
        GL.TexCoord2(1, 0); GL.Vertex(p1);
        GL.TexCoord2(1, 1); GL.Vertex(p2);
        GL.TexCoord2(0, 1); GL.Vertex(p3);

        GL.End();

        GL.PopMatrix();
    }

    public void SaveSlicePNG(string fileName = null)
    {
        if (pressureTexture == null)
        {
            Debug.LogWarning("Slice texture is missing.");
            return;
        }

        var readback = new Texture2D(
            pressureTexture.width,
            pressureTexture.height,
            TextureFormat.RFloat,
            false
        );

        RenderTexture previous = RenderTexture.active;
        RenderTexture.active = pressureTexture;

        readback.ReadPixels(
            new Rect(0, 0, pressureTexture.width, pressureTexture.height),
            0,
            0
        );
        readback.Apply();

        RenderTexture.active = previous;

        var pixels = readback.GetPixels();
        var output = new Texture2D(
            pressureTexture.width,
            pressureTexture.height,
            TextureFormat.RGBA32,
            false
        );

        float maxValue = 0f;

        for (int i = 0; i < pixels.Length; i++)
            maxValue = Mathf.Max(maxValue, pixels[i].r);

        maxValue = Mathf.Max(maxValue, 0.0001f);

        for (int i = 0; i < pixels.Length; i++)
        {
            float v = pixels[i].r / maxValue;
            output.SetPixel(i % pressureTexture.width, i / pressureTexture.width, new Color(v, v, v, 1f));
        }

        output.Apply();

        byte[] bytes = output.EncodeToPNG();

        if (string.IsNullOrEmpty(fileName))
            fileName = "slice_" + System.DateTime.Now.ToString("yyyyMMdd_HHmmss") + ".png";

        string path = System.IO.Path.Combine(
            System.Environment.GetFolderPath(System.Environment.SpecialFolder.Desktop),
            fileName
        );

        System.IO.File.WriteAllBytes(path, bytes);

        Destroy(readback);
        Destroy(output);

        Debug.Log("Saved slice: " + path);
    }

    public void SaveRenderedSlicePNG(int outputResolution = 1024, string fileName = null)
    {
        if (pressureTexture == null || previewMaterial == null)
            return;

        var outputRT = new RenderTexture(
            outputResolution,
            outputResolution,
            0,
            RenderTextureFormat.ARGB32
        );

        outputRT.enableRandomWrite = false;
        outputRT.filterMode = FilterMode.Bilinear;
        outputRT.wrapMode = TextureWrapMode.Clamp;
        outputRT.Create();

        previewMaterial.SetTexture("_PressureTexture", pressureTexture);
        previewMaterial.SetTexture("_MaxTexture", maxTexture);
        previewMaterial.SetFloat("_Opacity", 1f);

        Graphics.Blit(null, outputRT, previewMaterial);

        var tex = new Texture2D(
            outputResolution,
            outputResolution,
            TextureFormat.RGBA32,
            false
        );

        RenderTexture previous = RenderTexture.active;
        RenderTexture.active = outputRT;

        tex.ReadPixels(
            new Rect(0, 0, outputResolution, outputResolution),
            0,
            0
        );

        tex.Apply();

        RenderTexture.active = previous;

        byte[] bytes = tex.EncodeToPNG();

        if (string.IsNullOrEmpty(fileName))
            fileName = "slice_" + System.DateTime.Now.ToString("yyyyMMdd_HHmmss") + ".png";

        string path = System.IO.Path.Combine(
            System.Environment.GetFolderPath(System.Environment.SpecialFolder.Desktop),
            fileName
        );

        System.IO.File.WriteAllBytes(path, bytes);

        Destroy(tex);
        outputRT.Release();
        Destroy(outputRT);

        Debug.Log("Saved slice: " + path);
    }
}