PhysarumSimulationHandler.cs

using System.Collections;
using System.Collections.Generic;
using UnityEngine;
using System.Runtime.InteropServices;
using System.Linq;
using UnityEngine.UI;

/// <summary>
/// contain basic parameters of compute shader
/// </summary>
public class KernelParamsHandler
{
    // name of kernel
    public string _name { get; private set; }
    // index of kernel
    public int _index { get; private set; }
    // thread size x
    public uint _x { get; private set; }
    // thread size y
    public uint _y { get; private set; }
    // thread size z
    public uint _z { get; private set; }
    // thread size in signed int
    public Vector3Int _signedThreadSize { get; private set; }

    /// <summary>
    /// constructor
    /// </summary>
    /// <param name="computeShader"></param>
    /// <param name="kernelName"></param>
    public KernelParamsHandler(ComputeShader computeShader, string kernelName)
    {
        _name = kernelName;
        _index = computeShader.FindKernel(_name);

        uint x, y, z;
        computeShader.GetKernelThreadGroupSizes(_index, out x, out y, out z);
        _x = x;
        _y = y;
        _z = z;

        _signedThreadSize = new Vector3Int((int)_x, (int)_y, (int)_z);
    }
}

/// <summary>
/// physarum simulation parameters
/// </summary>
[System.Serializable]
public struct PhysarumSimulationParams
{
    // angle to turn (degrees)
    public float _turnAngle;

    // angle of sensor (degrees)
    public float _sensorAngle;

    // distance of sensors
    public float _sensorDistance;

    // distance to move
    public float _moveDistance;

    // deposit
    public float _depositValue;

    // decay
    public float _decayFactor;
}

public class PhysarumSimulationHandler : MonoBehaviour
{
    [SerializeField, Tooltip("Compute Shader to diffuse texture")]
    protected ComputeShader _computeShader;

    [SerializeField, Tooltip("resolution of the texture")]
    protected Vector2Int _textureResolution = new Vector2Int(512, 512);

    // render texture to paint
    protected RenderTexture _resultTexture;

    // name of render texture on the compute shader
    protected string _resultTextureName = "_resultTexture";

    // kernel to simulate diffusion of render texture
    protected KernelParamsHandler _simulationKernel;

    // name of the simulation kernel
    protected string _simulationKernelName;

    // kernel to reset render texture
    protected KernelParamsHandler _resetKernel;

    // name of the reset kernel
    protected string _resetKernelName = "Reset";

    [SerializeField, Tooltip("optional input as a texture")]
    protected Texture _optionalInputTexture;

    protected string _optionalInputTextureName = "_inputTexture";

    [SerializeField, Tooltip("mesh to check paint result")]
    protected MeshRenderer _testMesh;

    [SerializeField, Tooltip("ui image to check paint result")]
    protected RawImage _testUiImage;

    // whether this diffusion system is enabled or not
    protected bool _isEnabled;

    // kernel to initialize particles
    private KernelParamsHandler _initParticlesKernel;

    // name of _initParticlesKernel in the compute shader
    private string _initParticlesKernelName = "InitParticles";

    // kernel to initialize input
    private KernelParamsHandler _initInputKernel;

    // name of _initInputKernel in the compute shader
    private string _initInputKernelName = "InitInput";

    // name of _diffuseDecayKernel in the compute shader
    private string _simulateyKernelName = "Simulate";

    // kernel to simulate duffse and decay
    private KernelParamsHandler _diffuseDecayKernel;

    // name of _diffuseDecayKernel in the compute shader
    private string _diffuseDecayKernelName = "DiffuseDecay";

    [SerializeField, Tooltip("whether use input texture for initialization")]
    private bool _useInputTexture;

    // name of _useInputTexture in the compute shader
    private string _useInputTextureName = "_useInputTexture";

    [SerializeField, Tooltip("number of particles")]
    private int _particleNum = 100000;

