Skip to content

Instantly share code, notes, and snippets.

@brianasu

brianasu/SPHBurst.cs

Created July 7, 2021 12:03
Show Gist options
  • Save brianasu/6c314596c615f9900fe07afbd50588de to your computer and use it in GitHub Desktop.
Save brianasu/6c314596c615f9900fe07afbd50588de to your computer and use it in GitHub Desktop.
Download ZIP
Simple example of Smooth Particle Hydrodynamics using Unity Jobs and Burst
Raw
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);
}
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment