Merge pull request #225 from molstar/tubular-helices

Add tubularHelices parameter to Cartoon representation

CHANGELOG.md

@@ -6,6 +6,7 @@ Note that since we don't clearly distinguish between a public and private interf
 
 ## [Unreleased]
 
+- Add `tubularHelices` parameter to Cartoon representation
 
 ## [v2.1.0] - 2021-07-05
 

src/mol-model-props/computed/helix-orientation.ts

+/**
+ * Copyright (c) 2021 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author Alexander Rose <alexander.rose@weirdbyte.de>
+ */
+
+import { ParamDefinition as PD } from '../../mol-util/param-definition';
+import { Model } from '../../mol-model/structure';
+import { CustomPropertyDescriptor } from '../../mol-model/custom-property';
+import { CustomModelProperty } from '../common/custom-model-property';
+import { calcHelixOrientation, HelixOrientation } from './helix-orientation/helix-orientation';
+
+export const HelixOrientationParams = { };
+export type HelixOrientationParams = typeof HelixOrientationParams
+export type HelixOrientationProps = PD.Values<HelixOrientationParams>
+
+export type HelixOrientationValue = HelixOrientation;
+
+export const HelixOrientationProvider: CustomModelProperty.Provider<HelixOrientationParams, HelixOrientationValue> = CustomModelProperty.createProvider({
+    label: 'Helix Orientation',
+    descriptor: CustomPropertyDescriptor({
+        name: 'molstar_helix_orientation'
+    }),
+    type: 'dynamic',
+    defaultParams: {},
+    getParams: () => ({}),
+    isApplicable: (data: Model) => true,
+    obtain: async (ctx, data) => {
+        return { value: calcHelixOrientation(data) };
+    }
+});
\ No newline at end of file

src/mol-model-props/computed/helix-orientation/helix-orientation.ts

