From 486e8605750d56a3f4915044ee7285ae7674d77a Mon Sep 17 00:00:00 2001
From: David Sehnal <david.sehnal@gmail.com>
Date: Thu, 23 Aug 2018 18:00:49 +0200
Subject: [PATCH] mol-model/query expandProperty, some "query modifer" tweaks
---
src/mol-data/generic/unique-array.ts | 4 +-
.../structure/query/queries/filters.ts | 4 ++
.../structure/query/queries/modifiers.ts | 61 ++++++++++++++++++-
src/mol-model/structure/query/selection.ts | 1 +
4 files changed, 66 insertions(+), 4 deletions(-)
diff --git a/src/mol-data/generic/unique-array.ts b/src/mol-data/generic/unique-array.ts
index dad669395..e19d77fa4 100644
--- a/src/mol-data/generic/unique-array.ts
+++ b/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: [] };
}
diff --git a/src/mol-model/structure/query/queries/filters.ts b/src/mol-model/structure/query/queries/filters.ts
index 60c47c067..85989cdda 100644
--- a/src/mol-model/structure/query/queries/filters.ts
+++ b/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();
}
diff --git a/src/mol-model/structure/query/queries/modifiers.ts b/src/mol-model/structure/query/queries/modifiers.ts
index c249fa266..b2540584c 100644
--- a/src/mol-model/structure/query/queries/modifiers.ts
+++ b/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
\ No newline at end of file
+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"
\ 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 650c0ae04..bd258621e 100644
--- a/src/mol-model/structure/query/selection.ts
+++ b/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)) {
--
GitLab