From ca0754572783d9aed1d5c178f16a03e9ce656a71 Mon Sep 17 00:00:00 2001
From: David Sehnal <david.sehnal@gmail.com>
Date: Tue, 17 Jul 2018 21:15:39 +0200
Subject: [PATCH] wip, mol-model/queries
---
src/mol-model/structure/query/context.ts | 22 ++--
.../structure/query/queries/combinators.ts | 17 +++
.../structure/query/queries/filters.ts | 104 ++++++++++++++----
src/mol-model/structure/query/selection.ts | 18 +++
4 files changed, 125 insertions(+), 36 deletions(-)
diff --git a/src/mol-model/structure/query/context.ts b/src/mol-model/structure/query/context.ts
index cefbc8346..846e84377 100644
--- a/src/mol-model/structure/query/context.ts
+++ b/src/mol-model/structure/query/context.ts
@@ -12,34 +12,32 @@ export interface QueryContextView {
}
export class QueryContext implements QueryContextView {
- private currentStack: StructureElement[] = [];
- private currentLocked = false;
+ private currentElementStack: StructureElement[] = [];
+ private currentStructureStack: Structure[] = [];
readonly inputStructure: Structure;
/** Current element */
readonly element: StructureElement = StructureElement.create();
- readonly currentStructure: Structure = void 0 as any;
+ currentStructure: Structure = void 0 as any;
pushCurrentElement(): StructureElement {
- this.currentStack[this.currentStack.length] = this.element;
+ this.currentElementStack[this.currentElementStack.length] = this.element;
(this.element as StructureElement) = StructureElement.create();
return this.element;
}
popCurrentElement() {
- (this.element as StructureElement) = this.currentStack.pop()!;
+ (this.element as StructureElement) = this.currentElementStack.pop()!;
}
- lockCurrentStructure(structure: Structure) {
- if (this.currentLocked) throw new Error('Current structure already locked.');
- this.currentLocked = true;
- (this.currentStructure as Structure) = structure;
+ pushCurrentStructure() {
+ if (this.currentStructure) this.currentStructureStack.push(this.currentStructure);
}
- unlockCurrentStructure() {
- this.currentLocked = false;
- (this.currentStructure as any) = void 0;
+ popCurrentStructure() {
+ if (this.currentStructureStack.length) (this.currentStructure as Structure) = this.currentStructureStack.pop()!;
+ else (this.currentStructure as Structure) = void 0 as any;
}
constructor(structure: Structure) {
diff --git a/src/mol-model/structure/query/queries/combinators.ts b/src/mol-model/structure/query/queries/combinators.ts
index de7200ac1..7bcf70406 100644
--- a/src/mol-model/structure/query/queries/combinators.ts
+++ b/src/mol-model/structure/query/queries/combinators.ts
@@ -3,3 +3,20 @@
*
* @author David Sehnal <david.sehnal@gmail.com>
*/
+
+import { StructureQuery } from '../query';
+import { StructureSelection } from '../selection';
+
+export function merge(queries: ArrayLike<StructureQuery>): StructureQuery {
+ return ctx => {
+ const ret = StructureSelection.UniqueBuilder(ctx.inputStructure);
+ for (let i = 0; i < queries.length; i++) {
+ StructureSelection.forEach(queries[i](ctx), s => {
+ ret.add(s);
+ });
+ }
+ return ret.getSelection();
+ }
+}
+
+// TODO: intersect, distanceCluster
\ No newline at end of file
diff --git a/src/mol-model/structure/query/queries/filters.ts b/src/mol-model/structure/query/queries/filters.ts
index 3bf99a41f..6d7363123 100644
--- a/src/mol-model/structure/query/queries/filters.ts
+++ b/src/mol-model/structure/query/queries/filters.ts
@@ -4,37 +4,93 @@
* @author David Sehnal <david.sehnal@gmail.com>
*/
-import { SortedArray } from 'mol-data/int';
-import { Structure } from '../../structure';
-import { QueryPredicate } from '../context';
+import { isSuperset } from 'mol-util/set';
+import { Unit } from '../../structure';
+import { QueryContext, QueryFn, QueryPredicate } from '../context';
import { StructureQuery } from '../query';
import { StructureSelection } from '../selection';
+import { structureAreIntersecting } from '../utils/structure';
export function pick(query: StructureQuery, pred: QueryPredicate): StructureQuery {
return ctx => {
const sel = query(ctx);
+ const ret = StructureSelection.LinearBuilder(ctx.inputStructure);
+ ctx.pushCurrentElement();
+ StructureSelection.forEach(sel, s => {
+ ctx.currentStructure = s;
+ if (pred(ctx)) ret.add(s);
+ });
+ ctx.popCurrentStructure();
+ return ret.getSelection();
+ };
+}
- if (StructureSelection.isSingleton(sel)) {
- const ret = StructureSelection.LinearBuilder(ctx.inputStructure);
- for (const unit of ctx.inputStructure.units) {
- const { elements } = unit;
- for (let i = 0, _i = elements.length; i < _i; i++) {
- // TODO: optimize this somehow???
- const s = Structure.create([unit.getChild(SortedArray.ofSingleton(elements[i]))]);
- ctx.lockCurrentStructure(s);
- if (pred(ctx)) ret.add(s);
- ctx.unlockCurrentStructure();
- }
- }
- return ret.getSelection();
- } else {
- const ret = StructureSelection.LinearBuilder(ctx.inputStructure);
- for (const s of sel.structures) {
- ctx.lockCurrentStructure(s);
- if (pred(ctx)) ret.add(s);
- ctx.unlockCurrentStructure();
- }
- return ret.getSelection();
+export interface UnitTypeProperties { atomic?: QueryFn, coarse?: QueryFn }
+
+export function getCurrentStructureProperties(ctx: QueryContext, props: UnitTypeProperties, set: Set<any>) {
+ const { units } = ctx.currentStructure;
+ const l = ctx.pushCurrentElement();
+
+ for (const unit of units) {
+ l.unit = unit;
+ const elements = unit.elements;
+
+ let fn;
+ if (Unit.isAtomic(unit)) fn = props.atomic;
+ else fn = props.coarse;
+ if (!fn) continue;
+
+ for (let j = 0, _j = elements.length; j < _j; j++) {
+ l.element = elements[j];
+ set.add(fn(ctx));
}
+ }
+ ctx.popCurrentElement();
+ return set;
+}
+
+function getSelectionProperties(ctx: QueryContext, query: StructureQuery, props: UnitTypeProperties) {
+ const set = new Set();
+
+ const sel = query(ctx);
+ ctx.pushCurrentElement();
+ StructureSelection.forEach(sel, s => {
+ ctx.currentStructure = s;
+ getCurrentStructureProperties(ctx, props, set);
+ });
+ ctx.popCurrentElement();
+ return set;
+}
+
+export function withSameAtomProperties(query: StructureQuery, propertySource: StructureQuery, props: UnitTypeProperties): StructureQuery {
+ return ctx => {
+ const sel = query(ctx);
+ const propSet = getSelectionProperties(ctx, propertySource, props);
+
+ const ret = StructureSelection.LinearBuilder(ctx.inputStructure);
+ ctx.pushCurrentStructure();
+ StructureSelection.forEach(sel, s => {
+ ctx.currentStructure = s;
+ const currentProps = getCurrentStructureProperties(ctx, props, new Set());
+ if (isSuperset(currentProps, propSet)) {
+ ret.add(s);
+ }
+ });
+ ctx.popCurrentStructure();
+ return ret.getSelection();
};
}
+
+export function areIntersectedBy(query: StructureQuery, by: StructureQuery): StructureQuery {
+ return ctx => {
+ const mask = StructureSelection.unionStructure(by(ctx));
+ const ret = StructureSelection.LinearBuilder(ctx.inputStructure);
+
+ StructureSelection.forEach(query(ctx), s => {
+ if (structureAreIntersecting(mask, s)) ret.add(s);
+ });
+ return ret.getSelection();
+ };
+}
+
+// TODO: within, isConnectedTo
\ No newline at end of file
diff --git a/src/mol-model/structure/query/selection.ts b/src/mol-model/structure/query/selection.ts
index d2d55553e..a92cf93aa 100644
--- a/src/mol-model/structure/query/selection.ts
+++ b/src/mol-model/structure/query/selection.ts
@@ -101,6 +101,24 @@ namespace StructureSelection {
export function LinearBuilder(structure: Structure): Builder { return new LinearBuilderImpl(structure); }
export function UniqueBuilder(structure: Structure): Builder { return new HashBuilderImpl(structure); }
+ export function forEach(sel: StructureSelection, fn: (s: Structure, i: number) => void) {
+ let idx = 0;
+ if (StructureSelection.isSingleton(sel)) {
+ for (const unit of sel.structure.units) {
+ const { elements } = unit;
+ for (let i = 0, _i = elements.length; i < _i; i++) {
+ // TODO: optimize this somehow???
+ const s = Structure.create([unit.getChild(SortedArray.ofSingleton(elements[i]))]);
+ fn(s, idx++);
+ }
+ }
+ } else {
+ for (const s of sel.structures) {
+ fn(s, idx++);
+ }
+ }
+ }
+
// TODO: spatial lookup
}
--
GitLab