mol-model/query expandProperty, some "query modifer" tweaks

src/mol-data/generic/unique-array.ts

@@ -4,13 +4,13 @@
  * @author David Sehnal <david.sehnal@gmail.com>
  */
 
-interface UniqueArray<K, T> {
+interface UniqueArray<K, T = K> {
     keys: Set<K>,
     array: T[]
 }
 
 namespace UniqueArray {
-    export function create<K, T>(): UniqueArray<K, T> {
+    export function create<K, T = K>(): UniqueArray<K, T> {
         return { keys: new Set<K>(), array: [] };
     }
 

src/mol-model/structure/query/queries/filters.ts

@@ -177,12 +177,14 @@ function withinMaxRadius({ queryCtx, selection, target, maxRadius, invert, eleme
     const targetStructure = StructureSelection.unionStructure(target);
     const ret = StructureSelection.LinearBuilder(queryCtx.inputStructure);
 
+    queryCtx.pushCurrentElement();
     StructureSelection.forEach(selection, (s, sI) => {
         let withinRadius = checkStructureMaxRadiusDistance(queryCtx, targetStructure, s, maxRadius, elementRadius);
         if (invert) withinRadius = !withinRadius;
         if (withinRadius) ret.add(s);
         if (sI % 10 === 0) queryCtx.throwIfTimedOut();
     });
+    queryCtx.popCurrentElement();
 
     return ret.getSelection();
 }
@@ -191,12 +193,14 @@ function withinMinMaxRadius({ queryCtx, selection, target, minRadius, maxRadius,
     const targetStructure = StructureSelection.unionStructure(target);
     const ret = StructureSelection.LinearBuilder(queryCtx.inputStructure);
 
+    queryCtx.pushCurrentElement();
     StructureSelection.forEach(selection, (s, sI) => {
         let withinRadius = checkStructureMinMaxDistance(queryCtx, targetStructure, s, minRadius, maxRadius, elementRadius);
         if (invert) withinRadius = !withinRadius;
         if (withinRadius) ret.add(s);
         if (sI % 10 === 0) queryCtx.throwIfTimedOut();
     });
+    queryCtx.popCurrentElement();
 
     return ret.getSelection();
 }

src/mol-model/structure/query/queries/modifiers.ts

@@ -10,8 +10,10 @@ import { StructureQuery } from '../query';
 import { StructureSelection } from '../selection';
 import { UniqueStructuresBuilder } from '../utils/builders';
 import { StructureUniqueSubsetBuilder } from '../../structure/util/unique-subset-builder';
-import { QueryContext } from '../context';
+import { QueryContext, QueryFn } from '../context';
 import { structureIntersect, structureSubtract } from '../utils/structure-set';
+import { UniqueArray } from 'mol-data/generic';
+import { StructureSubsetBuilder } from '../../structure/util/subset-builder';
 
 function getWholeResidues(ctx: QueryContext, source: Structure, structure: Structure) {
     const builder = source.subsetBuilder(true);
@@ -128,6 +130,7 @@ export function intersectBy(query: StructureQuery, by: StructureQuery): Structur
         StructureSelection.forEach(selection, (s, sI) => {
             const ii = structureIntersect(unionBy, s);
             if (ii.elementCount !== 0) ret.add(ii);
+            if (sI % 50 === 0) ctx.throwIfTimedOut();
         });
 
         return ret.getSelection();
@@ -147,6 +150,7 @@ export function exceptBy(query: StructureQuery, by: StructureQuery): StructureQu
         StructureSelection.forEach(selection, (s, sI) => {
             const diff = structureSubtract(s, subtractBy);
             if (diff.elementCount !== 0) ret.add(diff);
+            if (sI % 50 === 0) ctx.throwIfTimedOut();
         });
 
         return ret.getSelection();
@@ -161,4 +165,57 @@ export function union(query: StructureQuery): StructureQuery {
     };
 }
 
-// TODO: unionBy (skip this one?), cluster, includeConnected, includeSurroundings with "radii", expandProperty
+export function expandProperty(query: StructureQuery, property: QueryFn): StructureQuery {
+    return ctx => {
+        const src = query(ctx);
+        const propertyToStructureIndexMap = new Map<any, UniqueArray<number>>();
+
+        const builders: StructureSubsetBuilder[] = [];
+        ctx.pushCurrentElement();
+        StructureSelection.forEach(src, (s, sI) => {
+            for (const unit of s.units) {
+                ctx.element.unit = unit;
+                const elements = unit.elements;
+                for (let i = 0, _i = elements.length; i < _i; i++) {
+                    ctx.element.element = elements[i];
+                    const p = property(ctx);
+                    let arr: UniqueArray<number>;
+                    if (propertyToStructureIndexMap.has(p)) arr = propertyToStructureIndexMap.get(p)!;
+                    else {
+                        arr = UniqueArray.create<number>();
+                        propertyToStructureIndexMap.set(p, arr);
+                    }
+                    UniqueArray.add(arr, sI, sI);
+                }
+            }
+            builders[sI] = ctx.inputStructure.subsetBuilder(true);
+
+            if (sI % 10 === 0) ctx.throwIfTimedOut();
+        });
+
+        for (const unit of ctx.inputStructure.units) {
+            ctx.element.unit = unit;
+            const elements = unit.elements;
+            for (let i = 0, _i = elements.length; i < _i; i++) {
+                ctx.element.element = elements[i];
+                const p = property(ctx);
+                if (!propertyToStructureIndexMap.has(p)) continue;
+
+                const indices = propertyToStructureIndexMap.get(p)!.array;
+
+                for (let _sI = 0, __sI = indices.length; _sI < __sI; _sI++) {
+                    builders[indices[i]].addToUnit(unit.id, elements[i]);
+                }
+            }
+        }
+
+        ctx.popCurrentElement();
+
+        const ret = StructureSelection.UniqueBuilder(ctx.inputStructure);
+        for (const b of builders) ret.add(b.getStructure());
+
+        return ret.getSelection();
+    };
+}
+
+// TODO: unionBy (skip this one?), cluster, includeConnected, includeSurroundings with "radii"

src/mol-model/structure/query/selection.ts

@@ -101,6 +101,7 @@ namespace StructureSelection {
     export function LinearBuilder(structure: Structure): Builder { return new LinearBuilderImpl(structure); }
     export function UniqueBuilder(structure: Structure): Builder { return new HashBuilderImpl(structure); }
 
+    // TODO: build timeout checking into this?
     export function forEach(sel: StructureSelection, fn: (s: Structure, i: number) => void) {
         let idx = 0;
         if (StructureSelection.isSingleton(sel)) {