wip, mol-model query

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

@@ -15,9 +15,11 @@ export interface QueryContextView {
 export class QueryContext implements QueryContextView {
     private currentElementStack: StructureElement[] = [];
     private currentStructureStack: Structure[] = [];
+    private inputStructureStack: Structure[] = [];
+
     private timeCreated = now();
-    private timeoutMs: number;
 
+    readonly timeoutMs: number;
     readonly inputStructure: Structure;
 
     /** Current element */
@@ -43,6 +45,16 @@ export class QueryContext implements QueryContextView {
         else (this.currentStructure as Structure) = void 0 as any;
     }
 
+    pushInputStructure(structure: Structure) {
+        this.inputStructureStack.push(this.inputStructure);
+        (this.inputStructure as any) = structure;
+    }
+
+    popInputStructure() {
+        if (this.inputStructureStack.length === 0) throw new Error('Must push before pop.');
+        (this.inputStructure as any) = this.inputStructureStack.pop();
+    }
+
     throwIfTimedOut() {
         if (this.timeoutMs === 0) return;
         if (now() - this.timeCreated > this.timeoutMs) {
@@ -50,8 +62,6 @@ export class QueryContext implements QueryContextView {
         }
     }
 
-    // todo timeout
-
     constructor(structure: Structure, timeoutMs = 0) {
         this.inputStructure = structure;
         this.timeoutMs = timeoutMs;

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

@@ -7,6 +7,8 @@
 import { StructureQuery } from '../query';
 import { StructureSelection } from '../selection';
 import { none } from './generators';
+import { HashSet } from 'mol-data/generic';
+import { Structure } from '../../structure';
 
 export function merge(queries: ArrayLike<StructureQuery>): StructureQuery {
     if (queries.length === 0) {
@@ -26,4 +28,40 @@ export function merge(queries: ArrayLike<StructureQuery>): StructureQuery {
     }
 }
 
-// TODO: intersect, distanceCluster
+export function intersect(queries: ArrayLike<StructureQuery>): StructureQuery {
+    if (queries.length === 0) {
+        return none;
+    } else if (queries.length === 1) {
+        return queries[0];
+    }
+    return ctx => {
+        const selections: StructureSelection[] = [];
+        for (let i = 0; i < queries.length; i++) selections.push(queries[i](ctx));
+        let pivotIndex = 0, pivotLength = StructureSelection.structureCount(selections[0]);
+        for (let i = 1; i < selections.length; i++) {
+            const len = StructureSelection.structureCount(selections[i]);
+            if (len < pivotLength) {
+                pivotIndex = i;
+                pivotLength = len;
+            }
+        }
+
+        ctx.throwIfTimedOut();
+        const pivotSet = HashSet<Structure>(s => s.hashCode, Structure.areEqual);
+        StructureSelection.forEach(selections[pivotIndex], s => pivotSet.add(s));
+
+        const ret = StructureSelection.UniqueBuilder(ctx.inputStructure);
+
+        for (let pI = 0; pI < selections.length; pI++) {
+            if (pI === pivotIndex) continue;
+            StructureSelection.forEach(selections[pI], s => {
+                if (pivotSet.has(s)) ret.add(s);
+            });
+            ctx.throwIfTimedOut();
+        }
+
+        return ret.getSelection();
+    };
+}
+
+// TODO: distanceCluster

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

@@ -10,6 +10,7 @@ import { QueryContext, QueryFn, QueryPredicate } from '../context';
 import { StructureQuery } from '../query';
 import { StructureSelection } from '../selection';
 import { structureAreIntersecting } from '../utils/structure';
+import { Vec3 } from 'mol-math/linear-algebra';
 
 export function pick(query: StructureQuery, pred: QueryPredicate): StructureQuery {
     return ctx => {
@@ -101,4 +102,93 @@ export function areIntersectedBy(query: StructureQuery, by: StructureQuery): Str
     };
 }
 
-// TODO: within, isConnectedTo
+export interface WithinParams {
+    query: StructureQuery,
+    target: StructureQuery,
+    minRadius?: number,
+    maxRadius: number,
+    atomRadius?: QueryFn<number>,
+    invert?: boolean
+}
+
+function _zeroRadius(ctx: QueryContext) { return 0; }
+
+export function within(params: WithinParams): StructureQuery {
+    return queryCtx => {
+        const ctx: WithinContext = {
+            queryCtx,
+            selection: params.query(queryCtx),
+            target: params.target(queryCtx),
+            maxRadius: params.maxRadius,
+            minRadius: params.minRadius ? params.minRadius : 0,
+            atomRadius: params.atomRadius || _zeroRadius,
+            invert: !!params.invert,
+        }
+
+        if (ctx.minRadius === 0 && ctx.atomRadius === _zeroRadius) {
+            return withinMaxRadius(ctx);
+        } else {
+            // TODO
+            throw 'not implemented';
+            // return withinMinMaxRadius(ctx);
+        }
+    }
+}
+
+interface WithinContext {
+    queryCtx: QueryContext,
+    selection: StructureSelection,
+    target: StructureSelection,
+    minRadius: number,
+    maxRadius: number,
+    invert: boolean,
+    atomRadius: QueryFn<number>
+}
+function withinMaxRadius({ queryCtx, selection, target, maxRadius, invert }: WithinContext) {
+    const targetLookup = StructureSelection.unionStructure(target).lookup3d;
+    const ret = StructureSelection.LinearBuilder(queryCtx.inputStructure);
+
+    const pos = Vec3.zero();
+    StructureSelection.forEach(selection, (s, sI) => {
+        const { units } = s;
+
+        let withinRadius = false;
+        for (let i = 0, _i = units.length; i < _i; i++) {
+            const unit = units[i];
+            const { elements, conformation } = unit;
+
+            switch (unit.kind) {
+                case Unit.Kind.Atomic:
+                // TODO: assign radius to gaussian elements?
+                case Unit.Kind.Gaussians:
+                    for (let i = 0, _i = elements.length; i < _i; i++) {
+                        conformation.position(elements[i], pos);
+                        if (targetLookup.check(pos[0], pos[1], pos[2], maxRadius)) {
+                            withinRadius = true;
+                            break;
+                        }
+                    }
+                    break;
+                case Unit.Kind.Spheres:
+                    const radius = unit.coarseConformation.radius;
+                    for (let i = 0, _i = elements.length; i < _i; i++) {
+                        conformation.position(elements[i], pos);
+                        if (targetLookup.check(pos[0], pos[1], pos[2], maxRadius + radius[elements[i]])) {
+                            withinRadius = true;
+                            break;
+                        }
+                    }
+                    break;
+            }
+            if (withinRadius) break;
+        }
+        if (invert) withinRadius = !withinRadius;
+        if (withinRadius) ret.add(s);
+        if (sI % 10 === 0) queryCtx.throwIfTimedOut();
+    });
+
+    return ret.getSelection();
+}
+
+
+// TODO: isConnectedTo

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

@@ -7,9 +7,14 @@
 import { StructureQuery } from '../query'
 import { StructureSelection } from '../selection'
 import { Unit, StructureProperties as P } from '../../structure'
-import { Segmentation } from 'mol-data/int'
+import { Segmentation, SortedArray } from 'mol-data/int'
 import { LinearGroupingBuilder } from '../utils/builders';
-import { QueryPredicate, QueryFn, QueryContextView } from '../context';
+import { QueryPredicate, QueryFn, QueryContextView, QueryContext } from '../context';
+import { UnitRing } from '../../structure/unit/rings';
+import Structure from '../../structure/structure';
+import { ElementIndex } from '../../model';
+import { UniqueArray } from 'mol-data/generic';
+import { structureSubtract } from '../utils/structure';
 
 export const none: StructureQuery = ctx => StructureSelection.Sequence(ctx.inputStructure, []);
 export const all: StructureQuery = ctx => StructureSelection.Singletons(ctx.inputStructure, ctx.inputStructure);
@@ -157,4 +162,63 @@ function atomGroupsGrouped({ entityTest, chainTest, residueTest, atomTest, group
         ctx.popCurrentElement();
         return builder.getSelection();
     };
+}
+
+function getRingStructure(unit: Unit.Atomic, ring: UnitRing) {
+    const elements = new Int32Array(ring.length) as any as ElementIndex[];
+    for (let i = 0, _i = ring.length; i < _i; i++) elements[i] = unit.elements[ring[i]];
+    return Structure.create([unit.getChild(SortedArray.ofSortedArray(elements))])
+}
+
+export function rings(fingerprints?: ArrayLike<UnitRing.Fingerprint>): StructureQuery {
+    return ctx => {
+        const { units } = ctx.inputStructure;
+        let ret = StructureSelection.LinearBuilder(ctx.inputStructure);
+
+        if (!fingerprints || fingerprints.length === 0) {
+            for (const u of units) {
+                if (!Unit.isAtomic(u)) continue;
+
+                for (const r of u.rings.all) {
+                    ret.add(getRingStructure(u, r));
+                }
+            }
+        } else {
+            const uniqueFps = UniqueArray.create<UnitRing.Fingerprint, UnitRing.Fingerprint>();
+            for (let i = 0; i < fingerprints.length; i++) UniqueArray.add(uniqueFps, fingerprints[i], fingerprints[i]);
+
+            for (const u of units) {
+                if (!Unit.isAtomic(u)) continue;
+
+                const rings = u.rings;
+                for (const fp of uniqueFps.array) {
+                    if (!rings.byFingerprint.has(fp)) continue;
+                    for (const r of rings.byFingerprint.get(fp)!) {
+                        ret.add(getRingStructure(u, rings.all[r]));
+                    }
+                }
+            }
+        }
+
+        return ret.getSelection();
+    }
+}
+
+export function querySelection(selection: StructureQuery, query: StructureQuery, inComplement: boolean = false): StructureQuery {
+    return ctx => {
+        const targetSel = selection(ctx);
+        if (StructureSelection.structureCount(targetSel) === 0) return targetSel;
+
+        const target = inComplement
+        ? structureSubtract(ctx.inputStructure, StructureSelection.unionStructure(targetSel))
+        : StructureSelection.unionStructure(targetSel);
+
+        if (target.elementCount === 0) return StructureSelection.Empty(ctx.inputStructure);
+        ctx.throwIfTimedOut();
+
+        ctx.pushInputStructure(target);
+        const result = query(ctx);
+        ctx.popInputStructure();
+        return StructureSelection.withInputStructure(result, ctx.inputStructure);
+    }
 }

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

@@ -59,6 +59,7 @@ export function wholeResidues(query: StructureQuery, isFlat: boolean): Structure
 
 export interface IncludeSurroundingsParams {
     radius: number,
+    // TODO
     // atomRadius?: Element.Property<number>,
     wholeResidues?: boolean
 }
@@ -96,4 +97,6 @@ export function includeSurroundings(query: StructureQuery, params: IncludeSurrou
             return builder.getSelection();
         }
     };
-}
+}
+
+// TODO: queryEach, intersectBy, exceptBy, unionBy, union, cluster, includeConnected

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

@@ -119,7 +119,12 @@ namespace StructureSelection {
         }
     }
 
-    // TODO: spatial lookup
+    export function withInputStructure(selection: StructureSelection, structure: Structure) {
+        if (isSingleton(selection)) return Singletons(structure, selection.structure);
+        return Sequence(structure, selection.structures);
+    }
+
+    // TODO: spatial lookup?
 }
 
 export { StructureSelection }

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