    // name of _particleNum in the compute shader
    private string _particleNumName = "_particleNum";

    // struct of a particle
    private struct Particle
    {
        public Vector2 position;
        public Vector2 direction;
    }

    // buffer for particles
    private ComputeBuffer _particles;

    // name of particles buffer in the compute shader
    private string _particlesName = "_particles";

    [SerializeField, Tooltip("simulation parameters")]
    private PhysarumSimulationParams[] _simulationParams = new PhysarumSimulationParams[2];

    // name of _sensorRRot in compute shader
    private string _sensorRRotName = "_sensorRRot";

    // name of _sensorLRot in compute shader
    private string _sensorLRotName = "_sensorLRot";

    // name of _turnRRot in compute shader
    private string _turnRRotName = "_turnRRot";

    // name of _turnLRot in compute shader
    private string _turnLRotName = "_turnLRot";

    // name of _sensorAngle in compute shader
    private string _sensorAngleName = "_sensorAngle";

    // name of _turnAngle in compute shader
    private string _turnAngleName = "_turnAngle";

    // name of _moveDistance in compute shader
    private string _moveDistanceName = "_moveDistance";

    // name of _sensorDistance in compute shader
    private string _sensorDistanceName = "_sensorDistance";

    // name of _depositValue in compute shader
    private string _depositValueName = "_depositValue";

    // name of _decayFactor in compute shader
    private string _decayFactorName = "_decayFactor";

    private Particle[] _debugParticles;

    [SerializeField, Tooltip("particle index to show for debugging")]
    private int _debugParticleIndex;

    /// <summary>
    /// initialization
    /// </summary>
    private void Start()
    {
        _debugParticles = new Particle[_particleNum];
        InitTexture();
        InitKenel();
        InitSimulation();
        SetTextureOnUI();
        SetTextureOnMesh();
        //ResetSimulation();
    }

    /// <summary>
    /// update simulation
    /// </summary>
    private void Update()
    {
        if (Input.GetKeyDown(KeyCode.Escape))
        {
            ResetSimulation();
        }

        if (Input.GetKeyDown(KeyCode.Space))
        {
            _isEnabled = !_isEnabled;
        }

        if (_isEnabled) 
        {
            Simulate();
        }

        if (Input.GetKeyDown(KeyCode.A))
        {
            Debug.Log("pos: " + _debugParticles[_debugParticleIndex].position + " / dir: " + _debugParticles[_debugParticleIndex].direction);
        }
    }

    private void OnDestroy()
    {
        _resultTexture.Release();

        if (_optionalInputTexture is RenderTexture)
        {
            RenderTexture noiseTexture = _optionalInputTexture as RenderTexture;
            noiseTexture.Release();
        }
        
        _particles.Dispose();
    }

    /// <summary>
    /// initialize render texture
    /// </summary>
    private void InitTexture()
    {
        _resultTexture = new RenderTexture(_textureResolution.x, _textureResolution.y, 0, RenderTextureFormat.ARGB32);
        _resultTexture.enableRandomWrite = true;
        _resultTexture.Create();
    }

    /// <summary>
    /// initialize diffuse kernel
    /// </summary>
    private void InitSimulationKernel(string diffuseKernelName)
    {
        _simulationKernelName = diffuseKernelName;
        _simulationKernel = new KernelParamsHandler(_computeShader, _simulationKernelName);
        _computeShader.SetTexture(_simulationKernel._index, _resultTextureName, _resultTexture);
    }

    /// <summary>
    /// initialize reset kernel
    /// </summary>
    private void InitResetKernel()
    {
        _resetKernel = new KernelParamsHandler(_computeShader, _resetKernelName);
        _computeShader.SetTexture(_resetKernel._index, _resultTextureName, _resultTexture);
    }

    /// <summary>
    /// initialize kernel
    /// </summary>
    private void InitKenel()
    {
        // main simulation; sense, rotate, move, deposit
        InitSimulationKernel(_simulateyKernelName);
        

        // reset
        InitResetKernel();

        // initialize particles
        _initParticlesKernel = new KernelParamsHandler(_computeShader, _initParticlesKernelName);

        // initialize input
        _initInputKernel = new KernelParamsHandler(_computeShader, _initInputKernelName);
        _computeShader.SetBool(_useInputTextureName, _useInputTexture);

        // diffuse and decay
        _diffuseDecayKernel = new KernelParamsHandler(_computeShader, _diffuseDecayKernelName);
        _computeShader.SetTexture(_diffuseDecayKernel._index, _resultTextureName, _resultTexture);

        if (_useInputTexture)
        {
            
        }
        else
        {
            _initInputKernel = new KernelParamsHandler(_computeShader, _initInputKernelName);
        }

        
    }


    /// <summary>
    /// initialize simulation
    /// </summary>
    private void InitSimulation()
    {
        // set parameters
        SetSimulationParams();

        // initialize particles
        _particles = new ComputeBuffer(_particleNum, Marshal.SizeOf(typeof(Particle)));
        
        Particle defaultParticle = new Particle();
        defaultParticle.position = Vector2.zero;
        defaultParticle.direction = Vector2.up;

        _particles.SetData(Enumerable.Repeat<Particle>(defaultParticle, _particleNum).ToArray());
        _computeShader.SetBuffer(_initParticlesKernel._index, _particlesName, _particles);
        _computeShader.SetBuffer(_simulationKernel._index, _particlesName, _particles);

        _computeShader.SetInt(_particleNumName, _particleNum);
        _computeShader.SetTexture(_initParticlesKernel._index, _resultTextureName, _resultTexture);
        _computeShader.Dispatch(_initParticlesKernel._index, _particleNum / _initParticlesKernel._signedThreadSize.x, _initParticlesKernel._signedThreadSize.y, _initParticlesKernel._signedThreadSize.z);
        
        _particles.GetData(_debugParticles);
        Debug.Log("pos: " + _debugParticles[_debugParticleIndex].position + " / dir: " + _debugParticles[_debugParticleIndex].direction);

        // initialize input
        _computeShader.SetBool(_useInputTextureName, _useInputTexture);
        if (_useInputTexture)
        {
            // set input texture
            _computeShader.SetTexture(_initInputKernel._index, _optionalInputTextureName, _optionalInputTexture);
            _computeShader.SetTexture(_simulationKernel._index, _optionalInputTextureName, _optionalInputTexture);
            _computeShader.SetTexture(_diffuseDecayKernel._index, _optionalInputTextureName, _optionalInputTexture);
        }
        else
        {
            // initialize with noise in the compute shader
            RenderTexture noiseInputTexture = new RenderTexture(_resultTexture);
            noiseInputTexture.enableRandomWrite = true;
            noiseInputTexture.Create();
            _optionalInputTexture = noiseInputTexture;
            _computeShader.SetTexture(_initInputKernel._index, _optionalInputTextureName, _optionalInputTexture);
        }

        _computeShader.SetTexture(_initInputKernel._index, _resultTextureName, _resultTexture);
        _computeShader.Dispatch(_initInputKernel._index, _textureResolution.x / _initInputKernel._signedThreadSize.x, _textureResolution.y / _initInputKernel._signedThreadSize.y, _initInputKernel._signedThreadSize.z);

    }

    /// <summary>
    /// set simulation parameters
    /// </summary>
    private void SetSimulationParams()
    {
        // set rotation matrix for sensor and turn left/right
        /*
        Matrix4x4 sensorLRot = DegreeToRotationMatrix(_simulationParams[0]._sensorAngle);
        _computeShader.SetMatrix(_sensorRRotName, sensorLRot);

        Matrix4x4 sensorRRot = DegreeToRotationMatrix(-_simulationParams[0]._sensorAngle);
        _computeShader.SetMatrix(_sensorRRotName, sensorRRot);

        Matrix4x4 turnLRot = DegreeToRotationMatrix(_simulationParams[0]._turnAngle);
        _computeShader.SetMatrix(_sensorRRotName, turnLRot);

        Matrix4x4 turnRRot = DegreeToRotationMatrix(-_simulationParams[0]._turnAngle);
        _computeShader.SetMatrix(_sensorRRotName, turnRRot);
        */

        for(int i = 0; i < _simulationParams.Length; i++)
        {
            // set sensor angle
            _computeShader.SetFloat(_sensorAngleName + i.ToString(), _simulationParams[i]._sensorAngle);

            // set turn angle
            _computeShader.SetFloat(_turnAngleName + i.ToString(), _simulationParams[i]._turnAngle);

            // set move distance
            _computeShader.SetFloat(_moveDistanceName + i.ToString(), _simulationParams[i]._moveDistance);

            // set sensor distance
            _computeShader.SetFloat(_sensorDistanceName + i.ToString(), _simulationParams[i]._sensorAngle);

            // set deposit value
            _computeShader.SetFloat(_depositValueName + i.ToString(), _simulationParams[i]._depositValue);

            // set decay factor
            _computeShader.SetFloat(_decayFactorName + i.ToString(), _simulationParams[i]._decayFactor);
        }
        


    }

    /// <summary>
    /// create rotation matrix from degrees
    /// </summary>
    /// <param name="angleInDegree"></param>
    /// <returns></returns>
    private Matrix4x4 DegreeToRotationMatrix(float angleInDegree)
    {
        Matrix4x4 rotationMatrix = Matrix4x4.identity;

        rotationMatrix.SetRow(0, new Vector4(
            Mathf.Cos(angleInDegree * Mathf.Deg2Rad),
            -Mathf.Sin(angleInDegree * Mathf.Deg2Rad),
            0,
            0
        ));
        rotationMatrix.SetRow(1, new Vector4(
            Mathf.Sin(angleInDegree * Mathf.Deg2Rad),
            Mathf.Cos(angleInDegree * Mathf.Deg2Rad),
            0,
            0
        ));

        return rotationMatrix;
    }

    /// <summary>
    /// execute simulation
    /// </summary>
    private void Simulate()
    {
        SetSimulationParams();


        if (_useInputTexture)
        {
            _computeShader.SetTexture(_simulationKernel._index, _optionalInputTextureName, _optionalInputTexture);
            _computeShader.SetTexture(_diffuseDecayKernel._index, _optionalInputTextureName, _optionalInputTexture);
        }

        _computeShader.Dispatch(_simulationKernel._index, _particleNum / _simulationKernel._signedThreadSize.x, _simulationKernel._signedThreadSize.y, _simulationKernel._signedThreadSize.z);
        _particles.GetData(_debugParticles);
        _computeShader.Dispatch(_diffuseDecayKernel._index, _textureResolution.x / _diffuseDecayKernel._signedThreadSize.x, _textureResolution.y / _diffuseDecayKernel._signedThreadSize.y, _diffuseDecayKernel._signedThreadSize.z);

    }

    /// <summary>
    /// reset simulation and texture
    /// </summary>
    public virtual void ResetSimulation()
    {
        _computeShader.Dispatch(_resetKernel._index, _textureResolution.x / _resetKernel._signedThreadSize.x, _textureResolution.y / _resetKernel._signedThreadSize.y, _resetKernel._signedThreadSize.z);
    }

    /// <summary>
    /// assign paint texture on the test mesh
    /// </summary>
    private void SetTextureOnMesh()
    {
        if (!_testMesh) { return; }
        _testMesh.material.SetTexture("_MainTex", _resultTexture);
    }

    /// <summary>
    /// assign paint texture on the test ui
    /// </summary>
    private void SetTextureOnUI()
    {
        if (!_testUiImage) { return; }
        _testUiImage.texture = _resultTexture;
    }
}