diff --git a/src/mol-model/structure/query/queries/modifiers.ts b/src/mol-model/structure/query/queries/modifiers.ts index b2540584ce7eddcdfc2f416fe929ddfa5880c1ac..2a6641adb0e5bb5375dc9471a0f59dcb1a9d2980 100644 --- a/src/mol-model/structure/query/queries/modifiers.ts +++ b/src/mol-model/structure/query/queries/modifiers.ts @@ -14,6 +14,7 @@ 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'; +import StructureElement from '../../structure/element'; function getWholeResidues(ctx: QueryContext, source: Structure, structure: Structure) { const builder = source.subsetBuilder(true); @@ -59,8 +60,7 @@ export function wholeResidues(query: StructureQuery): StructureQuery { export interface IncludeSurroundingsParams { radius: number, - // TODO - // atomRadius?: Element.Property<number>, + elementRadius?: QueryFn<number>, wholeResidues?: boolean } @@ -82,9 +82,88 @@ function getIncludeSurroundings(ctx: QueryContext, source: Structure, structure: return !!params.wholeResidues ? getWholeResidues(ctx, source, builder.getStructure()) : builder.getStructure(); } +interface IncludeSurroundingsParamsWithRadius extends IncludeSurroundingsParams { + elementRadius: QueryFn<number>, + elementRadiusClosure: StructureElement.Property<number>, + sourceMaxRadius: number +} + +function getIncludeSurroundingsWithRadius(ctx: QueryContext, source: Structure, structure: Structure, params: IncludeSurroundingsParamsWithRadius) { + const builder = new StructureUniqueSubsetBuilder(source); + const lookup = source.lookup3d; + const { elementRadius, elementRadiusClosure, sourceMaxRadius, radius } = params; + + ctx.pushCurrentElement(); + for (const unit of structure.units) { + ctx.element.unit = unit; + const { x, y, z } = unit.conformation; + const elements = unit.elements; + + for (let i = 0, _i = elements.length; i < _i; i++) { + const e = elements[i]; + ctx.element.element = e; + const eRadius = elementRadius(ctx); + lookup.findIntoBuilderWithRadius(x(e), y(e), z(e), eRadius, sourceMaxRadius, radius, elementRadiusClosure, builder); + } + + ctx.throwIfTimedOut(); + } + + ctx.popCurrentElement(); + return !!params.wholeResidues ? getWholeResidues(ctx, source, builder.getStructure()) : builder.getStructure(); +} + +function createElementRadiusFn(ctx: QueryContext, eRadius: QueryFn<number>): StructureElement.Property<number> { + return e => { + ctx.element.unit = e.unit; + ctx.element.element = e.element; + return eRadius(ctx); + } +} + +function findStructureRadius(ctx: QueryContext, eRadius: QueryFn<number>) { + let r = 0; + 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++) { + const e = elements[i]; + ctx.element.element = e; + const eR = eRadius(ctx); + if (eR > r) r = eR; + } + + } + ctx.throwIfTimedOut(); + return r; +} + export function includeSurroundings(query: StructureQuery, params: IncludeSurroundingsParams): StructureQuery { return ctx => { const inner = query(ctx); + + if (params.elementRadius) { + const prms: IncludeSurroundingsParamsWithRadius = { + ...params, + elementRadius: params.elementRadius, + elementRadiusClosure: createElementRadiusFn(ctx, params.elementRadius), + sourceMaxRadius: findStructureRadius(ctx, params.elementRadius) + }; + + if (StructureSelection.isSingleton(inner)) { + const surr = getIncludeSurroundingsWithRadius(ctx, ctx.inputStructure, inner.structure, prms); + const ret = StructureSelection.Singletons(ctx.inputStructure, surr); + return ret; + } else { + const builder = new UniqueStructuresBuilder(ctx.inputStructure); + for (const s of inner.structures) { + builder.add(getIncludeSurroundingsWithRadius(ctx, ctx.inputStructure, s, prms)); + } + return builder.getSelection(); + } + } + if (StructureSelection.isSingleton(inner)) { const surr = getIncludeSurroundings(ctx, ctx.inputStructure, inner.structure, params); const ret = StructureSelection.Singletons(ctx.inputStructure, surr); @@ -218,4 +297,4 @@ export function expandProperty(query: StructureQuery, property: QueryFn): Struct }; } -// TODO: unionBy (skip this one?), cluster, includeConnected, includeSurroundings with "radii" \ No newline at end of file +// TODO: unionBy (skip this one?), cluster, includeConnected \ No newline at end of file diff --git a/src/mol-model/structure/structure/util/lookup3d.ts b/src/mol-model/structure/structure/util/lookup3d.ts index 19187af2f868d1d5d1dac11abfbcafc2be2b57bc..2427ba51e2d0bb44a4200fe4b0384a63f8669664 100644 --- a/src/mol-model/structure/structure/util/lookup3d.ts +++ b/src/mol-model/structure/structure/util/lookup3d.ts @@ -10,6 +10,7 @@ import { Vec3 } from 'mol-math/linear-algebra'; import { computeStructureBoundary } from './boundary'; import { OrderedSet } from 'mol-data/int'; import { StructureUniqueSubsetBuilder } from './unique-subset-builder'; +import StructureElement from '../element'; export class StructureLookup3D { private unitLookup: Lookup3D; @@ -66,6 +67,39 @@ export class StructureLookup3D { } } + findIntoBuilderWithRadius(x: number, y: number, z: number, pivotR: number, maxRadius: number, radius: number, eRadius: StructureElement.Property<number>, builder: StructureUniqueSubsetBuilder) { + const { units } = this.structure; + const closeUnits = this.unitLookup.find(x, y, z, radius); + if (closeUnits.count === 0) return; + + const se = StructureElement.create(); + const queryRadius = pivotR + maxRadius + radius; + + for (let t = 0, _t = closeUnits.count; t < _t; t++) { + const unit = units[closeUnits.indices[t]]; + Vec3.set(this.pivot, x, y, z); + if (!unit.conformation.operator.isIdentity) { + Vec3.transformMat4(this.pivot, this.pivot, unit.conformation.operator.inverse); + } + const unitLookup = unit.lookup3d; + const groupResult = unitLookup.find(this.pivot[0], this.pivot[1], this.pivot[2], queryRadius); + if (groupResult.count === 0) continue; + + const elements = unit.elements; + se.unit = unit; + builder.beginUnit(unit.id); + for (let j = 0, _j = groupResult.count; j < _j; j++) { + se.element = elements[groupResult.indices[j]]; + const rr = eRadius(se); + if (Math.sqrt(groupResult.squaredDistances[j]) - pivotR - rr > radius) continue; + builder.addElement(elements[groupResult.indices[j]]); + } + builder.commitUnit(); + } + } + + + check(x: number, y: number, z: number, radius: number): boolean { const { units } = this.structure; const closeUnits = this.unitLookup.find(x, y, z, radius);