more specialized location-iterators

src/mol-geo/representation/structure/visual/nucleotide-block-mesh.ts

@@ -9,14 +9,15 @@ import { UnitsVisual } from '..';
 import { RuntimeContext } from 'mol-task'
 import { Mesh } from '../../../mesh/mesh';
 import { MeshBuilder } from '../../../mesh/mesh-builder';
-import { getElementLoci, markElement, StructureElementIterator } from './util/element';
+import { getElementLoci, markElement } from './util/element';
 import { Vec3, Mat4 } from 'mol-math/linear-algebra';
-import { Segmentation, SortedArray } from 'mol-data/int';
+import { Segmentation } from 'mol-data/int';
 import { MoleculeType, isNucleic, isPurinBase, isPyrimidineBase } from 'mol-model/structure/model/types';
 import { getElementIndexForAtomId, getElementIndexForAtomRole } from 'mol-model/structure/util';
 import { DefaultUnitsMeshProps, UnitsMeshVisual } from '../units-visual';
 import { addCylinder } from '../../../mesh/builder/cylinder';
 import { Box } from '../../../primitive/box';
+import { NucleotideLocationIterator } from './util/nucleotide';
 
 const p1 = Vec3.zero()
 const p2 = Vec3.zero()
@@ -84,7 +85,7 @@ async function createNucleotideBlockMesh(ctx: RuntimeContext, unit: Unit, props:
 
                 if (idx5 !== -1 && idx6 !== -1) {
                     pos(idx5, p5); pos(idx6, p6)
-                    builder.setGroup(SortedArray.findPredecessorIndex(elements, idx6))
+                    builder.setGroup(i)
                     addCylinder(builder, p5, p6, 1, { radiusTop: 0.2, radiusBottom: 0.2 })
                     if (idx1 !== -1 && idx2 !== -1 && idx3 !== -1 && idx4 !== -1) {
                         pos(idx1, p1); pos(idx2, p2); pos(idx3, p3); pos(idx4, p4);
@@ -98,7 +99,6 @@ async function createNucleotideBlockMesh(ctx: RuntimeContext, unit: Unit, props:
                         builder.add(t, box)
                     }
                 }
-            }
 
                 if (i % 10000 === 0 && ctx.shouldUpdate) {
                     await ctx.update({ message: 'Nucleotide block mesh', current: i });
@@ -106,6 +106,7 @@ async function createNucleotideBlockMesh(ctx: RuntimeContext, unit: Unit, props:
                 ++i
             }
         }
+    }
 
     return builder.getMesh()
 }
