Simple example of Smooth Particle Hydrodynamics using Unity Jobs and Burst

SPHBurst.cs

// Smooth Particle Hydrodynamics using burst
// Requires native collections which is removed from unity 2020/2019. This has to be manually installed
// https://forum.unity.com/threads/visibility-changes-for-preview-packages-in-2020-1.910880/
// Burst and Mathematics needed
// @brianasu

using System.Collections.Generic;
using Unity.Burst;
using Unity.Collections;
using Unity.Jobs;
using Unity.Mathematics;
using UnityEngine;

public class SPHBurst : MonoBehaviour
{
    public const uint MAX_NEIGHBOUR = 32; // max number of neighbors to consider. larger number is more accurate

    [SerializeField]
    private float3 _gravity = new float3(0.0f, 3000 * -9.8f, 0.0f);
    [SerializeField]
    private float _restDensity = 1000.0f; // rest density
    [SerializeField]
    private float _gasConstant = 2000.0f; // const for equation of state
    [SerializeField]
    private float _radius = 2.5f; // radius
    [SerializeField]
    private float _mass = 2.0f; // assume all particles have the same mass
    [SerializeField]
    private float _viscosity = 250.0f; // viscosity constant
    [SerializeField]
    private float _timeStep = 0.0005f;
    [SerializeField]
    private float _boundDamping = -0.5f;
    [SerializeField]
    private Bounds _bounds = new Bounds(Vector3.zero, Vector3.one * 128);
    [SerializeField]
    private int _iterations = 1;

    // smoothing kernels defined in Müller and their gradients
    private float _poly6;
    private float _spikyGrad;
    private float _viscLap;

    private NativeArray<float3> _positions;
    private NativeArray<float3> _velocities;
    private NativeArray<float3> _forces;

    private NativeArray<float2> _densityPressures;
    private NativeMultiHashMap<int, int> _gridHash;
    private JobHandle _handleIntegrate;

    private void OnDestroy()
    {
        _handleIntegrate.Complete();
        if (_positions.IsCreated)
        {
            _positions.Dispose();
            _velocities.Dispose();
            _forces.Dispose();
            _densityPressures.Dispose();
            _gridHash.Dispose();
        }
    }

    private void InitSPH()
    {
        _poly6 = 315.0f / (65.0f * Mathf.PI * Mathf.Pow(_radius, 9.0f));
        _spikyGrad = -45.0f / (Mathf.PI * Mathf.Pow(_radius, 6.0f));
        _viscLap = 45.0f / (Mathf.PI * Mathf.Pow(_radius, 6.0f));
        OnDestroy();

        var tmpParticles = new List<float3>();
        for (int x = 20; x < 50; x++)
        {
            for (int z = 20; z < 50; z++)
            {
                for (int y = 20; y < 50; y++)
                {
                    var pos = new float3(x / 50f, y / 50f, z / 50f);
                    pos.x *= _bounds.size.x;
                    pos.y *= _bounds.size.y;
                    pos.z *= _bounds.size.z;
                    pos.x += UnityEngine.Random.value * 0.01f;
                    pos.y += UnityEngine.Random.value * 0.5f;
                    pos.z += UnityEngine.Random.value * 0.5f;
                    pos -= (float3)_bounds.extents;
                    pos += (float3)_bounds.center;
                    pos += new float3(1, 1, 1) * _radius * 0.5f;
                    tmpParticles.Add(new float3(pos.x, pos.y, pos.z));
                }
            }
        }

        Debug.Log(tmpParticles.Count + " particles");
        _positions = new NativeArray<float3>(tmpParticles.Count, Allocator.Persistent);
        _velocities = new NativeArray<float3>(tmpParticles.Count, Allocator.Persistent);
        _forces = new NativeArray<float3>(tmpParticles.Count, Allocator.Persistent);
        _densityPressures = new NativeArray<float2>(tmpParticles.Count, Allocator.Persistent);
        var voxelSize = GetVoxelSize();
        _gridHash = new NativeMultiHashMap<int, int>(voxelSize.x * voxelSize.y * voxelSize.z, Allocator.Persistent);
        _positions.CopyFrom(tmpParticles.ToArray());
    }

