├── LICENSE.meta ├── README.md.meta ├── package.json.meta ├── KNN.asmdef.meta ├── KnnJobs.cs.meta ├── MinMaxHeap.cs.meta ├── QueryNode.cs.meta ├── KDNode.cs.meta ├── KDBounds.cs.meta ├── KnnContainer.cs.meta ├── package.json ├── KNN.asmdef ├── LICENSE ├── QueryNode.cs ├── KDBounds.cs ├── KDNode.cs ├── README.md ├── KnnJobs.cs ├── MinMaxHeap.cs └── KnnContainer.cs /LICENSE.meta: -------------------------------------------------------------------------------- 1 | fileFormatVersion: 2 2 | guid: 7d7cd75ac774c71449596e8d8abaeafa 3 | DefaultImporter: 4 | externalObjects: {} 5 | userData: 6 | assetBundleName: 7 | assetBundleVariant: 8 | -------------------------------------------------------------------------------- /README.md.meta: -------------------------------------------------------------------------------- 1 | fileFormatVersion: 2 2 | guid: c665b892543e1e34586e514a7b054872 3 | TextScriptImporter: 4 | externalObjects: {} 5 | userData: 6 | assetBundleName: 7 | assetBundleVariant: 8 | -------------------------------------------------------------------------------- /package.json.meta: -------------------------------------------------------------------------------- 1 | fileFormatVersion: 2 2 | guid: ee7fb34896e71334dbb05e79876ce0ef 3 | PackageManifestImporter: 4 | externalObjects: {} 5 | userData: 6 | assetBundleName: 7 | assetBundleVariant: 8 | -------------------------------------------------------------------------------- /KNN.asmdef.meta: -------------------------------------------------------------------------------- 1 | fileFormatVersion: 2 2 | guid: 93dfb46a4941698438e1e2dce7f5e021 3 | AssemblyDefinitionImporter: 4 | externalObjects: {} 5 | userData: 6 | assetBundleName: 7 | assetBundleVariant: 8 | -------------------------------------------------------------------------------- /KnnJobs.cs.meta: -------------------------------------------------------------------------------- 1 | fileFormatVersion: 2 2 | guid: d6d4892e1b2ae3d4aa785b2f2802a0f8 3 | MonoImporter: 4 | externalObjects: {} 5 | serializedVersion: 2 6 | defaultReferences: [] 7 | executionOrder: 0 8 | icon: {instanceID: 0} 9 | userData: 10 | assetBundleName: 11 | assetBundleVariant: 12 | -------------------------------------------------------------------------------- /MinMaxHeap.cs.meta: -------------------------------------------------------------------------------- 1 | fileFormatVersion: 2 2 | guid: 354799470fae8d74d8f977dd8db040cf 3 | MonoImporter: 4 | externalObjects: {} 5 | serializedVersion: 2 6 | defaultReferences: [] 7 | executionOrder: 0 8 | icon: {instanceID: 0} 9 | userData: 10 | assetBundleName: 11 | assetBundleVariant: 12 | -------------------------------------------------------------------------------- /QueryNode.cs.meta: -------------------------------------------------------------------------------- 1 | fileFormatVersion: 2 2 | guid: cbe8b45a46d50254494605e93da9d3ed 3 | timeCreated: 1521907406 4 | licenseType: Free 5 | MonoImporter: 6 | externalObjects: {} 7 | serializedVersion: 2 8 | defaultReferences: [] 9 | executionOrder: 0 10 | icon: {instanceID: 0} 11 | userData: 12 | assetBundleName: 13 | assetBundleVariant: 14 | -------------------------------------------------------------------------------- /KDNode.cs.meta: -------------------------------------------------------------------------------- 1 | fileFormatVersion: 2 2 | guid: c4440ff733ae6bc488c26393873a5950 3 | timeCreated: 1520605573 4 | licenseType: Free 5 | MonoImporter: 6 | externalObjects: {} 7 | serializedVersion: 2 8 | defaultReferences: [] 9 | executionOrder: 0 10 | icon: {instanceID: 0} 11 | userData: 12 | assetBundleName: 13 | assetBundleVariant: 14 | -------------------------------------------------------------------------------- /KDBounds.cs.meta: -------------------------------------------------------------------------------- 1 | fileFormatVersion: 2 2 | guid: 1b92db3ff3e67ee43a15eda23aea87a3 3 | timeCreated: 1521907406 4 | licenseType: Free 5 | MonoImporter: 6 | externalObjects: {} 7 | serializedVersion: 2 8 | defaultReferences: [] 9 | executionOrder: 0 10 | icon: {instanceID: 0} 11 | userData: 12 | assetBundleName: 13 | assetBundleVariant: 14 | -------------------------------------------------------------------------------- /KnnContainer.cs.meta: -------------------------------------------------------------------------------- 1 | fileFormatVersion: 2 2 | guid: 437ce4903d14e2f4182fad0f2ac92678 3 | timeCreated: 1520696434 4 | licenseType: Free 5 | MonoImporter: 6 | externalObjects: {} 7 | serializedVersion: 2 8 | defaultReferences: [] 9 | executionOrder: 0 10 | icon: {instanceID: 0} 11 | userData: 12 | assetBundleName: 13 | assetBundleVariant: 14 | -------------------------------------------------------------------------------- /package.json: -------------------------------------------------------------------------------- 1 | { 2 | "name": "com.akb.knn", 3 | "displayName": "KNN", 4 | "version": "1.0.1", 5 | "unity": "2019.2", 6 | "description": "Super fast K-nearest neighbour search using Burst and Unity's job system", 7 | "keywords": [ 8 | "knn", 9 | "unity" 10 | ], 11 | "category": "Unity", 12 | "dependencies": { 13 | "com.unity.burst": "1.1.2", 14 | "com.unity.collections": "0.1.1-preview", 15 | "com.unity.mathematics": "1.1.0" 16 | } 17 | } 18 | -------------------------------------------------------------------------------- /KNN.asmdef: -------------------------------------------------------------------------------- 1 | { 2 | "name": "KNN", 3 | "references": [ 4 | "GUID:d8b63aba1907145bea998dd612889d6b", 5 | "GUID:e0cd26848372d4e5c891c569017e11f1", 6 | "GUID:8a2eafa29b15f444eb6d74f94a930e1d", 7 | "GUID:2665a8d13d1b3f18800f46e256720795" 8 | ], 9 | "includePlatforms": [], 10 | "excludePlatforms": [], 11 | "allowUnsafeCode": true, 12 | "overrideReferences": false, 13 | "precompiledReferences": [], 14 | "autoReferenced": true, 15 | "defineConstraints": [], 16 | "versionDefines": [] 17 | } -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2019 Arthur Brussee 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /QueryNode.cs: -------------------------------------------------------------------------------- 1 | //MIT License 2 | // 3 | //Copyright(c) 2018 Vili Volčini / viliwonka 4 | // 5 | //Permission is hereby granted, free of charge, to any person obtaining a copy 6 | //of this software and associated documentation files (the "Software"), to deal 7 | //in the Software without restriction, including without limitation the rights 8 | //to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | //copies of the Software, and to permit persons to whom the Software is 10 | //furnished to do so, subject to the following conditions: 11 | // 12 | //The above copyright notice and this permission notice shall be included in all 13 | //copies or substantial portions of the Software. 14 | // 15 | //THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | //IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | //FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | //AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | //LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | //OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | //SOFTWARE. 22 | // 23 | // Modifed 2019 Arthur Brussee 24 | 25 | using Unity.Mathematics; 26 | 27 | namespace KNN.Internal { 28 | public struct QueryNode { 29 | public int NodeIndex; 30 | public float3 TempClosestPoint; 31 | public float Distance; 32 | } 33 | } -------------------------------------------------------------------------------- /KDBounds.cs: -------------------------------------------------------------------------------- 1 | //MIT License 2 | // 3 | //Copyright(c) 2018 Vili Volčini / viliwonka 4 | // 5 | //Permission is hereby granted, free of charge, to any person obtaining a copy 6 | //of this software and associated documentation files (the "Software"), to deal 7 | //in the Software without restriction, including without limitation the rights 8 | //to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | //copies of the Software, and to permit persons to whom the Software is 10 | //furnished to do so, subject to the following conditions: 11 | // 12 | //The above copyright notice and this permission notice shall be included in all 13 | //copies or substantial portions of the Software. 14 | // 15 | //THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | //IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | //FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | //AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | //LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | //OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | //SOFTWARE. 22 | // 23 | // Modifed 2019 Arthur Brussee 24 | 25 | using Unity.Mathematics; 26 | 27 | namespace KNN.Internal { 28 | public struct KdNodeBounds { 29 | public float3 Min; 30 | public float3 Max; 31 | 32 | public float3 Size => Max - Min; 33 | 34 | public float3 ClosestPoint(float3 point) { 35 | return math.clamp(point, Min, Max); 36 | } 37 | } 38 | } -------------------------------------------------------------------------------- /KDNode.cs: -------------------------------------------------------------------------------- 1 | //MIT License 2 | // 3 | //Copyright(c) 2018 Vili Volčini / viliwonka 4 | // 5 | //Permission is hereby granted, free of charge, to any person obtaining a copy 6 | //of this software and associated documentation files (the "Software"), to deal 7 | //in the Software without restriction, including without limitation the rights 8 | //to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | //copies of the Software, and to permit persons to whom the Software is 10 | //furnished to do so, subject to the following conditions: 11 | // 12 | //The above copyright notice and this permission notice shall be included in all 13 | //copies or substantial portions of the Software. 14 | // 15 | //THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | //IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | //FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | //AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | //LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | //OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | //SOFTWARE. 22 | // 23 | // Modifed 2019 Arthur Brussee 24 | 25 | namespace KNN.Internal { 26 | public struct KdNode { 27 | public KdNodeBounds Bounds; 28 | 29 | public int Start; 30 | public int End; 31 | 32 | public int PartitionAxis; 33 | public float PartitionCoordinate; 34 | 35 | public int NegativeChildIndex; 36 | public int PositiveChildIndex; 37 | 38 | public int Count => End - Start; 39 | public bool Leaf => PartitionAxis == -1; 40 | } 41 | } -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | ![Demo Image](http://g2f.nl/0jx82rw) 2 | 3 | # KNN 4 | 5 | Simple K-Nearest Neighbour library for Unity, using the 'Dots' technology stack (the Burst compiler and Unity's job system). It uses K-D trees to speed up queries. 6 | 7 | It is totally free from managed allocations and can run multi-threaded. The Burst compiler heavily vectorizes the searching code. 8 | 9 | The implementation is a heavily modified version of a KD-Tree by viliwonka: https://github.com/viliwonka/KDTree . It only includes a sub-set of functionality however. 10 | 11 | As a rough benchmark, here are perf numbers on a i7-770hq, querying the k=10 nearest neighbours, 100.000 times, in 100.000 points 12 | 13 | ![Performance comparison](http://g2f.nl/04y9z7g) 14 | 15 | This shows an approximately ~130x speedup over the original implementation! Note that this is very much apples to oranges though (eg. this is multi-threaded, the original implementation was not); I don't mean to pick on the original at all, but does hopefully it shows this implementation is really fast! 16 | 17 | # API Overview 18 | 19 | ```C# 20 | // First let's create a random point cloud 21 | var points = new NativeArray(100000, Allocator.Persistent); 22 | var rand = new Random(123456); 23 | for (int i = 0; i < points.Length; ++i) { 24 | points[i] = rand.NextFloat3(); 25 | } 26 | 27 | // Number of neighbours we want to query 28 | const int kNeighbours = 10; 29 | float3 queryPosition = float3.zero; 30 | 31 | // Create a container that accelerates querying for neighbours. 32 | // The 2nd argument indicates whether we want to build the tree straight away or not 33 | // Let's hold off on building it a little bit 34 | var knnContainer = new KnnContainer(points, false, Allocator.TempJob); 35 | 36 | // Whenever your point cloud changes, you can make a job to rebuild the container: 37 | var rebuildJob = new KnnRebuildJob(knnContainer); 38 | rebuildJob.Schedule().Complete(); 39 | 40 | // Most basic usage: 41 | // Get 10 nearest neighbours as indices into our points array! 42 | // This is NOT burst accelerated yet! Unity need to implement compiling delegates with Burst 43 | var result = new NativeArray(kNeighbours, Allocator.TempJob); 44 | knnContainer.QueryKNearest(queryPosition, result); 45 | 46 | // The result array at this point contains indices into the points array with the nearest neighbours! 47 | Profiler.BeginSample("Simple Query"); 48 | // Get a job to do the query. 49 | var queryJob = new QueryKNearestJob(knnContainer, queryPosition, result); 50 | 51 | // And just run immediately on the main thread for now. This uses Burst! 52 | queryJob.Schedule().Complete(); 53 | Profiler.EndSample(); 54 | 55 | // Or maybe we want to query neighbours for multiple points. 56 | const int queryPoints = 100000; 57 | 58 | // Keep an array of neighbour indices of all points 59 | var results = new NativeArray(queryPoints * kNeighbours, Allocator.TempJob); 60 | 61 | // Query at a few random points 62 | var queryPositions = new NativeArray(queryPoints, Allocator.TempJob); 63 | for (int i = 0; i < queryPoints; ++i) { 64 | queryPositions[i] = rand.NextFloat3() * 0.1f; 65 | } 66 | 67 | // Fire up job to get results for all points 68 | var batchQueryJob = new QueryKNearestBatchJob(knnContainer, queryPositions, results); 69 | 70 | // And just run immediately now. This will run on multiple threads! 71 | batchQueryJob.ScheduleBatch(queryPositions.Length, queryPositions.Length / 32).Complete(); 72 | 73 | // Or maybe we're interested in a range around eacht query point 74 | var queryRangeResult = new NativeList(Allocator.TempJob); 75 | var rangeQueryJob = new QueryRangeJob(knnContainer, queryPosition, 2.0f, queryRangeResult); 76 | 77 | // Store a list of particles in range 78 | var rangeResults = new NativeArray(queryPoints, Allocator.TempJob); 79 | 80 | // And just run immediately on the main thread for now. This uses Burst! 81 | rangeQueryJob.Schedule().Complete(); 82 | 83 | // Unfortunately, for batch range queries we do need to decide upfront the maximum nr. of neighbours we allow 84 | // This is due to limitation on allocations within a job. 85 | for (int i = 0; i < rangeResults.Length; ++i) { 86 | rangeResults[i] = new RangeQueryResult(128, Allocator.TempJob); 87 | } 88 | 89 | // Fire up job to get results for all points 90 | var batchRange = new QueryRangeBatchJob(knnContainer, queryPositions, 2.0f, rangeResults); 91 | 92 | // And just run immediately now. This will run on multiple threads! 93 | batchRange.ScheduleBatch(queryPositions.Length, queryPositions.Length / 32).Complete(); 94 | 95 | // Now the results array contains all the neighbours! 96 | queryRangeResult.Dispose(); 97 | foreach (var r in rangeResults) { 98 | r.Dispose(); 99 | } 100 | rangeResults.Dispose(); 101 | knnContainer.Dispose(); 102 | queryPositions.Dispose(); 103 | results.Dispose(); 104 | points.Dispose(); 105 | result.Dispose(); 106 | ``` 107 | 108 | # Demo 109 | 110 | The demo folder contains 2 demos: 111 | 112 | - KnnApiDemo.cs, Illustrates various API usages from a basic to advanced level 113 | - KnnVisualizationDemo.cs: Illustrates a real scene where particles are colored based on their neighbourhood information. Shows how to dynamically rebuild your container and other advanced API usages 114 | 115 | 116 | ## Installation 117 | 118 | The project was made as a Unity Package. Just add the git URL to your package manifest and Unity should install it. 119 | -------------------------------------------------------------------------------- /KnnJobs.cs: -------------------------------------------------------------------------------- 1 | using System; 2 | using System.Runtime.CompilerServices; 3 | using Unity.Burst; 4 | using Unity.Collections; 5 | using Unity.Collections.LowLevel.Unsafe; 6 | using Unity.Jobs; 7 | using Unity.Mathematics; 8 | using UnityEngine; 9 | using UnsafeUtilityEx = KNN.Internal.UnsafeUtilityEx; 10 | 11 | namespace KNN.Jobs { 12 | [BurstCompile(CompileSynchronously = true)] 13 | public struct QueryKNearestJob : IJob { 14 | [ReadOnly] KnnContainer m_container; 15 | [WriteOnly] NativeSlice m_result; 16 | 17 | float3 m_queryPosition; 18 | 19 | public QueryKNearestJob(KnnContainer container, float3 queryPosition, NativeSlice result) { 20 | m_result = result; 21 | m_queryPosition = queryPosition; 22 | m_container = container; 23 | } 24 | 25 | void IJob.Execute() { 26 | m_container.QueryKNearest(m_queryPosition, m_result); 27 | } 28 | } 29 | 30 | [BurstCompile(CompileSynchronously = true)] 31 | public struct QueryRangeJob : IJob { 32 | [ReadOnly] KnnContainer m_container; 33 | [WriteOnly] NativeList m_result; 34 | 35 | float m_range; 36 | float3 m_queryPosition; 37 | 38 | public QueryRangeJob(KnnContainer container, float3 queryPosition, float range, NativeList result) { 39 | m_result = result; 40 | m_range = range; 41 | m_queryPosition = queryPosition; 42 | m_container = container; 43 | } 44 | 45 | void IJob.Execute() { 46 | m_container.QueryRange(m_queryPosition, m_range, m_result); 47 | } 48 | } 49 | 50 | 51 | [BurstCompile(CompileSynchronously = true)] 52 | public struct QueryKNearestBatchJob : IJobParallelForBatch { 53 | [ReadOnly] KnnContainer m_container; 54 | [ReadOnly] NativeSlice m_queryPositions; 55 | 56 | // Unity really doesn't like it when we write to the same underlying array 57 | // Even if slices don't overlap... So we're just being dangerous here 58 | [NativeDisableParallelForRestriction, NativeDisableContainerSafetyRestriction] 59 | NativeSlice m_results; 60 | 61 | int m_k; 62 | 63 | public QueryKNearestBatchJob(KnnContainer container, NativeArray queryPositions, NativeSlice results) { 64 | m_container = container; 65 | m_queryPositions = queryPositions; 66 | m_results = results; 67 | 68 | #if ENABLE_UNITY_COLLECTIONS_CHECKS 69 | if (queryPositions.Length == 0 || results.Length % queryPositions.Length != 0) { 70 | Debug.LogError("Make sure your results array is a multiple in length of your querypositions array!"); 71 | } 72 | #endif 73 | 74 | m_k = results.Length / queryPositions.Length; 75 | } 76 | 77 | public void Execute(int startIndex, int count) { 78 | // Write results to proper slice! 79 | for (int index = startIndex; index < startIndex + count; ++index) { 80 | NativeSlice resultsSlice = m_results.Slice(index * m_k, m_k); 81 | m_container.QueryKNearest(m_queryPositions[index], resultsSlice); 82 | } 83 | } 84 | } 85 | 86 | public unsafe struct RangeQueryResult { 87 | public int Length; 88 | 89 | int* m_indices; 90 | int m_capacity; 91 | 92 | Allocator m_allocator; 93 | 94 | public int this[int index] { 95 | get { 96 | if (index >= m_capacity) { 97 | throw new IndexOutOfRangeException(); 98 | } 99 | 100 | return UnsafeUtility.ReadArrayElement(m_indices, index); 101 | } 102 | } 103 | 104 | public RangeQueryResult(int maxCount, Allocator allocator) { 105 | m_capacity = maxCount; 106 | m_indices = UnsafeUtilityEx.AllocArray(m_capacity, allocator); 107 | Length = 0; 108 | m_allocator = allocator; 109 | } 110 | 111 | public void SetResults(NativeList result) { 112 | UnsafeUtility.MemCpy(m_indices, result.GetUnsafePtr(), Mathf.Min(m_capacity, result.Length) * sizeof(int)); 113 | Length = Mathf.Min(m_capacity, result.Length); 114 | } 115 | 116 | public void Dispose() { 117 | UnsafeUtility.Free(m_indices, m_allocator); 118 | } 119 | } 120 | 121 | 122 | [BurstCompile(CompileSynchronously = true)] 123 | public struct QueryRangeBatchJob : IJobParallelForBatch { 124 | [ReadOnly] KnnContainer m_container; 125 | [ReadOnly] NativeSlice m_queryPositions; 126 | 127 | float m_range; 128 | 129 | public NativeArray Results; 130 | 131 | public QueryRangeBatchJob(KnnContainer container, NativeArray queryPositions, float range, NativeArray results) { 132 | m_container = container; 133 | m_queryPositions = queryPositions; 134 | m_range = range; 135 | Results = results; 136 | } 137 | 138 | public void Execute(int startIndex, int count) { 139 | // Write results to proper slice! 140 | for (int index = startIndex; index < startIndex + count; ++index) { 141 | var tempList = new NativeList(Allocator.Temp); 142 | m_container.QueryRange(m_queryPositions[index], m_range, tempList); 143 | 144 | var result = Results[index]; 145 | result.SetResults(tempList); 146 | 147 | Results[index] = result; 148 | } 149 | } 150 | } 151 | 152 | [BurstCompile(CompileSynchronously = true)] 153 | public struct KnnRebuildJob : IJob { 154 | KnnContainer m_container; 155 | 156 | public KnnRebuildJob(KnnContainer container) { 157 | m_container = container; 158 | } 159 | 160 | void IJob.Execute() { 161 | m_container.Rebuild(); 162 | } 163 | } 164 | } -------------------------------------------------------------------------------- /MinMaxHeap.cs: -------------------------------------------------------------------------------- 1 | //MIT License 2 | // 3 | //Copyright(c) 2018 Vili Volčini / viliwonka 4 | // 5 | //Permission is hereby granted, free of charge, to any person obtaining a copy 6 | //of this software and associated documentation files (the "Software"), to deal 7 | //in the Software without restriction, including without limitation the rights 8 | //to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | //copies of the Software, and to permit persons to whom the Software is 10 | //furnished to do so, subject to the following conditions: 11 | // 12 | //The above copyright notice and this permission notice shall be included in all 13 | //copies or substantial portions of the Software. 14 | // 15 | //THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | //IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | //FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | //AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | //LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | //OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | //SOFTWARE. 22 | // 23 | // Modifed 2019 Arthur Brussee 24 | 25 | using System; 26 | using Unity.Collections; 27 | using Unity.Collections.LowLevel.Unsafe; 28 | 29 | namespace KNN.Internal { 30 | public static class HeapUtils { 31 | public static int Parent(int index) { 32 | return index / 2; 33 | } 34 | 35 | public static int Left(int index) { 36 | return index * 2; 37 | } 38 | 39 | public static int Right(int index) { 40 | return index * 2 + 1; 41 | } 42 | } 43 | 44 | // Sorted heap with a self balancing tree 45 | // Can act as either a min or max heap 46 | public unsafe struct MinMaxHeap : IDisposable where T : unmanaged { 47 | [NativeDisableContainerSafetyRestriction] 48 | T* keys; //objects 49 | 50 | [NativeDisableContainerSafetyRestriction] 51 | float* values; 52 | 53 | public int Count; 54 | int m_capacity; 55 | 56 | public float HeadValue => values[1]; 57 | T HeadKey => keys[1]; 58 | 59 | public bool IsFull => Count == m_capacity; 60 | 61 | Allocator m_allocator; 62 | 63 | public MinMaxHeap(int startCapacity, Allocator allocator) { 64 | Count = 0; 65 | m_allocator = allocator; 66 | 67 | // Now alloc starting arrays 68 | m_capacity = startCapacity; 69 | values = UnsafeUtilityEx.AllocArray(startCapacity + 1, m_allocator); 70 | keys = UnsafeUtilityEx.AllocArray(startCapacity + 1, m_allocator); 71 | } 72 | 73 | void Swap(int indexA, int indexB) { 74 | float tempVal = values[indexA]; 75 | values[indexA] = values[indexB]; 76 | values[indexB] = tempVal; 77 | 78 | T tempKey = keys[indexA]; 79 | keys[indexA] = keys[indexB]; 80 | keys[indexB] = tempKey; 81 | } 82 | 83 | public void Dispose() { 84 | UnsafeUtility.Free(values, m_allocator); 85 | UnsafeUtility.Free(keys, m_allocator); 86 | values = null; 87 | keys = null; 88 | } 89 | 90 | public void Resize(int newSize) { 91 | // Allocate more space 92 | var newValues = UnsafeUtilityEx.AllocArray(newSize + 1, m_allocator); 93 | var newKeys = UnsafeUtilityEx.AllocArray(newSize + 1, m_allocator); 94 | 95 | // Copy over old arrays 96 | UnsafeUtility.MemCpy(newValues, values, (m_capacity + 1) * sizeof(int)); 97 | UnsafeUtility.MemCpy(newKeys, keys, (m_capacity + 1) * sizeof(int)); 98 | 99 | // Get rid of old arrays 100 | Dispose(); 101 | 102 | // And now use old arrays 103 | values = newValues; 104 | keys = newKeys; 105 | m_capacity = newSize; 106 | } 107 | 108 | // bubble down, MaxHeap version 109 | void BubbleDownMax(int index) { 110 | int l = HeapUtils.Left(index); 111 | int r = HeapUtils.Right(index); 112 | 113 | // bubbling down, 2 kids 114 | while (r <= Count) { 115 | // if heap property is violated between index and Left child 116 | if (values[index] < values[l]) { 117 | if (values[l] < values[r]) { 118 | Swap(index, r); // left has bigger priority 119 | index = r; 120 | } else { 121 | Swap(index, l); // right has bigger priority 122 | index = l; 123 | } 124 | } else { 125 | // if heap property is violated between index and R 126 | if (values[index] < values[r]) { 127 | Swap(index, r); 128 | index = r; 129 | } else { 130 | index = l; 131 | l = HeapUtils.Left(index); 132 | break; 133 | } 134 | } 135 | 136 | l = HeapUtils.Left(index); 137 | r = HeapUtils.Right(index); 138 | } 139 | 140 | // only left & last children available to test and swap 141 | if (l <= Count && values[index] < values[l]) { 142 | Swap(index, l); 143 | } 144 | } 145 | 146 | void BubbleDownMin(int index) { 147 | int l = HeapUtils.Left(index); 148 | int r = HeapUtils.Right(index); 149 | 150 | // bubbling down, 2 kids 151 | while (r <= Count) { 152 | // if heap property is violated between index and Left child 153 | if (values[index] > values[l]) { 154 | if (values[l] > values[r]) { 155 | Swap(index, r); // right has smaller priority 156 | index = r; 157 | } else { 158 | Swap(index, l); // left has smaller priority 159 | index = l; 160 | } 161 | } else { 162 | // if heap property is violated between index and R 163 | if (values[index] > values[r]) { 164 | Swap(index, r); 165 | index = r; 166 | } else { 167 | index = l; 168 | l = HeapUtils.Left(index); 169 | break; 170 | } 171 | } 172 | 173 | l = HeapUtils.Left(index); 174 | r = HeapUtils.Right(index); 175 | } 176 | 177 | // only left & last children available to test and swap 178 | if (l <= Count && values[index] > values[l]) { 179 | Swap(index, l); 180 | } 181 | } 182 | 183 | void BubbleUpMax(int index) { 184 | int p = HeapUtils.Parent(index); 185 | 186 | //swap, until Heap property isn't violated anymore 187 | while (p > 0 && values[p] < values[index]) { 188 | Swap(p, index); 189 | index = p; 190 | p = HeapUtils.Parent(index); 191 | } 192 | } 193 | 194 | void BubbleUpMin(int index) { 195 | int p = HeapUtils.Parent(index); 196 | 197 | //swap, until Heap property isn't violated anymore 198 | while (p > 0 && values[p] > values[index]) { 199 | Swap(p, index); 200 | index = p; 201 | p = HeapUtils.Parent(index); 202 | } 203 | } 204 | 205 | public void PushObjMax(T key, float val) { 206 | // if heap full 207 | if (Count == m_capacity) { 208 | // if Heads priority is smaller than input priority, then ignore that item 209 | if (HeadValue > val) { 210 | values[1] = val; // remove top element 211 | keys[1] = key; 212 | BubbleDownMax(1); // bubble it down 213 | } 214 | } 215 | else { 216 | Count++; 217 | values[Count] = val; 218 | keys[Count] = key; 219 | BubbleUpMax(Count); 220 | } 221 | } 222 | 223 | public void PushObjMin(T key, float val) { 224 | // if heap full 225 | if (Count == m_capacity) { 226 | // if Heads priority is smaller than input priority, then ignore that item 227 | if (HeadValue < val) { 228 | values[1] = val; // remove top element 229 | keys[1] = key; 230 | BubbleDownMin(1); // bubble it down 231 | } 232 | } 233 | else { 234 | Count++; 235 | values[Count] = val; 236 | keys[Count] = key; 237 | BubbleUpMin(Count); 238 | } 239 | } 240 | 241 | T PopHeadObj() { 242 | T result = HeadKey; 243 | 244 | values[1] = values[Count]; 245 | keys[1] = keys[Count]; 246 | Count--; 247 | 248 | return result; 249 | } 250 | 251 | public T PopObjMax() { 252 | T result = PopHeadObj(); 253 | BubbleDownMax(1); 254 | return result; 255 | } 256 | 257 | public T PopObjMin() { 258 | T result = PopHeadObj(); 259 | BubbleDownMin(1); 260 | return result; 261 | } 262 | } 263 | } -------------------------------------------------------------------------------- /KnnContainer.cs: -------------------------------------------------------------------------------- 1 | //MIT License 2 | // 3 | //Copyright(c) 2018 Vili Volčini / viliwonka 4 | // 5 | //Permission is hereby granted, free of charge, to any person obtaining a copy 6 | //of this software and associated documentation files (the "Software"), to deal 7 | //in the Software without restriction, including without limitation the rights 8 | //to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | //copies of the Software, and to permit persons to whom the Software is 10 | //furnished to do so, subject to the following conditions: 11 | // 12 | //The above copyright notice and this permission notice shall be included in all 13 | //copies or substantial portions of the Software. 14 | // 15 | //THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | //IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | //FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | //AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | //LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | //OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | //SOFTWARE. 22 | // 23 | // Modifed 2019 Arthur Brussee 24 | 25 | using System; 26 | using KNN.Internal; 27 | using KNN.Jobs; 28 | using Unity.Collections; 29 | using Unity.Collections.LowLevel.Unsafe; 30 | using Unity.Jobs; 31 | using Unity.Mathematics; 32 | 33 | namespace KNN.Internal { 34 | public static unsafe class UnsafeUtilityEx { 35 | public static T* AllocArray(int length, Allocator allocator) where T : unmanaged { 36 | return (T*)UnsafeUtility.Malloc(length * UnsafeUtility.SizeOf(), UnsafeUtility.AlignOf(), allocator); 37 | } 38 | } 39 | } 40 | 41 | namespace KNN { 42 | [NativeContainerSupportsDeallocateOnJobCompletion, NativeContainer, System.Diagnostics.DebuggerDisplay("Length = {Points.Length}")] 43 | public struct KnnContainer : IDisposable { 44 | // We manage safety by our own sentinel. Disable unity's safety system for internal caches / arrays 45 | [NativeDisableContainerSafetyRestriction] 46 | public NativeArray Points; 47 | 48 | [NativeDisableContainerSafetyRestriction] 49 | NativeArray m_permutation; 50 | 51 | [NativeDisableContainerSafetyRestriction] 52 | NativeList m_nodes; 53 | 54 | [NativeDisableContainerSafetyRestriction] 55 | NativeArray m_rootNodeIndex; 56 | 57 | [NativeDisableContainerSafetyRestriction] 58 | NativeQueue m_buildQueue; 59 | 60 | KdNode RootNode => m_nodes[m_rootNodeIndex[0]]; 61 | 62 | #if ENABLE_UNITY_COLLECTIONS_CHECKS 63 | // Note: MUST be named m_Safey, m_DisposeSentinel exactly 64 | // ReSharper disable once InconsistentNaming 65 | internal AtomicSafetyHandle m_Safety; 66 | [NativeSetClassTypeToNullOnSchedule] 67 | // ReSharper disable once InconsistentNaming 68 | internal DisposeSentinel m_DisposeSentinel; 69 | #endif 70 | 71 | const int c_maxPointsPerLeafNode = 64; 72 | 73 | public struct KnnQueryTemp : IDisposable { 74 | public MinMaxHeap MaxHeap; 75 | public MinMaxHeap MinHeap; 76 | 77 | public static KnnQueryTemp Create(int kCapacity) { 78 | KnnQueryTemp temp; 79 | temp.MaxHeap = new MinMaxHeap(kCapacity, Allocator.Temp); 80 | 81 | // Min heap keeps track of current stack. 82 | // The max stack depth is the tree depth 83 | // The tree depth is log_c(nodes) 84 | // Let's assume people have a tree at most 32 deep (which equals 2^32 * c_maxPointsPerLeafNode ~ 2^39 nodes) 85 | // There are left/right nodes -> 64 max on stack at any given time 86 | temp.MinHeap = new MinMaxHeap(64, Allocator.Temp); 87 | return temp; 88 | } 89 | 90 | public void PushQueryNode(int index, float3 closestPoint, float3 queryPosition) { 91 | float lengthsq = math.lengthsq(closestPoint - queryPosition); 92 | 93 | MinHeap.PushObjMin(new QueryNode { 94 | NodeIndex = index, 95 | TempClosestPoint = closestPoint, 96 | Distance = lengthsq 97 | }, lengthsq); 98 | } 99 | 100 | public void Dispose() { 101 | MaxHeap.Dispose(); 102 | MinHeap.Dispose(); 103 | } 104 | } 105 | 106 | public KnnContainer(NativeArray points, bool buildNow, Allocator allocator) { 107 | int nodeCountEstimate = 4 * (int) math.ceil(points.Length / (float) c_maxPointsPerLeafNode + 1) + 1; 108 | Points = points; 109 | 110 | // Both arrays are filled in as we go, so start with uninitialized mem 111 | m_nodes = new NativeList(nodeCountEstimate, allocator); 112 | 113 | // Dumb way to create an int* essentially.. 114 | m_permutation = new NativeArray(points.Length, allocator, NativeArrayOptions.UninitializedMemory); 115 | m_rootNodeIndex = new NativeArray(new[] {-1}, allocator); 116 | m_buildQueue = new NativeQueue(allocator); 117 | 118 | #if ENABLE_UNITY_COLLECTIONS_CHECKS 119 | if (allocator <= Allocator.None) { 120 | throw new ArgumentException("Allocator must be Temp, TempJob or Persistent", nameof(allocator)); 121 | } 122 | 123 | if (points.Length <= 0) { 124 | throw new ArgumentOutOfRangeException(nameof(points), "Input points length must be >= 0"); 125 | } 126 | 127 | DisposeSentinel.Create(out m_Safety, out m_DisposeSentinel, 0, allocator); 128 | #endif 129 | 130 | if (buildNow) { 131 | var rebuild = new KnnRebuildJob(this); 132 | rebuild.Schedule().Complete(); 133 | } 134 | } 135 | 136 | public void Rebuild() { 137 | #if ENABLE_UNITY_COLLECTIONS_CHECKS 138 | AtomicSafetyHandle.CheckWriteAndThrow(m_Safety); 139 | #endif 140 | 141 | m_nodes.Clear(); 142 | 143 | for (int i = 0; i < m_permutation.Length; ++i) { 144 | m_permutation[i] = i; 145 | } 146 | 147 | int rootNode = GetKdNode(MakeBounds(), 0, Points.Length); 148 | 149 | m_rootNodeIndex[0] = rootNode; 150 | m_buildQueue.Enqueue(rootNode); 151 | 152 | while (m_buildQueue.Count > 0) { 153 | int index = m_buildQueue.Dequeue(); 154 | SplitNode(index, out int posNodeIndex, out int negNodeIndex); 155 | 156 | if (m_nodes[negNodeIndex].Count > c_maxPointsPerLeafNode) { 157 | m_buildQueue.Enqueue(posNodeIndex); 158 | } 159 | 160 | if (m_nodes[posNodeIndex].Count > c_maxPointsPerLeafNode) { 161 | m_buildQueue.Enqueue(negNodeIndex); 162 | } 163 | } 164 | } 165 | 166 | public void Dispose() { 167 | #if ENABLE_UNITY_COLLECTIONS_CHECKS 168 | DisposeSentinel.Dispose(ref m_Safety, ref m_DisposeSentinel); 169 | #endif 170 | 171 | m_permutation.Dispose(); 172 | m_nodes.Dispose(); 173 | m_rootNodeIndex.Dispose(); 174 | m_buildQueue.Dispose(); 175 | } 176 | 177 | public JobHandle Dispose(JobHandle job) { 178 | #if ENABLE_UNITY_COLLECTIONS_CHECKS 179 | job.Complete(); 180 | DisposeSentinel.Dispose(ref m_Safety, ref m_DisposeSentinel); 181 | #endif 182 | job = m_permutation.Dispose(job); 183 | job = m_nodes.Dispose(job); 184 | job = m_rootNodeIndex.Dispose(job); 185 | job = m_buildQueue.Dispose(job); 186 | return job; 187 | } 188 | 189 | int GetKdNode(KdNodeBounds bounds, int start, int end) { 190 | m_nodes.Add(new KdNode { 191 | Bounds = bounds, 192 | Start = start, 193 | End = end, 194 | PartitionAxis = -1, 195 | PartitionCoordinate = 0.0f, 196 | PositiveChildIndex = -1, 197 | NegativeChildIndex = -1 198 | }); 199 | 200 | return m_nodes.Length - 1; 201 | } 202 | 203 | /// 204 | /// For calculating root node bounds 205 | /// 206 | /// Boundary of all Vector3 points 207 | KdNodeBounds MakeBounds() { 208 | var max = new float3(float.MinValue, float.MinValue, float.MinValue); 209 | var min = new float3(float.MaxValue, float.MaxValue, float.MaxValue); 210 | int even = Points.Length & ~1; // calculate even Length 211 | 212 | // min, max calculations 213 | // 3n/2 calculations instead of 2n 214 | for (int i0 = 0; i0 < even; i0 += 2) { 215 | int i1 = i0 + 1; 216 | 217 | // X Coords 218 | if (Points[i0].x > Points[i1].x) { 219 | // i0 is bigger, i1 is smaller 220 | if (Points[i1].x < min.x) { 221 | min.x = Points[i1].x; 222 | } 223 | 224 | if (Points[i0].x > max.x) { 225 | max.x = Points[i0].x; 226 | } 227 | } else { 228 | // i1 is smaller, i0 is bigger 229 | if (Points[i0].x < min.x) { 230 | min.x = Points[i0].x; 231 | } 232 | 233 | if (Points[i1].x > max.x) { 234 | max.x = Points[i1].x; 235 | } 236 | } 237 | 238 | // Y Coords 239 | if (Points[i0].y > Points[i1].y) { 240 | // i0 is bigger, i1 is smaller 241 | if (Points[i1].y < min.y) { 242 | min.y = Points[i1].y; 243 | } 244 | 245 | if (Points[i0].y > max.y) { 246 | max.y = Points[i0].y; 247 | } 248 | } else { 249 | // i1 is smaller, i0 is bigger 250 | if (Points[i0].y < min.y) { 251 | min.y = Points[i0].y; 252 | } 253 | 254 | if (Points[i1].y > max.y) { 255 | max.y = Points[i1].y; 256 | } 257 | } 258 | 259 | // Z Coords 260 | if (Points[i0].z > Points[i1].z) { 261 | // i0 is bigger, i1 is smaller 262 | if (Points[i1].z < min.z) { 263 | min.z = Points[i1].z; 264 | } 265 | 266 | if (Points[i0].z > max.z) { 267 | max.z = Points[i0].z; 268 | } 269 | } else { 270 | // i1 is smaller, i0 is bigger 271 | if (Points[i0].z < min.z) { 272 | min.z = Points[i0].z; 273 | } 274 | 275 | if (Points[i1].z > max.z) { 276 | max.z = Points[i1].z; 277 | } 278 | } 279 | } 280 | 281 | // if array was odd, calculate also min/max for the last element 282 | if (even != Points.Length) { 283 | // X 284 | if (min.x > Points[even].x) { 285 | min.x = Points[even].x; 286 | } 287 | 288 | if (max.x < Points[even].x) { 289 | max.x = Points[even].x; 290 | } 291 | 292 | // Y 293 | if (min.y > Points[even].y) { 294 | min.y = Points[even].y; 295 | } 296 | 297 | if (max.y < Points[even].y) { 298 | max.y = Points[even].y; 299 | } 300 | 301 | // Z 302 | if (min.z > Points[even].z) { 303 | min.z = Points[even].z; 304 | } 305 | 306 | if (max.z < Points[even].z) { 307 | max.z = Points[even].z; 308 | } 309 | } 310 | 311 | var b = new KdNodeBounds(); 312 | b.Min = min; 313 | b.Max = max; 314 | return b; 315 | } 316 | 317 | // TODO: When multiple points overlap exactly this function breaks. 318 | /// 319 | /// Recursive splitting procedure 320 | /// 321 | void SplitNode(int parentIndex, out int posNodeIndex, out int negNodeIndex) { 322 | KdNode parent = m_nodes[parentIndex]; 323 | 324 | // center of bounding box 325 | KdNodeBounds parentBounds = parent.Bounds; 326 | float3 parentBoundsSize = parentBounds.Size; 327 | 328 | // Find axis where bounds are largest 329 | int splitAxis = 0; 330 | float axisSize = parentBoundsSize.x; 331 | 332 | if (axisSize < parentBoundsSize.y) { 333 | splitAxis = 1; 334 | axisSize = parentBoundsSize.y; 335 | } 336 | 337 | if (axisSize < parentBoundsSize.z) { 338 | splitAxis = 2; 339 | } 340 | 341 | // Our axis min-max bounds 342 | float boundsStart = parentBounds.Min[splitAxis]; 343 | float boundsEnd = parentBounds.Max[splitAxis]; 344 | 345 | // Calculate the spiting coords 346 | float splitPivot = CalculatePivot(parent.Start, parent.End, boundsStart, boundsEnd, splitAxis); 347 | 348 | // 'Spiting' array to two sub arrays 349 | int splittingIndex = Partition(parent.Start, parent.End, splitPivot, splitAxis); 350 | 351 | // Negative / Left node 352 | float3 negMax = parentBounds.Max; 353 | negMax[splitAxis] = splitPivot; 354 | 355 | var bounds = parentBounds; 356 | bounds.Max = negMax; 357 | negNodeIndex = GetKdNode(bounds, parent.Start, splittingIndex); 358 | 359 | parent.PartitionAxis = splitAxis; 360 | parent.PartitionCoordinate = splitPivot; 361 | 362 | // Positive / Right node 363 | float3 posMin = parentBounds.Min; 364 | posMin[splitAxis] = splitPivot; 365 | 366 | bounds = parentBounds; 367 | bounds.Min = posMin; 368 | posNodeIndex = GetKdNode(bounds, splittingIndex, parent.End); 369 | 370 | parent.NegativeChildIndex = negNodeIndex; 371 | parent.PositiveChildIndex = posNodeIndex; 372 | 373 | // Write back node to array to update those values 374 | m_nodes[parentIndex] = parent; 375 | } 376 | 377 | /// 378 | /// Sliding midpoint splitting pivot calculation 379 | /// 1. First splits node to two equal parts (midPoint) 380 | /// 2. Checks if elements are in both sides of splitted bounds 381 | /// 3a. If they are, just return midPoint 382 | /// 3b. If they are not, then points are only on left or right bound. 383 | /// 4. Move the splitting pivot so that it shrinks part with points completely (calculate min or max dependent) and return. 384 | /// 385 | float CalculatePivot(int start, int end, float boundsStart, float boundsEnd, int axis) { 386 | //! sliding midpoint rule 387 | float midPoint = (boundsStart + boundsEnd) / 2.0f; 388 | 389 | bool negative = false; 390 | bool positive = false; 391 | 392 | float negMax = float.MinValue; 393 | float posMin = float.MaxValue; 394 | 395 | // this for loop section is used both for sorted and unsorted data 396 | for (int i = start; i < end; i++) { 397 | float val = Points[m_permutation[i]][axis]; 398 | 399 | if (val < midPoint) { 400 | negative = true; 401 | } else { 402 | positive = true; 403 | } 404 | 405 | if (negative && positive) { 406 | return midPoint; 407 | } 408 | } 409 | 410 | if (negative) { 411 | for (int i = start; i < end; i++) { 412 | float val = Points[m_permutation[i]][axis]; 413 | 414 | if (negMax < val) { 415 | negMax = val; 416 | } 417 | } 418 | 419 | return negMax; 420 | } 421 | 422 | for (int i = start; i < end; i++) { 423 | float val = Points[m_permutation[i]][axis]; 424 | 425 | if (posMin > val) { 426 | posMin = val; 427 | } 428 | } 429 | 430 | return posMin; 431 | } 432 | 433 | /// 434 | /// Similar to Hoare partitioning algorithm (used in Quick Sort) 435 | /// Modification: pivot is not left-most element but is instead argument of function 436 | /// Calculates splitting index and partially sorts elements (swaps them until they are on correct side - depending on pivot) 437 | /// Complexity: O(n) 438 | /// 439 | /// Start index 440 | /// End index 441 | /// Pivot that decides boundary between left and right 442 | /// Axis of this pivoting 443 | /// 444 | /// Returns splitting index that subdivides array into 2 smaller arrays 445 | /// left = [start, pivot), 446 | /// right = [pivot, end) 447 | /// 448 | int Partition(int start, int end, float partitionPivot, int axis) { 449 | // note: increasing right pointer is actually decreasing! 450 | int lp = start - 1; // left pointer (negative side) 451 | int rp = end; // right pointer (positive side) 452 | 453 | while (true) { 454 | do { 455 | // move from left to the right until "out of bounds" value is found 456 | lp++; 457 | } while (lp < rp && Points[m_permutation[lp]][axis] < partitionPivot); 458 | 459 | do { 460 | // move from right to the left until "out of bounds" value found 461 | rp--; 462 | } while (lp < rp && Points[m_permutation[rp]][axis] >= partitionPivot); 463 | 464 | if (lp < rp) { 465 | // swap 466 | int temp = m_permutation[lp]; 467 | m_permutation[lp] = m_permutation[rp]; 468 | m_permutation[rp] = temp; 469 | } else { 470 | return lp; 471 | } 472 | } 473 | } 474 | 475 | public void QueryRange(float3 queryPosition, float radius, NativeList result) { 476 | #if ENABLE_UNITY_COLLECTIONS_CHECKS 477 | AtomicSafetyHandle.CheckReadAndThrow(m_Safety); 478 | #endif 479 | 480 | // Start with a temp of some size. This will be resized dynamically 481 | var temp = KnnQueryTemp.Create(32); 482 | 483 | // Biggest Smallest Squared Radius 484 | float bssr = radius * radius; 485 | float3 rootClosestPoint = RootNode.Bounds.ClosestPoint(queryPosition); 486 | 487 | temp.PushQueryNode(m_rootNodeIndex[0], rootClosestPoint, queryPosition); 488 | 489 | while (temp.MinHeap.Count > 0) { 490 | QueryNode queryNode = temp.MinHeap.PopObjMin(); 491 | 492 | if (queryNode.Distance > bssr) { 493 | continue; 494 | } 495 | 496 | KdNode node = m_nodes[queryNode.NodeIndex]; 497 | 498 | if (!node.Leaf) { 499 | int partitionAxis = node.PartitionAxis; 500 | float partitionCoord = node.PartitionCoordinate; 501 | float3 tempClosestPoint = queryNode.TempClosestPoint; 502 | 503 | if (tempClosestPoint[partitionAxis] - partitionCoord < 0) { 504 | // we already know we are on the side of negative bound/node, 505 | // so we don't need to test for distance 506 | // push to stack for later querying 507 | temp.PushQueryNode(node.NegativeChildIndex, tempClosestPoint, queryPosition); 508 | 509 | // project the tempClosestPoint to other bound 510 | tempClosestPoint[partitionAxis] = partitionCoord; 511 | 512 | if (node.Count != 0) { 513 | temp.PushQueryNode(node.PositiveChildIndex, tempClosestPoint, queryPosition); 514 | } 515 | } 516 | else { 517 | // we already know we are on the side of positive bound/node, 518 | // so we don't need to test for distance 519 | // push to stack for later querying 520 | temp.PushQueryNode(node.PositiveChildIndex, tempClosestPoint, queryPosition); 521 | 522 | // project the tempClosestPoint to other bound 523 | tempClosestPoint[partitionAxis] = partitionCoord; 524 | 525 | if (node.Count != 0) { 526 | temp.PushQueryNode(node.NegativeChildIndex, tempClosestPoint, queryPosition); 527 | } 528 | } 529 | } else { 530 | for (int i = node.Start; i < node.End; i++) { 531 | int index = m_permutation[i]; 532 | float sqrDist = math.lengthsq(Points[index] - queryPosition); 533 | 534 | if (sqrDist <= bssr) { 535 | // Unlike the k-query we want to keep _all_ objects in range 536 | // So resize the heap when pushing this node 537 | if (temp.MaxHeap.IsFull) { 538 | temp.MaxHeap.Resize(temp.MaxHeap.Count * 2); 539 | } 540 | 541 | temp.MaxHeap.PushObjMax(index, sqrDist); 542 | } 543 | } 544 | } 545 | } 546 | 547 | while (temp.MaxHeap.Count > 0) { 548 | result.Add(temp.MaxHeap.PopObjMax()); 549 | } 550 | 551 | temp.Dispose(); 552 | } 553 | 554 | public void QueryKNearest(float3 queryPosition, NativeSlice result) { 555 | #if ENABLE_UNITY_COLLECTIONS_CHECKS 556 | AtomicSafetyHandle.CheckReadAndThrow(m_Safety); 557 | #endif 558 | 559 | var temp = KnnQueryTemp.Create(result.Length); 560 | int k = result.Length; 561 | 562 | // Biggest Smallest Squared Radius 563 | float bssr = float.PositiveInfinity; 564 | float3 rootClosestPoint = RootNode.Bounds.ClosestPoint(queryPosition); 565 | 566 | temp.PushQueryNode(m_rootNodeIndex[0], rootClosestPoint, queryPosition); 567 | 568 | while (temp.MinHeap.Count > 0) { 569 | QueryNode queryNode = temp.MinHeap.PopObjMin(); 570 | 571 | if (queryNode.Distance > bssr) { 572 | continue; 573 | } 574 | 575 | KdNode node = m_nodes[queryNode.NodeIndex]; 576 | 577 | if (!node.Leaf) { 578 | int partitionAxis = node.PartitionAxis; 579 | float partitionCoord = node.PartitionCoordinate; 580 | float3 tempClosestPoint = queryNode.TempClosestPoint; 581 | 582 | if (tempClosestPoint[partitionAxis] - partitionCoord < 0) { 583 | // we already know we are on the side of negative bound/node, 584 | // so we don't need to test for distance 585 | // push to stack for later querying 586 | temp.PushQueryNode(node.NegativeChildIndex, tempClosestPoint, queryPosition); 587 | 588 | // project the tempClosestPoint to other bound 589 | tempClosestPoint[partitionAxis] = partitionCoord; 590 | 591 | if (node.Count != 0) { 592 | temp.PushQueryNode(node.PositiveChildIndex, tempClosestPoint, queryPosition); 593 | } 594 | } else { 595 | // we already know we are on the side of positive bound/node, 596 | // so we don't need to test for distance 597 | // push to stack for later querying 598 | temp.PushQueryNode(node.PositiveChildIndex, tempClosestPoint, queryPosition); 599 | 600 | // project the tempClosestPoint to other bound 601 | tempClosestPoint[partitionAxis] = partitionCoord; 602 | 603 | if (node.Count != 0) { 604 | temp.PushQueryNode(node.NegativeChildIndex, tempClosestPoint, queryPosition); 605 | } 606 | } 607 | } else { 608 | for (int i = node.Start; i < node.End; i++) { 609 | int index = m_permutation[i]; 610 | float sqrDist = math.lengthsq(Points[index] - queryPosition); 611 | 612 | if (sqrDist <= bssr) { 613 | temp.MaxHeap.PushObjMax(index, sqrDist); 614 | 615 | if (temp.MaxHeap.Count == k) { 616 | bssr = temp.MaxHeap.HeadValue; 617 | } 618 | } 619 | } 620 | } 621 | } 622 | 623 | // Astero Change: AR-156 624 | // Only add results up to the size of the heap. If the result set is looking for 625 | // more results than the heap has, fill it with -1 to indicate as such. 626 | var originalCount = temp.MaxHeap.Count; 627 | for (int i = 0; i < k; i++) 628 | { 629 | if (i < originalCount) 630 | { 631 | result[i] = temp.MaxHeap.PopObjMax(); 632 | } 633 | else 634 | { 635 | result[i] = -1; 636 | } 637 | } 638 | 639 | temp.Dispose(); 640 | } 641 | } 642 | } 643 | --------------------------------------------------------------------------------