Parallel Variable Radii Poisson Sampler in 3D

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;
    }
}