@@ -21,6 +21,7 @@ import StructureProperties from './properties';
 import { ResidueIndex } from '../model/indexing';
 import { Carbohydrates } from './carbohydrates/data';
 import { computeCarbohydrates } from './carbohydrates/compute';
+import { Vec3 } from 'mol-math/linear-algebra';
 
 class Structure {
     readonly unitMap: IntMap<Unit>;
@@ -321,6 +322,81 @@ namespace Structure {
         sortArray(uniqueResidues.array);
         return uniqueResidues.array;
     }
+
+    const distVec = Vec3.zero();
+    function atomicOrGaussianDistance(u: Unit.Atomic | Unit.Gaussians, p: Vec3, r: number) {
+        const { elements, conformation } = u;
+        let minD = Number.MAX_VALUE;
+        for (let i = 0, _i = elements.length; i < _i; i++) {
+            const d = Vec3.distance(p, conformation.position(elements[i], distVec)) - r;
+            if (d < minD) minD = d;
+        }
+        return minD;
+    }
+
+    function sphereDistance(u: Unit.Spheres, p: Vec3, r: number) {
+        const { elements, conformation } = u;
+        const radius = u.coarseConformation.radius;
+        let minD = Number.MAX_VALUE;
+        for (let i = 0, _i = elements.length; i < _i; i++) {
+            const d = Vec3.distance(p, conformation.position(elements[i], distVec)) - r - radius[elements[i]];
+            if (d < minD) minD = d;
+        }
+        return minD;
+    }
+
+    export function minDistanceToPoint(s: Structure, point: Vec3, radius: number) {
+        const { units } = s;
+        let minD = Number.MAX_VALUE, d = 0;
+        for (let i = 0, _i = units.length; i < _i; i++) {
+            const unit = units[i];
+            switch (unit.kind) {
+                case Unit.Kind.Atomic:
+                // TODO: assign radius to gaussian elements?
+                case Unit.Kind.Gaussians:
+                    d = atomicOrGaussianDistance(unit, point, radius);
+                    break;
+                case Unit.Kind.Spheres:
+                    d = sphereDistance(unit, point, radius);
+                    break;
+            }
+            if (d < minD) minD = d;
+        }
+        return minD;
+    }
+
+    const distPivot = Vec3.zero();
+    export function distance(a: Structure, b: Structure) {
+        if (a.elementCount === 0 || b.elementCount === 0) return 0;
+
+        const { units } = a;
+        let minD = Number.MAX_VALUE, d = 0;
+
+        for (let i = 0, _i = units.length; i < _i; i++) {
+            const unit = units[i];
+            const { elements, conformation } = unit;
+
+            switch (unit.kind) {
+                case Unit.Kind.Atomic:
+                // TODO: assign radius to gaussian elements?
+                case Unit.Kind.Gaussians:
+                    for (let i = 0, _i = elements.length; i < _i; i++) {
+                        const d = minDistanceToPoint(b, conformation.position(elements[i], distPivot), 0);
+                        if (d < minD) minD = d;
+                    }
+                    break;
+                case Unit.Kind.Spheres:
+                    const radius = unit.coarseConformation.radius;
+                    for (let i = 0, _i = elements.length; i < _i; i++) {
+                        const d = minDistanceToPoint(b, conformation.position(elements[i], distPivot), radius[elements[i]]);
+                        if (d < minD) minD = d;
+                    }
+                    break;
+            }
+            if (d < minD) minD = d;
+        }
+        return minD;
+    }
 }
 
 export default Structure