+/**
+ * Copyright (c) 2021 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author Alexander Rose <alexander.rose@weirdbyte.de>
+ */
+
+import { ElementIndex } from '../../../mol-model/structure';
+import { Segmentation } from '../../../mol-data/int/segmentation';
+import { SortedRanges } from '../../../mol-data/int/sorted-ranges';
+import { OrderedSet } from '../../../mol-data/int';
+import { Model } from '../../../mol-model/structure/model';
+import { Vec3 } from '../../../mol-math/linear-algebra';
+
+export interface HelixOrientation {
+    centers: ArrayLike<number>
+}
+
+/** Usees same definition as GROMACS' helixorient */
+export function calcHelixOrientation(model: Model): HelixOrientation {
+    const { x, y, z } = model.atomicConformation;
+    const { polymerType, traceElementIndex } = model.atomicHierarchy.derived.residue;
+    const n = polymerType.length;
+
+    const elements = OrderedSet.ofBounds(0, model.atomicConformation.atomId.rowCount) as OrderedSet<ElementIndex>;
+    const polymerIt = SortedRanges.transientSegments(model.atomicRanges.polymerRanges, elements);
+    const residueIt = Segmentation.transientSegments(model.atomicHierarchy.residueAtomSegments, elements);
+
+    const centers = new Float32Array(n * 3);
+    const axes = new Float32Array(n * 3);
+
+    let i = 0;
+    let j = -1;
+    let s = -1;
+
+    const a1 = Vec3();
+    const a2 = Vec3();
+    const a3 = Vec3();
+    const a4 = Vec3();
+
+    const r12 = Vec3();
+    const r23 = Vec3();
+    const r34 = Vec3();
+
+    const v1 = Vec3();
+    const v2 = Vec3();
+    const vt = Vec3();
+
+    const diff13 = Vec3();
+    const diff24 = Vec3();
+
+    const axis = Vec3();
+    const prevAxis = Vec3();
+
+    while (polymerIt.hasNext) {
+        const ps = polymerIt.move();
+        residueIt.setSegment(ps);
+        i = -1;
+        s = -1;
+        while (residueIt.hasNext) {
+            i += 1;
+            const { index } = residueIt.move();
+            if (i === 0) s = index;
+
+            j = (index - 2);
+            const j3 = j * 3;
+
+            Vec3.copy(a1, a2);
+            Vec3.copy(a2, a3);
+            Vec3.copy(a3, a4);
+
+            const eI = traceElementIndex[index];
+            Vec3.set(a4, x[eI], y[eI], z[eI]);
+
+            if (i < 3) continue;
+
+            Vec3.sub(r12, a2, a1);
+            Vec3.sub(r23, a3, a2);
+            Vec3.sub(r34, a4, a3);
+
+            Vec3.sub(diff13, r12, r23);
+            Vec3.sub(diff24, r23, r34);
+
+            Vec3.cross(axis, diff13, diff24);
+            Vec3.normalize(axis, axis);
+            Vec3.toArray(axis, axes, j3);
+
+            const tmp = Math.cos(Vec3.angle(diff13, diff24));
+
+            const diff13Length = Vec3.magnitude(diff13);
+            const diff24Length = Vec3.magnitude(diff24);
+
+            const r = (
+                Math.sqrt(diff24Length * diff13Length) /
+                // clamp, to avoid numerical instabilities for when
+                // angle between diff13 and diff24 is close to 0
+                Math.max(2.0, 2.0 * (1.0 - tmp))
+            );
+
+            Vec3.scale(v1, diff13, r / diff13Length);
+            Vec3.sub(v1, a2, v1);
+            Vec3.toArray(v1, centers, j3);
+
+            Vec3.scale(v2, diff24, r / diff24Length);
+            Vec3.sub(v2, a3, v2);
+            Vec3.toArray(v2, centers, j3 + 3);
+
+            Vec3.copy(prevAxis, axis);
+        }
+
+        // calc axis as dir of second and third center pos
+        // project first trace atom onto axis to get first center pos
+        const s3 = s * 3;
+        Vec3.fromArray(v1, centers, s3 + 3);
+        Vec3.fromArray(v2, centers, s3 + 6);
+        Vec3.normalize(axis, Vec3.sub(axis, v1, v2));
+        const sI = traceElementIndex[s];
+        Vec3.set(a1, x[sI], y[sI], z[sI]);
+        Vec3.copy(vt, a1);
+        Vec3.projectPointOnVector(vt, vt, axis, v1);
+        Vec3.toArray(vt, centers, s3);
+
+        // calc axis as dir of n-1 and n-2 center pos
+        // project last traceAtom onto axis to get last center pos
+        const e = j + 2;
+        const e3 = e * 3;
+        Vec3.fromArray(v1, centers, e3 - 3);
+        Vec3.fromArray(v2, centers, e3 - 6);
+        Vec3.normalize(axis, Vec3.sub(axis, v1, v2));
+        const eI = traceElementIndex[e];
+        Vec3.set(a1, x[eI], y[eI], z[eI]);Vec3.copy(vt, a1);
+        Vec3.projectPointOnVector(vt, vt, axis, v1);
+        Vec3.toArray(vt, centers, e3);
+    }
+
+    return {
+        centers
+    };
+}

src/mol-repr/structure/representation/cartoon.ts

 /**
- * Copyright (c) 2018 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ * Copyright (c) 2018-2021 mol* contributors, licensed under MIT, See LICENSE file for more info.
  *
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
  */
@@ -17,6 +17,7 @@ import { PolymerGapParams, PolymerGapVisual } from '../visual/polymer-gap-cylind
 import { PolymerTraceParams, PolymerTraceVisual } from '../visual/polymer-trace-mesh';
 import { SecondaryStructureProvider } from '../../../mol-model-props/computed/secondary-structure';
 import { CustomProperty } from '../../../mol-model-props/common/custom-property';
