mol-plugin: optimized sequence marking

src/mol-data/int/impl/ordered-set.ts

@@ -288,6 +288,30 @@ export function forEach(set: OrderedSetImpl, f: (value: number, i: number, ctx:
     return ctx;
 }
 
+export function forEachSegment(set: OrderedSetImpl, segment: (v: number) => number, f: (value: number, segIndex: number, ctx: any) => void, ctx: any) {
+    if (I.is(set)) {
+        let sI = 0;
+        for (let i = I.min(set), _i = I.max(set); i <= _i; i++) {
+            const s = segment(i);
+            let endI = i + 1;
+            while (endI < _i && segment(endI) === s) endI++;
+            i = endI - 1;
+            f(s, sI, ctx);
+            sI++;
+        }
+    } else {
+        let sI = 0;
+        for (let i = 0, _i = set.length; i < _i; i++) {
+            const s = segment(set[i]);
+            let endI = i + 1;
+            while (endI < _i && segment(set[endI]) === s) endI++;
+            i = endI - 1;
+            f(s, sI, ctx);
+            sI++;
+        }
+    }
+    return ctx;
+}
 
 export function indexedIntersect(idxA: OrderedSetImpl, a: S, b: S): OrderedSetImpl {
     if (a === b) return idxA;

src/mol-data/int/ordered-set.ts

@@ -57,6 +57,10 @@ namespace OrderedSet {
         return Base.forEach(set as any, f as any, ctx);
     }
 
+    export function forEachSegment<T extends number, S extends number, Ctx>(set: OrderedSet<T>, segment: (v: T) => S, f: (v: S, sI: number, ctx: Ctx) => void, ctx?: Ctx): Ctx {
+        return Base.forEachSegment(set as any, segment as any, f as any, ctx);
+    }
+
     export function isInterval<T extends number = number>(set: OrderedSet<T>): set is Interval<T> {
         return Interval.is(set);
     }

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

@@ -43,6 +43,8 @@ export function Loci(structure: Structure, elements: ArrayLike<{ unit: Unit, ind
 }
 
 export namespace Loci {
+    export type Element = Loci['elements'][0]
+
     export function is(x: any): x is Loci {
         return !!x && x.kind === 'element-loci';
     }

src/mol-plugin/ui/sequence/chain.ts

@@ -9,6 +9,7 @@ import { SequenceWrapper, StructureUnit } from './wrapper';
 import { OrderedSet, Interval } from '../../../mol-data/int';
 import { Loci } from '../../../mol-model/loci';
 import { ColorNames } from '../../../mol-util/color/names';
+import { MarkerAction, applyMarkerAction } from '../../../mol-util/marker-action';
 
 export class ChainSequenceWrapper extends SequenceWrapper<StructureUnit> {
     private label: string
@@ -22,7 +23,7 @@ export class ChainSequenceWrapper extends SequenceWrapper<StructureUnit> {
         return ColorNames.black
     }
 
-    eachResidue(loci: Loci, apply: (set: OrderedSet) => boolean) {
+    mark(loci: Loci, action: MarkerAction) {
         let changed = false
         const { structure } = this.data
         if (StructureElement.Loci.is(loci)) {
@@ -33,14 +34,14 @@ export class ChainSequenceWrapper extends SequenceWrapper<StructureUnit> {
                 const indices = this.unitIndices.get(e.unit.id)
                 if (indices) {
                     if (OrderedSet.isSubset(indices, e.indices)) {
-                        if (apply(Interval.ofSingleton(0))) changed = true
+                        if (applyMarkerAction(this.markerArray, Interval.ofSingleton(0), action)) changed = true
                     }
                 }
             }
         } else if (Structure.isLoci(loci)) {
             if (!Structure.areRootsEquivalent(loci.structure, structure)) return false
 
-            if (apply(Interval.ofSingleton(0))) changed = true
+            if (applyMarkerAction(this.markerArray, Interval.ofSingleton(0), action)) changed = true
         }
         return changed
     }

src/mol-plugin/ui/sequence/element.ts

@@ -9,6 +9,7 @@ import { SequenceWrapper, StructureUnit } from './wrapper';
 import { OrderedSet, Interval } from '../../../mol-data/int';
 import { Loci } from '../../../mol-model/loci';
 import { ColorNames } from '../../../mol-util/color/names';
+import { MarkerAction, applyMarkerAction } from '../../../mol-util/marker-action';
 
 export class ElementSequenceWrapper extends SequenceWrapper<StructureUnit> {
     private unitIndices: Map<number, Interval<StructureElement.UnitIndex>>
@@ -20,7 +21,7 @@ export class ElementSequenceWrapper extends SequenceWrapper<StructureUnit> {
         return ColorNames.black
     }
 
-    eachResidue(loci: Loci, apply: (set: OrderedSet) => boolean) {
+    mark(loci: Loci, action: MarkerAction) {
         let changed = false
         const { structure, units } = this.data
         if (StructureElement.Loci.is(loci)) {
@@ -31,7 +32,7 @@ export class ElementSequenceWrapper extends SequenceWrapper<StructureUnit> {
                 const indices = this.unitIndices.get(e.unit.id)
                 if (indices) {
                     if (OrderedSet.isSubset(indices, e.indices)) {
-                        if (apply(e.indices)) changed = true
+                        if (applyMarkerAction(this.markerArray, e.indices, action)) changed = true
                     }
                 }
             }
@@ -40,7 +41,7 @@ export class ElementSequenceWrapper extends SequenceWrapper<StructureUnit> {
 
             for (let i = 0, il = units.length; i < il; ++i) {
                 const indices = this.unitIndices.get(units[i].id)!
-                if (apply(indices)) changed = true
+                if (applyMarkerAction(this.markerArray, indices, action)) changed = true
             }
         }
         return changed

src/mol-plugin/ui/sequence/hetero.ts

@@ -9,6 +9,7 @@ import { SequenceWrapper, StructureUnit } from './wrapper';
 import { OrderedSet, Segmentation, Interval, SortedArray } from '../../../mol-data/int';
 import { Loci } from '../../../mol-model/loci';
 import { ColorNames } from '../../../mol-util/color/names';
+import { MarkerAction, applyMarkerAction } from '../../../mol-util/marker-action';
 
 export class HeteroSequenceWrapper extends SequenceWrapper<StructureUnit> {
     private readonly unitMap: Map<number, Unit>
@@ -24,7 +25,7 @@ export class HeteroSequenceWrapper extends SequenceWrapper<StructureUnit> {
         return ColorNames.black
     }
 
-    eachResidue(loci: Loci, apply: (set: OrderedSet) => boolean) {
+    mark(loci: Loci, action: MarkerAction) {
         let changed = false
         const { structure } = this.data
         if (StructureElement.Loci.is(loci)) {
@@ -37,14 +38,14 @@ export class HeteroSequenceWrapper extends SequenceWrapper<StructureUnit> {
                     const { index: residueIndex } = e.unit.model.atomicHierarchy.residueAtomSegments
                     OrderedSet.forEach(e.indices, v => {
                         const seqIdx = this.sequenceIndices.get(residueIndex[unit.elements[v]])
-                        if (seqIdx !== undefined && apply(Interval.ofSingleton(seqIdx))) changed = true
+                        if (seqIdx !== undefined && applyMarkerAction(this.markerArray, Interval.ofSingleton(seqIdx), action)) changed = true
                     })
                 }
             }
         } else if (Structure.isLoci(loci)) {
             if (!Structure.areRootsEquivalent(loci.structure, structure)) return false
 
-            if (apply(Interval.ofBounds(0, this.length))) changed = true
+            if (applyMarkerAction(this.markerArray, Interval.ofBounds(0, this.length), action)) changed = true
         }
         return changed
     }

src/mol-plugin/ui/sequence/polymer.ts

@@ -4,13 +4,14 @@
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
  */
 
-import { StructureSelection, StructureQuery, Structure, Queries, StructureProperties as SP, StructureElement, Unit, ElementIndex } from '../../../mol-model/structure';
+import { StructureSelection, StructureQuery, Structure, Queries, StructureProperties as SP, StructureElement, Unit } from '../../../mol-model/structure';
 import { SequenceWrapper, StructureUnit } from './wrapper';
 import { OrderedSet, Interval, SortedArray } from '../../../mol-data/int';
 import { Loci } from '../../../mol-model/loci';
 import { Sequence } from '../../../mol-model/sequence';
 import { MissingResidues } from '../../../mol-model/structure/model/properties/common';
 import { ColorNames } from '../../../mol-util/color/names';
+import { MarkerAction, applyMarkerAction, applyMarkerActionAtPosition } from '../../../mol-util/marker-action';
 
 export class PolymerSequenceWrapper extends SequenceWrapper<StructureUnit> {
     private readonly unitMap: Map<number, Unit>
@@ -28,13 +29,14 @@ export class PolymerSequenceWrapper extends SequenceWrapper<StructureUnit> {
     residueLabel(seqIdx: number) {
         return this.sequence.label.value(seqIdx)
     }
+
     residueColor(seqIdx: number) {
         return this.missing.has(this.modelNum, this.asymId, this.seqId(seqIdx))
             ? ColorNames.grey
             : ColorNames.black
     }
 
-    eachResidue(loci: Loci, apply: (set: OrderedSet) => boolean) {
+    mark(loci: Loci, action: MarkerAction): boolean {
         let changed = false
         const { structure } = this.data
         if (StructureElement.Loci.is(loci)) {
@@ -43,16 +45,17 @@ export class PolymerSequenceWrapper extends SequenceWrapper<StructureUnit> {
 
             const { offset } = this.sequence
             for (const e of loci.elements) {
-                if (this.unitMap.has(e.unit.id)) {
-                    OrderedSet.forEach(e.indices, v => {
-                        if (apply(getSeqIndices(e.unit, e.unit.elements[v], offset))) changed = true
-                    })
+                if (!this.unitMap.has(e.unit.id)) continue;
+
+                if (Unit.isAtomic(e.unit)) {
+                    changed = applyMarkerAtomic(e, action, this.markerArray, offset) || changed;
+                } else {
+                    changed = applyMarkerCoarse(e, action, this.markerArray, offset) || changed;
                 }
             }
         } else if (Structure.isLoci(loci)) {
             if (!Structure.areRootsEquivalent(loci.structure, structure)) return false
-
-            if (apply(this.observed)) changed = true
+            if (applyMarkerAction(this.markerArray, this.observed, action)) changed = true
         }
         return changed
     }
@@ -109,22 +112,29 @@ function createResidueQuery(chainGroupId: number, operatorName: string, label_se
     });
 }
 
-function getSeqIndices(unit: Unit, element: ElementIndex, offset: number): Interval {
-    const { model } = unit
-    switch (unit.kind) {
-        case Unit.Kind.Atomic:
-            const residueIndex = model.atomicHierarchy.residueAtomSegments.index[element]
-            const seqId = model.atomicHierarchy.residues.label_seq_id.value(residueIndex)
-            return Interval.ofSingleton(seqId - 1 - offset)
-        case Unit.Kind.Spheres:
-            return Interval.ofRange(
-                model.coarseHierarchy.spheres.seq_id_begin.value(element) - 1 - offset,
-                model.coarseHierarchy.spheres.seq_id_end.value(element) - 1 - offset
-            )
-        case Unit.Kind.Gaussians:
-            return Interval.ofRange(
-                model.coarseHierarchy.gaussians.seq_id_begin.value(element) - 1 - offset,
-                model.coarseHierarchy.gaussians.seq_id_end.value(element) - 1 - offset
-            )
-    }
+function applyMarkerAtomic(e: StructureElement.Loci.Element, action: MarkerAction, markerArray: Uint8Array, offset: number) {
+    const { model, elements } = e.unit;
+    const { index } = model.atomicHierarchy.residueAtomSegments;
+    const { label_seq_id } = model.atomicHierarchy.residues;
+
+    let changed = false;
+    OrderedSet.forEachSegment(e.indices, i => index[elements[i]], rI => {
+        const seqId = label_seq_id.value(rI)
+        changed = applyMarkerActionAtPosition(markerArray, seqId - 1 - offset, action) || changed;
+    });
+    return changed;
+}
+
+function applyMarkerCoarse(e: StructureElement.Loci.Element, action: MarkerAction, markerArray: Uint8Array, offset: number) {
+    const { model, elements } = e.unit;
+    const begin = Unit.isSpheres(e.unit) ? model.coarseHierarchy.spheres.seq_id_begin : model.coarseHierarchy.gaussians.seq_id_begin;
+    const end = Unit.isSpheres(e.unit) ? model.coarseHierarchy.spheres.seq_id_end : model.coarseHierarchy.gaussians.seq_id_end;
+
+    let changed = false;
+    OrderedSet.forEach(e.indices, i => {
+        for (let s = begin.value(elements[i]) - 1 - offset, e = end.value(elements[i]) - 1 - offset; s <= e; s++) {
+            changed = applyMarkerActionAtPosition(markerArray, s, action) || changed;
+        }
+    });
+    return changed;
 }

src/mol-plugin/ui/sequence/wrapper.ts

@@ -4,7 +4,7 @@
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
  */
 
-import { OrderedSet, Interval } from '../../../mol-data/int';
+import { Interval } from '../../../mol-data/int';
 import { Loci, isEveryLoci } from '../../../mol-model/loci';
 import { MarkerAction, applyMarkerAction } from '../../../mol-util/marker-action';
 import { StructureElement, Structure, Unit } from '../../../mol-model/structure';
@@ -18,16 +18,15 @@ abstract class SequenceWrapper<D> {
     abstract residueLabel(seqIdx: number): string
     abstract residueColor(seqIdx: number): Color
 
-    abstract eachResidue(loci: Loci, apply: (set: OrderedSet) => boolean): boolean
     abstract getLoci(seqIdx: number): StructureElement.Loci
 
+    abstract mark(loci: Loci, action: MarkerAction): boolean;
+
     markResidue(loci: Loci, action: MarkerAction) {
         if (isEveryLoci(loci)) {
             return applyMarkerAction(this.markerArray, Interval.ofLength(this.length), action)
         } else {
-            return this.eachResidue(loci, (set: OrderedSet) => {
-                return applyMarkerAction(this.markerArray, set, action)
-            })
+            return this.mark(loci, action);
         }
     }
 

src/mol-repr/structure/visual/util/element.ts

@@ -94,7 +94,7 @@ export function eachElement(loci: Loci, structureGroup: StructureGroup, apply: (
                 for (let i = 0, _i = e.indices.length; i < _i; i++) {
                     const start = e.indices[i];
                     let endI = i + 1;
-                    while (endI < _i && e.indices[i] === start) endI++;
+                    while (endI < _i && e.indices[endI] === start) endI++;
                     i = endI - 1;
                     const end = e.indices[i];
                     changed = apply(Interval.ofRange(start, end)) || changed;

src/mol-repr/structure/visual/util/polymer.ts

@@ -138,6 +138,7 @@ export function eachAtomicUnitTracedElement(intervalOffset: number, groupSize: n
     if (Interval.is(e.indices)) {
         if (Interval.start(e.indices) === 0 && Interval.end(e.indices) === e.unit.elements.length) {
             // full unit here
+            console.log('full unit');
             changed = apply(Interval.ofBounds(intervalOffset, intervalOffset + groupSize)) || changed;
         } else {
             let r1 = resIndex[elements[Interval.min(e.indices)]];
@@ -147,10 +148,11 @@ export function eachAtomicUnitTracedElement(intervalOffset: number, groupSize: n
     } else {
         const { indices } = e;
 
+        
         for (let i = 0, _i = indices.length; i < _i; i++) {
             const r1 = resIndex[elements[indices[i]]];
             let r2 = r1;
-
+            
             let endI = i + 1;
             while (endI < _i) {
                 const _r = resIndex[elements[indices[endI]]];
@@ -158,6 +160,7 @@ export function eachAtomicUnitTracedElement(intervalOffset: number, groupSize: n
                 r2 = _r;
                 endI++;
             }
+            console.log('indices', i === _i);
             i = endI - 1;
             changed = tryApplyResidueInterval(tracedElements, traceElementIndex, apply, r1, r2) || changed;
         }
@@ -194,7 +197,7 @@ export function eachPolymerElement(loci: Loci, structureGroup: StructureGroup, a
                 for (let i = 0, _i = e.indices.length; i < _i; i++) {
                     const start = e.indices[i];
                     let endI = i + 1;
-                    while (endI < _i && e.indices[i] === start) endI++;
+                    while (endI < _i && e.indices[endI] === start) endI++;
                     i = endI - 1;
                     const end = e.indices[i];
                     changed = apply(Interval.ofRange(start, end)) || changed;

src/mol-util/marker-action.ts

@@ -15,7 +15,7 @@ export enum MarkerAction {
     Clear
 }
 
-function applyAction(array: Uint8Array, i: number, action: MarkerAction) {
+export function applyMarkerActionAtPosition(array: Uint8Array, i: number, action: MarkerAction) {
     let v = array[i];
     switch (action) {
         case MarkerAction.Highlight:
@@ -54,11 +54,11 @@ export function applyMarkerAction(array: Uint8Array, set: OrderedSet, action: Ma
     let changed = false;
     if (Interval.is(set)) {
         for (let i = Interval.start(set), _i = Interval.end(set); i < _i; i++) {
-            changed = applyAction(array, i, action) || changed;
+            changed = applyMarkerActionAtPosition(array, i, action) || changed;
         }
     } else {
         for (let i = 0, _i = set.length; i < _i; i++) {
-            changed = applyAction(array, set[i], action) || changed;
+            changed = applyMarkerActionAtPosition(array, set[i], action) || changed;
         }
     }
     return changed;