    // add all particles to a voxel grid for neighbour lookup
    [BurstCompile]
    public struct InjectGrid : IJob
    {
        [WriteOnly]
        public NativeMultiHashMap<int, int> gridHash;
        [ReadOnly]
        public NativeArray<float3> positions;
        [ReadOnly]
        public int3 voxelSize;
        [ReadOnly]
        public float3 boundsCenter;
        [ReadOnly]
        public float3 boundsExtents;
        [ReadOnly]
        public float3 boundsSize;
        public void Execute()
        {
            gridHash.Clear();
            for (var index = 0; index < positions.Length; index++)
            {
                var p = positions[index];
                var boundPos = p + boundsExtents - boundsCenter;
                var normalizedBoundsPos = boundPos / boundsSize;
                var idx = (int3)math.floor(normalizedBoundsPos * voxelSize);
                var flatIndex = idx.x + idx.y * voxelSize.x + idx.z * (voxelSize.x * voxelSize.y);
                gridHash.Add(flatIndex, index);
            }
        }
    }

    [BurstCompile]
    public struct ComputeDensityPressureJob : IJobParallelFor
    {
        [WriteOnly]
        public NativeArray<float2> densityPressures;

        [ReadOnly]
        public NativeArray<float3> positions;
        [ReadOnly]
        public float hSquared;
        [ReadOnly]
        public float mass;
        [ReadOnly]
        public float poly6;
        [ReadOnly]
        public float gasConst;
        [ReadOnly]
        public float restDensity;
        [ReadOnly]
        public NativeMultiHashMap<int, int> gridHash;
        [ReadOnly]
        public int3 voxelSize;
        [ReadOnly]
        public float3 boundsCenter;
        [ReadOnly]
        public float3 boundsExtents;
        [ReadOnly]
        public float3 boundsSize;

        public void Execute(int index)
        {
            var myPosition = positions[index];
            var density = 0f;

            var boundPos = myPosition + boundsExtents - boundsCenter;
            var normalizedBoundsPos = boundPos / boundsSize;
            var idx = (int3)math.floor(normalizedBoundsPos * voxelSize);
            int total = voxelSize.x * voxelSize.y * voxelSize.z;
            var massPoly6 = mass * poly6;
            var voxelSliceSize = voxelSize.x * voxelSize.y;
            var voxelWidth = voxelSize.x;
            uint count = 0;

            for (var xi = -1; xi <= 1; xi++)
            {
                for (var yi = -1; yi <= 1; yi++)
                {
                    for (var zi = -1; zi <= 1; zi++)
                    {
                        var flatIndex = (idx.x + xi) + (idx.y + yi) * voxelWidth + (idx.z + zi) * voxelSliceSize;
                        if (flatIndex >= 0 && flatIndex < total)
                        {
                            var values = gridHash.GetValuesForKey(flatIndex);
                            while (values.MoveNext())
                            {
                                var otherPosition = positions[values.Current];
                                var rij = otherPosition - myPosition;
                                var lenSquared = math.lengthsq(rij);
                                if (lenSquared < hSquared)
                                {
                                    density += massPoly6 * math.pow(math.max(0, hSquared - lenSquared), 3.0f);
                                    if (count++ > MAX_NEIGHBOUR)
                                    {
                                        break;
                                    }
                                }
                            }
                        }
                    }
                }
            }

            densityPressures[index] = new float2(density, gasConst * (density - restDensity));
        }
    }

    [BurstCompile]
    public struct ComputeForceJob : IJobParallelFor
    {
        [WriteOnly]
        public NativeArray<float3> forces;