@@ -119,7 +120,7 @@ export function NucleotideBlockVisual(): UnitsVisual<NucleotideBlockProps> {
     return UnitsMeshVisual<NucleotideBlockProps>({
         defaultProps: DefaultNucleotideBlockProps,
         createMesh: createNucleotideBlockMesh,
-        createLocationIterator: StructureElementIterator.fromGroup,
+        createLocationIterator: NucleotideLocationIterator.fromGroup,
         getLoci: getElementLoci,
         mark: markElement,
         setUpdateState: () => {}

src/mol-geo/representation/structure/visual/polymer-direction-wedge.ts

@@ -7,15 +7,14 @@
 import { Unit } from 'mol-model/structure';
 import { UnitsVisual } from '..';
 import { RuntimeContext } from 'mol-task'
-import { markElement, getElementLoci, StructureElementIterator } from './util/element';
+import { markElement, getElementLoci } from './util/element';
 import { Mesh } from '../../../mesh/mesh';
 import { MeshBuilder } from '../../../mesh/mesh-builder';
-import { getPolymerElementCount, PolymerTraceIterator, createCurveSegmentState, interpolateCurveSegment } from './util/polymer';
+import { getPolymerElementCount, PolymerTraceIterator, createCurveSegmentState, interpolateCurveSegment, PolymerLocationIterator } from './util/polymer';
 import { Vec3, Mat4 } from 'mol-math/linear-algebra';
 import { SecondaryStructureType, MoleculeType } from 'mol-model/structure/model/types';
 import { DefaultUnitsMeshProps, UnitsMeshVisual } from '../units-visual';
 import { SizeThemeProps, SizeTheme } from 'mol-view/theme/size';
-import { OrderedSet } from 'mol-data/int';
 import { Wedge } from '../../../primitive/wedge';
 
 const t = Mat4.identity()
@@ -51,7 +50,7 @@ async function createPolymerDirectionWedgeMesh(ctx: RuntimeContext, unit: Unit,
     const polymerTraceIt = PolymerTraceIterator(unit)
     while (polymerTraceIt.hasNext) {
         const v = polymerTraceIt.move()
-        builder.setGroup(OrderedSet.findPredecessorIndex(unit.elements, v.center.element))
+        builder.setGroup(i)
 
         const isNucleic = v.moleculeType === MoleculeType.DNA || v.moleculeType === MoleculeType.RNA
         const isSheet = SecondaryStructureType.is(v.secStrucType, SecondaryStructureType.Flag.Beta)
@@ -95,7 +94,7 @@ export function PolymerDirectionVisual(): UnitsVisual<PolymerDirectionProps> {
     return UnitsMeshVisual<PolymerDirectionProps>({
         defaultProps: DefaultPolymerDirectionProps,
         createMesh: createPolymerDirectionWedgeMesh,
-        createLocationIterator: StructureElementIterator.fromGroup,
+        createLocationIterator: PolymerLocationIterator.fromGroup,
         getLoci: getElementLoci,
         mark: markElement,
         setUpdateState: () => {}

src/mol-geo/representation/structure/visual/polymer-gap-cylinder.ts

@@ -9,8 +9,8 @@ import { UnitsVisual } from '..';
 import { RuntimeContext } from 'mol-task'
 import { Mesh } from '../../../mesh/mesh';
 import { MeshBuilder } from '../../../mesh/mesh-builder';
-import { getPolymerGapCount, PolymerGapIterator } from './util/polymer';
-import { getElementLoci, markElement, StructureElementIterator } from './util/element';
+import { getPolymerGapCount, PolymerGapIterator, PolymerGapLocationIterator } from './util/polymer';
+import { getElementLoci, markElement } from './util/element';
 import { Vec3 } from 'mol-math/linear-algebra';
 import { UnitsMeshVisual, DefaultUnitsMeshProps } from '../units-visual';
 import { SizeThemeProps, SizeTheme } from 'mol-view/theme/size';
@@ -35,11 +35,9 @@ async function createPolymerGapCylinderMesh(ctx: RuntimeContext, unit: Unit, pro
     const vertexCountEstimate = segmentCount * radialSegments * 2 * polymerGapCount * 2
     const builder = MeshBuilder.create(vertexCountEstimate, vertexCountEstimate / 10, mesh)
 
-    const { elements } = unit
     const pos = unit.conformation.invariantPosition
     const pA = Vec3.zero()
     const pB = Vec3.zero()
-    const l = StructureElement.create(unit)
     const cylinderProps: CylinderProps = {
         radiusTop: 1, radiusBottom: 1, topCap: true, bottomCap: true, radialSegments
     }
@@ -49,30 +47,26 @@ async function createPolymerGapCylinderMesh(ctx: RuntimeContext, unit: Unit, pro
     while (polymerGapIt.hasNext) {
         const { centerA, centerB } = polymerGapIt.move()
         if (centerA.element === centerB.element) {
-            builder.setGroup(centerA.element)
-            pos(elements[centerA.element], pA)
+            builder.setGroup(i)
+            pos(centerA.element, pA)
             addSphere(builder, pA, 0.6, 0)
         } else {
-            const elmA = elements[centerA.element]
-            const elmB = elements[centerB.element]
-            pos(elmA, pA)
-            pos(elmB, pB)
+            pos(centerA.element, pA)
+            pos(centerB.element, pB)
 
-            l.element = elmA
-            cylinderProps.radiusTop = cylinderProps.radiusBottom = sizeTheme.size(l)
-            builder.setGroup(centerA.element)
+            cylinderProps.radiusTop = cylinderProps.radiusBottom = sizeTheme.size(centerA)
+            builder.setGroup(i)
             addFixedCountDashedCylinder(builder, pA, pB, 0.5, segmentCount, cylinderProps)
 
-            l.element = elmB
-            cylinderProps.radiusTop = cylinderProps.radiusBottom = sizeTheme.size(l)
-            builder.setGroup(centerB.element)
+            cylinderProps.radiusTop = cylinderProps.radiusBottom = sizeTheme.size(centerB)
+            builder.setGroup(i + 1)
             addFixedCountDashedCylinder(builder, pB, pA, 0.5, segmentCount, cylinderProps)
         }
 
         if (i % 10000 === 0 && ctx.shouldUpdate) {
             await ctx.update({ message: 'Gap mesh', current: i, max: polymerGapCount });
         }
-        ++i
+        i += 2
     }
 
     return builder.getMesh()
@@ -88,7 +82,7 @@ export function PolymerGapVisual(): UnitsVisual<PolymerGapProps> {
     return UnitsMeshVisual<PolymerGapProps>({
         defaultProps: DefaultPolymerGapProps,
         createMesh: createPolymerGapCylinderMesh,
-        createLocationIterator: StructureElementIterator.fromGroup,
+        createLocationIterator: PolymerGapLocationIterator.fromGroup,
         getLoci: getElementLoci,
         mark: markElement,
         setUpdateState: () => {}

src/mol-geo/representation/structure/visual/polymer-trace-mesh.ts

@@ -7,14 +7,13 @@
 import { Unit } from 'mol-model/structure';
 import { UnitsVisual, MeshUpdateState } from '..';
 import { RuntimeContext } from 'mol-task'
-import { markElement, getElementLoci, StructureElementIterator } from './util/element';
+import { markElement, getElementLoci } from './util/element';
 import { Mesh } from '../../../mesh/mesh';
 import { MeshBuilder } from '../../../mesh/mesh-builder';
-import { getPolymerElementCount, PolymerTraceIterator, createCurveSegmentState, interpolateCurveSegment } from './util/polymer';
+import { getPolymerElementCount, PolymerTraceIterator, createCurveSegmentState, interpolateCurveSegment, PolymerLocationIterator } from './util/polymer';
 import { SecondaryStructureType, isNucleic } from 'mol-model/structure/model/types';
 import { UnitsMeshVisual, DefaultUnitsMeshProps } from '../units-visual';
 import { SizeThemeProps, SizeTheme } from 'mol-view/theme/size';
-import { OrderedSet } from 'mol-data/int';
 import { addSheet } from '../../../mesh/builder/sheet';
 import { addTube } from '../../../mesh/builder/tube';
 
@@ -45,7 +44,7 @@ async function createPolymerTraceMesh(ctx: RuntimeContext, unit: Unit, props: Po
     const polymerTraceIt = PolymerTraceIterator(unit)
     while (polymerTraceIt.hasNext) {
         const v = polymerTraceIt.move()
-        builder.setGroup(OrderedSet.findPredecessorIndex(unit.elements, v.center.element))
+        builder.setGroup(i)
 
         const isNucleicType = isNucleic(v.moleculeType)
         const isSheet = SecondaryStructureType.is(v.secStrucType, SecondaryStructureType.Flag.Beta)
@@ -95,7 +94,7 @@ export function PolymerTraceVisual(): UnitsVisual<PolymerTraceProps> {
     return UnitsMeshVisual<PolymerTraceProps>({
         defaultProps: DefaultPolymerTraceProps,
         createMesh: createPolymerTraceMesh,
-        createLocationIterator: StructureElementIterator.fromGroup,
+        createLocationIterator: PolymerLocationIterator.fromGroup,
         getLoci: getElementLoci,
         mark: markElement,
         setUpdateState: (state: MeshUpdateState, newProps: PolymerTraceProps, currentProps: PolymerTraceProps) => {

src/mol-geo/representation/structure/visual/util/nucleotide.ts

+/**
+ * Copyright (c) 2018 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author Alexander Rose <alexander.rose@weirdbyte.de>
+ */
+import { Unit, ElementIndex, StructureElement } from 'mol-model/structure';
+import { LocationIterator } from '../../../../util/location-iterator';
+import { Segmentation } from 'mol-data/int';
+import { isNucleic, MoleculeType } from 'mol-model/structure/model/types';
+import { getElementIndexForAtomRole } from 'mol-model/structure/util';
+
+export function getNucleotideElementIndices(unit: Unit) {
+    const indices: ElementIndex[] = []
+    const { elements, model } = unit
+    const { chemicalComponentMap } = model.properties
+    const { chainAtomSegments, residueAtomSegments, residues } = model.atomicHierarchy
+    const { label_comp_id } = residues
+    const chainIt = Segmentation.transientSegments(chainAtomSegments, elements)
+    const residueIt = Segmentation.transientSegments(residueAtomSegments, elements)
+    while (chainIt.hasNext) {
+        residueIt.setSegment(chainIt.move());
+
+        while (residueIt.hasNext) {
+            const { index } = residueIt.move();
+            const cc = chemicalComponentMap.get(label_comp_id.value(index))
+            const moleculeType = cc ? cc.moleculeType : MoleculeType.unknown
+
+            if (isNucleic(moleculeType)) {
+                const elementIndex = getElementIndexForAtomRole(model, index, 'trace')
+                indices.push(elementIndex === -1 ? residueAtomSegments.offsets[index] : elementIndex)
+            }
+        }
+    }
+    return indices
+}
+
+export namespace NucleotideLocationIterator {
+    export function fromGroup(group: Unit.SymmetryGroup): LocationIterator {
+        const nucleotideElementIndices = getNucleotideElementIndices(group.units[0])
+        const groupCount = nucleotideElementIndices.length
+        const instanceCount = group.units.length
+        const location = StructureElement.create()
+        const getLocation = (groupIndex: number, instanceIndex: number) => {
+            const unit = group.units[instanceIndex]
+            location.unit = unit
+            location.element = nucleotideElementIndices[groupIndex]
+            return location
+        }
+        return LocationIterator(groupCount, instanceCount, getLocation)
+    }
+}
\ No newline at end of file

src/mol-geo/representation/structure/visual/util/polymer.ts

@@ -4,9 +4,12 @@
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
  */
 
-import { Unit, ElementIndex } from 'mol-model/structure';
+import { Unit, ElementIndex, StructureElement } from 'mol-model/structure';
 import { Segmentation, OrderedSet, Interval } from 'mol-data/int';
 import SortedRanges from 'mol-data/int/sorted-ranges';
+import { LocationIterator } from '../../../../util/location-iterator';
+import { getElementIndexForAtomRole } from 'mol-model/structure/util';
+import { PolymerGapIterator } from './polymer/gap-iterator';
 
 export * from './polymer/backbone-iterator'
 export * from './polymer/gap-iterator'
@@ -57,6 +60,54 @@ export function getPolymerElementCount(unit: Unit) {
     return count
 }
 
+export function getPolymerElementIndices(unit: Unit) {
+    const indices: ElementIndex[] = []
+    const { elements, model } = unit
+    const { residueAtomSegments } = unit.model.atomicHierarchy
+    const polymerIt = SortedRanges.transientSegments(getPolymerRanges(unit), elements)
+    switch (unit.kind) {
+        case Unit.Kind.Atomic:
+            const residueIt = Segmentation.transientSegments(residueAtomSegments, elements)
+            while (polymerIt.hasNext) {
+                const polymerSegment = polymerIt.move()
+                residueIt.setSegment(polymerSegment)
+                while (residueIt.hasNext) {
+                    const residueSegment = residueIt.move()
+                    const { start, end, index } = residueSegment
+                    if (OrderedSet.areIntersecting(Interval.ofBounds(elements[start], elements[end - 1]), elements)) {
+                        const elementIndex = getElementIndexForAtomRole(model, index, 'trace')
+                        indices.push(elementIndex === -1 ? residueAtomSegments.offsets[index] : elementIndex)
+                    }
+                }
+            }
+            break
+        case Unit.Kind.Spheres:
+        case Unit.Kind.Gaussians:
+            while (polymerIt.hasNext) {
+                const { start, end } = polymerIt.move()
+                for (let i = start; i < end; ++i) { indices.push(elements[i]) }
+            }
+            break
+    }
+    return indices
+}
+
+export namespace PolymerLocationIterator {
+    export function fromGroup(group: Unit.SymmetryGroup): LocationIterator {
+        const polymerElementIndices = getPolymerElementIndices(group.units[0])
+        const groupCount = polymerElementIndices.length
+        const instanceCount = group.units.length
+        const location = StructureElement.create()
+        const getLocation = (groupIndex: number, instanceIndex: number) => {
+            const unit = group.units[instanceIndex]
+            location.unit = unit
+            location.element = polymerElementIndices[groupIndex]
+            return location
+        }
+        return LocationIterator(groupCount, instanceCount, getLocation)
+    }
+}
+
 export function getPolymerGapCount(unit: Unit) {
     let count = 0
     const { elements } = unit
@@ -67,3 +118,29 @@ export function getPolymerGapCount(unit: Unit) {
     }
     return count
 }
+
+export function getPolymerGapElementIndices(unit: Unit) {
+    const indices: ElementIndex[] = []
+    const polymerGapIt = PolymerGapIterator(unit)
+    while (polymerGapIt.hasNext) {
+        const { centerA, centerB } = polymerGapIt.move()
+        indices.push(centerA.element, centerB.element)
+    }
+    return indices
+}
+
+export namespace PolymerGapLocationIterator {
+    export function fromGroup(group: Unit.SymmetryGroup): LocationIterator {
+        const polymerGapElementIndices = getPolymerGapElementIndices(group.units[0])
+        const groupCount = polymerGapElementIndices.length
+        const instanceCount = group.units.length
+        const location = StructureElement.create()
+        const getLocation = (groupIndex: number, instanceIndex: number) => {
+            const unit = group.units[instanceIndex]
+            location.unit = unit
+            location.element = polymerGapElementIndices[groupIndex]
+            return location
+        }
+        return LocationIterator(groupCount, instanceCount, getLocation)
+    }
+}
\ No newline at end of file

src/mol-geo/representation/structure/visual/util/polymer/gap-iterator.ts

@@ -5,7 +5,6 @@
  */
 
 import { Unit, StructureElement, ElementIndex, ResidueIndex } from 'mol-model/structure';
-import { SortedArray } from 'mol-data/int';
 import { AtomRole } from 'mol-model/structure/model/types';
 import Iterator from 'mol-data/iterator';
 import SortedRanges from 'mol-data/int/sorted-ranges';
@@ -40,13 +39,8 @@ export class AtomicPolymerGapIterator implements Iterator<PolymerGapPair> {
     hasNext: boolean = false;
 
     private getElementIndex(residueIndex: ResidueIndex, atomRole: AtomRole) {
-        const index = getElementIndexForAtomRole(this.unit.model, residueIndex, atomRole)
-        // TODO handle case when it returns -1
-        const elementIndex = SortedArray.indexOf(this.unit.elements, index) as ElementIndex
-        if (elementIndex === -1) {
-            console.log('-1', residueIndex, atomRole, index)
-        }
-        return elementIndex === -1 ? 0 as ElementIndex : elementIndex
+        const elementIndex = getElementIndexForAtomRole(this.unit.model, residueIndex, atomRole)
+        return elementIndex === -1 ? this.unit.model.atomicHierarchy.residueAtomSegments.offsets[residueIndex] : elementIndex
     }
 
     move() {

src/mol-geo/util/color-data.ts

@@ -96,9 +96,7 @@ export function createGroupColor(locationIt: LocationIterator, color: LocationCo
 export function createGroupInstanceColor(locationIt: LocationIterator, color: LocationColor, colorData?: ColorData): ColorData {
     const { groupCount, instanceCount } = locationIt
     const count = instanceCount * groupCount
-    console.log(count, instanceCount, groupCount)
     const colors = colorData && colorData.tColor.ref.value.array.length >= count * 3 ? colorData.tColor.ref.value : createTextureImage(count, 3)
-    console.log(colors.array.length / 3, count)
     while (locationIt.hasNext) {
         const { location, isSecondary, index } = locationIt.move()
         Color.toArray(color(location, isSecondary), colors.array, index * 3)

src/mol-view/theme/color/unit-index.ts

@@ -20,7 +20,6 @@ export function UnitIndexColorTheme(props: ColorThemeProps): ColorTheme {
 
         color = (location: Location): Color => {
             if (StructureElement.isLocation(location)) {
-                // console.log(location.unit.id, Unit.findUnitById(location.unit.id, units), units.length, units)
                 return scale.color(Unit.findUnitById(location.unit.id, units))
             } else if (Link.isLocation(location)) {
                 return scale.color(Unit.findUnitById(location.aUnit.id, units))