molstar
Version:
A comprehensive macromolecular library.
342 lines (341 loc) • 17 kB
JavaScript
/**
* Copyright (c) 2017-2022 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 { __assign } from "tslib";
import { UniqueArray } from '../../../../mol-data/generic';
import { Segmentation, SortedArray } from '../../../../mol-data/int';
import { StructureProperties as P, Unit } from '../../structure';
import { Structure } from '../../structure/structure';
import { StructureSelection } from '../selection';
import { LinearGroupingBuilder } from '../utils/builders';
import { structureSubtract } from '../utils/structure-set';
export var none = function (ctx) { return StructureSelection.Sequence(ctx.inputStructure, []); };
export var all = function (ctx) { return StructureSelection.Singletons(ctx.inputStructure, ctx.inputStructure); };
export function residues(params) { return atoms(__assign(__assign({}, params), { groupBy: function (ctx) { return P.residue.key(ctx.element); } })); }
export function chains(params) { return atoms(__assign(__assign({}, params), { groupBy: function (ctx) { return P.chain.key(ctx.element); } })); }
function _true(ctx) { return true; }
function _zero(ctx) { return 0; }
export function atoms(params) {
if (!params || (!params.atomTest && !params.residueTest && !params.chainTest && !params.entityTest && !params.unitTest && !params.groupBy))
return all;
if (!!params.atomTest && !params.residueTest && !params.chainTest && !params.entityTest && !params.unitTest && !params.groupBy)
return atomGroupsLinear(params.atomTest);
var normalized = {
unitTest: params.unitTest || _true,
entityTest: params.entityTest || _true,
chainTest: params.chainTest || _true,
residueTest: params.residueTest || _true,
atomTest: params.atomTest || _true,
groupBy: params.groupBy || _zero,
};
if (!params.groupBy)
return atomGroupsSegmented(normalized);
return atomGroupsGrouped(normalized);
}
function atomGroupsLinear(atomTest) {
return function query_atomGroupsLinear(ctx) {
var inputStructure = ctx.inputStructure;
var units = inputStructure.units;
var l = ctx.pushCurrentElement();
var builder = inputStructure.subsetBuilder(true);
l.structure = inputStructure;
for (var _a = 0, units_1 = units; _a < units_1.length; _a++) {
var unit = units_1[_a];
l.unit = unit;
var elements = unit.elements;
builder.beginUnit(unit.id);
for (var j = 0, _j = elements.length; j < _j; j++) {
l.element = elements[j];
if (atomTest(ctx))
builder.addElement(l.element);
}
builder.commitUnit();
ctx.throwIfTimedOut();
}
ctx.popCurrentElement();
return StructureSelection.Singletons(inputStructure, builder.getStructure());
};
}
function atomGroupsSegmented(_a) {
var unitTest = _a.unitTest, entityTest = _a.entityTest, chainTest = _a.chainTest, residueTest = _a.residueTest, atomTest = _a.atomTest;
return function query_atomGroupsSegmented(ctx) {
var inputStructure = ctx.inputStructure;
var units = inputStructure.units;
var l = ctx.pushCurrentElement();
var builder = inputStructure.subsetBuilder(true);
var chainLevel = residueTest === _true && atomTest === _true;
var residueLevel = atomTest === _true;
l.structure = inputStructure;
for (var _a = 0, units_2 = units; _a < units_2.length; _a++) {
var unit = units_2[_a];
l.unit = unit;
if (!unitTest(ctx))
continue;
var elements = unit.elements, model = unit.model;
builder.beginUnit(unit.id);
if (unit.kind === 0 /* Unit.Kind.Atomic */) {
var chainsIt = Segmentation.transientSegments(unit.model.atomicHierarchy.chainAtomSegments, elements);
var residuesIt = Segmentation.transientSegments(unit.model.atomicHierarchy.residueAtomSegments, elements);
while (chainsIt.hasNext) {
var chainSegment = chainsIt.move();
l.element = elements[chainSegment.start];
// test entity and chain
if (!entityTest(ctx) || !chainTest(ctx))
continue;
if (chainLevel) {
builder.addElementRange(elements, chainSegment.start, chainSegment.end);
continue;
}
residuesIt.setSegment(chainSegment);
while (residuesIt.hasNext) {
var residueSegment = residuesIt.move();
l.element = elements[residueSegment.start];
// test residue
if (!residueTest(ctx))
continue;
if (residueLevel) {
builder.addElementRange(elements, residueSegment.start, residueSegment.end);
continue;
}
for (var j = residueSegment.start, _j = residueSegment.end; j < _j; j++) {
l.element = elements[j];
// test atom
if (atomTest(ctx)) {
builder.addElement(l.element);
}
}
}
}
}
else {
var chainElementSegments = (1 /* Unit.Kind.Spheres */ ? model.coarseHierarchy.spheres : model.coarseHierarchy.gaussians).chainElementSegments;
var chainsIt = Segmentation.transientSegments(chainElementSegments, elements);
while (chainsIt.hasNext) {
var chainSegment = chainsIt.move();
l.element = elements[chainSegment.start];
// test entity and chain
if (!entityTest(ctx) || !chainTest(ctx))
continue;
if (chainLevel) {
builder.addElementRange(elements, chainSegment.start, chainSegment.end);
continue;
}
for (var j = chainSegment.start, _j = chainSegment.end; j < _j; j++) {
l.element = elements[j];
// test residue/coarse element
if (residueTest(ctx)) {
builder.addElement(l.element);
}
}
}
}
builder.commitUnit();
ctx.throwIfTimedOut();
}
ctx.popCurrentElement();
return StructureSelection.Singletons(inputStructure, builder.getStructure());
};
}
function atomGroupsGrouped(_a) {
var unitTest = _a.unitTest, entityTest = _a.entityTest, chainTest = _a.chainTest, residueTest = _a.residueTest, atomTest = _a.atomTest, groupBy = _a.groupBy;
return function query_atomGroupsGrouped(ctx) {
var inputStructure = ctx.inputStructure;
var units = inputStructure.units;
var l = ctx.pushCurrentElement();
var builder = new LinearGroupingBuilder(inputStructure);
l.structure = inputStructure;
for (var _a = 0, units_3 = units; _a < units_3.length; _a++) {
var unit = units_3[_a];
l.unit = unit;
if (!unitTest(ctx))
continue;
var elements = unit.elements, model = unit.model;
if (unit.kind === 0 /* Unit.Kind.Atomic */) {
var chainsIt = Segmentation.transientSegments(model.atomicHierarchy.chainAtomSegments, elements);
var residuesIt = Segmentation.transientSegments(model.atomicHierarchy.residueAtomSegments, elements);
while (chainsIt.hasNext) {
var chainSegment = chainsIt.move();
l.element = elements[chainSegment.start];
// test entity and chain
if (!entityTest(ctx) || !chainTest(ctx))
continue;
residuesIt.setSegment(chainSegment);
while (residuesIt.hasNext) {
var residueSegment = residuesIt.move();
l.element = elements[residueSegment.start];
// test residue
if (!residueTest(ctx))
continue;
for (var j = residueSegment.start, _j = residueSegment.end; j < _j; j++) {
l.element = elements[j];
// test atom
if (atomTest(ctx)) {
builder.add(groupBy(ctx), unit.id, l.element);
}
}
}
}
}
else {
var chainElementSegments = (1 /* Unit.Kind.Spheres */ ? model.coarseHierarchy.spheres : model.coarseHierarchy.gaussians).chainElementSegments;
var chainsIt = Segmentation.transientSegments(chainElementSegments, elements);
while (chainsIt.hasNext) {
var chainSegment = chainsIt.move();
l.element = elements[chainSegment.start];
// test entity and chain
if (!entityTest(ctx) || !chainTest(ctx))
continue;
for (var j = chainSegment.start, _j = chainSegment.end; j < _j; j++) {
l.element = elements[j];
// test residue/coarse element
if (residueTest(ctx)) {
builder.add(groupBy(ctx), unit.id, l.element);
}
}
}
}
ctx.throwIfTimedOut();
}
ctx.popCurrentElement();
return builder.getSelection();
};
}
function getRingStructure(unit, ring, inputStructure) {
var elements = new Int32Array(ring.length);
for (var i = 0, _i = ring.length; i < _i; i++)
elements[i] = unit.elements[ring[i]];
return Structure.create([unit.getChild(SortedArray.ofSortedArray(elements))], { parent: inputStructure });
}
export function rings(fingerprints, onlyAromatic) {
return function query_rings(ctx) {
var units = ctx.inputStructure.units;
var ret = StructureSelection.LinearBuilder(ctx.inputStructure);
if (!fingerprints || fingerprints.length === 0) {
for (var _a = 0, units_4 = units; _a < units_4.length; _a++) {
var u = units_4[_a];
if (!Unit.isAtomic(u))
continue;
if (onlyAromatic) {
for (var _b = 0, _c = u.rings.aromaticRings; _b < _c.length; _b++) {
var r = _c[_b];
ret.add(getRingStructure(u, u.rings.all[r], ctx.inputStructure));
}
}
else {
for (var _d = 0, _e = u.rings.all; _d < _e.length; _d++) {
var r = _e[_d];
ret.add(getRingStructure(u, r, ctx.inputStructure));
}
}
}
}
else {
var uniqueFps = UniqueArray.create();
for (var i = 0; i < fingerprints.length; i++)
UniqueArray.add(uniqueFps, fingerprints[i], fingerprints[i]);
for (var _f = 0, units_5 = units; _f < units_5.length; _f++) {
var u = units_5[_f];
if (!Unit.isAtomic(u))
continue;
var rings_1 = u.rings;
for (var _g = 0, _h = uniqueFps.array; _g < _h.length; _g++) {
var fp = _h[_g];
if (!rings_1.byFingerprint.has(fp))
continue;
for (var _k = 0, _l = rings_1.byFingerprint.get(fp); _k < _l.length; _k++) {
var r = _l[_k];
if (onlyAromatic && !rings_1.aromaticRings.includes(r))
continue;
ret.add(getRingStructure(u, rings_1.all[r], ctx.inputStructure));
}
}
}
}
return ret.getSelection();
};
}
export function querySelection(selection, query, inComplement) {
if (inComplement === void 0) { inComplement = false; }
return function query_querySelection(ctx) {
var targetSel = selection(ctx);
if (StructureSelection.structureCount(targetSel) === 0)
return targetSel;
var target = inComplement
? structureSubtract(ctx.inputStructure, StructureSelection.unionStructure(targetSel))
: StructureSelection.unionStructure(targetSel);
if (target.elementCount === 0)
return StructureSelection.Empty(ctx.inputStructure);
ctx.throwIfTimedOut();
ctx.pushInputStructure(target);
var result = query(ctx);
ctx.popInputStructure();
return StructureSelection.withInputStructure(result, ctx.inputStructure);
};
}
export function bondedAtomicPairs(bondTest) {
return function query_bondedAtomicPairs(ctx) {
var structure = ctx.inputStructure;
var interBonds = structure.interUnitBonds;
// Note: each bond is called twice, that's why we need the unique builder.
var ret = StructureSelection.UniqueBuilder(ctx.inputStructure);
ctx.pushCurrentBond();
var atomicBond = ctx.atomicBond;
atomicBond.setTestFn(bondTest);
atomicBond.setStructure(structure);
// Process intra unit bonds
for (var _a = 0, _b = structure.units; _a < _b.length; _a++) {
var unit = _b[_a];
if (unit.kind !== 0 /* Unit.Kind.Atomic */)
continue;
var _c = unit.bonds, intraBondOffset = _c.offset, intraBondB = _c.b, _d = _c.edgeProps, flags = _d.flags, order = _d.order, key = _d.key;
atomicBond.a.unit = unit;
atomicBond.b.unit = unit;
for (var i = 0, _i = unit.elements.length; i < _i; i++) {
atomicBond.aIndex = i;
atomicBond.a.element = unit.elements[i];
// check intra unit bonds
for (var lI = intraBondOffset[i], _lI = intraBondOffset[i + 1]; lI < _lI; lI++) {
atomicBond.bIndex = intraBondB[lI];
atomicBond.b.element = unit.elements[intraBondB[lI]];
atomicBond.type = flags[lI];
atomicBond.order = order[lI];
atomicBond.key = key[lI];
// No need to "swap test" because each bond direction will be visited eventually.
if (atomicBond.test(ctx, false)) {
var b = structure.subsetBuilder(false);
b.beginUnit(unit.id);
b.addElement(atomicBond.a.element);
b.addElement(atomicBond.b.element);
b.commitUnit();
ret.add(b.getStructure());
}
}
}
}
// Process inter unit bonds
for (var _e = 0, _f = interBonds.edges; _e < _f.length; _e++) {
var bond = _f[_e];
atomicBond.a.unit = structure.unitMap.get(bond.unitA);
atomicBond.a.element = atomicBond.a.unit.elements[bond.indexA];
atomicBond.aIndex = bond.indexA;
atomicBond.b.unit = structure.unitMap.get(bond.unitB);
atomicBond.b.element = atomicBond.b.unit.elements[bond.indexB];
atomicBond.bIndex = bond.indexB;
atomicBond.order = bond.props.order;
atomicBond.type = bond.props.flag;
atomicBond.key = bond.props.key;
// No need to "swap test" because each bond direction will be visited eventually.
if (atomicBond.test(ctx, false)) {
var b = structure.subsetBuilder(false);
b.addToUnit(atomicBond.a.unit.id, atomicBond.a.element);
b.addToUnit(atomicBond.b.unit.id, atomicBond.b.element);
ret.add(b.getStructure());
}
}
ctx.popCurrentBond();
return ret.getSelection();
};
}