diff --git a/src/mol-data/generic/unique-array.ts b/src/mol-data/generic/unique-array.ts index dad6693959d1d47d7bf66a50f2456bbc0b341ee8..e19d77fa44b42d6f3f9688caae21a7889c5194d7 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 60c47c06796abf4e4b9f62527e8a00ee30f9bdcd..85989cdda5145bb3807ec47a82dbe689480897c9 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 c249fa266d1d6d5c22a37cbd597e466e8b452250..b2540584ce7eddcdfc2f416fe929ddfa5880c1ac 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 650c0ae04e23685be2d770b73672df57d7755b74..bd258621e8cc202239fb0d5a2c92b801c7047a06 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)) {