+import { HelixOrientationProvider } from '../../../mol-model-props/computed/helix-orientation';
 
 const CartoonVisuals = {
     'polymer-trace': (ctx: RepresentationContext, getParams: RepresentationParamsGetter<Structure, PolymerTraceParams>) => UnitsRepresentation('Polymer trace mesh', ctx, getParams, PolymerTraceVisual),
@@ -67,7 +68,17 @@ export const CartoonRepresentationProvider = StructureRepresentationProvider({
     defaultSizeTheme: { name: 'uniform' },
     isApplicable: (structure: Structure) => structure.polymerResidueCount > 0,
     ensureCustomProperties: {
-        attach: (ctx: CustomProperty.Context, structure: Structure) => SecondaryStructureProvider.attach(ctx, structure, void 0, true),
-        detach: (data) => SecondaryStructureProvider.ref(data, false)
+        attach: async (ctx: CustomProperty.Context, structure: Structure) => {
+            await SecondaryStructureProvider.attach(ctx, structure, void 0, true);
+            for (const m of structure.models) {
+                await HelixOrientationProvider.attach(ctx, m, void 0, true);
+            }
+        },
+        detach: (data) => {
+            SecondaryStructureProvider.ref(data, false);
+            for (const m of data.models) {
+                HelixOrientationProvider.ref(m, false);
+            }
+        }
     }
 });
\ No newline at end of file

src/mol-repr/structure/visual/polymer-trace-mesh.ts

 /**
- * Copyright (c) 2018-2019 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ * Copyright (c) 2018-2021 mol* contributors, licensed under MIT, See LICENSE file for more info.
  *
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
  */
@@ -27,6 +27,7 @@ export const PolymerTraceMeshParams = {
     sizeFactor: PD.Numeric(0.2, { min: 0, max: 10, step: 0.01 }),
     aspectRatio: PD.Numeric(5, { min: 0.1, max: 10, step: 0.1 }),
     arrowFactor: PD.Numeric(1.5, { min: 0, max: 3, step: 0.1 }),
+    tubularHelices: PD.Boolean(false),
     detail: PD.Numeric(0, { min: 0, max: 3, step: 1 }, BaseGeometry.CustomQualityParamInfo),
     linearSegments: PD.Numeric(8, { min: 1, max: 48, step: 1 }, BaseGeometry.CustomQualityParamInfo),
     radialSegments: PD.Numeric(16, { min: 2, max: 56, step: 2 }, BaseGeometry.CustomQualityParamInfo)
@@ -40,7 +41,7 @@ function createPolymerTraceMesh(ctx: VisualContext, unit: Unit, structure: Struc
     const polymerElementCount = unit.polymerElements.length;
 
     if (!polymerElementCount) return Mesh.createEmpty(mesh);
-    const { sizeFactor, detail, linearSegments, radialSegments, aspectRatio, arrowFactor } = props;
+    const { sizeFactor, detail, linearSegments, radialSegments, aspectRatio, arrowFactor, tubularHelices } = props;
 
     const vertexCount = linearSegments * radialSegments * polymerElementCount + (radialSegments + 1) * polymerElementCount * 2;
     const builderState = MeshBuilder.createState(vertexCount, vertexCount / 10, mesh);
@@ -50,7 +51,7 @@ function createPolymerTraceMesh(ctx: VisualContext, unit: Unit, structure: Struc
     const { curvePoints, normalVectors, binormalVectors, widthValues, heightValues } = state;
 
     let i = 0;
-    const polymerTraceIt = PolymerTraceIterator(unit, structure);
+    const polymerTraceIt = PolymerTraceIterator(unit, structure, { ignoreSecondaryStructure: false, useHelixOrientation: tubularHelices });
     while (polymerTraceIt.hasNext) {
         const v = polymerTraceIt.move();
         builderState.currentGroup = i;
@@ -58,7 +59,7 @@ function createPolymerTraceMesh(ctx: VisualContext, unit: Unit, structure: Struc
         const isNucleicType = isNucleic(v.moleculeType);
         const isSheet = SecondaryStructureType.is(v.secStrucType, SecondaryStructureType.Flag.Beta);
         const isHelix = SecondaryStructureType.is(v.secStrucType, SecondaryStructureType.Flag.Helix);
-        const tension = isHelix ? HelixTension : StandardTension;
+        const tension = isHelix && !tubularHelices ? HelixTension : StandardTension;
         const shift = isNucleicType ? NucleicShift : StandardShift;
 
         interpolateCurveSegment(state, v, tension, shift);
@@ -112,9 +113,19 @@ function createPolymerTraceMesh(ctx: VisualContext, unit: Unit, structure: Struc
         } else {
             let h0: number, h1: number, h2: number;
             if (isHelix && !v.isCoarseBackbone) {
-                h0 = w0 * aspectRatio;
-                h1 = w1 * aspectRatio;
-                h2 = w2 * aspectRatio;
+                if (tubularHelices) {
+                    w0 *= aspectRatio * 1.5;
+                    w1 *= aspectRatio * 1.5;
+                    w2 *= aspectRatio * 1.5;
+
+                    h0 = w0;
+                    h1 = w1;
+                    h2 = w2;
+                } else {
+                    h0 = w0 * aspectRatio;
+                    h1 = w1 * aspectRatio;
+                    h2 = w2 * aspectRatio;
+                }
             } else if (isNucleicType && !v.isCoarseBackbone) {
                 h0 = w0 * aspectRatio;
                 h1 = w1 * aspectRatio;
@@ -172,6 +183,7 @@ export function PolymerTraceVisual(materialId: number): UnitsVisual<PolymerTrace
         setUpdateState: (state: VisualUpdateState, newProps: PD.Values<PolymerTraceParams>, currentProps: PD.Values<PolymerTraceParams>, newTheme: Theme, currentTheme: Theme, newStructureGroup: StructureGroup, currentStructureGroup: StructureGroup) => {
             state.createGeometry = (
                 newProps.sizeFactor !== currentProps.sizeFactor ||
+                newProps.tubularHelices !== currentProps.tubularHelices ||
                 newProps.detail !== currentProps.detail ||
                 newProps.linearSegments !== currentProps.linearSegments ||
                 newProps.radialSegments !== currentProps.radialSegments ||

src/mol-repr/structure/visual/polymer-tube-mesh.ts

@@ -46,7 +46,7 @@ function createPolymerTubeMesh(ctx: VisualContext, unit: Unit, structure: Struct
     const { curvePoints, normalVectors, binormalVectors, widthValues, heightValues } = state;
 
     let i = 0;
-    const polymerTraceIt = PolymerTraceIterator(unit, structure, true);
+    const polymerTraceIt = PolymerTraceIterator(unit, structure, { ignoreSecondaryStructure: true });
     while (polymerTraceIt.hasNext) {
         const v = polymerTraceIt.move();
         builderState.currentGroup = i;

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

 /**
- * Copyright (c) 2018-2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ * Copyright (c) 2018-2021 mol* contributors, licensed under MIT, See LICENSE file for more info.
  *
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
  */
@@ -14,14 +14,31 @@ import { CoarseSphereConformation, CoarseGaussianConformation } from '../../../.
 import { getPolymerRanges } from '../polymer';
 import { AtomicConformation } from '../../../../../mol-model/structure/model/properties/atomic';
 import { SecondaryStructureProvider } from '../../../../../mol-model-props/computed/secondary-structure';
+import { HelixOrientationProvider } from '../../../../../mol-model-props/computed/helix-orientation';
+import { SecondaryStructure } from '../../../../../mol-model/structure/model/properties/seconday-structure';
+
+function isHelixSS(ss: SecondaryStructureType.Flag) {
+    return SecondaryStructureType.is(ss, SecondaryStructureType.Flag.Helix);
+}
+
+function isSheetSS(ss: SecondaryStructureType.Flag) {
+    return SecondaryStructureType.is(ss, SecondaryStructureType.Flag.Beta);
+}
+
+//
+
+type PolymerTraceIteratorOptions = {
+    ignoreSecondaryStructure?: boolean,
+    useHelixOrientation?: boolean
+}
 
 /**
  * Iterates over individual residues/coarse elements in polymers of a unit while
  * providing information about the neighbourhood in the underlying model for drawing splines
  */
-export function PolymerTraceIterator(unit: Unit, structure: Structure, ignoreSecondaryStructure = false): Iterator<PolymerTraceElement> {
+export function PolymerTraceIterator(unit: Unit, structure: Structure, options: PolymerTraceIteratorOptions = {}): Iterator<PolymerTraceElement> {
     switch (unit.kind) {
-        case Unit.Kind.Atomic: return new AtomicPolymerTraceIterator(unit, structure, ignoreSecondaryStructure);
+        case Unit.Kind.Atomic: return new AtomicPolymerTraceIterator(unit, structure, options);
         case Unit.Kind.Spheres:
         case Unit.Kind.Gaussians:
             return new CoarsePolymerTraceIterator(unit, structure);
@@ -91,6 +108,8 @@ export class AtomicPolymerTraceIterator implements Iterator<PolymerTraceElement>
     private directionToElementIndex: ArrayLike<ElementIndex | -1>
     private moleculeType: ArrayLike<MoleculeType>
     private atomicConformation: AtomicConformation
+    private secondaryStructure: SecondaryStructure | undefined
+    private helixOrientationCenters: ArrayLike<number> | undefined
 
     private p0 = Vec3()
     private p1 = Vec3()
@@ -107,7 +126,7 @@ export class AtomicPolymerTraceIterator implements Iterator<PolymerTraceElement>
 
     hasNext: boolean = false;
 
-    private pos(target: Vec3, index: number) {
+    private atomicPos(target: Vec3, index: ElementIndex | -1) {
         if (index !== -1) {
             target[0] = this.atomicConformation.x[index];
             target[1] = this.atomicConformation.y[index];
@@ -115,6 +134,15 @@ export class AtomicPolymerTraceIterator implements Iterator<PolymerTraceElement>
         }
     }
 
+    private pos(target: Vec3, residueIndex: ResidueIndex, ss: SecondaryStructureType.Flag) {
+        const index = this.traceElementIndex[residueIndex];
+        if (this.helixOrientationCenters && isHelixSS(ss)) {
+            Vec3.fromArray(target, this.helixOrientationCenters, residueIndex * 3);
+        } else {
+            this.atomicPos(target, index);
+        }
+    }
+
     private updateResidueSegmentRange(polymerSegment: Segmentation.Segment<number>) {
         const { index } = this.residueAtomSegments;
         this.residueSegmentMin = index[this.polymerRanges[polymerSegment.index * 2]];
@@ -140,23 +168,31 @@ export class AtomicPolymerTraceIterator implements Iterator<PolymerTraceElement>
         return residueIndex as ResidueIndex;
     }
 
-    private getSecStruc: (residueIndex: ResidueIndex) => SecondaryStructureType.Flag
+    private getSecStruc(residueIndex: ResidueIndex): SecondaryStructureType.Flag {
+        if (this.secondaryStructure) {
+            const { type, getIndex } = this.secondaryStructure;
+            const ss = type[getIndex(residueIndex)];
+            // normalize helix-type
+            return isHelixSS(ss) ? SecondaryStructureType.Flag.Helix : ss;
+        } else {
+            return SecStrucTypeNA;
+        }
+    }
 
-    private setControlPoint(out: Vec3, p1: Vec3, p2: Vec3, p3: Vec3, residueIndex: ResidueIndex) {
-        const ss =  this.getSecStruc(residueIndex);
-        if (SecondaryStructureType.is(ss, SecondaryStructureType.Flag.Beta)) {
+    private setControlPoint(out: Vec3, p1: Vec3, p2: Vec3, p3: Vec3, ss: SecondaryStructureType.Flag) {
+        if (isSheetSS(ss) || (this.helixOrientationCenters && isHelixSS(ss))) {
             Vec3.scale(out, Vec3.add(out, p1, Vec3.add(out, p3, Vec3.add(out, p2, p2))), 1 / 4);
         } else {
             Vec3.copy(out, p2);
         }
     }
 
-    private setFromToVector(out: Vec3, residueIndex: ResidueIndex) {
-        if (this.value.isCoarseBackbone) {
+    private setFromToVector(out: Vec3, residueIndex: ResidueIndex, ss: SecondaryStructureType.Flag) {
+        if (this.value.isCoarseBackbone || (this.helixOrientationCenters && isHelixSS(ss))) {
             Vec3.set(out, 1, 0, 0);
         } else {
-            this.pos(tmpVecA, this.directionFromElementIndex[residueIndex]);
-            this.pos(tmpVecB, this.directionToElementIndex[residueIndex]);
+            this.atomicPos(tmpVecA, this.directionFromElementIndex[residueIndex]);
+            this.atomicPos(tmpVecB, this.directionToElementIndex[residueIndex]);
             Vec3.sub(out, tmpVecB, tmpVecA);
         }
     }
@@ -203,7 +239,7 @@ export class AtomicPolymerTraceIterator implements Iterator<PolymerTraceElement>
 
             this.prevCoarseBackbone = this.currCoarseBackbone;
             this.currCoarseBackbone = this.nextCoarseBackbone;
-            this.nextCoarseBackbone = residueIndex === residueIndexNext1 ? false : (this.directionFromElementIndex[residueIndexNext1] === -1 || this.directionToElementIndex[residueIndexNext1] === -1);
+            this.nextCoarseBackbone = this.directionFromElementIndex[residueIndexNext1] === -1 || this.directionToElementIndex[residueIndexNext1] === -1;
 
             value.secStrucType = this.currSecStrucType;
             value.secStrucFirst = this.prevSecStrucType !== this.currSecStrucType;
@@ -214,7 +250,6 @@ export class AtomicPolymerTraceIterator implements Iterator<PolymerTraceElement>
             value.first = residueIndex === this.residueSegmentMin;
             value.last = residueIndex === this.residueSegmentMax;
             value.moleculeType = this.moleculeType[residueIndex];
-            value.isCoarseBackbone = this.directionFromElementIndex[residueIndex] === -1 || this.directionToElementIndex[residueIndex] === -1;
 
             value.initial = residueIndex === residueIndexPrev1;
             value.final = residueIndex === residueIndexNext1;
@@ -223,53 +258,132 @@ export class AtomicPolymerTraceIterator implements Iterator<PolymerTraceElement>
             value.center.element = this.traceElementIndex[residueIndex];
             value.centerNext.element = this.traceElementIndex[residueIndexNext1];
 
-            this.pos(this.p0, this.traceElementIndex[residueIndexPrev3]);
-            this.pos(this.p1, this.traceElementIndex[residueIndexPrev2]);
-            this.pos(this.p2, this.traceElementIndex[residueIndexPrev1]);
-            this.pos(this.p3, this.traceElementIndex[residueIndex]);
-            this.pos(this.p4, this.traceElementIndex[residueIndexNext1]);
-            this.pos(this.p5, this.traceElementIndex[residueIndexNext2]);
-            this.pos(this.p6, this.traceElementIndex[residueIndexNext3]);
+            const ssPrev3 = this.getSecStruc(residueIndexPrev3);
+            const ssPrev2 = this.getSecStruc(residueIndexPrev2);
+            const ssPrev1 = this.getSecStruc(residueIndexPrev1);
+            const ss = this.getSecStruc(residueIndex);
+            const ssNext1 = this.getSecStruc(residueIndexNext1);
+            const ssNext2 = this.getSecStruc(residueIndexNext2);
+            const ssNext3 = this.getSecStruc(residueIndexNext3);
+
+            this.pos(this.p0, residueIndexPrev3, ssPrev3);
+            this.pos(this.p1, residueIndexPrev2, ssPrev2);
+            this.pos(this.p2, residueIndexPrev1, ssPrev1);
+            this.pos(this.p3, residueIndex, ss);
+            this.pos(this.p4, residueIndexNext1, ssNext1);
+            this.pos(this.p5, residueIndexNext2, ssNext2);
+            this.pos(this.p6, residueIndexNext3, ssNext3);
+
+            const isHelixPrev3 = isHelixSS(ssPrev3);
+            const isHelixPrev2 = isHelixSS(ssPrev2);
+            const isHelixPrev1 = isHelixSS(ssPrev1);
+            const isHelix = isHelixSS(ss);
+            const isHelixNext1 = isHelixSS(ssNext1);
+            const isHelixNext2 = isHelixSS(ssNext2);
+            const isHelixNext3 = isHelixSS(ssNext3);
+
+            // handle positions for tubular helices
+            if (this.helixOrientationCenters) {
+                if (isHelix !== isHelixPrev1) {
+                    if (isHelix) {
+                        Vec3.copy(this.p0, this.p3);
+                        Vec3.copy(this.p1, this.p3);
+                        Vec3.copy(this.p2, this.p3);
+                    } else if(isHelixPrev1) {
+                        Vec3.scale(tmpDir, Vec3.sub(tmpDir, this.p2, this.p3), 2);
+                        Vec3.add(this.p2, this.p3, tmpDir);
+                        Vec3.add(this.p1, this.p2, tmpDir);
+                        Vec3.add(this.p0, this.p1, tmpDir);
+                    }
+                } else if (isHelix !== isHelixPrev2) {
+                    if (isHelix) {
+                        Vec3.copy(this.p0, this.p2);
+                        Vec3.copy(this.p1, this.p2);
+                    } else if(isHelixPrev2) {
+                        Vec3.scale(tmpDir, Vec3.sub(tmpDir, this.p1, this.p2), 2);
+                        Vec3.add(this.p1, this.p2, tmpDir);
+                        Vec3.add(this.p0, this.p1, tmpDir);
+                    }
+                } else if (isHelix !== isHelixPrev3) {
+                    if (isHelix) {
+                        Vec3.copy(this.p0, this.p1);
+                    } else if(isHelixPrev3) {
+                        Vec3.scale(tmpDir, Vec3.sub(tmpDir, this.p0, this.p1), 2);
+                        Vec3.add(this.p0, this.p1, tmpDir);
+                    }
+                }
+
+                if (isHelix !== isHelixNext1) {
+                    if (isHelix) {
+                        Vec3.copy(this.p4, this.p3);
+                        Vec3.copy(this.p5, this.p3);
+                        Vec3.copy(this.p6, this.p3);
+                    } else if(isHelixNext1) {
+                        Vec3.scale(tmpDir, Vec3.sub(tmpDir, this.p4, this.p3), 2);
+                        Vec3.add(this.p4, this.p3, tmpDir);
+                        Vec3.add(this.p5, this.p4, tmpDir);
+                        Vec3.add(this.p6, this.p5, tmpDir);
+                    }
+                } else if (isHelix !== isHelixNext2) {
+                    if (isHelix) {
+                        Vec3.copy(this.p5, this.p4);
+                        Vec3.copy(this.p6, this.p4);
+                    } else if(isHelixNext2) {
+                        Vec3.scale(tmpDir, Vec3.sub(tmpDir, this.p5, this.p4), 2);
+                        Vec3.add(this.p5, this.p4, tmpDir);
+                        Vec3.add(this.p6, this.p5, tmpDir);
+                    }
+                } else if (isHelix !== isHelixNext3) {
+                    if (isHelix) {
+                        Vec3.copy(this.p6, this.p5);
+                    } else if(isHelixNext3) {
+                        Vec3.scale(tmpDir, Vec3.sub(tmpDir, this.p6, this.p5), 2);
+                        Vec3.add(this.p6, this.p5, tmpDir);
+                    }
+                }
+            }
+
+            this.setFromToVector(this.d01, residueIndexPrev1, ssPrev1);
+            this.setFromToVector(this.d12, residueIndex, ss);
+            this.setFromToVector(this.d23, residueIndexNext1, ssNext1);
+            this.setFromToVector(this.d34, residueIndexNext2, ssNext2);
 
-            this.setFromToVector(this.d01, residueIndexPrev1);
-            this.setFromToVector(this.d12, residueIndex);
-            this.setFromToVector(this.d23, residueIndexNext1);
-            this.setFromToVector(this.d34, residueIndexNext2);
+            const helixFlag = isHelix && this.helixOrientationCenters;
 
             // extend termini
-            const f = 0.5;
-            if (residueIndex === residueIndexPrev1) {
-                Vec3.scale(tmpDir, Vec3.sub(tmpDir, this.p3, this.p4), f);
+            const f = 1.5;
+            if (residueIndex === residueIndexPrev1 || (ss !== ssPrev1 && helixFlag)) {
+                Vec3.setMagnitude(tmpDir, Vec3.sub(tmpDir, this.p3, this.p4), f);
                 Vec3.add(this.p2, this.p3, tmpDir);
                 Vec3.add(this.p1, this.p2, tmpDir);
                 Vec3.add(this.p0, this.p1, tmpDir);
-            } else if (residueIndexPrev1 === residueIndexPrev2) {
-                Vec3.scale(tmpDir, Vec3.sub(tmpDir, this.p2, this.p3), f);
+            } else if (residueIndexPrev1 === residueIndexPrev2 || (ss !== ssPrev2 && helixFlag)) {
+                Vec3.setMagnitude(tmpDir, Vec3.sub(tmpDir, this.p2, this.p3), f);
                 Vec3.add(this.p1, this.p2, tmpDir);
                 Vec3.add(this.p0, this.p1, tmpDir);
-            } else if (residueIndexPrev2 === residueIndexPrev3) {
-                Vec3.scale(tmpDir, Vec3.sub(tmpDir, this.p1, this.p2), f);
+            } else if (residueIndexPrev2 === residueIndexPrev3 || (ss !== ssPrev3 && helixFlag)) {
+                Vec3.setMagnitude(tmpDir, Vec3.sub(tmpDir, this.p1, this.p2), f);
                 Vec3.add(this.p0, this.p1, tmpDir);
             }
-            if (residueIndex === residueIndexNext1) {
-                Vec3.scale(tmpDir, Vec3.sub(tmpDir, this.p3, this.p2), f);
+            if (residueIndex === residueIndexNext1 || (ss !== ssNext1 && helixFlag)) {
+                Vec3.setMagnitude(tmpDir, Vec3.sub(tmpDir, this.p3, this.p2), f);
                 Vec3.add(this.p4, this.p3, tmpDir);
                 Vec3.add(this.p5, this.p4, tmpDir);
                 Vec3.add(this.p6, this.p5, tmpDir);
-            } else if (residueIndexNext1 === residueIndexNext2) {
-                Vec3.scale(tmpDir, Vec3.sub(tmpDir, this.p4, this.p3), f);
+            } else if (residueIndexNext1 === residueIndexNext2 || (ss !== ssNext2 && helixFlag)) {
+                Vec3.setMagnitude(tmpDir, Vec3.sub(tmpDir, this.p4, this.p3), f);
                 Vec3.add(this.p5, this.p4, tmpDir);
                 Vec3.add(this.p6, this.p5, tmpDir);
-            } else if (residueIndexNext2 === residueIndexNext3) {
-                Vec3.scale(tmpDir, Vec3.sub(tmpDir, this.p5, this.p4), f);
+            } else if (residueIndexNext2 === residueIndexNext3 || (ss !== ssNext3 && helixFlag)) {
+                Vec3.setMagnitude(tmpDir, Vec3.sub(tmpDir, this.p5, this.p4), f);
                 Vec3.add(this.p6, this.p5, tmpDir);
             }
 
-            this.setControlPoint(value.p0, this.p0, this.p1, this.p2, residueIndexPrev2);
-            this.setControlPoint(value.p1, this.p1, this.p2, this.p3, residueIndexPrev1);
-            this.setControlPoint(value.p2, this.p2, this.p3, this.p4, residueIndex);
-            this.setControlPoint(value.p3, this.p3, this.p4, this.p5, residueIndexNext1);
-            this.setControlPoint(value.p4, this.p4, this.p5, this.p6, residueIndexNext2);
+            this.setControlPoint(value.p0, this.p0, this.p1, this.p2, ssPrev2);
+            this.setControlPoint(value.p1, this.p1, this.p2, this.p3, ssPrev1);
+            this.setControlPoint(value.p2, this.p2, this.p3, this.p4, ss);
+            this.setControlPoint(value.p3, this.p3, this.p4, this.p5, ssNext1);
+            this.setControlPoint(value.p4, this.p4, this.p5, this.p6, ssNext2);
 
             this.setDirection(value.d12, this.d01, this.d12, this.d23);
             this.setDirection(value.d23, this.d12, this.d23, this.d34);
@@ -284,7 +398,7 @@ export class AtomicPolymerTraceIterator implements Iterator<PolymerTraceElement>
         return this.value;
     }
 
-    constructor(private unit: Unit.Atomic, structure: Structure, ignoreSecondaryStructure = false) {
+    constructor(private unit: Unit.Atomic, structure: Structure, options: PolymerTraceIteratorOptions = {}) {
         this.atomicConformation = unit.model.atomicConformation;
         this.residueAtomSegments = unit.model.atomicHierarchy.residueAtomSegments;
         this.polymerRanges = unit.model.atomicRanges.polymerRanges;
@@ -298,12 +412,15 @@ export class AtomicPolymerTraceIterator implements Iterator<PolymerTraceElement>
         this.value = createPolymerTraceElement(structure, unit);
         this.hasNext = this.residueIt.hasNext && this.polymerIt.hasNext;
 
-        const secondaryStructure = !ignoreSecondaryStructure && SecondaryStructureProvider.get(structure).value?.get(unit.invariantId);
-        if (secondaryStructure) {
-            const { type, getIndex } = secondaryStructure;
-            this.getSecStruc = (residueIndex: ResidueIndex) => type[getIndex(residueIndex as ResidueIndex)];
-        } else {
-            this.getSecStruc = (residueIndex: ResidueIndex) => SecStrucTypeNA;
+        if (!options.ignoreSecondaryStructure) {
+            this.secondaryStructure = SecondaryStructureProvider.get(structure).value?.get(unit.invariantId);
+        }
+
+        if (options.useHelixOrientation) {
+            const helixOrientation = HelixOrientationProvider.get(unit.model).value;
+            if (!helixOrientation) throw new Error('missing helix-orientation');
+
+            this.helixOrientationCenters = helixOrientation.centers;
         }
     }
 }