UNPKG

molstar

Version:

A comprehensive macromolecular library.

294 lines (293 loc) 11.4 kB
/** * Copyright (c) 2018-2023 mol* contributors, licensed under MIT, See LICENSE file for more info. * * @author David Sehnal <david.sehnal@gmail.com> * @author Alexander Rose <alexander.rose@weirdbyte.de> */ import { SetUtils } from '../../../../mol-util/set'; import { Unit } from '../../structure'; import { StructureSelection } from '../selection'; import { structureAreIntersecting } from '../utils/structure-set'; import { Vec3 } from '../../../../mol-math/linear-algebra'; import { checkStructureMaxRadiusDistance, checkStructureMinMaxDistance } from '../utils/structure-distance'; import { Structure } from '../../structure/structure'; import { SortedArray } from '../../../../mol-data/int'; export function pick(query, pred) { return ctx => { const sel = query(ctx); const ret = StructureSelection.LinearBuilder(ctx.inputStructure); ctx.pushCurrentElement(); StructureSelection.forEach(sel, (s, i) => { ctx.currentStructure = s; if (pred(ctx)) ret.add(s); if (i % 100) ctx.throwIfTimedOut(); }); ctx.popCurrentStructure(); return ret.getSelection(); }; } export function first(query) { return ctx => { const sel = query(ctx); const ret = StructureSelection.LinearBuilder(ctx.inputStructure); if (sel.kind === 'singletons') { if (sel.structure.elementCount > 0) { const u = sel.structure.units[0]; const s = Structure.create([u.getChild(SortedArray.ofSingleton(u.elements[0]))], { parent: ctx.inputStructure }); ret.add(s); } } else { if (sel.structures.length > 0) { ret.add(sel.structures[0]); } } return ret.getSelection(); }; } export function getCurrentStructureProperties(ctx, props, set) { const { units } = ctx.currentStructure; const l = ctx.pushCurrentElement(); l.structure = ctx.currentStructure; for (const unit of units) { l.unit = unit; const elements = unit.elements; const fn = props; // 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.throwIfTimedOut(); } ctx.popCurrentElement(); return set; } function getSelectionProperties(ctx, query, props) { const set = new Set(); const sel = query(ctx); ctx.pushCurrentElement(); StructureSelection.forEach(sel, (s, i) => { ctx.currentStructure = s; getCurrentStructureProperties(ctx, props, set); if (i % 10) ctx.throwIfTimedOut(); }); ctx.popCurrentElement(); return set; } export function withSameAtomProperties(query, propertySource, props) { return ctx => { const sel = query(ctx); const propSet = getSelectionProperties(ctx, propertySource, props); const ret = StructureSelection.LinearBuilder(ctx.inputStructure); ctx.pushCurrentStructure(); StructureSelection.forEach(sel, (s, i) => { ctx.currentStructure = s; const currentProps = getCurrentStructureProperties(ctx, props, new Set()); if (SetUtils.isSuperset(propSet, currentProps)) { ret.add(s); } if (i % 10) ctx.throwIfTimedOut(); }); ctx.popCurrentStructure(); return ret.getSelection(); }; } export function areIntersectedBy(query, by) { return ctx => { const mask = StructureSelection.unionStructure(by(ctx)); const ret = StructureSelection.LinearBuilder(ctx.inputStructure); StructureSelection.forEach(query(ctx), (s, i) => { if (structureAreIntersecting(mask, s)) ret.add(s); if (i % 10) ctx.throwIfTimedOut(); }); return ret.getSelection(); }; } export function within(params) { return queryCtx => { const ctx = { queryCtx, selection: params.query(queryCtx), target: params.target(queryCtx), maxRadius: params.maxRadius, minRadius: params.minRadius ? Math.max(0, params.minRadius) : 0, elementRadius: params.elementRadius, invert: !!params.invert, }; if (ctx.minRadius === 0 && typeof params.minRadius === 'undefined') { return withinMaxRadiusLookup(ctx); } else if (ctx.minRadius === 0) { return withinMaxRadius(ctx); } else { return withinMinMaxRadius(ctx); } }; } function withinMaxRadiusLookup({ queryCtx, selection, target, maxRadius, invert }) { const targetLookup = StructureSelection.unionStructure(target).lookup3d; const ret = StructureSelection.LinearBuilder(queryCtx.inputStructure); const pos = Vec3.zero(); StructureSelection.forEach(selection, (s, sI) => { const { units } = s; let withinRadius = false; for (let i = 0, _i = units.length; i < _i; i++) { const unit = units[i]; const { elements, conformation: c } = unit; for (let i = 0, _i = elements.length; i < _i; i++) { const e = elements[i]; c.position(e, pos); if (targetLookup.check(pos[0], pos[1], pos[2], maxRadius + c.r(e))) { withinRadius = true; break; } } if (withinRadius) break; } if (invert) withinRadius = !withinRadius; if (withinRadius) ret.add(s); if (sI % 10 === 0) queryCtx.throwIfTimedOut(); }); return ret.getSelection(); } function withinMaxRadius({ queryCtx, selection, target, maxRadius, invert, elementRadius }) { 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(); } function withinMinMaxRadius({ queryCtx, selection, target, minRadius, maxRadius, invert, elementRadius }) { 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(); } function checkConnected(ctx, structure) { const { queryCtx, input, target, disjunct } = ctx; const atomicBond = queryCtx.atomicBond; const interBonds = input.interUnitBonds; atomicBond.setStructure(input); for (const unit of structure.units) { if (!Unit.isAtomic(unit)) continue; const inputUnit = input.unitMap.get(unit.id); const { offset, b, edgeProps: { flags, order, key } } = inputUnit.bonds; const bondedUnits = interBonds.getConnectedUnits(unit.id); const buCount = bondedUnits.length; const srcElements = unit.elements; const inputElements = inputUnit.elements; for (let i = 0, _i = srcElements.length; i < _i; i++) { const inputIndex = SortedArray.indexOf(inputElements, srcElements[i]); atomicBond.a.unit = inputUnit; atomicBond.b.unit = inputUnit; // tElement.unit = unit; for (let l = offset[inputIndex], _l = offset[inputIndex + 1]; l < _l; l++) { // tElement.element = inputElements[b[l]]; atomicBond.b.element = inputUnit.elements[b[l]]; if (disjunct && SortedArray.has(unit.elements, atomicBond.b.element)) continue; if (!target.hasElement(atomicBond.b)) continue; atomicBond.aIndex = inputIndex; atomicBond.a.element = srcElements[i]; atomicBond.bIndex = b[l]; atomicBond.type = flags[l]; atomicBond.order = order[l]; atomicBond.key = key[l]; if (atomicBond.test(queryCtx, true)) return true; } for (let li = 0; li < buCount; li++) { const lu = bondedUnits[li]; const bUnit = input.unitMap.get(lu.unitB); const bElements = bUnit.elements; const bonds = lu.getEdges(inputIndex); for (let bi = 0, _bi = bonds.length; bi < _bi; bi++) { const bond = bonds[bi]; atomicBond.b.unit = bUnit; atomicBond.b.element = bElements[bond.indexB]; if (!target.hasElement(atomicBond.b)) continue; if (disjunct && structure.hasElement(atomicBond.b)) continue; atomicBond.a.unit = inputUnit; atomicBond.aIndex = inputIndex; atomicBond.a.element = srcElements[i]; atomicBond.bIndex = bond.indexB; atomicBond.type = bond.props.flag; atomicBond.order = bond.props.order; atomicBond.key = bond.props.key; if (atomicBond.test(queryCtx, true)) return true; } } } } return false; } export function isConnectedTo({ query, target, disjunct, invert, bondTest }) { return ctx => { const targetSel = target(ctx); if (StructureSelection.isEmpty(targetSel)) return targetSel; const selection = query(ctx); if (StructureSelection.isEmpty(selection)) return selection; const connCtx = { queryCtx: ctx, input: ctx.inputStructure, disjunct, target: StructureSelection.unionStructure(targetSel) }; const ret = StructureSelection.LinearBuilder(ctx.inputStructure); ctx.pushCurrentBond(); ctx.atomicBond.setTestFn(bondTest); StructureSelection.forEach(selection, (s, sI) => { if (checkConnected(connCtx, s)) { ret.add(s); } else if (invert) { ret.add(s); } if (sI % 5 === 0) ctx.throwIfTimedOut(); }); ctx.popCurrentBond(); return ret.getSelection(); }; }