Updated carbohydrate detection in mol-model

src/mol-model/structure/structure/carbohydrates/compute.ts

@@ -4,45 +4,57 @@
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
  */
 
-import Unit from '../unit';
+import { Segmentation } from 'mol-data/int';
+import { combinations } from 'mol-data/util/combination';
+import { areConnected } from 'mol-math/graph';
+import { Vec3 } from 'mol-math/linear-algebra';
+import PrincipalAxes from 'mol-math/linear-algebra/matrix/principal-axes';
+import { fillSerial } from 'mol-util/array';
 import { ResidueIndex } from '../../model';
-import { Interval, Segmentation } from 'mol-data/int';
+import { ElementSymbol, MoleculeType } from '../../model/types';
+import { getMoleculeType, getPositionMatrix } from '../../util';
+import StructureElement from '../element';
 import Structure from '../structure';
-import { Carbohydrates, CarbohydrateLink, CarbohydrateTerminalLink, CarbohydrateElement } from './data';
+import Unit from '../unit';
 import { SaccharideNameMap, UnknownSaccharideComponent } from './constants';
-import { Vec3 } from 'mol-math/linear-algebra';
-import { getMoleculeType, getPositionMatrix } from '../../util';
-import { MoleculeType, ElementSymbol } from '../../model/types';
-import { areConnected } from 'mol-math/graph';
-import { combinations } from 'mol-data/util/combination';
-import { fillSerial } from 'mol-util/array';
-import PrincipalAxes from 'mol-math/linear-algebra/matrix/principal-axes';
+import { CarbohydrateElement, CarbohydrateLink, Carbohydrates, CarbohydrateTerminalLink } from './data';
 
 function getResidueIndex(elementIndex: number, unit: Unit.Atomic) {
     return unit.model.atomicHierarchy.residueAtomSegments.index[unit.elements[elementIndex]]
 }
 
-function getRingIndices(unit: Unit.Atomic, rI: ResidueIndex) {
-    const { offsets } = unit.model.atomicHierarchy.residueAtomSegments
-    const { elements } = unit
-    const interval = Interval.ofBounds(offsets[rI], offsets[rI + 1])
-    const rings: number[] = []
-    rings.push(...unit.rings.byFingerprint.get('C-C-C-C-C-O') || [])
-    rings.push(...unit.rings.byFingerprint.get('C-C-C-C-O') || [])
-    const sugarRings: ReadonlyArray<number>[] = []
-    for (let i = 0, il = rings.length; i < il; ++i) {
-        let withinIntervalCount = 0
-        const ring = unit.rings.all[rings[i]]
-        for (let j = 0, jl = ring.length; j < jl; ++j) {
-            if (Interval.has(interval, elements[ring[j]])) ++withinIntervalCount
+function sugarResidueIdx(unit: Unit.Atomic, ring: ArrayLike<StructureElement.UnitIndex>): ResidueIndex {
+    const { elements } = unit;
+    const residueIndex = unit.model.atomicHierarchy.residueAtomSegments.index;
+    const idx = residueIndex[elements[ring[0]]];
+    for (let rI = 1, _rI = ring.length; rI < _rI; rI++) {
+        if (idx !== residueIndex[elements[ring[rI]]]) return -1 as ResidueIndex;
+    }
+    return idx;
+}
+
+function addSugarRings(unit: Unit.Atomic, fp: string, sugarResidues: Map<ResidueIndex, number[]>) {
+    const rings = unit.rings;
+    const byFp = rings.byFingerprint.get(fp);
+    if (!byFp) return;
+    for (const r of byFp) {
+        const idx = sugarResidueIdx(unit, rings.all[r]);
+        if (idx >= 0) {
+            if (sugarResidues.has(idx)) sugarResidues.get(idx)!.push(r);
+            else sugarResidues.set(idx, [r]);
         }
-        if (withinIntervalCount === ring.length) sugarRings.push(ring)
     }
-    return sugarRings
+}
+
+function getSugarRingIndices(unit: Unit.Atomic) {
+    const sugarResidues = new Map<ResidueIndex, number[]>();
+    addSugarRings(unit, 'C-C-C-C-C-O', sugarResidues);
+    addSugarRings(unit, 'C-C-C-C-O', sugarResidues);
+    return sugarResidues;
 }
 
 const C = ElementSymbol('C')
-function getDirection(direction: Vec3, unit: Unit.Atomic, indices: ReadonlyArray<number>, center: Vec3) {
+function getDirection(direction: Vec3, unit: Unit.Atomic, indices: ReadonlyArray<StructureElement.UnitIndex>, center: Vec3) {
     let indexC1 = -1, indexC1X = -1, indexC = -1
     const { elements } = unit
     const { position } = unit.conformation
@@ -108,6 +120,8 @@ export function computeCarbohydrates(structure: Structure): Carbohydrates {
         const chainIt = Segmentation.transientSegments(chainAtomSegments, unit.elements)
         const residueIt = Segmentation.transientSegments(residueAtomSegments, unit.elements)
 
+        let sugarResidueMap: Map<ResidueIndex, number[]> | undefined = void 0;
+
         while (chainIt.hasNext) {
             residueIt.setSegment(chainIt.move());
 
@@ -119,14 +133,27 @@ export function computeCarbohydrates(structure: Structure): Carbohydrates {
                     if (getMoleculeType(unit.model, residueIndex) !== MoleculeType.saccharide) continue
                 }
 
-                const sugarRings = getRingIndices(unit, residueIndex)
+                if (!sugarResidueMap) {
+                    sugarResidueMap = getSugarRingIndices(unit);
+                }
+
+                const sugarRings = sugarResidueMap.get(residueIndex);
+
+                if (!sugarRings || !sugarRings.length) {
+                    console.warn(`No ring found for carbohydrate on residue with index ${residueIndex}, unit ${unit.id}. Residue skipped.`);
+                    continue;
+                }
+
+                const rings = unit.rings;
                 const ringElements: number[] = []
 
                 for (let j = 0, jl = sugarRings.length; j < jl; ++j) {
-                    const pa = new PrincipalAxes(getPositionMatrix(unit, sugarRings[j]))
+                    const ringAtoms = rings.all[sugarRings[j]];
+
+                    const pa = new PrincipalAxes(getPositionMatrix(unit, ringAtoms))
                     const center = Vec3.copy(Vec3.zero(), pa.center)
                     const normal = Vec3.copy(Vec3.zero(), pa.normVecC)
-                    const direction = getDirection(Vec3.zero(), unit, sugarRings[j], center)
+                    const direction = getDirection(Vec3.zero(), unit, ringAtoms, center)
                     Vec3.orthogonalize(direction, normal, direction)
 
                     const elementIndex = elements.length
@@ -138,8 +165,9 @@ export function computeCarbohydrates(structure: Structure): Carbohydrates {
                 // add carbohydrate links induced by intra-residue bonds
                 const ringCombinations = combinations(fillSerial(new Array(sugarRings.length)), 2)
                 for (let j = 0, jl = ringCombinations.length; j < jl; ++j) {
-                    const rc = ringCombinations[j]
-                    if (areConnected(sugarRings[rc[0]], sugarRings[rc[1]], unit.links, 2)) {
+                    const rc = ringCombinations[j];
+                    const r0 = rings.all[sugarRings[rc[0]]], r1 = rings.all[sugarRings[rc[1]]];
+                    if (areConnected(r0, r1, unit.links, 2)) {
                         // fix both directions as it is unlcear where the C1 atom is
                         fixLinkDirection(ringElements[rc[0]], ringElements[rc[1]])
                         fixLinkDirection(ringElements[rc[1]], ringElements[rc[0]])
@@ -153,10 +181,6 @@ export function computeCarbohydrates(structure: Structure): Carbohydrates {
                         })
                     }
                 }
-
-                if (!sugarRings.length) {
-                    console.warn('No ring found for carbohydrate')
-                }
             }
         }
     }