        [ReadOnly]
        public NativeArray<float3> positions;
        [ReadOnly]
        public NativeArray<float3> velocities;
        [ReadOnly]
        public NativeArray<float2> densityPressures;
        [ReadOnly]
        public float radius;
        [ReadOnly]
        public float mass;
        [ReadOnly]
        public float spikyGrad;
        [ReadOnly]
        public float visc;
        [ReadOnly]
        public float viscLap;
        [ReadOnly]
        public float3 gravity;
        [ReadOnly]
        public NativeMultiHashMap<int, int> gridHash;
        [ReadOnly]
        public int3 voxelSize;
        [ReadOnly]
        public float3 boundsCenter;
        [ReadOnly]
        public float3 boundsExtents;
        [ReadOnly]
        public float3 boundsSize;

        public void Execute(int index)
        {
            var myPosition = positions[index];
            var myVelocity = velocities[index];
            var myPressure = densityPressures[index].x;

            var fpress = float3.zero;
            var fvisc = float3.zero;

            var boundPos = myPosition + boundsExtents - boundsCenter;
            var normalizedBoundsPos = boundPos / boundsSize;
            var voxelIdx = (int3)math.floor(normalizedBoundsPos * voxelSize);
            var total = voxelSize.x * voxelSize.y * voxelSize.z;
            var voxelSliceSize = voxelSize.x * voxelSize.y;
            var voxelWidth = voxelSize.x;

            uint count = 0;
            var radiusSquared = radius * radius;
            var viscMass = visc * mass;

            for (var xi = -1; xi <= 1; xi++)
            {
                for (var yi = -1; yi <= 1; yi++)
                {
                    for (var zi = -1; zi <= 1; zi++)
                    {
                        var flatIndex = (voxelIdx.x + xi) + (voxelIdx.y + yi) * voxelWidth + (voxelIdx.z + zi) * voxelSliceSize;
                        if (flatIndex >= 0 && flatIndex < total)
                        {
                            var values = gridHash.GetValuesForKey(flatIndex);
                            while (values.MoveNext())
                            {
                                var otherIdx = values.Current;
                                if (index == otherIdx)
                                {
                                    continue;
                                }

                                var rij = positions[otherIdx] - myPosition;
                                var lenSquared = math.lengthsq(rij);

                                if (lenSquared < radiusSquared)
                                {
                                    var len = math.sqrt(lenSquared);
                                    var dist = math.max(0, radius - len);
                                    var otherDensityPressure = densityPressures[otherIdx];
                                    fpress += -math.normalize(rij) * mass * (myPressure + otherDensityPressure.y) / (2.0f * otherDensityPressure.x) * spikyGrad * (dist * dist);
                                    fvisc += viscMass * (velocities[otherIdx] - myVelocity) / otherDensityPressure.x * viscLap * dist;
                                    if (count++ > MAX_NEIGHBOUR)
                                    {
                                        break;
                                    }
                                }
                            }
                        }
                    }
                }
            }
            var fgrav = gravity * densityPressures[index].x;
            forces[index] = fpress + fvisc + fgrav;
        }
    }

    [BurstCompile]
    public struct IntegrateJob : IJobParallelFor
    {        
        public NativeArray<float3> positions;
        public NativeArray<float3> velocities;

        [ReadOnly]
        public NativeArray<float2> densityPressures;
        [ReadOnly]
        public NativeArray<float3> forces;
        [ReadOnly]
        public float dt;
        [ReadOnly]
        public float radius;
        [ReadOnly]
        public float size;
        [ReadOnly]
        public float boundDamping;

