mol-model: count chains in Loci stats

src/mol-model/structure/structure/element/stats.ts

@@ -2,22 +2,26 @@
  * Copyright (c) 2019 mol* contributors, licensed under MIT, See LICENSE file for more info.
  *
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
+ * @author David Sehnal <david.sehnal@gmail.com>
  */
 
 import { OrderedSet } from '../../../../mol-data/int';
 import Unit from '../unit';
 import { Loci } from './loci';
 import { Location } from './location';
+import { ChainIndex } from '../../model/indexing';
 
 export interface Stats {
     elementCount: number
     conformationCount: number
     residueCount: number
+    chainCount: number
     unitCount: number
 
     firstElementLoc: Location
     firstConformationLoc: Location
     firstResidueLoc: Location
+    firstChainLoc: Location
     firstUnitLoc: Location
 }
 
@@ -27,15 +31,22 @@ export namespace Stats {
             elementCount: 0,
             conformationCount: 0,
             residueCount: 0,
+            chainCount: 0,
             unitCount: 0,
 
             firstElementLoc: Location.create(),
             firstConformationLoc: Location.create(),
             firstResidueLoc: Location.create(),
+            firstChainLoc: Location.create(),
             firstUnitLoc: Location.create(),
         }
     }
 
+    function addCountHelper<K>(map: Map<K, number>, key: K, inc: number) {
+        const count = map.get(key) || 0
+        map.set(key, count + inc)
+    }
+
     function handleElement(stats: Stats, element: Loci['elements'][0]) {
         const { indices, unit } = element
         const { elements } = unit
@@ -43,10 +54,6 @@ export namespace Stats {
 
         const lociResidueAltIdCounts = new Map<string, number>()
         const residueAltIdCounts = new Map<string, number>()
-        const addCount = (map: Map<string, number>, altId: string) => {
-            const count = map.get(altId) || 0
-            map.set(altId, count + 1)
-        }
 
         if (size > 0) {
             Location.set(stats.firstElementLoc, unit, elements[OrderedSet.start(indices)])
@@ -72,13 +79,13 @@ export namespace Stats {
                     let j = 0
                     const eI = elements[OrderedSet.getAt(indices, i)]
                     const rI = index[eI]
-                    addCount(lociResidueAltIdCounts, label_alt_id.value(eI))
+                    addCountHelper(lociResidueAltIdCounts, label_alt_id.value(eI), 1)
                     ++i
                     ++j
                     while (i < size) {
                         const eI = elements[OrderedSet.getAt(indices, i)]
                         if (index[eI] !== rI) break
-                        addCount(lociResidueAltIdCounts, label_alt_id.value(eI))
+                        addCountHelper(lociResidueAltIdCounts, label_alt_id.value(eI), 1)
                         ++i
                         ++j
                     }
@@ -93,7 +100,7 @@ export namespace Stats {
                         // partial residue
                         residueAltIdCounts.clear()
                         for (let l = offsets[rI], _l = offsets[rI + 1]; l < _l; ++l) {
-                            addCount(residueAltIdCounts, label_alt_id.value(l))
+                            addCountHelper(residueAltIdCounts, label_alt_id.value(l), 1)
                         }
                         // check if shared atom count match
                         if (residueAltIdCounts.get('') === lociResidueAltIdCounts.get('')) {
@@ -125,10 +132,180 @@ export namespace Stats {
         }
     }
 
+    function handleUnitChainsSimple(stats: Stats, element: Loci['elements'][0]) {
+        const { indices, unit } = element;
+        const size = OrderedSet.size(indices);
+        if (size === 0) return;
+
+        const { elements } = unit;
+
+        if (!Unit.Traits.is(unit.traits, Unit.Trait.MultiChain)) {
+            if (size === elements.length) {
+                stats.chainCount += 1;
+                if (stats.chainCount === 1) {
+                    Location.set(stats.firstChainLoc, unit, elements[OrderedSet.start(indices)]);
+                }
+            }
+            return;
+        }
+
+        const segments = Unit.isAtomic(unit)
+            ? unit.model.atomicHierarchy.chainAtomSegments
+            : Unit.isSpheres(unit)
+            ? unit.model.coarseHierarchy.spheres.chainElementSegments
+            : Unit.isGaussians(unit)
+            ? unit.model.coarseHierarchy.gaussians.chainElementSegments
+            : void 0;
+
+        if (!segments) {
+            console.warn('StructureElement loci stats: unknown unit type');
+            return;
+        }
+
+        const { index, offsets } = segments;
+        let i = 0;
+        while (i < size) {
+            let j = 0
+            const eI = elements[OrderedSet.getAt(indices, i)];
+            const cI = index[eI];
+            ++i;
+            ++j;
+            while (i < size) {
+                const eI = elements[OrderedSet.getAt(indices, i)];
+                if (index[eI] !== cI) break;
+                ++i;
+                ++j;
+            }
+
+            if (offsets[cI + 1] - offsets[cI] === j) {
+                // full chain
+                stats.chainCount += 1;
+                if (stats.chainCount === 1) {
+                    Location.set(stats.firstChainLoc, unit, elements[offsets[cI]]);
+                }
+            }
+        }
+    }
+
+    function handleUnitChainsPartitioned(stats: Stats, lociElements: Loci['elements'], start: number, end: number) {
+        let element = lociElements[start];
+
+        // all the elements have the same model since they are part of the same group so this is ok.
+        const segments = Unit.isAtomic(element.unit)
+            ? element.unit.model.atomicHierarchy.chainAtomSegments
+            : Unit.isSpheres(element.unit)
+            ? element.unit.model.coarseHierarchy.spheres.chainElementSegments
+            : Unit.isGaussians(element.unit)
+            ? element.unit.model.coarseHierarchy.gaussians.chainElementSegments
+            : void 0;
+
+        if (!segments) {
+            console.warn('StructureElement loci stats: unknown unit type');
+            return;
+        }
+
+        const { index, offsets } = segments;
+
+        const chainCounts = new Map<ChainIndex, number>();
+
+        for (let elIndex = start; elIndex < end; elIndex++) {
+            element = lociElements[elIndex];
+
+            const { indices, unit } = element
+            const size = OrderedSet.size(indices);
+            if (size === 0) continue;
+
+            const { elements } = unit;
+
+            if (!Unit.Traits.is(unit.traits, Unit.Trait.MultiChain)) {
+                const eI = elements[OrderedSet.start(indices)];
+                addCountHelper(chainCounts, index[eI], elements.length);
+                continue;
+            }
+
+            let i = 0;
+            while (i < size) {
+                let j = 0
+                const eI = elements[OrderedSet.getAt(indices, i)];
+                const cI = index[eI];
+                ++i;
+                ++j;
+                while (i < size) {
+                    const eI = elements[OrderedSet.getAt(indices, i)];
+                    if (index[eI] !== cI) break;
+                    ++i;
+                    ++j;
+                }
+
+                addCountHelper(chainCounts, cI, j);
+            }
+        }
+
+        let firstCI = -1;
+        chainCounts.forEach((count, cI) => {
+            if (offsets[cI + 1] - offsets[cI] === count) {
+                // full chain
+                stats.chainCount += 1
+                if (stats.chainCount === 1) {
+                    firstCI = cI;
+                }
+            }
+        });
+
+        if (firstCI < 0) return;
+
+        for (let elIndex = start; elIndex < end; elIndex++) {
+            element = lociElements[elIndex];
+
+            const { indices, unit } = element
+            const size = OrderedSet.size(indices);
+            if (size === 0) continue;
+
+            const { elements } = unit;
+
+            let i = 0;
+            while (i < size) {
+                const eI = elements[OrderedSet.getAt(indices, i)];
+                const cI = index[eI];
+                if (cI === firstCI) {
+                    Location.set(stats.firstChainLoc, unit, elements[offsets[cI]]);
+                    return;
+                }
+            }
+        }
+    }
+
     export function ofLoci(loci: Loci) {
         const stats = create()
+        let hasPartitions = false;
         if (!Loci.isEmpty(loci)) {
-            for (const e of loci.elements) handleElement(stats, e)
+            for (const e of loci.elements) {
+                handleElement(stats, e)
+                if (!Unit.Traits.is(e.unit.traits, Unit.Trait.Patitioned)) {
+                    handleUnitChainsSimple(stats, e);
+                } else {
+                    hasPartitions = true;
+                }
+            }
+        }
+
+        if (hasPartitions) {
+            for (let i = 0, len = loci.elements.length; i < len; i++) {
+                const e = loci.elements[i];
+                if (!Unit.Traits.is(e.unit.traits, Unit.Trait.Patitioned)) continue;
+
+                const start = i;
+                while (i < len && loci.elements[i].unit.chainGroupId === e.unit.chainGroupId) {
+                    i++;
+                }
+                const end = i;
+                i--;
+                if (end - start === 1) {
+                    handleUnitChainsSimple(stats, e);
+                } else {
+                    handleUnitChainsPartitioned(stats, loci.elements, start, end);
+                }
+            }
         }
         return stats
     }

src/mol-theme/label.ts

@@ -52,29 +52,29 @@ function getResidueCount(unit: Unit.Atomic) {
 }
 
 export function structureElementStatsLabel(stats: StructureElement.Stats, countsOnly = false): string {
-    const { unitCount, residueCount, conformationCount, elementCount } = stats
+    const { chainCount, residueCount, conformationCount, elementCount } = stats
 
-    if (!countsOnly && elementCount === 1 && residueCount === 0 && unitCount === 0) {
+    if (!countsOnly && elementCount === 1 && residueCount === 0 && chainCount === 0) {
         return elementLabel(stats.firstElementLoc, { granularity: 'element' })
-    } else if (!countsOnly && elementCount === 0 && residueCount === 1 && unitCount === 0) {
+    } else if (!countsOnly && elementCount === 0 && residueCount === 1 && chainCount === 0) {
         return elementLabel(stats.firstResidueLoc, { granularity: 'residue' })
-    } else if (!countsOnly && elementCount === 0 && residueCount === 0 && unitCount === 1) {
-        const { unit } = stats.firstUnitLoc
+    } else if (!countsOnly && elementCount === 0 && residueCount === 0 && chainCount === 1) {
+        const { unit } = stats.firstChainLoc
         const granularity = (Unit.isAtomic(unit) && getResidueCount(unit) === 1) ? 'residue' : 'chain'
-        return elementLabel(stats.firstUnitLoc, { granularity })
+        return elementLabel(stats.firstChainLoc, { granularity })
     } else if (!countsOnly) {
         const label: string[] = []
         let hidePrefix = false;
-        if (unitCount > 0) {
-            label.push(unitCount === 1 ? elementLabel(stats.firstUnitLoc, { granularity: 'chain' }) : otherLabel(unitCount, stats.firstElementLoc, 'chain', false))
+        if (chainCount > 0) {
+            label.push(chainCount === 1 ? elementLabel(stats.firstChainLoc, { granularity: 'chain' }) : otherLabel(chainCount, stats.firstChainLoc || stats.firstElementLoc, 'chain', false))
             hidePrefix = true;
         }
         if (residueCount > 0) {
-            label.push(residueCount === 1 ? elementLabel(stats.firstResidueLoc, { granularity: 'residue', hidePrefix }) : otherLabel(residueCount, stats.firstElementLoc, 'residue', hidePrefix))
+            label.push(residueCount === 1 ? elementLabel(stats.firstResidueLoc, { granularity: 'residue', hidePrefix }) : otherLabel(residueCount, stats.firstResidueLoc || stats.firstElementLoc, 'residue', hidePrefix))
             hidePrefix = true;
         }
         if (conformationCount > 0) {
-            label.push(conformationCount === 1 ? elementLabel(stats.firstConformationLoc, { granularity: 'conformation', hidePrefix }) : otherLabel(conformationCount, stats.firstElementLoc, 'conformation', hidePrefix))
+            label.push(conformationCount === 1 ? elementLabel(stats.firstConformationLoc, { granularity: 'conformation', hidePrefix }) : otherLabel(conformationCount, stats.firstConformationLoc || stats.firstElementLoc, 'conformation', hidePrefix))
             hidePrefix = true;
         }
         if (elementCount > 0) {
@@ -83,7 +83,7 @@ export function structureElementStatsLabel(stats: StructureElement.Stats, counts
         return label.join('<small> + </small>')
     } else {
         const label: string[] = []
-        if (unitCount > 0) label.push(countLabel(unitCount, 'Chain'))
+        if (chainCount > 0) label.push(countLabel(chainCount, 'Chain'))
         if (residueCount > 0) label.push(countLabel(residueCount, 'Residue'))
         if (conformationCount > 0) label.push(countLabel(conformationCount, 'Conformation'))
         if (elementCount > 0) label.push(countLabel(elementCount, 'Element'))