Skip to content

Instantly share code, notes, and snippets.

@JevinJ

JevinJ/PoissonSampling3D.cs

Last active September 4, 2020 19:36
Show Gist options
  • Save JevinJ/431e034396a7090da5bb636f8886070a to your computer and use it in GitHub Desktop.
Save JevinJ/431e034396a7090da5bb636f8886070a to your computer and use it in GitHub Desktop.
Download ZIP
Parallel Variable Radii Poisson Sampler in 3D
Raw
PoissonSampling3D.cs
//TODO, parallel sort of points, refactoring.
/// <summary>
/// Variable radii poisson disk sampling in 3D.
/// Based on: graphics.cs.umass.edu/pubs/sa_2010.pdf
/// </summary>
public static class PoissonSampling3D {
struct PointDataComparer : IComparer<float4> {
public AABB GridBounds;
public int3 CellsPerAxis;
public int Compare(float4 a, float4 b) {
int3 cellId = GetCellID(GridBounds, CellsPerAxis, a.xyz);
int3 otherCellId = GetCellID(GridBounds, CellsPerAxis, b.xyz);
int xDiff = cellId.x.CompareTo(otherCellId.x);
if(xDiff == 0) {
int yDiff = cellId.y.CompareTo(otherCellId.y);
if(yDiff == 0) {
return cellId.z.CompareTo(otherCellId.z);
}
return yDiff;
}
return xDiff;
}
}
static int3 GetCellID(AABB gridBounds, int3 cellsPerAxis, float3 position) {
float3 adjustment = -gridBounds.Center + gridBounds.Extents;
return (int3)(math.floor(cellsPerAxis * ((position + adjustment) / gridBounds.Size)));
}
static NativeArray<int3> GetPhaseGroups() {
int groupSize = 3;
var phaseGroups = new NativeArray<int3>((int)math.pow(groupSize, groupSize), Allocator.TempJob);
int i = 0;
for(int x = 0; x < groupSize; x++) {
for(int y = 0; y < groupSize; y++) {
for(int z = 0; z < groupSize; z++) {
phaseGroups[i] = new int3(x, y, z);
i++;
}
}
}
return phaseGroups;
}
/// <summary>
/// Pack points and their weight into a float4, where xyz is the point, x is the weight.
/// </summary>
[BurstCompile]
struct PackPointDataJob : IJobParallelFor {
[ReadOnly] public NativeArray<float3> Points;
[ReadOnly] public NativeArray<float> PointWeights;
public NativeArray<float4> PointDataOut;
public void Execute(int index) {
PointDataOut[index] = new float4(Points[index], PointWeights[index]);
}
}
[BurstCompatible]
struct InitCellDataJob : IJob {
public int TotalCells;
public int Trials;
public NativeHashMap<int3, int> CellFirstPointIndices;
public NativeMultiHashMap<int3, int> CellSampleIndices;
public void Execute() {
CellFirstPointIndices.Capacity = TotalCells;
CellSampleIndices.Capacity = TotalCells * Trials;
}
}
/// <summary>
/// Sort points by their cell id
/// </summary>
[BurstCompile]
struct SortPointDataJob : IJob {
public AABB GridBounds;
public int3 CellsPerAxis;
public NativeArray<float4> PointData;
public void Execute() {
PointData.Sort(new PointDataComparer {
GridBounds = GridBounds,
CellsPerAxis = CellsPerAxis
});
}
}
/// <summary>
/// Given an array sorted by cell id, finds the indices of the first point having a cell id.
/// </summary>
[BurstCompile]
struct FindFirstPointIndicesWithCellIDJob : IJobParallelFor {
[ReadOnly] public int3 CellsPerAxis;
[ReadOnly] public AABB GridBounds;
[ReadOnly, NativeDisableParallelForRestriction]
public NativeArray<float4> PointDatas;
public NativeHashMap<int3, int>.ParallelWriter CellFirstPointIndices;
public void Execute(int index) {
int3 cellId = GetCellID(GridBounds, CellsPerAxis, PointDatas[index].xyz);
if(index > 0) {
int3 previousCellId = GetCellID(GridBounds, CellsPerAxis, PointDatas[index - 1].xyz);
if(previousCellId.Equals(cellId)) {
return;
}
}
CellFirstPointIndices.TryAdd(cellId, index);
}
}
[BurstCompile]
struct AddCellsToPhaseGroupJob : IJobParallelFor {
[ReadOnly] public NativeArray<int3> Cells;
public NativeMultiHashMap<int3, int3>.ParallelWriter PhaseGroupCellsOut;
public void Execute(int index) {
int3 cellId = Cells[index];
int3 groupId = cellId % 3;
PhaseGroupCellsOut.Add(groupId, cellId);
}
}
[BurstCompile]
struct PoissonSampleJob : IJobParallelFor {
public AABB GridBounds;
public int3 CellsPerAxis;
public int Trial;
[ReadOnly] public NativeList<int3> CellsToProcess;
[ReadOnly] public NativeArray<float4> PointDatas;
[ReadOnly] public NativeHashMap<int3, int> CellFirstPointIndices;
[ReadOnly] public NativeMultiHashMap<int3, int> CellSampleIndices;
public NativeStream.Writer CellSampleIndexUpdates;
public void Execute(int index) {
int3 cellId = CellsToProcess[index];
int pointIndex = CellFirstPointIndices[cellId] + Trial;
if(pointIndex >= PointDatas.Length) {
return;
}
float3 point = PointDatas[pointIndex].xyz;
float pointWeight = PointDatas[pointIndex].w;
int3 pointCellId = GetCellID(GridBounds, CellsPerAxis, point);
if(!pointCellId.Equals(cellId)) {
return;
}
int3 startCell = cellId - 1;
for(int x = 0; x < 3; x++) {
for(int y = 0; y < 3; y++) {
for(int z = 0; z < 3; z++) {
int3 otherCellId = startCell + new int3(x, y, z);
if(CellSampleIndices.TryGetFirstValue(otherCellId, out int otherCellSampleIndex, out var it)) {
do {
float3 otherSamplePoint = PointDatas[otherCellSampleIndex].xyz;
float otherSampleWeight = PointDatas[otherCellSampleIndex].w;
float threshold = math.max(otherSampleWeight, pointWeight);
if(math.distance(point, otherSamplePoint) < threshold) {
return;
}
} while(CellSampleIndices.TryGetNextValue(out otherCellSampleIndex, ref it));
}
}
}
}
int4 packedSampleData = new int4(cellId, pointIndex);
CellSampleIndexUpdates.BeginForEachIndex(index);
CellSampleIndexUpdates.Write(packedSampleData);
CellSampleIndexUpdates.EndForEachIndex();
}
}
/// <summary>
/// Update cells with new sample indices found after a sampler iteration.
/// </summary>
[BurstCompile]
struct UpdateCellSamplesJob : IJobParallelFor {
public NativeStream.Reader Updates;
public NativeMultiHashMap<int3, int>.ParallelWriter CellSampleIndices;
public void Execute(int index) {
Updates.BeginForEachIndex(index);
while(Updates.RemainingItemCount > 0) {
int4 value = Updates.Read<int4>();
int3 cellId = value.xyz;
int pointIndex = value.w;
CellSampleIndices.Add(cellId, pointIndex);
}
Updates.EndForEachIndex();
}
}
[BurstCompile]
struct AddPointsToOutputJob : IJobParallelFor {
[ReadOnly] public NativeArray<int3> Cells;
[ReadOnly] public NativeArray<float4> PointDatas;
[ReadOnly] public NativeMultiHashMap<int3, int> CellSampleIndices;
public NativeList<float3>.ParallelWriter PointsOut;
public void Execute(int index) {
int3 cellId = Cells[index];
var indices = CellSampleIndices.GetValuesForKey(cellId);
do {
float3 point = PointDatas[indices.Current].xyz;
PointsOut.AddNoResize(point);
} while(indices.MoveNext());
}
}
/// <summary>
/// Variable radii poisson disk sampling in 3D.
/// Based on: graphics.cs.umass.edu/pubs/sa_2010.pdf
/// </summary>
/// <param name="points">Set of points to sample.</param>
/// <param name="pointWeights">Weights of points in range [0f, 1f], must be same length as <paramref name="points"/></param>
/// <param name="outPoints">Resulting point set.</param>
/// <param name="k">Number of trials or "iteration" for the sampler. Higher values are more accurate, but slower.</param>
/// <param name="maxR">The minimum distance of a point to others, if the point has a weight of 1f.</param>
/// <returns>The JobHandle of the sampler job.</returns>
public static JobHandle Calculate(
NativeArray<float3> points, NativeArray<float> pointWeights, out NativeList<float3> outPoints,
int k = 30, float maxR = 1f) {
float cellSize = maxR / math.sqrt(3);
float3 minPoint = new float3(float.MaxValue);
float3 maxPoint = new float3(float.MinValue);
foreach(var point in points) {
minPoint = math.min(minPoint, point);
maxPoint = math.max(maxPoint, point);
}
float3 center = (maxPoint + minPoint) / 2;
var gridBounds = new AABB {
Center = center,
Extents = math.max(cellSize * math.ceil((maxPoint - center) / cellSize), new int3(1))
};
int3 cellsPerAxis = (int3)(gridBounds.Size / cellSize);
int totalCells = cellsPerAxis.x * cellsPerAxis.y * cellsPerAxis.z;
var pointData = new NativeArray<float4>(points.Length, Allocator.TempJob);
var convertToPointDataJob = new PackPointDataJob {
Points = points,
PointWeights = pointWeights,
PointDataOut = pointData
}.Schedule(points.Length, innerloopBatchCount:8192);
var sortPointDataJob = new SortPointDataJob {
GridBounds = gridBounds,
CellsPerAxis = cellsPerAxis,
PointData = pointData
}.Schedule(dependsOn:convertToPointDataJob);
var cellFirstPointIndices = new NativeHashMap<int3, int>(totalCells, Allocator.TempJob);
var cellSampleIndices = new NativeMultiHashMap<int3, int>(totalCells * k, Allocator.TempJob);
var initializeCellDataJob = new InitCellDataJob {
TotalCells = totalCells,
Trials = k,
CellFirstPointIndices = cellFirstPointIndices,
CellSampleIndices = cellSampleIndices
}.Schedule(sortPointDataJob);
var findFirstPointIndicesJob = new FindFirstPointIndicesWithCellIDJob {
GridBounds = gridBounds,
CellsPerAxis = cellsPerAxis,
PointDatas = pointData,
CellFirstPointIndices = cellFirstPointIndices.AsParallelWriter()
}.Schedule(points.Length, innerloopBatchCount:1024, dependsOn:initializeCellDataJob);
findFirstPointIndicesJob.Complete();
NativeArray<int3> uniqueCellIds = cellFirstPointIndices.GetKeyArray(Allocator.TempJob);
// indices into points array
var phaseGroupCells = new NativeMultiHashMap<int3, int3>(uniqueCellIds.Length, Allocator.TempJob);
var getCellsInPhaseGroupJob = new AddCellsToPhaseGroupJob {
Cells = uniqueCellIds,
PhaseGroupCellsOut = phaseGroupCells.AsParallelWriter()
}.Schedule(uniqueCellIds.Length, innerloopBatchCount:1024, dependsOn:findFirstPointIndicesJob);
getCellsInPhaseGroupJob.Complete();
var phaseGroupCellArrays = new List<NativeList<int3>>();
using(NativeArray<int3> phaseGroups = GetPhaseGroups()) {
foreach(int3 phaseGroup in phaseGroups) {
var cellIds = phaseGroupCells.GetValuesForKey(phaseGroup);
var cells = new NativeList<int3>(Allocator.TempJob);
foreach(int3 cell in cellIds) {
cells.Add(cell);
}
phaseGroupCellArrays.Add(cells);
}
}
JobHandle poissonHandle = getCellsInPhaseGroupJob;
for(int trial = 1; trial < k; trial++) {
foreach(NativeList<int3> cellsToProcess in phaseGroupCellArrays) {
var cellSampleUpdates = new NativeStream(cellsToProcess.Length, Allocator.TempJob);
var poissonSampleJob = new PoissonSampleJob {
GridBounds = gridBounds,
CellsPerAxis = cellsPerAxis,
Trial = trial,
CellsToProcess = cellsToProcess,
PointDatas = pointData,
CellFirstPointIndices = cellFirstPointIndices,
CellSampleIndices = cellSampleIndices,
CellSampleIndexUpdates = cellSampleUpdates.AsWriter()
}.Schedule(cellsToProcess.Length, innerloopBatchCount:32, dependsOn:poissonHandle);
var updateCellSampleIndicesJob = new UpdateCellSamplesJob {
Updates = cellSampleUpdates.AsReader(),
CellSampleIndices = cellSampleIndices.AsParallelWriter()
}.Schedule(cellsToProcess.Length, innerloopBatchCount:32, dependsOn:poissonSampleJob);
poissonHandle = cellSampleUpdates.Dispose(updateCellSampleIndicesJob);
}
}
outPoints = new NativeList<float3>(points.Length, Allocator.TempJob);
var addPointsToOutputJob = new AddPointsToOutputJob {
Cells = uniqueCellIds,
PointDatas = pointData,
PointsOut = outPoints.AsParallelWriter(),
CellSampleIndices = cellSampleIndices
}.Schedule(uniqueCellIds.Length, innerloopBatchCount:128, dependsOn:poissonHandle);
JobHandle finalHandle = addPointsToOutputJob;
foreach(var i in phaseGroupCellArrays) {
finalHandle = i.Dispose(finalHandle);
}
finalHandle = phaseGroupCells.Dispose(finalHandle);
return finalHandle;
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment