modifiers.ts 26 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690
  1. /**
  2. * Copyright (c) 2017-2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
  3. *
  4. * @author David Sehnal <david.sehnal@gmail.com>
  5. */
  6. import { Segmentation, SortedArray } from '../../../../mol-data/int';
  7. import { Structure, Unit } from '../../structure';
  8. import { StructureQuery } from '../query';
  9. import { StructureSelection } from '../selection';
  10. import { UniqueStructuresBuilder } from '../utils/builders';
  11. import { StructureUniqueSubsetBuilder } from '../../structure/util/unique-subset-builder';
  12. import { QueryContext, QueryFn } from '../context';
  13. import { structureIntersect, structureSubtract, structureUnion } from '../utils/structure-set';
  14. import { UniqueArray } from '../../../../mol-data/generic';
  15. import { StructureSubsetBuilder } from '../../structure/util/subset-builder';
  16. import { StructureElement } from '../../structure/element';
  17. import { MmcifFormat } from '../../../../mol-model-formats/structure/mmcif';
  18. import { ResidueSet, ResidueSetEntry } from '../../model/properties/utils/residue-set';
  19. import { StructureProperties } from '../../structure/properties';
  20. import { arraySetAdd } from '../../../../mol-util/array';
  21. function getWholeResidues(ctx: QueryContext, source: Structure, structure: Structure) {
  22. const builder = source.subsetBuilder(true);
  23. for (const unit of structure.units) {
  24. if (unit.kind !== Unit.Kind.Atomic) {
  25. // just copy non-atomic units.
  26. builder.setUnit(unit.id, unit.elements);
  27. continue;
  28. }
  29. const { residueAtomSegments } = unit.model.atomicHierarchy;
  30. const sourceElements = source.unitMap.get(unit.id).elements;
  31. const elements = unit.elements;
  32. builder.beginUnit(unit.id);
  33. const residuesIt = Segmentation.transientSegments(residueAtomSegments, elements);
  34. while (residuesIt.hasNext) {
  35. const rI = residuesIt.move().index;
  36. for (let j = residueAtomSegments.offsets[rI], _j = residueAtomSegments.offsets[rI + 1]; j < _j; j++) {
  37. if (SortedArray.has(sourceElements, j)) builder.addElement(j);
  38. }
  39. }
  40. builder.commitUnit();
  41. ctx.throwIfTimedOut();
  42. }
  43. return builder.getStructure();
  44. }
  45. export function wholeResidues(query: StructureQuery): StructureQuery {
  46. return function query_wholeResidues(ctx) {
  47. const inner = query(ctx);
  48. if (StructureSelection.isSingleton(inner)) {
  49. return StructureSelection.Singletons(ctx.inputStructure, getWholeResidues(ctx, ctx.inputStructure, inner.structure));
  50. } else {
  51. const builder = new UniqueStructuresBuilder(ctx.inputStructure);
  52. for (const s of inner.structures) {
  53. builder.add(getWholeResidues(ctx, ctx.inputStructure, s));
  54. }
  55. return builder.getSelection();
  56. }
  57. };
  58. }
  59. export interface IncludeSurroundingsParams {
  60. radius: number,
  61. elementRadius?: QueryFn<number>,
  62. wholeResidues?: boolean
  63. }
  64. function getIncludeSurroundings(ctx: QueryContext, source: Structure, structure: Structure, params: IncludeSurroundingsParams) {
  65. const builder = new StructureUniqueSubsetBuilder(source);
  66. const lookup = source.lookup3d;
  67. const r = params.radius;
  68. for (const unit of structure.units) {
  69. const { x, y, z } = unit.conformation;
  70. const elements = unit.elements;
  71. for (let i = 0, _i = elements.length; i < _i; i++) {
  72. const e = elements[i];
  73. lookup.findIntoBuilder(x(e), y(e), z(e), r, builder);
  74. }
  75. ctx.throwIfTimedOut();
  76. }
  77. return !!params.wholeResidues ? getWholeResidues(ctx, source, builder.getStructure()) : builder.getStructure();
  78. }
  79. interface IncludeSurroundingsParamsWithRadius extends IncludeSurroundingsParams {
  80. elementRadius: QueryFn<number>,
  81. elementRadiusClosure: StructureElement.Property<number>,
  82. sourceMaxRadius: number
  83. }
  84. function getIncludeSurroundingsWithRadius(ctx: QueryContext, source: Structure, structure: Structure, params: IncludeSurroundingsParamsWithRadius) {
  85. const builder = new StructureUniqueSubsetBuilder(source);
  86. const lookup = source.lookup3d;
  87. const { elementRadius, elementRadiusClosure, sourceMaxRadius, radius } = params;
  88. ctx.pushCurrentElement();
  89. ctx.element.structure = structure;
  90. for (const unit of structure.units) {
  91. ctx.element.unit = unit;
  92. const { x, y, z } = unit.conformation;
  93. const elements = unit.elements;
  94. for (let i = 0, _i = elements.length; i < _i; i++) {
  95. const e = elements[i];
  96. ctx.element.element = e;
  97. const eRadius = elementRadius(ctx);
  98. lookup.findIntoBuilderWithRadius(x(e), y(e), z(e), eRadius, sourceMaxRadius, radius, elementRadiusClosure, builder);
  99. }
  100. ctx.throwIfTimedOut();
  101. }
  102. ctx.popCurrentElement();
  103. return !!params.wholeResidues ? getWholeResidues(ctx, source, builder.getStructure()) : builder.getStructure();
  104. }
  105. function createElementRadiusFn(ctx: QueryContext, eRadius: QueryFn<number>): StructureElement.Property<number> {
  106. return e => {
  107. ctx.element.structure = e.structure;
  108. ctx.element.unit = e.unit;
  109. ctx.element.element = e.element;
  110. return eRadius(ctx);
  111. };
  112. }
  113. function findStructureRadius(ctx: QueryContext, eRadius: QueryFn<number>) {
  114. let r = 0;
  115. ctx.element.structure = ctx.inputStructure;
  116. for (const unit of ctx.inputStructure.units) {
  117. ctx.element.unit = unit;
  118. const elements = unit.elements;
  119. for (let i = 0, _i = elements.length; i < _i; i++) {
  120. const e = elements[i];
  121. ctx.element.element = e;
  122. const eR = eRadius(ctx);
  123. if (eR > r) r = eR;
  124. }
  125. }
  126. ctx.throwIfTimedOut();
  127. return r;
  128. }
  129. export function includeSurroundings(query: StructureQuery, params: IncludeSurroundingsParams): StructureQuery {
  130. return function query_includeSurroundings(ctx) {
  131. const inner = query(ctx);
  132. if (params.elementRadius) {
  133. const prms: IncludeSurroundingsParamsWithRadius = {
  134. ...params,
  135. elementRadius: params.elementRadius,
  136. elementRadiusClosure: createElementRadiusFn(ctx, params.elementRadius),
  137. sourceMaxRadius: findStructureRadius(ctx, params.elementRadius)
  138. };
  139. if (StructureSelection.isSingleton(inner)) {
  140. const surr = getIncludeSurroundingsWithRadius(ctx, ctx.inputStructure, inner.structure, prms);
  141. const ret = StructureSelection.Singletons(ctx.inputStructure, surr);
  142. return ret;
  143. } else {
  144. const builder = new UniqueStructuresBuilder(ctx.inputStructure);
  145. for (const s of inner.structures) {
  146. builder.add(getIncludeSurroundingsWithRadius(ctx, ctx.inputStructure, s, prms));
  147. }
  148. return builder.getSelection();
  149. }
  150. }
  151. if (StructureSelection.isSingleton(inner)) {
  152. const surr = getIncludeSurroundings(ctx, ctx.inputStructure, inner.structure, params);
  153. const ret = StructureSelection.Singletons(ctx.inputStructure, surr);
  154. return ret;
  155. } else {
  156. const builder = new UniqueStructuresBuilder(ctx.inputStructure);
  157. for (const s of inner.structures) {
  158. builder.add(getIncludeSurroundings(ctx, ctx.inputStructure, s, params));
  159. }
  160. return builder.getSelection();
  161. }
  162. };
  163. }
  164. export function querySelection(selection: StructureQuery, query: StructureQuery): StructureQuery {
  165. return function query_querySelection(ctx) {
  166. const targetSel = selection(ctx);
  167. if (StructureSelection.structureCount(targetSel) === 0) return targetSel;
  168. const ret = StructureSelection.UniqueBuilder(ctx.inputStructure);
  169. const add = (s: Structure) => ret.add(s);
  170. StructureSelection.forEach(targetSel, (s, sI) => {
  171. ctx.pushInputStructure(s);
  172. StructureSelection.forEach(query(ctx), add);
  173. ctx.popInputStructure();
  174. if (sI % 10 === 0) ctx.throwIfTimedOut();
  175. });
  176. return ret.getSelection();
  177. };
  178. }
  179. export function intersectBy(query: StructureQuery, by: StructureQuery): StructureQuery {
  180. return function query_intersectBy(ctx) {
  181. const selection = query(ctx);
  182. if (StructureSelection.structureCount(selection) === 0) return selection;
  183. const bySel = by(ctx);
  184. if (StructureSelection.structureCount(bySel) === 0) return StructureSelection.Empty(ctx.inputStructure);
  185. const unionBy = StructureSelection.unionStructure(bySel);
  186. const ret = StructureSelection.UniqueBuilder(ctx.inputStructure);
  187. StructureSelection.forEach(selection, (s, sI) => {
  188. const ii = structureIntersect(unionBy, s);
  189. if (ii.elementCount !== 0) ret.add(ii);
  190. if (sI % 50 === 0) ctx.throwIfTimedOut();
  191. });
  192. return ret.getSelection();
  193. };
  194. }
  195. export function exceptBy(query: StructureQuery, by: StructureQuery): StructureQuery {
  196. return function query_exceptBy(ctx) {
  197. const selection = query(ctx);
  198. if (StructureSelection.structureCount(selection) === 0) return selection;
  199. const bySel = by(ctx);
  200. if (StructureSelection.structureCount(bySel) === 0) return selection;
  201. const subtractBy = StructureSelection.unionStructure(bySel);
  202. const ret = StructureSelection.UniqueBuilder(ctx.inputStructure);
  203. StructureSelection.forEach(selection, (s, sI) => {
  204. const diff = structureSubtract(s, subtractBy);
  205. if (diff.elementCount !== 0) ret.add(diff);
  206. if (sI % 50 === 0) ctx.throwIfTimedOut();
  207. });
  208. return ret.getSelection();
  209. };
  210. }
  211. export function union(query: StructureQuery): StructureQuery {
  212. return function query_union(ctx) {
  213. const ret = StructureSelection.LinearBuilder(ctx.inputStructure);
  214. ret.add(StructureSelection.unionStructure(query(ctx)));
  215. return ret.getSelection();
  216. };
  217. }
  218. export function expandProperty(query: StructureQuery, property: QueryFn): StructureQuery {
  219. return function query_expandProperty(ctx) {
  220. const src = query(ctx);
  221. const propertyToStructureIndexMap = new Map<any, UniqueArray<number>>();
  222. const builders: StructureSubsetBuilder[] = [];
  223. ctx.pushCurrentElement();
  224. StructureSelection.forEach(src, (s, sI) => {
  225. ctx.element.structure = s;
  226. for (const unit of s.units) {
  227. ctx.element.unit = unit;
  228. const elements = unit.elements;
  229. for (let i = 0, _i = elements.length; i < _i; i++) {
  230. ctx.element.element = elements[i];
  231. const p = property(ctx);
  232. let arr: UniqueArray<number>;
  233. if (propertyToStructureIndexMap.has(p)) arr = propertyToStructureIndexMap.get(p)!;
  234. else {
  235. arr = UniqueArray.create<number>();
  236. propertyToStructureIndexMap.set(p, arr);
  237. }
  238. UniqueArray.add(arr, sI, sI);
  239. }
  240. }
  241. builders[sI] = ctx.inputStructure.subsetBuilder(true);
  242. if (sI % 10 === 0) ctx.throwIfTimedOut();
  243. });
  244. ctx.element.structure = ctx.inputStructure;
  245. for (const unit of ctx.inputStructure.units) {
  246. ctx.element.unit = unit;
  247. const elements = unit.elements;
  248. for (let i = 0, _i = elements.length; i < _i; i++) {
  249. ctx.element.element = elements[i];
  250. const p = property(ctx);
  251. if (!propertyToStructureIndexMap.has(p)) continue;
  252. const indices = propertyToStructureIndexMap.get(p)!.array;
  253. for (let _sI = 0, __sI = indices.length; _sI < __sI; _sI++) {
  254. builders[indices[_sI]].addToUnit(unit.id, elements[i]);
  255. }
  256. }
  257. }
  258. ctx.popCurrentElement();
  259. const ret = StructureSelection.UniqueBuilder(ctx.inputStructure);
  260. for (const b of builders) ret.add(b.getStructure());
  261. return ret.getSelection();
  262. };
  263. }
  264. export interface IncludeConnectedParams {
  265. query: StructureQuery,
  266. bondTest?: QueryFn<boolean>,
  267. layerCount: number,
  268. wholeResidues: boolean,
  269. fixedPoint: boolean
  270. }
  271. export function includeConnected({ query, layerCount, wholeResidues, bondTest, fixedPoint }: IncludeConnectedParams): StructureQuery {
  272. const lc = Math.max(layerCount, 0);
  273. return function query_includeConnected(ctx) {
  274. const builder = StructureSelection.UniqueBuilder(ctx.inputStructure);
  275. const src = query(ctx);
  276. ctx.pushCurrentBond();
  277. ctx.atomicBond.setTestFn(bondTest);
  278. StructureSelection.forEach(src, (s, sI) => {
  279. let incl = s;
  280. if (fixedPoint) {
  281. while (true) {
  282. const prevCount = incl.elementCount;
  283. incl = includeConnectedStep(ctx, wholeResidues, incl);
  284. if (incl.elementCount === prevCount) break;
  285. }
  286. } else {
  287. for (let i = 0; i < lc; i++) {
  288. incl = includeConnectedStep(ctx, wholeResidues, incl);
  289. }
  290. }
  291. builder.add(incl);
  292. if (sI % 10 === 0) ctx.throwIfTimedOut();
  293. });
  294. ctx.popCurrentBond();
  295. return builder.getSelection();
  296. };
  297. }
  298. function includeConnectedStep(ctx: QueryContext, wholeResidues: boolean, structure: Structure) {
  299. const expanded = expandConnected(ctx, structure);
  300. if (wholeResidues) return getWholeResidues(ctx, ctx.inputStructure, expanded);
  301. return expanded;
  302. }
  303. function expandConnected(ctx: QueryContext, structure: Structure) {
  304. const inputStructure = ctx.inputStructure;
  305. const interBonds = inputStructure.interUnitBonds;
  306. const builder = new StructureUniqueSubsetBuilder(inputStructure);
  307. const atomicBond = ctx.atomicBond;
  308. // Process intra unit bonds
  309. for (const unit of structure.units) {
  310. if (unit.kind !== Unit.Kind.Atomic) {
  311. // add the whole unit
  312. builder.beginUnit(unit.id);
  313. for (let i = 0, _i = unit.elements.length; i < _i; i++) {
  314. builder.addElement(unit.elements[i]);
  315. }
  316. builder.commitUnit();
  317. continue;
  318. }
  319. const inputUnitA = inputStructure.unitMap.get(unit.id) as Unit.Atomic;
  320. const { offset: intraBondOffset, b: intraBondB, edgeProps: { flags, order } } = inputUnitA.bonds;
  321. atomicBond.setStructure(inputStructure);
  322. // Process intra unit bonds
  323. atomicBond.a.unit = inputUnitA;
  324. atomicBond.b.unit = inputUnitA;
  325. for (let i = 0, _i = unit.elements.length; i < _i; i++) {
  326. // add the current element
  327. builder.addToUnit(unit.id, unit.elements[i]);
  328. const aIndex = SortedArray.indexOf(inputUnitA.elements, unit.elements[i]) as StructureElement.UnitIndex;
  329. // check intra unit bonds
  330. for (let lI = intraBondOffset[aIndex], _lI = intraBondOffset[aIndex + 1]; lI < _lI; lI++) {
  331. const bIndex = intraBondB[lI] as StructureElement.UnitIndex;
  332. const bElement = inputUnitA.elements[bIndex];
  333. // Check if the element is already present:
  334. if (SortedArray.has(unit.elements, bElement) || builder.has(unit.id, bElement)) continue;
  335. atomicBond.aIndex = aIndex;
  336. atomicBond.a.element = unit.elements[i];
  337. atomicBond.bIndex = bIndex;
  338. atomicBond.b.element = bElement;
  339. atomicBond.type = flags[lI];
  340. atomicBond.order = order[lI];
  341. if (atomicBond.test(ctx, true)) {
  342. builder.addToUnit(unit.id, bElement);
  343. }
  344. }
  345. }
  346. // Process inter unit bonds
  347. for (const bondedUnit of interBonds.getConnectedUnits(inputUnitA.id)) {
  348. const currentUnitB = structure.unitMap.get(bondedUnit.unitB);
  349. const inputUnitB = inputStructure.unitMap.get(bondedUnit.unitB) as Unit.Atomic;
  350. for (const aI of bondedUnit.connectedIndices) {
  351. // check if the element is in the expanded structure
  352. if (!SortedArray.has(unit.elements, inputUnitA.elements[aI])) continue;
  353. for (const bond of bondedUnit.getEdges(aI)) {
  354. const bElement = inputUnitB.elements[bond.indexB];
  355. // Check if the element is already present:
  356. if ((currentUnitB && SortedArray.has(currentUnitB.elements, bElement)) || builder.has(bondedUnit.unitB, bElement)) continue;
  357. atomicBond.a.unit = inputUnitA;
  358. atomicBond.aIndex = aI;
  359. atomicBond.a.element = inputUnitA.elements[aI];
  360. atomicBond.b.unit = inputUnitB;
  361. atomicBond.bIndex = bond.indexB;
  362. atomicBond.b.element = bElement;
  363. atomicBond.type = bond.props.flag;
  364. atomicBond.order = bond.props.order;
  365. if (atomicBond.test(ctx, true)) {
  366. builder.addToUnit(bondedUnit.unitB, bElement);
  367. }
  368. }
  369. }
  370. }
  371. }
  372. return builder.getStructure();
  373. }
  374. export interface SurroundingLigandsParams {
  375. query: StructureQuery,
  376. radius: number,
  377. includeWater: boolean
  378. }
  379. /**
  380. * Includes expanded surrounding ligands based on radius from the source, struct_conn entries & pdbx_molecule entries.
  381. */
  382. export function surroundingLigands({ query, radius, includeWater }: SurroundingLigandsParams): StructureQuery {
  383. return function query_surroundingLigands(ctx) {
  384. const inner = StructureSelection.unionStructure(query(ctx));
  385. const surroundings = getWholeResidues(ctx, ctx.inputStructure, getIncludeSurroundings(ctx, ctx.inputStructure, inner, { radius }));
  386. const prd = getPrdAsymIdx(ctx.inputStructure);
  387. const graph = getStructConnInfo(ctx.inputStructure);
  388. const l = StructureElement.Location.create(surroundings);
  389. const includedPrdChains = new Map<string, string[]>();
  390. const componentResidues = new ResidueSet({ checkOperator: true });
  391. for (const unit of surroundings.units) {
  392. if (unit.kind !== Unit.Kind.Atomic) continue;
  393. l.unit = unit;
  394. const { elements } = unit;
  395. const chainsIt = Segmentation.transientSegments(unit.model.atomicHierarchy.chainAtomSegments, elements);
  396. const residuesIt = Segmentation.transientSegments(unit.model.atomicHierarchy.residueAtomSegments, elements);
  397. while (chainsIt.hasNext) {
  398. const chainSegment = chainsIt.move();
  399. l.element = elements[chainSegment.start];
  400. const asym_id = StructureProperties.chain.label_asym_id(l);
  401. const op_name = StructureProperties.unit.operator_name(l);
  402. // check for PRD molecules
  403. if (prd.has(asym_id)) {
  404. if (includedPrdChains.has(asym_id)) {
  405. arraySetAdd(includedPrdChains.get(asym_id)!, op_name);
  406. } else {
  407. includedPrdChains.set(asym_id, [op_name]);
  408. }
  409. continue;
  410. }
  411. const entityType = StructureProperties.entity.type(l);
  412. // test entity and chain
  413. if (entityType === 'water' || entityType === 'polymer') continue;
  414. residuesIt.setSegment(chainSegment);
  415. while (residuesIt.hasNext) {
  416. const residueSegment = residuesIt.move();
  417. l.element = elements[residueSegment.start];
  418. graph.addComponent(ResidueSet.getEntryFromLocation(l), componentResidues);
  419. }
  420. }
  421. ctx.throwIfTimedOut();
  422. }
  423. // assemble the core structure
  424. const builder = ctx.inputStructure.subsetBuilder(true);
  425. for (const unit of ctx.inputStructure.units) {
  426. if (unit.kind !== Unit.Kind.Atomic) continue;
  427. l.unit = unit;
  428. const { elements } = unit;
  429. const chainsIt = Segmentation.transientSegments(unit.model.atomicHierarchy.chainAtomSegments, elements);
  430. const residuesIt = Segmentation.transientSegments(unit.model.atomicHierarchy.residueAtomSegments, elements);
  431. builder.beginUnit(unit.id);
  432. while (chainsIt.hasNext) {
  433. const chainSegment = chainsIt.move();
  434. l.element = elements[chainSegment.start];
  435. const asym_id = StructureProperties.chain.label_asym_id(l);
  436. const op_name = StructureProperties.unit.operator_name(l);
  437. if (includedPrdChains.has(asym_id) && includedPrdChains.get(asym_id)!.indexOf(op_name) >= 0) {
  438. builder.addElementRange(elements, chainSegment.start, chainSegment.end);
  439. continue;
  440. }
  441. if (!componentResidues.hasLabelAsymId(asym_id)) {
  442. continue;
  443. }
  444. residuesIt.setSegment(chainSegment);
  445. while (residuesIt.hasNext) {
  446. const residueSegment = residuesIt.move();
  447. l.element = elements[residueSegment.start];
  448. if (!componentResidues.has(l)) continue;
  449. builder.addElementRange(elements, residueSegment.start, residueSegment.end);
  450. }
  451. }
  452. builder.commitUnit();
  453. ctx.throwIfTimedOut();
  454. }
  455. const components = structureUnion(ctx.inputStructure, [builder.getStructure(), inner]);
  456. // add water
  457. if (includeWater) {
  458. const finalBuilder = new StructureUniqueSubsetBuilder(ctx.inputStructure);
  459. const lookup = ctx.inputStructure.lookup3d;
  460. for (const unit of components.units) {
  461. const { x, y, z } = unit.conformation;
  462. const elements = unit.elements;
  463. for (let i = 0, _i = elements.length; i < _i; i++) {
  464. const e = elements[i];
  465. lookup.findIntoBuilderIf(x(e), y(e), z(e), radius, finalBuilder, testIsWater);
  466. finalBuilder.addToUnit(unit.id, e);
  467. }
  468. ctx.throwIfTimedOut();
  469. }
  470. return StructureSelection.Sequence(ctx.inputStructure, [finalBuilder.getStructure()]);
  471. } else {
  472. return StructureSelection.Sequence(ctx.inputStructure, [components]);
  473. }
  474. };
  475. }
  476. const _entity_type = StructureProperties.entity.type;
  477. function testIsWater(l: StructureElement.Location) {
  478. return _entity_type(l) === 'water';
  479. }
  480. function getPrdAsymIdx(structure: Structure) {
  481. const model = structure.models[0];
  482. const ids = new Set<string>();
  483. if (!MmcifFormat.is(model.sourceData)) return ids;
  484. const { _rowCount, asym_id } = model.sourceData.data.db.pdbx_molecule;
  485. for (let i = 0; i < _rowCount; i++) {
  486. ids.add(asym_id.value(i));
  487. }
  488. return ids;
  489. }
  490. function getStructConnInfo(structure: Structure) {
  491. const model = structure.models[0];
  492. const graph = new StructConnGraph();
  493. if (!MmcifFormat.is(model.sourceData)) return graph;
  494. const struct_conn = model.sourceData.data.db.struct_conn;
  495. const { conn_type_id } = struct_conn;
  496. const { ptnr1_label_asym_id, ptnr1_label_comp_id, ptnr1_label_seq_id, ptnr1_symmetry, pdbx_ptnr1_label_alt_id, pdbx_ptnr1_PDB_ins_code } = struct_conn;
  497. const { ptnr2_label_asym_id, ptnr2_label_comp_id, ptnr2_label_seq_id, ptnr2_symmetry, pdbx_ptnr2_label_alt_id, pdbx_ptnr2_PDB_ins_code } = struct_conn;
  498. for (let i = 0; i < struct_conn._rowCount; i++) {
  499. const bondType = conn_type_id.value(i);
  500. if (bondType !== 'covale' && bondType !== 'metalc') continue;
  501. const a: ResidueSetEntry = {
  502. label_asym_id: ptnr1_label_asym_id.value(i),
  503. label_comp_id: ptnr1_label_comp_id.value(i),
  504. label_seq_id: ptnr1_label_seq_id.value(i),
  505. label_alt_id: pdbx_ptnr1_label_alt_id.value(i),
  506. ins_code: pdbx_ptnr1_PDB_ins_code.value(i),
  507. operator_name: ptnr1_symmetry.value(i) ?? '1_555'
  508. };
  509. const b: ResidueSetEntry = {
  510. label_asym_id: ptnr2_label_asym_id.value(i),
  511. label_comp_id: ptnr2_label_comp_id.value(i),
  512. label_seq_id: ptnr2_label_seq_id.value(i),
  513. label_alt_id: pdbx_ptnr2_label_alt_id.value(i),
  514. ins_code: pdbx_ptnr2_PDB_ins_code.value(i),
  515. operator_name: ptnr2_symmetry.value(i) ?? '1_555'
  516. };
  517. graph.addEdge(a, b);
  518. }
  519. return graph;
  520. }
  521. class StructConnGraph {
  522. vertices = new Map<string, ResidueSetEntry>();
  523. edges = new Map<string, string[]>();
  524. private addVertex(e: ResidueSetEntry, label: string) {
  525. if (this.vertices.has(label)) return;
  526. this.vertices.set(label, e);
  527. this.edges.set(label, []);
  528. }
  529. addEdge(a: ResidueSetEntry, b: ResidueSetEntry) {
  530. const al = ResidueSet.getLabel(a);
  531. const bl = ResidueSet.getLabel(b);
  532. this.addVertex(a, al);
  533. this.addVertex(b, bl);
  534. arraySetAdd(this.edges.get(al)!, bl);
  535. arraySetAdd(this.edges.get(bl)!, al);
  536. }
  537. addComponent(start: ResidueSetEntry, set: ResidueSet) {
  538. const startLabel = ResidueSet.getLabel(start);
  539. if (!this.vertices.has(startLabel)) {
  540. set.add(start);
  541. return;
  542. }
  543. const visited = new Set<string>();
  544. const added = new Set<string>();
  545. const stack = [startLabel];
  546. added.add(startLabel);
  547. set.add(start);
  548. while (stack.length > 0) {
  549. const a = stack.pop()!;
  550. visited.add(a);
  551. const u = this.vertices.get(a)!;
  552. for (const b of this.edges.get(a)!) {
  553. if (visited.has(b)) continue;
  554. stack.push(b);
  555. if (added.has(b)) continue;
  556. added.add(b);
  557. const v = this.vertices.get(b)!;
  558. if (u.operator_name === v.operator_name) {
  559. set.add({ ...v, operator_name: start.operator_name });
  560. } else {
  561. set.add(v);
  562. }
  563. }
  564. }
  565. }
  566. }
  567. // TODO: unionBy (skip this one?), cluster