polymer iterator cleanup

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

@@ -23,6 +23,7 @@ import { createMeshValues, updateMeshValues, updateRenderableState, createRender
 import { MeshBuilder } from '../../../shape/mesh-builder';
 import { getPolymerElementCount, PolymerBackboneIterator } from './util/polymer';
 import { getElementLoci, markElement } from './util/element';
+import { Vec3 } from 'mol-math/linear-algebra';
 
 async function createPolymerBackboneCylinderMesh(ctx: RuntimeContext, unit: Unit, mesh?: Mesh) {
     const polymerElementCount = getPolymerElementCount(unit)
@@ -32,15 +33,22 @@ async function createPolymerBackboneCylinderMesh(ctx: RuntimeContext, unit: Unit
     // TODO better vertex count estimates
     const builder = MeshBuilder.create(polymerElementCount * 30, polymerElementCount * 30 / 2, mesh)
 
+    const { elements } = unit
+    const pos = unit.conformation.invariantPosition
+    const pA = Vec3.zero()
+    const pB = Vec3.zero()
+
     let i = 0
     const polymerBackboneIt = PolymerBackboneIterator(unit)
     while (polymerBackboneIt.hasNext) {
-        const { indexA, indexB, posA, posB } = polymerBackboneIt.move()
-        builder.setId(indexA)
         // TODO size theme
-        builder.addCylinder(posA, posB, 0.5, { radiusTop: 0.2, radiusBottom: 0.2 })
-        builder.setId(indexB)
-        builder.addCylinder(posB, posA, 0.5, { radiusTop: 0.2, radiusBottom: 0.2 })
+        const { centerA, centerB } = polymerBackboneIt.move()
+        pos(centerA.element, pA)
+        pos(centerB.element, pB)
+        builder.setId(elements[centerA.element])
+        builder.addCylinder(pA, pB, 0.5, { radiusTop: 0.2, radiusBottom: 0.2 })
+        builder.setId(elements[centerB.element])
+        builder.addCylinder(pB, pA, 0.5, { radiusTop: 0.2, radiusBottom: 0.2 })
 
         if (i % 10000 === 0 && ctx.shouldUpdate) {
             await ctx.update({ message: 'Backbone mesh', current: i, max: polymerElementCount });

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

@@ -73,7 +73,7 @@ function interpolateNormals(controlPoints: Helpers.NumberArray, tangentVectors:
         //     console.warn(i, 'flip compared to prev', radToDeg(Math.acos(Vec3.dot(prevNormal, normalVec))))
         // }
         Vec3.copy(prevNormal, normalVec)
-        
+
         Vec3.normalize(binormalVec, Vec3.cross(binormalVec, tangentVec, normalVec))
         Vec3.toArray(binormalVec, binormalVectors, i * 3)
     }
@@ -84,6 +84,8 @@ async function createPolymerTraceMesh(ctx: RuntimeContext, unit: Unit, mesh?: Me
     console.log('polymerElementCount', polymerElementCount)
     if (!polymerElementCount) return Mesh.createEmpty(mesh)
 
+    const { elements } = unit
+
     // TODO better vertex count estimates
     const builder = MeshBuilder.create(polymerElementCount * 30, polymerElementCount * 30 / 2, mesh)
     const linearSegments = 12
@@ -108,7 +110,7 @@ async function createPolymerTraceMesh(ctx: RuntimeContext, unit: Unit, mesh?: Me
     const polymerTraceIt = PolymerTraceIterator(unit)
     while (polymerTraceIt.hasNext) {
         const v = polymerTraceIt.move()
-        builder.setId(v.index)
+        builder.setId(elements[v.center.element])
 
         for (let j = 0; j <= linearSegments; ++j) {
             const t = j * 1.0 / linearSegments;

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

@@ -4,15 +4,14 @@
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
  */
 
-import { Unit, StructureElement, StructureProperties, Model, ElementIndex } from 'mol-model/structure';
+import { Unit, StructureElement, Model, ElementIndex, ResidueIndex } from 'mol-model/structure';
 import { Segmentation, OrderedSet, Interval } from 'mol-data/int';
 import { MoleculeType, SecondaryStructureType } from 'mol-model/structure/model/types';
 import Iterator from 'mol-data/iterator';
 import { Vec3 } from 'mol-math/linear-algebra';
-import { SymmetryOperator } from 'mol-math/geometry';
 import SortedRanges from 'mol-data/int/sorted-ranges';
 
-export function getPolymerRanges(unit: Unit) {
+export function getPolymerRanges(unit: Unit): SortedRanges<ElementIndex> {
     switch (unit.kind) {
         case Unit.Kind.Atomic: return unit.model.atomicHierarchy.polymerRanges
         case Unit.Kind.Spheres: return unit.model.coarseHierarchy.spheres.polymerRanges
@@ -48,78 +47,49 @@ export function getPolymerElementCount(unit: Unit) {
     return count
 }
 
-function getTraceName(l: StructureElement) {
-    const compId = StructureProperties.residue.label_comp_id(l)
-    const chemCompMap = l.unit.model.properties.chemicalComponentMap
-    const cc = chemCompMap.get(compId)
-    const moleculeType = cc ? cc.moleculeType : MoleculeType.unknown
-    let traceName = ''
-    if (moleculeType === MoleculeType.protein) {
-        traceName = 'CA'
-    } else if (moleculeType === MoleculeType.DNA || moleculeType === MoleculeType.RNA) {
-        traceName = 'P'
-    }
-    return traceName
-}
-
-function setTraceElement(l: StructureElement, residueSegment: Segmentation.Segment) {
-    const elements = l.unit.elements
-    l.element = elements[residueSegment.start]
-    const traceName = getTraceName(l)
-
-    for (let j = residueSegment.start, _j = residueSegment.end; j < _j; j++) {
-        l.element = elements[j];
-        if (StructureProperties.atom.label_atom_id(l) === traceName) return j
+function getResidueTypeAtomId(moleculeType: MoleculeType, atomType: 'trace' | 'direction') {
+    switch (moleculeType) {
+        case MoleculeType.protein:
+            switch (atomType) {
+                case 'trace': return 'CA'
+                case 'direction': return 'O'
+            }
+            break
+        case MoleculeType.RNA:
+            switch (atomType) {
+                case 'trace': return 'C4\''
+                case 'direction': return 'C3\''
+            }
+            break
+        case MoleculeType.DNA:
+            switch (atomType) {
+                case 'trace': return 'C3\''
+                case 'direction': return 'O4\''
+            }
+            break
     }
-    return residueSegment.end - 1
+    return ''
 }
 
-
-function getTraceName2(model: Model, residueModelIndex: number) {
-    const compId = model.atomicHierarchy.residues.label_comp_id.value(residueModelIndex)
+function getMoleculeType(model: Model, residueIndex: number) {
+    const compId = model.atomicHierarchy.residues.label_comp_id.value(residueIndex)
     const chemCompMap = model.properties.chemicalComponentMap
     const cc = chemCompMap.get(compId)
-    const moleculeType = cc ? cc.moleculeType : MoleculeType.unknown
-    let traceName = ''
-    if (moleculeType === MoleculeType.protein) {
-        traceName = 'CA'
-    } else if (moleculeType === MoleculeType.DNA) {
-        // traceName = 'P'
-        traceName = 'C3\''
-    } else if (moleculeType === MoleculeType.RNA) {
-        // traceName = 'P'
-        traceName = 'C4\''
-    }
-    return traceName
+    return cc ? cc.moleculeType : MoleculeType.unknown
 }
 
-// TODO fix type
-function getTraceElement2(model: Model, residueModelSegment: Segmentation.Segment<number>) {
-    const traceName = getTraceName2(model, residueModelSegment.index)
-
-    for (let j = residueModelSegment.start, _j = residueModelSegment.end; j < _j; j++) {
-        if (model.atomicHierarchy.atoms.label_atom_id.value(j) === traceName) return j
+function getElementIndexForAtomId(unit: Unit.Atomic, residueSegment: Segmentation.Segment, atomId: string) {
+    const elements = unit.elements
+    const { label_atom_id } = unit.model.atomicHierarchy.atoms
+    for (let j = residueSegment.start, _j = residueSegment.end; j < _j; j++) {
+        if (label_atom_id.value(elements[j]) === atomId) return j as ElementIndex
     }
-    console.log(`trace name element "${traceName}" not found`, { ...residueModelSegment })
-    return residueModelSegment.start
+    return residueSegment.end - 1 as ElementIndex
 }
 
-function getDirectionName2(model: Model, residueModelIndex: number) {
-    const compId = model.atomicHierarchy.residues.label_comp_id.value(residueModelIndex)
-    const chemCompMap = model.properties.chemicalComponentMap
-    const cc = chemCompMap.get(compId)
-    const moleculeType = cc ? cc.moleculeType : MoleculeType.unknown
-    let traceName = ''
-    if (moleculeType === MoleculeType.protein) {
-        traceName = 'O'
-    } else if (moleculeType === MoleculeType.DNA) {
-        traceName = 'O4\''
-    } else if (moleculeType === MoleculeType.RNA) {
-        traceName = 'C3\''
-    } else {
-        console.log('molecule type unknown', Number.isInteger(residueModelIndex) ? compId : residueModelIndex, chemCompMap)
-    }
-    return traceName
+function getResidueTypeAtomIdElementIndex(unit: Unit.Atomic, residueSegment: Segmentation.Segment, type: 'trace' | 'direction') {
+    const atomId = getResidueTypeAtomId(getMoleculeType(unit.model, residueSegment.index), type)
+    return getElementIndexForAtomId(unit, residueSegment, atomId)
 }
 
 // function residueLabel(model: Model, rI: number) {
@@ -130,18 +100,6 @@ function getDirectionName2(model: Model, residueModelIndex: number) {
 //     return `${label_asym_id.value(cI)} ${label_comp_id.value(rI)} ${label_seq_id.value(rI)}`
 // }
 
-// TODO fix type
-function getDirectionElement2(model: Model, residueModelSegment: Segmentation.Segment<number>) {
-    const traceName = getDirectionName2(model, residueModelSegment.index)
-
-    for (let j = residueModelSegment.start, _j = residueModelSegment.end; j < _j; j++) {
-        if (model.atomicHierarchy.atoms.label_atom_id.value(j) === traceName) return j
-    }
-    // console.log('direction name element not found', { ...residueModelSegment })
-    return residueModelSegment.start
-}
-
-
 /** Iterates over consecutive pairs of residues/coarse elements in polymers */
 export function PolymerBackboneIterator(unit: Unit): Iterator<PolymerBackbonePair> {
     switch (unit.kind) {
@@ -155,20 +113,12 @@ export function PolymerBackboneIterator(unit: Unit): Iterator<PolymerBackbonePai
 interface PolymerBackbonePair {
     centerA: StructureElement
     centerB: StructureElement
-    indexA: number
-    indexB: number
-    posA: Vec3
-    posB: Vec3
 }
 
 function createPolymerBackbonePair (unit: Unit) {
     return {
         centerA: StructureElement.create(unit),
         centerB: StructureElement.create(unit),
-        indexA: 0,
-        indexB: 0,
-        posA: Vec3.zero(),
-        posB: Vec3.zero()
     }
 }
 
@@ -176,29 +126,19 @@ const enum AtomicPolymerBackboneIteratorState { nextPolymer, firstResidue, nextR
 
 export class AtomicPolymerBackboneIterator<T extends number = number> implements Iterator<PolymerBackbonePair> {
     private value: PolymerBackbonePair
-
     private polymerIt: SortedRanges.Iterator<ElementIndex>
-    private residueIt: Segmentation.SegmentIterator<number> // TODO specific type
-    private polymerSegment: Segmentation.Segment<ElementIndex>
+    private residueIt: Segmentation.SegmentIterator<ResidueIndex>
     private state: AtomicPolymerBackboneIteratorState = AtomicPolymerBackboneIteratorState.nextPolymer
-    private pos: SymmetryOperator.CoordinateMapper
-
     hasNext: boolean = false;
 
     move() {
-        const { residueIt, polymerIt, value, pos } = this
-
         if (this.state === AtomicPolymerBackboneIteratorState.nextPolymer) {
-            while (polymerIt.hasNext) {
-                this.polymerSegment = polymerIt.move();
-                // console.log('polymerSegment', this.polymerSegment)
-                residueIt.setSegment(this.polymerSegment);
-
-                const residueSegment = residueIt.move();
-                // console.log('first residueSegment', residueSegment, residueIt.hasNext)
-                if (residueIt.hasNext) {
-                    value.indexB = setTraceElement(value.centerB, residueSegment)
-                    pos(value.centerB.element, value.posB)
+            while (this.polymerIt.hasNext) {
+                const residueSegment = this.polymerIt.move()
+                this.residueIt.setSegment(residueSegment);
+                if (this.residueIt.hasNext) {
+                    this.value.centerB.element = getResidueTypeAtomIdElementIndex(this.unit, residueSegment, 'trace')
+                    // setTraceElement(this.value.centerB, this.residueIt.move())
                     this.state = AtomicPolymerBackboneIteratorState.nextResidue
                     break
                 }
@@ -206,36 +146,23 @@ export class AtomicPolymerBackboneIterator<T extends number = number> implements
         }
 
         if (this.state === AtomicPolymerBackboneIteratorState.nextResidue) {
-            const residueSegment = residueIt.move();
-            // console.log('next residueSegment', residueSegment)
-            value.centerA.element = value.centerB.element
-            value.indexA = value.indexB
-            Vec3.copy(value.posA, value.posB)
-            value.indexB = setTraceElement(value.centerB, residueSegment)
-            pos(value.centerB.element, value.posB)
-
-            if (!residueIt.hasNext) {
+            this.value.centerA.element = this.value.centerB.element
+            this.value.centerB.element = getResidueTypeAtomIdElementIndex(this.unit, this.residueIt.move(), 'trace')
+            // setTraceElement(this.value.centerB, this.residueIt.move())
+            if (!this.residueIt.hasNext) {
                 // TODO need to advance to a polymer that has two or more residues (can't assume it has)
                 this.state = AtomicPolymerBackboneIteratorState.nextPolymer
             }
         }
 
-        this.hasNext = residueIt.hasNext || polymerIt.hasNext
-
-        // console.log('hasNext', this.hasNext)
-        // console.log('value', this.value)
-
+        this.hasNext = this.residueIt.hasNext || this.polymerIt.hasNext
         return this.value;
     }
 
-    constructor(unit: Unit.Atomic) {
-        const { residueAtomSegments } = unit.model.atomicHierarchy
-        // console.log('unit.elements', OrderedSet.toArray(unit.elements))
+    constructor(private unit: Unit.Atomic) {
         this.polymerIt = SortedRanges.transientSegments(getPolymerRanges(unit), unit.elements)
-        this.residueIt = Segmentation.transientSegments(residueAtomSegments, unit.elements)
-        this.pos = unit.conformation.invariantPosition
+        this.residueIt = Segmentation.transientSegments(unit.model.atomicHierarchy.residueAtomSegments, unit.elements)
         this.value = createPolymerBackbonePair(unit)
-
         this.hasNext = this.residueIt.hasNext && this.polymerIt.hasNext
     }
 }
@@ -244,29 +171,20 @@ const enum CoarsePolymerBackboneIteratorState { nextPolymer, nextElement }
 
 export class CoarsePolymerBackboneIterator<T extends number = number> implements Iterator<PolymerBackbonePair> {
     private value: PolymerBackbonePair
-
     private polymerIt: SortedRanges.Iterator<ElementIndex>
     private polymerSegment: Segmentation.Segment<ElementIndex>
     private state: CoarsePolymerBackboneIteratorState = CoarsePolymerBackboneIteratorState.nextPolymer
-    private pos: SymmetryOperator.CoordinateMapper
     private elementIndex: number
-
     hasNext: boolean = false;
 
     move() {
-        const { polymerIt, value, pos } = this
-
         if (this.state === CoarsePolymerBackboneIteratorState.nextPolymer) {
-            if (polymerIt.hasNext) {
-                this.polymerSegment = polymerIt.move();
-                console.log('polymer', this.polymerSegment)
+            if (this.polymerIt.hasNext) {
+                this.polymerSegment = this.polymerIt.move();
                 this.elementIndex = this.polymerSegment.start
-                this.elementIndex += 1
+                // this.elementIndex += 1
                 if (this.elementIndex + 1 < this.polymerSegment.end) {
-                    value.centerB.element = value.centerB.unit.elements[this.elementIndex]
-                    value.indexB = this.elementIndex
-                    pos(value.centerB.element, value.posB)
-
+                    this.value.centerB.element = this.value.centerB.unit.elements[this.elementIndex]
                     this.state = CoarsePolymerBackboneIteratorState.nextElement
                 } else {
                     this.state = CoarsePolymerBackboneIteratorState.nextPolymer
@@ -276,31 +194,20 @@ export class CoarsePolymerBackboneIterator<T extends number = number> implements
 
         if (this.state === CoarsePolymerBackboneIteratorState.nextElement) {
             this.elementIndex += 1
-            value.centerA.element = value.centerB.element
-            value.indexA = value.indexB
-            Vec3.copy(value.posA, value.posB)
-
-            value.centerB.element = value.centerB.unit.elements[this.elementIndex]
-            value.indexB = this.elementIndex
-            pos(value.centerB.element, value.posB)
-
+            this.value.centerA.element = this.value.centerB.element
+            this.value.centerB.element = this.value.centerB.unit.elements[this.elementIndex]
             if (this.elementIndex + 1 >= this.polymerSegment.end) {
                 this.state = CoarsePolymerBackboneIteratorState.nextPolymer
             }
         }
 
-        this.hasNext = this.elementIndex + 1 < this.polymerSegment.end || polymerIt.hasNext
+        this.hasNext = this.elementIndex + 1 < this.polymerSegment.end || this.polymerIt.hasNext
         return this.value;
     }
 
     constructor(unit: Unit.Spheres | Unit.Gaussians) {
         this.polymerIt = SortedRanges.transientSegments(getPolymerRanges(unit), unit.elements);
-
-        this.pos = unit.conformation.invariantPosition
         this.value = createPolymerBackbonePair(unit)
-
-        console.log('CoarsePolymerBackboneIterator', this.polymerIt.hasNext)
-
         this.hasNext = this.polymerIt.hasNext
     }
 }
@@ -323,37 +230,19 @@ export function PolymerTraceIterator(unit: Unit): Iterator<PolymerTraceElement>
 
 interface PolymerTraceElement {
     center: StructureElement
-    index: number
-    first: boolean
-    last: boolean
+    first: boolean, last: boolean
     secStrucType: SecondaryStructureType
-
-    t0: Vec3
-    t1: Vec3
-    t2: Vec3
-    t3: Vec3
-    t4: Vec3
-
-    d12: Vec3
-    d23: Vec3
+    t0: Vec3, t1: Vec3, t2: Vec3, t3: Vec3, t4: Vec3
+    d12: Vec3, d23: Vec3
 }
 
 function createPolymerTraceElement (unit: Unit): PolymerTraceElement {
     return {
         center: StructureElement.create(unit),
-        index: 0,
-        first: false,
-        last: false,
+        first: false, last: false,
         secStrucType: SecondaryStructureType.create(SecondaryStructureType.Flag.NA),
-
-        t0: Vec3.zero(),
-        t1: Vec3.zero(),
-        t2: Vec3.zero(),
-        t3: Vec3.zero(),
-        t4: Vec3.zero(),
-
-        d12: Vec3.zero(),
-        d23: Vec3.zero(),
+        t0: Vec3.zero(), t1: Vec3.zero(), t2: Vec3.zero(), t3: Vec3.zero(), t4: Vec3.zero(),
+        d12: Vec3.zero(), d23: Vec3.zero(),
     }
 }
 
@@ -372,17 +261,13 @@ function setSegment (outSegment: Segmentation.Segment<number>, index: number, se
 
 export class AtomicPolymerTraceIterator<T extends number = number> implements Iterator<PolymerTraceElement> {
     private value: PolymerTraceElement
-
     private polymerIt: SortedRanges.Iterator<ElementIndex>
-    private residueIt: Segmentation.SegmentIterator<number> // TODO specialize type
-    private residueSegmentMin: number
-    private residueSegmentMax: number
+    private residueIt: Segmentation.SegmentIterator<ResidueIndex>
+    private residueAtomSegmentMin: number
+    private residueAtomSegmentMax: number
     private state: AtomicPolymerTraceIteratorState = AtomicPolymerTraceIteratorState.nextPolymer
-    private residueSegments: Segmentation<ElementIndex>
-
-    private tmpSegment: Segmentation.Segment<number>
-
-    private unit: Unit.Atomic
+    private residueAtomSegments: Segmentation<ElementIndex, ResidueIndex>
+    private tmpSegment: Segmentation.Segment<ResidueIndex>
 
     hasNext: boolean = false;
 
@@ -396,8 +281,8 @@ export class AtomicPolymerTraceIterator<T extends number = number> implements It
         const { polymerRanges, residueAtomSegments } = this.unit.model.atomicHierarchy
         const sMin = polymerRanges[polymerSegment.index * 2]
         const sMax = polymerRanges[polymerSegment.index * 2 + 1]
-        this.residueSegmentMin = residueAtomSegments.index[sMin]
-        this.residueSegmentMax = residueAtomSegments.index[sMax]
+        this.residueAtomSegmentMin = residueAtomSegments.index[sMin]
+        this.residueAtomSegmentMax = residueAtomSegments.index[sMax]
     }
 
     move() {
@@ -417,32 +302,32 @@ export class AtomicPolymerTraceIterator<T extends number = number> implements It
         }
 
         if (this.state === AtomicPolymerTraceIteratorState.nextResidue) {
-            const { tmpSegment, residueSegments, residueSegmentMin, residueSegmentMax } = this
+            const { tmpSegment, residueAtomSegments, residueAtomSegmentMin, residueAtomSegmentMax } = this
             const residueSegment = residueIt.move();
             const resSegIdx = residueSegment.index
             // console.log(residueLabel(this.unit.model, resSegIdx), resSegIdx, this.unit.model.properties.secondaryStructure.type[resSegIdx])
-            value.index = setTraceElement(value.center, residueSegment)
+            value.center.element = getResidueTypeAtomIdElementIndex(this.unit, residueSegment, 'trace')
 
-            setSegment(tmpSegment, resSegIdx - 2, residueSegments, residueSegmentMin, residueSegmentMax)
-            this.pos(value.t0, getTraceElement2(this.unit.model, tmpSegment))
+            setSegment(tmpSegment, resSegIdx - 2, residueAtomSegments, residueAtomSegmentMin, residueAtomSegmentMax)
+            this.pos(value.t0, getResidueTypeAtomIdElementIndex(this.unit, tmpSegment, 'trace'))
 
-            setSegment(tmpSegment, resSegIdx - 1, residueSegments, residueSegmentMin, residueSegmentMax)
-            this.pos(value.t1, getTraceElement2(this.unit.model, tmpSegment))
-            this.pos(value.d12, getDirectionElement2(this.unit.model, tmpSegment))
+            setSegment(tmpSegment, resSegIdx - 1, residueAtomSegments, residueAtomSegmentMin, residueAtomSegmentMax)
+            this.pos(value.t1, getResidueTypeAtomIdElementIndex(this.unit, tmpSegment, 'trace'))
+            this.pos(value.d12, getResidueTypeAtomIdElementIndex(this.unit, tmpSegment, 'direction'))
 
-            setSegment(tmpSegment, resSegIdx, residueSegments, residueSegmentMin, residueSegmentMax)
+            setSegment(tmpSegment, resSegIdx, residueAtomSegments, residueAtomSegmentMin, residueAtomSegmentMax)
             value.secStrucType = this.unit.model.properties.secondaryStructure.type[resSegIdx]
-            this.pos(value.t2, getTraceElement2(this.unit.model, tmpSegment))
-            this.pos(value.d23, getDirectionElement2(this.unit.model, tmpSegment))
+            this.pos(value.t2, getResidueTypeAtomIdElementIndex(this.unit, tmpSegment, 'trace'))
+            this.pos(value.d23, getResidueTypeAtomIdElementIndex(this.unit, tmpSegment, 'direction'))
 
-            setSegment(tmpSegment, resSegIdx + 1, residueSegments, residueSegmentMin, residueSegmentMax)
-            this.pos(value.t3, getTraceElement2(this.unit.model, tmpSegment))
+            setSegment(tmpSegment, resSegIdx + 1, residueAtomSegments, residueAtomSegmentMin, residueAtomSegmentMax)
+            this.pos(value.t3, getResidueTypeAtomIdElementIndex(this.unit, tmpSegment, 'trace'))
 
-            setSegment(tmpSegment, resSegIdx + 2, residueSegments, residueSegmentMin, residueSegmentMax)
-            this.pos(value.t4, getTraceElement2(this.unit.model, tmpSegment))
+            setSegment(tmpSegment, resSegIdx + 2, residueAtomSegments, residueAtomSegmentMin, residueAtomSegmentMax)
+            this.pos(value.t4, getResidueTypeAtomIdElementIndex(this.unit, tmpSegment, 'trace'))
 
-            value.first = resSegIdx === residueSegmentMin
-            value.last = resSegIdx === residueSegmentMax
+            value.first = resSegIdx === residueAtomSegmentMin
+            value.last = resSegIdx === residueAtomSegmentMax
 
             if (!residueIt.hasNext) {
                 this.state = AtomicPolymerTraceIteratorState.nextPolymer
@@ -454,17 +339,13 @@ export class AtomicPolymerTraceIterator<T extends number = number> implements It
         return this.value;
     }
 
-    constructor(unit: Unit.Atomic) {
-        const { residueAtomSegments } = unit.model.atomicHierarchy
+    constructor(private unit: Unit.Atomic) {
+        this.residueAtomSegments = unit.model.atomicHierarchy.residueAtomSegments
         this.polymerIt = SortedRanges.transientSegments(getPolymerRanges(unit), unit.elements)
-        this.residueIt = Segmentation.transientSegments(residueAtomSegments, unit.elements);
-        this.residueSegments = residueAtomSegments
+        this.residueIt = Segmentation.transientSegments(this.residueAtomSegments, unit.elements);
         this.value = createPolymerTraceElement(unit)
+        this.tmpSegment = { index: 0 as ResidueIndex, start: 0 as ElementIndex, end: 0 as ElementIndex }
         this.hasNext = this.residueIt.hasNext && this.polymerIt.hasNext
-
-        this.tmpSegment = { index: 0, start: 0 as ElementIndex, end: 0 as ElementIndex }
-
-        this.unit = unit
     }
 }