        // TODO actually use the bounds
        public void Execute(int index)
        {
            var position = positions[index];
            var velocity = velocities[index];

            velocity += dt * (forces[index]) / densityPressures[index].x;
            position += dt * velocity;

            // enforce boundary conditions
            if (position.x - radius < -size)
            {
                velocity.x *= boundDamping;
                position.x = -size + radius;
            }
            if (position.x + radius > size)
            {
                velocity.x *= boundDamping;
                position.x = size - radius;
            }

            if (position.y - radius < -size)
            {
                velocity.y *= boundDamping;
                position.y = -size + radius;
            }
            if (position.y + radius > size)
            {
                velocity.y *= boundDamping;
                position.y = size - radius;
            }

            if (position.z - radius < -size)
            {
                velocity.z *= boundDamping;
                position.z = -size + radius;
            }

            if (position.z + radius > size)
            {
                velocity.z *= boundDamping;
                position.z = size - radius;
            }

            velocities[index] = velocity;
            positions[index] = position;
        }
    }

    private void OnGUI()
    {
        if (GUILayout.Button("Run"))
        {
            InitSPH();
        }
    }

    private void Update()
    {
        if (!_positions.IsCreated)
        {
            return;
        }

        for (var i = 0; i < _iterations; i++)
        {
            var voxelSize = GetVoxelSize();

            var injectJob = new InjectGrid()
            {
                positions = _positions,
                voxelSize = voxelSize,
                boundsCenter = _bounds.center,
                boundsExtents = _bounds.extents,
                boundsSize = _bounds.size,
                // output
                gridHash = _gridHash
            };

            var handleInjectGrid = i == 0 ? injectJob.Schedule() : injectJob.Schedule(_handleIntegrate);

            var computerDensityPressureJob = new ComputeDensityPressureJob()
            {
                positions = _positions,
                hSquared = _radius * _radius,
                mass = _mass,
                poly6 = _poly6,
                gasConst = _gasConstant,
                restDensity = _restDensity,

                gridHash = _gridHash,
                voxelSize = voxelSize,
                boundsCenter = _bounds.center,
                boundsExtents = _bounds.extents,
                boundsSize = _bounds.size,
                // outputs
                densityPressures = _densityPressures
            };

            var handleComputeDensityPressure = computerDensityPressureJob.Schedule(_positions.Length, 8, handleInjectGrid);
            var computeForceJob = new ComputeForceJob()
            {
                positions = _positions,
                velocities = _velocities,
                densityPressures = _densityPressures,
                radius = _radius,
                mass = _mass,
                spikyGrad = _spikyGrad,
                visc = _viscosity,
                viscLap = _viscLap,
                gravity = _gravity,

                gridHash = _gridHash,
                voxelSize = voxelSize,
                boundsCenter = _bounds.center,
                boundsExtents = _bounds.extents,
                boundsSize = _bounds.size,
                // outputs
                forces = _forces
            };

            var handleComputeForce = computeForceJob.Schedule(_positions.Length, 8, handleComputeDensityPressure);
            var integrateJob = new IntegrateJob()
            {
                densityPressures = _densityPressures,
                forces = _forces,
                dt = _timeStep,
                radius = _radius,
                size = 64,
                boundDamping = _boundDamping,

                //outputs
                positions = _positions,
                velocities = _velocities,
            };

            _handleIntegrate = integrateJob.Schedule(_positions.Length, 64, handleComputeForce);
        }
    }

    private void LateUpdate()
    {
        _handleIntegrate.Complete();
    }

    private int3 GetVoxelSize()
    {
        return (int3)math.floor(_bounds.size / _radius);
    }

    private void OnDrawGizmos()
    {
        var upVec = new float3(0, _radius, 0);
        Gizmos.DrawWireCube(_bounds.center, _bounds.size);
        for (var i = 0; i < _positions.Length; i++)
        {
            // Gizmos.DrawSphere(p, 0.5f * _radius); // sloooow
            Gizmos.color = Color.Lerp(Color.blue, Color.red, math.length(_forces[i]) / 100000.0f);
            Gizmos.DrawRay(_positions[i], math.normalize(_velocities[i]) * _radius);
        }
    }
}