Преглед изворни кода

Avoid using unnecessary extractMinimum from heap if k=1 on nearest search, add nearest method as unreleased.

giagitom пре 2 година
родитељ
комит
53df86c585

+ 1 - 0
CHANGELOG.md

@@ -6,6 +6,7 @@ Note that since we don't clearly distinguish between a public and private interf
 
 ## [Unreleased]
 
+- Add nearest method to lookup3d
 - Add mipmap-based blur for skybox backgrounds
 
 ## [v3.19.0] - 2022-10-01

+ 16 - 6
src/mol-math/geometry/lookup3d/grid.ts

@@ -404,13 +404,23 @@ function queryNearest<T extends number = number>(ctx: QueryContext, result: Resu
         }
         expandGrid = false;
         if (nextGCount === 0) {
-            while (!tmpHeapG.isEmpty() && (gridPointsFinished || tmpHeapG.findMinimum()!.key as unknown as number <= maxRange) && result.count < k) {
-                const node = tmpHeapG.extractMinimum();
-                const squaredDistance = node!.key, index = node!.value;
+            if (k === 1){
+              const node = tmpHeapG.findMinimum()
+              if (node) {
+                const { key: squaredDistance, value : index } = node!
+                //const squaredDistance = node!.key, index = node!.value;
                 Result.add(result, index as number, squaredDistance as number);
-                if (stopIf && !stop) {
-                    stop = stopIf(index, squaredDistance);
-                }
+                return true
+              }
+            } else {
+              while (!tmpHeapG.isEmpty() && (gridPointsFinished || tmpHeapG.findMinimum()!.key as number <= maxRange) && result.count < k) {
+                  const node = tmpHeapG.extractMinimum();
+                  const squaredDistance = node!.key, index = node!.value;
+                  Result.add(result, index as number, squaredDistance as number);
+                  if (stopIf && !stop) {
+                      stop = stopIf(index, squaredDistance);
+                  }
+              }
             }
             if (result.count >= k || stop || result.count >= indicesCount) return result.count > 0;
             expandGrid = true;

+ 15 - 8
src/mol-model/structure/structure/util/lookup3d.ts

@@ -53,11 +53,10 @@ export function StructureLookup3DResultContext(): StructureLookup3DResultContext
     return { result: StructureResult.create(), closeUnitsResult: Result.create(), unitGroupResult: Result.create() };
 }
 
-const tmpHeap = new FibonacciHeap();
-
 export class StructureLookup3D {
     private unitLookup: Lookup3D;
     private pivot = Vec3();
+    private heap = new FibonacciHeap();
 
     findUnitIndices(x: number, y: number, z: number, radius: number): Result<number> {
         return this.unitLookup.find(x, y, z, radius);
@@ -95,9 +94,9 @@ export class StructureLookup3D {
     }
 
     _nearest(x: number, y: number, z: number, k: number, ctx: StructureLookup3DResultContext): StructureResult {
-        const result = ctx.result;
+        const result = ctx.result, heap = this.heap;
         Result.reset(result);
-        tmpHeap.clear();
+        heap.clear();
         const { units } = this.structure;
         let elementsCount = 0;
         const closeUnits = this.unitLookup.nearest(x, y, z, units.length, (uid: number) => (elementsCount += units[uid].elements.length) >= k, ctx.closeUnitsResult); // sort units based on distance to the point
@@ -117,15 +116,23 @@ export class StructureLookup3D {
             totalCount += groupResult.count;
             maxDistResult = Math.max(maxDistResult, groupResult.squaredDistances[groupResult.count - 1]);
             for (let j = 0, _j = groupResult.count; j < _j; j++) {
-                tmpHeap.insert(groupResult.squaredDistances[j], { index: groupResult.indices[j], unit: unit });
+                heap.insert(groupResult.squaredDistances[j], { index: groupResult.indices[j], unit: unit });
             }
         }
-        while (!tmpHeap.isEmpty() && result.count < k) {
-            const node = tmpHeap.extractMinimum();
-            if (!node) throw new Error('Cannot extract minimum, should not happen');
+        if (k === 1){
+          const node = heap.findMinimum()
+          if (node) {
             const { key: squaredDistance } = node;
             const { unit, index } = node.value as { index: UnitIndex, unit: Unit };
             StructureResult.add(result, unit as Unit, index as UnitIndex, squaredDistance as number);
+          }
+        } else {
+          while (!heap.isEmpty() && result.count < k) {
+              const node = heap.extractMinimum();
+              const { key: squaredDistance } = node!;
+              const { unit, index } = node!.value as { index: UnitIndex, unit: Unit };
+              StructureResult.add(result, unit as Unit, index as UnitIndex, squaredDistance as number);
+          }
         }
         return result;
     }