diff --git a/src/mol-data/generic/unique-array.ts b/src/mol-data/generic/unique-array.ts
index e19d77fa44b42d6f3f9688caae21a7889c5194d7..8e02d5abe4ce37432f1b8ea5ce686aab496ab19d 100644
--- a/src/mol-data/generic/unique-array.ts
+++ b/src/mol-data/generic/unique-array.ts
@@ -20,6 +20,10 @@ namespace UniqueArray {
array[array.length] = value;
return true;
}
+
+ export function has<K, T>({ keys }: UniqueArray<K, T>, key: K) {
+ return keys.has(key);
+ }
}
export { UniqueArray }
\ No newline at end of file
diff --git a/src/mol-model/structure/query/context.ts b/src/mol-model/structure/query/context.ts
index e1621e676fe4c99082cc59ca5dd7b7af87ad131d..2c037526aabb0398c2523cdd0c6d5d5eb89f17da 100644
--- a/src/mol-model/structure/query/context.ts
+++ b/src/mol-model/structure/query/context.ts
@@ -9,6 +9,7 @@ import { now } from '../../../mol-util/now';
import { ElementIndex } from '../model';
import { LinkType } from '../model/types';
import { StructureSelection } from './selection';
+import { defaultLinkTest } from './queries/internal';
export interface QueryContextView {
readonly element: StructureElement.Location;
@@ -119,6 +120,35 @@ export class QueryContextLinkInfo<U extends Unit = Unit> {
type: LinkType = LinkType.Flag.None;
order: number = 0;
+ private testFn: QueryPredicate = defaultLinkTest;
+
+ setTestFn(fn?: QueryPredicate) {
+ this.testFn = fn || defaultLinkTest;
+ }
+
+ test(ctx: QueryContext, trySwap: boolean) {
+ if (this.testFn(ctx)) return true;
+ if (trySwap) {
+ this.swap();
+ return this.testFn(ctx);
+ }
+ return false;
+ }
+
+ private swap() {
+ const idxA = this.aIndex;
+ this.aIndex = this.bIndex;
+ this.bIndex = idxA;
+
+ const unitA = this.a.unit;
+ this.a.unit = this.b.unit;
+ this.b.unit = unitA;
+
+ const eA = this.a.element;
+ this.a.element = this.b.element;
+ this.b.element = eA;
+ }
+
get length() {
return StructureElement.Location.distance(this.a, this. b);
}
diff --git a/src/mol-model/structure/query/queries/filters.ts b/src/mol-model/structure/query/queries/filters.ts
index 2e9880aae1cee5518ee9dd6cfd0b8d8a5fe21698..08c9cde65a768cbd7ae0acd2b21ac5f00f15acd4 100644
--- a/src/mol-model/structure/query/queries/filters.ts
+++ b/src/mol-model/structure/query/queries/filters.ts
@@ -15,7 +15,6 @@ import { checkStructureMaxRadiusDistance, checkStructureMinMaxDistance } from '.
import Structure from '../../structure/structure';
import StructureElement from '../../structure/element';
import { SortedArray } from '../../../../mol-data/int';
-import { defaultLinkTest } from './internal';
export function pick(query: StructureQuery, pred: QueryPredicate): StructureQuery {
return ctx => {
@@ -231,13 +230,11 @@ function withinMinMaxRadius({ queryCtx, selection, target, minRadius, maxRadius,
interface IsConnectedToCtx {
queryCtx: QueryContext,
input: Structure,
- target: Structure,
- linkTest: QueryFn<boolean>,
- tElement: StructureElement.Location
+ target: Structure
}
function checkConnected(ctx: IsConnectedToCtx, structure: Structure) {
- const { queryCtx, input, target, linkTest, tElement } = ctx;
+ const { queryCtx, input, target } = ctx;
const atomicLink = queryCtx.atomicLink;
const interLinks = input.links;
@@ -258,36 +255,41 @@ function checkConnected(ctx: IsConnectedToCtx, structure: Structure) {
for (let i = 0 as StructureElement.UnitIndex, _i = srcElements.length; i < _i; i++) {
const inputIndex = SortedArray.indexOf(inputElements, srcElements[i]) as StructureElement.UnitIndex;
- atomicLink.aIndex = inputIndex;
- atomicLink.a.element = srcElements[i];
+ atomicLink.a.unit = inputUnit;
atomicLink.b.unit = inputUnit;
- tElement.unit = unit;
+ // tElement.unit = unit;
for (let l = offset[inputIndex], _l = offset[inputIndex + 1]; l < _l; l++) {
- tElement.element = inputElements[b[l]];
- if (!target.hasElement(tElement)) continue;
- atomicLink.bIndex = b[l] as StructureElement.UnitIndex;
+ // tElement.element = inputElements[b[l]];
atomicLink.b.element = inputUnit.elements[b[l]];
+ if (!target.hasElement(atomicLink.b)) continue;
+
+ atomicLink.aIndex = inputIndex;
+ atomicLink.a.element = srcElements[i];
+ atomicLink.bIndex = b[l] as StructureElement.UnitIndex;
atomicLink.type = flags[l];
atomicLink.order = order[l];
- if (linkTest(queryCtx)) return true;
+ if (atomicLink.test(queryCtx, true)) return true;
}
for (let li = 0; li < luCount; li++) {
const lu = linkedUnits[li];
- tElement.unit = lu.unitB;
- atomicLink.b.unit = lu.unitB;
const bElements = lu.unitB.elements;
const bonds = lu.getBonds(inputIndex);
for (let bi = 0, _bi = bonds.length; bi < _bi; bi++) {
const bond = bonds[bi];
- tElement.element = bElements[bond.indexB];
- if (!target.hasElement(tElement)) continue;
+ atomicLink.b.unit = lu.unitB;
+ atomicLink.b.element = bElements[bond.indexB];
+ if (!target.hasElement(atomicLink.b)) continue;
+
+ atomicLink.a.unit = inputUnit;
+ atomicLink.aIndex = inputIndex;
+ atomicLink.a.element = srcElements[i];
+
atomicLink.bIndex = bond.indexB;
- atomicLink.b.element = tElement.element;
atomicLink.type = bond.flag;
atomicLink.order = bond.order;
- if (linkTest(queryCtx)) return true;
+ if (atomicLink.test(queryCtx, true)) return true;
}
}
}
@@ -314,15 +316,17 @@ export function isConnectedTo({ query, target, disjunct, invert, linkTest }: IsC
const connCtx: IsConnectedToCtx = {
queryCtx: ctx,
input: ctx.inputStructure,
- target: StructureSelection.unionStructure(targetSel),
- linkTest: linkTest || defaultLinkTest,
- tElement: StructureElement.Location.create()
+ target: StructureSelection.unionStructure(targetSel)
}
const ret = StructureSelection.LinearBuilder(ctx.inputStructure);
ctx.pushCurrentLink();
+ ctx.atomicLink.setTestFn(linkTest);
+
StructureSelection.forEach(selection, (s, sI) => {
- if (checkConnected(connCtx, s)) ret.add(s);
+ if (checkConnected(connCtx, s)) {
+ ret.add(s);
+ }
if (sI % 5 === 0) ctx.throwIfTimedOut();
})
ctx.popCurrentLink();
diff --git a/src/mol-model/structure/query/queries/generators.ts b/src/mol-model/structure/query/queries/generators.ts
index 830fd516d266f9529bc1593849f59c3ea5893d48..74fe040c9bcee35d9e72da97de9b1bad1094345b 100644
--- a/src/mol-model/structure/query/queries/generators.ts
+++ b/src/mol-model/structure/query/queries/generators.ts
@@ -277,7 +277,6 @@ export function querySelection(selection: StructureQuery, query: StructureQuery,
}
export function linkedAtomicPairs(linkTest?: QueryPredicate): StructureQuery {
- const test = linkTest || ((ctx: any) => true);
return function query_linkedAtomicPairs(ctx) {
const structure = ctx.inputStructure;
@@ -287,6 +286,7 @@ export function linkedAtomicPairs(linkTest?: QueryPredicate): StructureQuery {
ctx.pushCurrentLink();
const atomicLink = ctx.atomicLink;
+ atomicLink.setTestFn(linkTest);
// Process intra unit links
for (const unit of structure.units) {
@@ -305,7 +305,8 @@ export function linkedAtomicPairs(linkTest?: QueryPredicate): StructureQuery {
atomicLink.b.element = unit.elements[intraLinkB[lI]];
atomicLink.type = flags[lI];
atomicLink.order = order[lI];
- if (test(ctx)) {
+ // No need to "swap test" because each bond direction will be visited eventually.
+ if (atomicLink.test(ctx, false)) {
const b = structure.subsetBuilder(false);
b.beginUnit(unit.id);
b.addElement(atomicLink.a.element);
@@ -328,7 +329,8 @@ export function linkedAtomicPairs(linkTest?: QueryPredicate): StructureQuery {
atomicLink.order = bond.order;
atomicLink.type = bond.flag;
- if (test(ctx)) {
+ // No need to "swap test" because each bond direction will be visited eventually.
+ if (atomicLink.test(ctx, false)) {
const b = structure.subsetBuilder(false);
b.addToUnit(atomicLink.a.unit.id, atomicLink.a.element);
b.addToUnit(atomicLink.b.unit.id, atomicLink.b.element);
diff --git a/src/mol-model/structure/query/queries/modifiers.ts b/src/mol-model/structure/query/queries/modifiers.ts
index 4f564a884df4b74e936ceda06ba85159aee9415c..553a92ee98f0e68ad8a6ac184c85c983e0d6d794 100644
--- a/src/mol-model/structure/query/queries/modifiers.ts
+++ b/src/mol-model/structure/query/queries/modifiers.ts
@@ -15,7 +15,6 @@ import { structureIntersect, structureSubtract } from '../utils/structure-set';
import { UniqueArray } from '../../../../mol-data/generic';
import { StructureSubsetBuilder } from '../../structure/util/subset-builder';
import StructureElement from '../../structure/element';
-import { defaultLinkTest } from './internal';
function getWholeResidues(ctx: QueryContext, source: Structure, structure: Structure) {
const builder = source.subsetBuilder(true);
@@ -305,16 +304,16 @@ export interface IncludeConnectedParams {
}
export function includeConnected({ query, layerCount, wholeResidues, linkTest }: IncludeConnectedParams): StructureQuery {
- const bt = linkTest || defaultLinkTest;
const lc = Math.max(layerCount, 0);
return function query_includeConnected(ctx) {
const builder = StructureSelection.UniqueBuilder(ctx.inputStructure);
const src = query(ctx);
ctx.pushCurrentLink();
+ ctx.atomicLink.setTestFn(linkTest);
StructureSelection.forEach(src, (s, sI) => {
let incl = s;
for (let i = 0; i < lc; i++) {
- incl = includeConnectedStep(ctx, bt, wholeResidues, incl);
+ incl = includeConnectedStep(ctx, wholeResidues, incl);
}
builder.add(incl);
if (sI % 10 === 0) ctx.throwIfTimedOut();
@@ -325,24 +324,24 @@ export function includeConnected({ query, layerCount, wholeResidues, linkTest }:
}
}
-function includeConnectedStep(ctx: QueryContext, linkTest: QueryFn<boolean>, wholeResidues: boolean, structure: Structure) {
- const expanded = expandConnected(ctx, structure, linkTest);
+function includeConnectedStep(ctx: QueryContext, wholeResidues: boolean, structure: Structure) {
+ const expanded = expandConnected(ctx, structure);
if (wholeResidues) return getWholeResidues(ctx, ctx.inputStructure, expanded);
return expanded;
}
-function expandConnected(ctx: QueryContext, structure: Structure, linkTest: QueryFn<boolean>) {
+function expandConnected(ctx: QueryContext, structure: Structure) {
const inputStructure = ctx.inputStructure;
const interLinks = inputStructure.links;
const builder = new StructureUniqueSubsetBuilder(inputStructure);
// Note: each link is visited twice so that link.atom-a and link.atom-b both get the "swapped values"
+ const visitedSourceUnits = new Set<number>();
const atomicLink = ctx.atomicLink;
// Process intra unit links
for (const unit of structure.units) {
-
if (unit.kind !== Unit.Kind.Atomic) {
// add the whole unit
builder.beginUnit(unit.id);
@@ -353,53 +352,71 @@ function expandConnected(ctx: QueryContext, structure: Structure, linkTest: Quer
continue;
}
- const inputUnit = inputStructure.unitMap.get(unit.id) as Unit.Atomic;
- const { offset: intraLinkOffset, b: intraLinkB, edgeProps: { flags, order } } = inputUnit.links;
+ const inputUnitA = inputStructure.unitMap.get(unit.id) as Unit.Atomic;
+ const { offset: intraLinkOffset, b: intraLinkB, edgeProps: { flags, order } } = inputUnitA.links;
// Process intra unit links
- atomicLink.a.unit = inputUnit;
- atomicLink.b.unit = inputUnit;
+ atomicLink.a.unit = inputUnitA;
+ atomicLink.b.unit = inputUnitA;
for (let i = 0, _i = unit.elements.length; i < _i; i++) {
// add the current element
builder.addToUnit(unit.id, unit.elements[i]);
- const srcIndex = SortedArray.indexOf(inputUnit.elements, unit.elements[i]);
- atomicLink.aIndex = srcIndex as StructureElement.UnitIndex;
- atomicLink.a.element = unit.elements[i];
+ const aIndex = SortedArray.indexOf(inputUnitA.elements, unit.elements[i]) as StructureElement.UnitIndex;
// check intra unit links
- for (let lI = intraLinkOffset[srcIndex], _lI = intraLinkOffset[srcIndex + 1]; lI < _lI; lI++) {
- atomicLink.bIndex = intraLinkB[lI] as StructureElement.UnitIndex;
- atomicLink.b.element = inputUnit.elements[intraLinkB[lI]];
+ for (let lI = intraLinkOffset[aIndex], _lI = intraLinkOffset[aIndex + 1]; lI < _lI; lI++) {
+ const bIndex = intraLinkB[lI] as StructureElement.UnitIndex;
+ const bElement = inputUnitA.elements[bIndex];
+
+ // Check if the element is already present:
+ if (SortedArray.has(unit.elements, bElement) || builder.has(unit.id, bElement)) continue;
+
+ atomicLink.aIndex = aIndex;
+ atomicLink.a.element = unit.elements[i];
+ atomicLink.bIndex = bIndex;
+ atomicLink.b.element = bElement;
atomicLink.type = flags[lI];
atomicLink.order = order[lI];
- if (linkTest(ctx)) {
- builder.addToUnit(unit.id, inputUnit.elements[intraLinkB[lI]]);
+
+ if (atomicLink.test(ctx, true)) {
+ builder.addToUnit(unit.id, bElement);
}
}
}
// Process inter unit links
- for (const linkedUnit of interLinks.getLinkedUnits(inputUnit)) {
- atomicLink.b.unit = linkedUnit.unitB;
+ for (const linkedUnit of interLinks.getLinkedUnits(inputUnitA)) {
+ if (visitedSourceUnits.has(linkedUnit.unitB.id)) continue;
+ const currentUnitB = structure.unitMap.get(linkedUnit.unitB.id);
+
for (const aI of linkedUnit.linkedElementIndices) {
// check if the element is in the expanded structure
- if (!SortedArray.has(unit.elements, inputUnit.elements[aI])) continue;
+ if (!SortedArray.has(unit.elements, inputUnitA.elements[aI])) continue;
- atomicLink.aIndex = aI;
- atomicLink.a.element = inputUnit.elements[aI];
for (const bond of linkedUnit.getBonds(aI)) {
+ const bElement = linkedUnit.unitB.elements[bond.indexB];
+
+ // Check if the element is already present:
+ if ((currentUnitB && SortedArray.has(currentUnitB.elements, bElement)) || builder.has(linkedUnit.unitB.id, bElement)) continue;
+
+ atomicLink.a.unit = inputUnitA;
+ atomicLink.aIndex = aI;
+ atomicLink.a.element = inputUnitA.elements[aI];
+ atomicLink.b.unit = linkedUnit.unitB;
atomicLink.bIndex = bond.indexB;
- // TODO: optimize the lookup?
- atomicLink.b.element = linkedUnit.unitB.elements[bond.indexB];
+ atomicLink.b.element = bElement;
atomicLink.type = bond.flag;
atomicLink.order = bond.order;
- if (linkTest(ctx)) {
- builder.addToUnit(linkedUnit.unitB.id, linkedUnit.unitB.elements[bond.indexB]);
+
+ if (atomicLink.test(ctx, true)) {
+ builder.addToUnit(linkedUnit.unitB.id, bElement);
}
}
}
}
+
+ visitedSourceUnits.add(inputUnitA.id);
}
return builder.getStructure();
diff --git a/src/mol-model/structure/structure/util/unique-subset-builder.ts b/src/mol-model/structure/structure/util/unique-subset-builder.ts
index 742afb255a41016fdb2f79d3e8ad14cd923b6459..21ebf6788f3e3ab062685525f5ae9b09620bd1da 100644
--- a/src/mol-model/structure/structure/util/unique-subset-builder.ts
+++ b/src/mol-model/structure/structure/util/unique-subset-builder.ts
@@ -34,6 +34,12 @@ export class StructureUniqueSubsetBuilder {
}
}
+ has(parentId: number, e: number) {
+ const unit = this.unitMap.get(parentId);
+ if (!unit) return false;
+ return UniqueArray.has(unit, e);
+ }
+
beginUnit(parentId: number) {
this.parentId = parentId;
if (this.unitMap.has(parentId)) {
diff --git a/src/mol-script/language/symbol-table/structure-query.ts b/src/mol-script/language/symbol-table/structure-query.ts
index 234ef399caaf51e59ae145a271190c1bf8339790..3d63fd44eb27c226afb3f7c5028456644e8cddbc 100644
--- a/src/mol-script/language/symbol-table/structure-query.ts
+++ b/src/mol-script/language/symbol-table/structure-query.ts
@@ -94,7 +94,7 @@ const generator = {
}), Types.ElementSelectionQuery, 'Return all atoms for which the tests are satisfied, grouped into sets.'),
linkedAtomicPairs: symbol(Arguments.Dictionary({
- 0: Argument(Type.Bool, { isOptional: true, defaultValue: true, description: 'Test each link with this predicate. Each link is visited twice with swapped atom order.' }),
+ 0: Argument(Type.Bool, { isOptional: true, defaultValue: 'true for covalent links' as any, description: 'Test each link with this predicate. Each link is visited twice with swapped atom order.' }),
// TODO: shoud we support this or just use queryEach to get similar behavior
// 'group-by': Argument(Type.Any, { isOptional: true, defaultValue: ``, description: 'Group the links using the privided value' }),
}), Types.ElementSelectionQuery, 'Return all pairs of atoms for which the test is satisfied.'),