add backbone repr

- atomistic and coarse units

CHANGELOG.md

@@ -6,6 +6,7 @@ Note that since we don't clearly distinguish between a public and private interf
 
 ## [Unreleased]
 
+- Add backbone representation
 
 ## [v2.0.7] - 2021-06-23
 

src/mol-repr/structure/registry.ts

@@ -1,5 +1,5 @@
 /**
- * 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>
  */
@@ -21,6 +21,7 @@ import { PuttyRepresentationProvider } from './representation/putty';
 import { SpacefillRepresentationProvider } from './representation/spacefill';
 import { LineRepresentationProvider } from './representation/line';
 import { GaussianVolumeRepresentationProvider } from './representation/gaussian-volume';
+import { BackboneRepresentationProvider } from './representation/backbone';
 
 export class StructureRepresentationRegistry extends RepresentationRegistry<Structure, StructureRepresentationState> {
     constructor() {
@@ -35,6 +36,7 @@ export class StructureRepresentationRegistry extends RepresentationRegistry<Stru
 export namespace StructureRepresentationRegistry {
     export const BuiltIn = {
         'cartoon': CartoonRepresentationProvider,
+        'backbone': BackboneRepresentationProvider,
         'ball-and-stick': BallAndStickRepresentationProvider,
         'carbohydrate': CarbohydrateRepresentationProvider,
         'ellipsoid': EllipsoidRepresentationProvider,

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

@@ -1,29 +1,57 @@
-// /**
-//  * Copyright (c) 2018 mol* contributors, licensed under MIT, See LICENSE file for more info.
-//  *
-//  * @author Alexander Rose <alexander.rose@weirdbyte.de>
-//  */
+/**
+ * Copyright (c) 2018-2021 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author Alexander Rose <alexander.rose@weirdbyte.de>
+ */
 
-// import { PolymerBackboneVisual, PolymerBackboneParams } from '../visual/polymer-backbone-cylinder';
-// import { ParamDefinition as PD } from 'mol-util/param-definition';
-// import { UnitsRepresentation } from '../units-representation';
-// import { StructureRepresentation } from '../representation';
-// import { Representation } from 'mol-repr/representation';
-// import { ThemeRegistryContext } from 'mol-theme/theme';
-// import { Structure } from 'mol-model/structure';
+import { PolymerBackboneCylinderVisual, PolymerBackboneCylinderParams } from '../visual/polymer-backbone-cylinder';
+import { ParamDefinition as PD } from '../../../mol-util/param-definition';
+import { UnitsRepresentation } from '../units-representation';
+import { StructureRepresentation, StructureRepresentationProvider, StructureRepresentationStateBuilder } from '../representation';
+import { Representation, RepresentationContext, RepresentationParamsGetter } from '../../../mol-repr/representation';
+import { ThemeRegistryContext } from '../../../mol-theme/theme';
+import { Structure } from '../../../mol-model/structure';
+import { PolymerBackboneSphereParams, PolymerBackboneSphereVisual } from '../visual/polymer-backbone-sphere';
+import { PolymerGapParams, PolymerGapVisual } from '../visual/polymer-gap-cylinder';
 
-// export const BackboneParams = {
-//     ...PolymerBackboneParams,
-// }
-// export function getBackboneParams(ctx: ThemeRegistryContext, structure: Structure) {
-//     return BackboneParams // TODO return copy
-// }
-// export type BackboneProps = PD.DefaultValues<typeof BackboneParams>
+const BackboneVisuals = {
+    'polymer-backbone-cylinder': (ctx: RepresentationContext, getParams: RepresentationParamsGetter<Structure, PolymerBackboneCylinderParams>) => UnitsRepresentation('Polymer backbone cylinder', ctx, getParams, PolymerBackboneCylinderVisual),
+    'polymer-backbone-sphere': (ctx: RepresentationContext, getParams: RepresentationParamsGetter<Structure, PolymerBackboneSphereParams>) => UnitsRepresentation('Polymer backbone sphere', ctx, getParams, PolymerBackboneSphereVisual),
+    'polymer-gap': (ctx: RepresentationContext, getParams: RepresentationParamsGetter<Structure, PolymerGapParams>) => UnitsRepresentation('Polymer gap cylinder', ctx, getParams, PolymerGapVisual),
+};
 
-// export type BackboneRepresentation = StructureRepresentation<BackboneProps>
+export const BackboneParams = {
+    ...PolymerBackboneSphereParams,
+    ...PolymerBackboneCylinderParams,
+    ...PolymerGapParams,
+    sizeAspectRatio: PD.Numeric(1, { min: 0.1, max: 3, step: 0.1 }),
+    visuals: PD.MultiSelect(['polymer-backbone-cylinder', 'polymer-backbone-sphere', 'polymer-gap'], PD.objectToOptions(BackboneVisuals))
+};
+export type BackboneParams = typeof BackboneParams
+export function getBackboneParams(ctx: ThemeRegistryContext, structure: Structure) {
+    const params = PD.clone(BackboneParams);
+    let hasGaps = false;
+    structure.units.forEach(u => {
+        if (!hasGaps && u.gapElements.length) hasGaps = true;
+    });
+    params.visuals.defaultValue = ['polymer-backbone-cylinder', 'polymer-backbone-sphere'];
+    if (hasGaps) params.visuals.defaultValue.push('polymer-gap');
+    return params;
+}
 
-// export function BackboneRepresentation(defaultProps: BackboneProps): BackboneRepresentation {
-//     return Representation.createMulti('Backbone', defaultProps, [
-//         UnitsRepresentation('Polymer backbone cylinder', defaultProps, PolymerBackboneVisual)
-//     ])
-// }
+export type BackboneRepresentation = StructureRepresentation<BackboneParams>
+export function BackboneRepresentation(ctx: RepresentationContext, getParams: RepresentationParamsGetter<Structure, BackboneParams>): BackboneRepresentation {
+    return Representation.createMulti('Backbone', ctx, getParams, StructureRepresentationStateBuilder, BackboneVisuals as unknown as Representation.Def<Structure, BackboneParams>);
+}
+
+export const BackboneRepresentationProvider = StructureRepresentationProvider({
+    name: 'backbone',
+    label: 'Backbone',
+    description: 'Displays polymer backbone with cylinders and spheres.',
+    factory: BackboneRepresentation,
+    getParams: getBackboneParams,
+    defaultValues: PD.getDefaultValues(BackboneParams),
+    defaultColorTheme: { name: 'chain-id' },
+    defaultSizeTheme: { name: 'uniform' },
+    isApplicable: (structure: Structure) => structure.polymerResidueCount > 0,
+});

src/mol-repr/structure/visual/polymer-backbone-cylinder.ts

@@ -1,34 +1,102 @@
 /**
- * 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>
  */
 
 import { ParamDefinition as PD } from '../../../mol-util/param-definition';
 import { VisualContext } from '../../visual';
-import { Unit, Structure } from '../../../mol-model/structure';
+import { Unit, Structure, StructureElement, ElementIndex } from '../../../mol-model/structure';
 import { Theme } from '../../../mol-theme/theme';
 import { Mesh } from '../../../mol-geo/geometry/mesh/mesh';
 import { MeshBuilder } from '../../../mol-geo/geometry/mesh/mesh-builder';
 import { Vec3 } from '../../../mol-math/linear-algebra';
 import { CylinderProps } from '../../../mol-geo/primitive/cylinder';
-import { PolymerBackboneIterator } from './util/polymer';
-import { OrderedSet } from '../../../mol-data/int';
+import { eachPolymerElement, getPolymerElementLoci, NucleicShift, PolymerLocationIterator, StandardShift } from './util/polymer';
 import { addCylinder } from '../../../mol-geo/geometry/mesh/builder/cylinder';
-import { UnitsMeshParams, UnitsVisual, UnitsMeshVisual } from '../units-visual';
-import { ElementIterator, getElementLoci, eachElement } from './util/element';
+import { UnitsMeshParams, UnitsVisual, UnitsMeshVisual, UnitsCylindersVisual, UnitsCylindersParams, StructureGroup } from '../units-visual';
 import { VisualUpdateState } from '../../util';
 import { BaseGeometry } from '../../../mol-geo/geometry/base';
 import { Sphere3D } from '../../../mol-math/geometry';
+import { isNucleic, MoleculeType } from '../../../mol-model/structure/model/types';
+import { WebGLContext } from '../../../mol-gl/webgl/context';
+import { Cylinders } from '../../../mol-geo/geometry/cylinders/cylinders';
+import { CylindersBuilder } from '../../../mol-geo/geometry/cylinders/cylinders-builder';
+import { eachPolymerBackboneLink } from './util/polymer/backbone';
+
+// avoiding namespace lookup improved performance in Chrome (Aug 2020)
+const v3scale = Vec3.scale;
+const v3add = Vec3.add;
+const v3sub = Vec3.sub;
 
 export const PolymerBackboneCylinderParams = {
+    ...UnitsMeshParams,
+    ...UnitsCylindersParams,
     sizeFactor: PD.Numeric(0.3, { min: 0, max: 10, step: 0.01 }),
     radialSegments: PD.Numeric(16, { min: 2, max: 56, step: 2 }, BaseGeometry.CustomQualityParamInfo),
+    tryUseImpostor: PD.Boolean(true),
 };
-export const DefaultPolymerBackboneCylinderProps = PD.getDefaultValues(PolymerBackboneCylinderParams);
-export type PolymerBackboneCylinderProps = typeof DefaultPolymerBackboneCylinderProps
+export type PolymerBackboneCylinderParams = typeof PolymerBackboneCylinderParams
+
+export function PolymerBackboneCylinderVisual(materialId: number, structure: Structure, props: PD.Values<PolymerBackboneCylinderParams>, webgl?: WebGLContext) {
+    return props.tryUseImpostor && webgl && webgl.extensions.fragDepth
+        ? PolymerBackboneCylinderImpostorVisual(materialId)
+        : PolymerBackboneCylinderMeshVisual(materialId);
+}
+
+interface PolymerBackboneCylinderProps {
+    radialSegments: number,
+    sizeFactor: number,
+}
+
+function createPolymerBackboneCylinderImpostor(ctx: VisualContext, unit: Unit, structure: Structure, theme: Theme, props: PolymerBackboneCylinderProps, cylinders?: Cylinders) {
+    const polymerElementCount = unit.polymerElements.length;
+    if (!polymerElementCount) return Cylinders.createEmpty(cylinders);
+
+    const cylindersCountEstimate = polymerElementCount * 2;
+    const builder = CylindersBuilder.create(cylindersCountEstimate, cylindersCountEstimate / 4, cylinders);
+
+    const pos = unit.conformation.invariantPosition;
+    const pA = Vec3();
+    const pB = Vec3();
+    const pM = Vec3();
+
+    const add = function(indexA: ElementIndex, indexB: ElementIndex, groupA: number, groupB: number, moleculeType: MoleculeType) {
+        pos(indexA, pA);
+        pos(indexB, pB);
+
+        const isNucleicType = isNucleic(moleculeType);
+        const shift = isNucleicType ? NucleicShift : StandardShift;
+
+        v3add(pM, pA, v3scale(pM, v3sub(pM, pB, pA), shift));
+        builder.add(pA[0], pA[1], pA[2], pM[0], pM[1], pM[2], 1, false, false, groupA);
+        builder.add(pM[0], pM[1], pM[2], pB[0], pB[1], pB[2], 1, false, false, groupB);
+    };
+
+    eachPolymerBackboneLink(unit, add);
+
+    const c = builder.getCylinders();
+
+    const sphere = Sphere3D.expand(Sphere3D(), unit.boundary.sphere, 1 * props.sizeFactor);
+    c.setBoundingSphere(sphere);
+
+    return c;
+}
+
+export function PolymerBackboneCylinderImpostorVisual(materialId: number): UnitsVisual<PolymerBackboneCylinderParams> {
+    return UnitsCylindersVisual<PolymerBackboneCylinderParams>({
+        defaultProps: PD.getDefaultValues(PolymerBackboneCylinderParams),
+        createGeometry: createPolymerBackboneCylinderImpostor,
+        createLocationIterator: PolymerLocationIterator.fromGroup,
+        getLoci: getPolymerElementLoci,
+        eachLocation: eachPolymerElement,
+        setUpdateState: (state: VisualUpdateState, newProps: PD.Values<PolymerBackboneCylinderParams>, currentProps: PD.Values<PolymerBackboneCylinderParams>) => { },
+        mustRecreate: (structureGroup: StructureGroup, props: PD.Values<PolymerBackboneCylinderParams>, webgl?: WebGLContext) => {
+            return !props.tryUseImpostor || !webgl;
+        }
+    }, materialId);
+}
 
-// TODO do group id based on polymer index not element index
 function createPolymerBackboneCylinderMesh(ctx: VisualContext, unit: Unit, structure: Structure, theme: Theme, props: PolymerBackboneCylinderProps, mesh?: Mesh) {
     const polymerElementCount = unit.polymerElements.length;
     if (!polymerElementCount) return Mesh.createEmpty(mesh);
@@ -38,26 +106,34 @@ function createPolymerBackboneCylinderMesh(ctx: VisualContext, unit: Unit, struc
     const vertexCountEstimate = radialSegments * 2 * polymerElementCount * 2;
     const builderState = MeshBuilder.createState(vertexCountEstimate, vertexCountEstimate / 10, mesh);
 
-    const { elements } = unit;
     const pos = unit.conformation.invariantPosition;
-    const pA = Vec3.zero();
-    const pB = Vec3.zero();
+    const pA = Vec3();
+    const pB = Vec3();
     const cylinderProps: CylinderProps = { radiusTop: 1, radiusBottom: 1, radialSegments };
 
-    const polymerBackboneIt = PolymerBackboneIterator(structure, unit);
-    while (polymerBackboneIt.hasNext) {
-        const { centerA, centerB } = polymerBackboneIt.move();
+    const centerA = StructureElement.Location.create(structure, unit);
+    const centerB = StructureElement.Location.create(structure, unit);
+
+    const add = function(indexA: ElementIndex, indexB: ElementIndex, groupA: number, groupB: number, moleculeType: MoleculeType) {
+        centerA.element = indexA;
+        centerB.element = indexB;
+
         pos(centerA.element, pA);
         pos(centerB.element, pB);
 
+        const isNucleicType = isNucleic(moleculeType);
+        const shift = isNucleicType ? NucleicShift : StandardShift;
+
         cylinderProps.radiusTop = cylinderProps.radiusBottom = theme.size.size(centerA) * sizeFactor;
-        builderState.currentGroup = OrderedSet.indexOf(elements, centerA.element);
-        addCylinder(builderState, pA, pB, 0.5, cylinderProps);
+        builderState.currentGroup = groupA;
+        addCylinder(builderState, pA, pB, shift, cylinderProps);
 
         cylinderProps.radiusTop = cylinderProps.radiusBottom = theme.size.size(centerB) * sizeFactor;
-        builderState.currentGroup = OrderedSet.indexOf(elements, centerB.element);
-        addCylinder(builderState, pB, pA, 0.5, cylinderProps);
-    }
+        builderState.currentGroup = groupB;
+        addCylinder(builderState, pB, pA, 1 - shift, cylinderProps);
+    };
+
+    eachPolymerBackboneLink(unit, add);
 
     const m = MeshBuilder.getMesh(builderState);
 
@@ -67,25 +143,21 @@ function createPolymerBackboneCylinderMesh(ctx: VisualContext, unit: Unit, struc
     return m;
 }
 
-export const PolymerBackboneParams = {
-    ...UnitsMeshParams,
-    ...PolymerBackboneCylinderParams,
-};
-export type PolymerBackboneParams = typeof PolymerBackboneParams
-
-export function PolymerBackboneVisual(materialId: number): UnitsVisual<PolymerBackboneParams> {
-    return UnitsMeshVisual<PolymerBackboneParams>({
-        defaultProps: PD.getDefaultValues(PolymerBackboneParams),
+export function PolymerBackboneCylinderMeshVisual(materialId: number): UnitsVisual<PolymerBackboneCylinderParams> {
+    return UnitsMeshVisual<PolymerBackboneCylinderParams>({
+        defaultProps: PD.getDefaultValues(PolymerBackboneCylinderParams),
         createGeometry: createPolymerBackboneCylinderMesh,
-        // TODO create a specialized location iterator
-        createLocationIterator: ElementIterator.fromGroup,
-        getLoci: getElementLoci,
-        eachLocation: eachElement,
-        setUpdateState: (state: VisualUpdateState, newProps: PD.Values<PolymerBackboneParams>, currentProps: PD.Values<PolymerBackboneParams>) => {
+        createLocationIterator: PolymerLocationIterator.fromGroup,
+        getLoci: getPolymerElementLoci,
+        eachLocation: eachPolymerElement,
+        setUpdateState: (state: VisualUpdateState, newProps: PD.Values<PolymerBackboneCylinderParams>, currentProps: PD.Values<PolymerBackboneCylinderParams>) => {
             state.createGeometry = (
                 newProps.sizeFactor !== currentProps.sizeFactor ||
                 newProps.radialSegments !== currentProps.radialSegments
             );
+        },
+        mustRecreate: (structureGroup: StructureGroup, props: PD.Values<PolymerBackboneCylinderParams>, webgl?: WebGLContext) => {
+            return props.tryUseImpostor && !!webgl;
         }
     }, materialId);
 }

src/mol-repr/structure/visual/polymer-backbone-sphere.ts

@@ -1 +1,131 @@
-// TODO
+/**
+ * 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 { VisualContext } from '../../visual';
+import { Unit, Structure, ElementIndex, StructureElement } from '../../../mol-model/structure';
+import { Theme } from '../../../mol-theme/theme';
+import { Mesh } from '../../../mol-geo/geometry/mesh/mesh';
+import { MeshBuilder } from '../../../mol-geo/geometry/mesh/mesh-builder';
+import { Vec3 } from '../../../mol-math/linear-algebra';
+import { eachPolymerElement, getPolymerElementLoci, PolymerLocationIterator } from './util/polymer';
+import { UnitsMeshParams, UnitsVisual, UnitsMeshVisual, UnitsSpheresVisual, UnitsSpheresParams, StructureGroup } from '../units-visual';
+import { VisualUpdateState } from '../../util';
+import { BaseGeometry } from '../../../mol-geo/geometry/base';
+import { Sphere3D } from '../../../mol-math/geometry';
+import { addSphere } from '../../../mol-geo/geometry/mesh/builder/sphere';
+import { sphereVertexCount } from '../../../mol-geo/primitive/sphere';
+import { WebGLContext } from '../../../mol-gl/webgl/context';
+import { Spheres } from '../../../mol-geo/geometry/spheres/spheres';
+import { SpheresBuilder } from '../../../mol-geo/geometry/spheres/spheres-builder';
+import { eachPolymerBackboneElement } from './util/polymer/backbone';
+
+export const PolymerBackboneSphereParams = {
+    ...UnitsMeshParams,
+    ...UnitsSpheresParams,
+    sizeFactor: PD.Numeric(0.3, { min: 0, max: 10, step: 0.01 }),
+    detail: PD.Numeric(0, { min: 0, max: 3, step: 1 }, BaseGeometry.CustomQualityParamInfo),
+    tryUseImpostor: PD.Boolean(true),
+};
+export type PolymerBackboneSphereParams = typeof PolymerBackboneSphereParams
+
+export function PolymerBackboneSphereVisual(materialId: number, structure: Structure, props: PD.Values<PolymerBackboneSphereParams>, webgl?: WebGLContext) {
+    return props.tryUseImpostor && webgl && webgl.extensions.fragDepth
+        ? PolymerBackboneSphereImpostorVisual(materialId)
+        : PolymerBackboneSphereMeshVisual(materialId);
+}
+
+interface PolymerBackboneSphereProps {
+    detail: number,
+    sizeFactor: number,
+}
+
+function createPolymerBackboneSphereImpostor(ctx: VisualContext, unit: Unit, structure: Structure, theme: Theme, props: PolymerBackboneSphereProps, spheres?: Spheres) {
+    const polymerElementCount = unit.polymerElements.length;
+    if (!polymerElementCount) return Spheres.createEmpty(spheres);
+
+    const builder = SpheresBuilder.create(polymerElementCount, polymerElementCount / 2, spheres);
+
+    const pos = unit.conformation.invariantPosition;
+    const p = Vec3();
+
+    const add = (index: ElementIndex, group: number) => {
+        pos(index, p);
+        builder.add(p[0], p[1], p[2], group);
+    };
+
+    eachPolymerBackboneElement(unit, add);
+
+    const s = builder.getSpheres();
+
+    const sphere = Sphere3D.expand(Sphere3D(), unit.boundary.sphere, 1 * props.sizeFactor);
+    s.setBoundingSphere(sphere);
+
+    return s;
+}
+
+export function PolymerBackboneSphereImpostorVisual(materialId: number): UnitsVisual<PolymerBackboneSphereParams> {
+    return UnitsSpheresVisual<PolymerBackboneSphereParams>({
+        defaultProps: PD.getDefaultValues(PolymerBackboneSphereParams),
+        createGeometry: createPolymerBackboneSphereImpostor,
+        createLocationIterator: PolymerLocationIterator.fromGroup,
+        getLoci: getPolymerElementLoci,
+        eachLocation: eachPolymerElement,
+        setUpdateState: (state: VisualUpdateState, newProps: PD.Values<PolymerBackboneSphereParams>, currentProps: PD.Values<PolymerBackboneSphereParams>) => { },
+        mustRecreate: (structureGroup: StructureGroup, props: PD.Values<PolymerBackboneSphereParams>, webgl?: WebGLContext) => {
+            return !props.tryUseImpostor || !webgl;
+        }
+    }, materialId);
+}
+
+function createPolymerBackboneSphereMesh(ctx: VisualContext, unit: Unit, structure: Structure, theme: Theme, props: PolymerBackboneSphereProps, mesh?: Mesh) {
+    const polymerElementCount = unit.polymerElements.length;
+    if (!polymerElementCount) return Mesh.createEmpty(mesh);
+
+    const { detail, sizeFactor } = props;
+
+    const vertexCount = polymerElementCount * sphereVertexCount(detail);
+    const builderState = MeshBuilder.createState(vertexCount, vertexCount / 2, mesh);
+
+    const pos = unit.conformation.invariantPosition;
+    const p = Vec3();
+    const center = StructureElement.Location.create(structure, unit);
+
+    const add = (index: ElementIndex, group: number) => {
+        center.element = index;
+        pos(center.element, p);
+        builderState.currentGroup = group;
+        addSphere(builderState, p, theme.size.size(center) * sizeFactor, detail);
+    };
+
+    eachPolymerBackboneElement(unit, add);
+
+    const m = MeshBuilder.getMesh(builderState);
+
+    const sphere = Sphere3D.expand(Sphere3D(), unit.boundary.sphere, 1 * props.sizeFactor);
+    m.setBoundingSphere(sphere);
+
+    return m;
+}
+
+export function PolymerBackboneSphereMeshVisual(materialId: number): UnitsVisual<PolymerBackboneSphereParams> {
+    return UnitsMeshVisual<PolymerBackboneSphereParams>({
+        defaultProps: PD.getDefaultValues(PolymerBackboneSphereParams),
+        createGeometry: createPolymerBackboneSphereMesh,
+        createLocationIterator: PolymerLocationIterator.fromGroup,
+        getLoci: getPolymerElementLoci,
+        eachLocation: eachPolymerElement,
+        setUpdateState: (state: VisualUpdateState, newProps: PD.Values<PolymerBackboneSphereParams>, currentProps: PD.Values<PolymerBackboneSphereParams>) => {
+            state.createGeometry = (
+                newProps.sizeFactor !== currentProps.sizeFactor ||
+                newProps.detail !== currentProps.detail
+            );
+        },
+        mustRecreate: (structureGroup: StructureGroup, props: PD.Values<PolymerBackboneSphereParams>, webgl?: WebGLContext) => {
+            return props.tryUseImpostor && !!webgl;
+        }
+    }, materialId);
+}

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

@@ -14,7 +14,7 @@ import { PickingId } from '../../../../mol-geo/geometry/picking';
 import { StructureGroup } from '../../../structure/units-visual';
 import { getResidueLoci } from './common';
 
-export * from './polymer/backbone-iterator';
+export * from './polymer/backbone';
 export * from './polymer/gap-iterator';
 export * from './polymer/trace-iterator';
 export * from './polymer/curve-segment';

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

@@ -1,134 +0,0 @@
-/**
- * Copyright (c) 2018 mol* contributors, licensed under MIT, See LICENSE file for more info.
- *
- * @author Alexander Rose <alexander.rose@weirdbyte.de>
- */
-
-import { Unit, Structure, StructureElement, ElementIndex, ResidueIndex } from '../../../../../mol-model/structure';
-import { Segmentation } from '../../../../../mol-data/int';
-import { Iterator } from '../../../../../mol-data/iterator';
-import { SortedRanges } from '../../../../../mol-data/int/sorted-ranges';
-import { getPolymerRanges } from '../polymer';
-
-/** Iterates over consecutive pairs of residues/coarse elements in polymers */
-export function PolymerBackboneIterator(structure: Structure, unit: Unit): Iterator<PolymerBackbonePair> {
-    switch (unit.kind) {
-        case Unit.Kind.Atomic: return new AtomicPolymerBackboneIterator(structure, unit);
-        case Unit.Kind.Spheres:
-        case Unit.Kind.Gaussians:
-            return new CoarsePolymerBackboneIterator(structure, unit);
-    }
-}
-
-interface PolymerBackbonePair {
-    centerA: StructureElement.Location
-    centerB: StructureElement.Location
-}
-
-function createPolymerBackbonePair (structure: Structure, unit: Unit) {
-    return {
-        centerA: StructureElement.Location.create(structure, unit),
-        centerB: StructureElement.Location.create(structure, unit),
-    };
-}
-
-const enum AtomicPolymerBackboneIteratorState { nextPolymer, firstResidue, nextResidue, cycle }
-
-export class AtomicPolymerBackboneIterator implements Iterator<PolymerBackbonePair> {
-    private traceElementIndex: ArrayLike<ElementIndex>
-    private value: PolymerBackbonePair
-    private polymerIt: SortedRanges.Iterator<ElementIndex, ResidueIndex>
-    private residueIt: Segmentation.SegmentIterator<ResidueIndex>
-    private state: AtomicPolymerBackboneIteratorState = AtomicPolymerBackboneIteratorState.nextPolymer
-    private residueSegment: Segmentation.Segment<ResidueIndex>
-    hasNext: boolean = false;
-
-    move() {
-        if (this.state === AtomicPolymerBackboneIteratorState.nextPolymer) {
-            while (this.polymerIt.hasNext) {
-                this.residueIt.setSegment(this.polymerIt.move());
-                if (this.residueIt.hasNext) {
-                    this.residueSegment = this.residueIt.move();
-                    this.value.centerB.element = this.traceElementIndex[this.residueSegment.index];
-                    this.state = AtomicPolymerBackboneIteratorState.nextResidue;
-                    break;
-                }
-            }
-        }
-
-        if (this.state === AtomicPolymerBackboneIteratorState.nextResidue) {
-            this.residueSegment = this.residueIt.move();
-            this.value.centerA.element = this.value.centerB.element;
-            this.value.centerB.element = this.traceElementIndex[this.residueSegment.index];
-            if (!this.residueIt.hasNext) {
-                if (this.unit.model.atomicRanges.cyclicPolymerMap.has(this.residueSegment.index)) {
-                    this.state = AtomicPolymerBackboneIteratorState.cycle;
-                } else {
-                    // TODO need to advance to a polymer that has two or more residues (can't assume it has)
-                    this.state = AtomicPolymerBackboneIteratorState.nextPolymer;
-                }
-            }
-        } else if (this.state === AtomicPolymerBackboneIteratorState.cycle) {
-            const { cyclicPolymerMap } = this.unit.model.atomicRanges;
-            this.value.centerA.element = this.value.centerB.element;
-            this.value.centerB.element = this.traceElementIndex[cyclicPolymerMap.get(this.residueSegment.index)!];
-            // TODO need to advance to a polymer that has two or more residues (can't assume it has)
-            this.state = AtomicPolymerBackboneIteratorState.nextPolymer;
-        }
-
-        this.hasNext = this.residueIt.hasNext || this.polymerIt.hasNext || this.state === AtomicPolymerBackboneIteratorState.cycle;
-        return this.value;
-    }
-
-    constructor(structure: Structure, private unit: Unit.Atomic) {
-        this.traceElementIndex = unit.model.atomicHierarchy.derived.residue.traceElementIndex as ArrayLike<ElementIndex>; // can assume it won't be -1 for polymer residues
-        this.polymerIt = SortedRanges.transientSegments(getPolymerRanges(unit), unit.elements);
-        this.residueIt = Segmentation.transientSegments(unit.model.atomicHierarchy.residueAtomSegments, unit.elements);
-        this.value = createPolymerBackbonePair(structure, unit);
-        this.hasNext = this.residueIt.hasNext && this.polymerIt.hasNext;
-    }
-}
-
-const enum CoarsePolymerBackboneIteratorState { nextPolymer, nextElement }
-
-export class CoarsePolymerBackboneIterator implements Iterator<PolymerBackbonePair> {
-    private value: PolymerBackbonePair
-    private polymerIt: SortedRanges.Iterator<ElementIndex, ResidueIndex>
-    private polymerSegment: Segmentation.Segment<ResidueIndex>
-    private state: CoarsePolymerBackboneIteratorState = CoarsePolymerBackboneIteratorState.nextPolymer
-    private elementIndex: number
-    hasNext: boolean = false;
-
-    move() {
-        if (this.state === CoarsePolymerBackboneIteratorState.nextPolymer) {
-            if (this.polymerIt.hasNext) {
-                this.polymerSegment = this.polymerIt.move();
-                this.elementIndex = this.polymerSegment.start;
-                if (this.elementIndex + 1 < this.polymerSegment.end) {
-                    this.value.centerB.element = this.unit.elements[this.elementIndex];
-                    this.state = CoarsePolymerBackboneIteratorState.nextElement;
-                } else {
-                    this.state = CoarsePolymerBackboneIteratorState.nextPolymer;
-                }
-            }
-        }
-
-        if (this.state === CoarsePolymerBackboneIteratorState.nextElement) {
-            this.elementIndex += 1;
-            this.value.centerA.element = this.value.centerB.element;
-            this.value.centerB.element = this.unit.elements[this.elementIndex];
-            if (this.elementIndex + 1 >= this.polymerSegment.end) {
-                this.state = CoarsePolymerBackboneIteratorState.nextPolymer;
-            }
-        }
-
-        this.hasNext = this.elementIndex + 1 < this.polymerSegment.end || this.polymerIt.hasNext;
-        return this.value;
-    }
-
-    constructor(structure: Structure, private unit: Unit.Spheres | Unit.Gaussians) {
-        this.polymerIt = SortedRanges.transientSegments(getPolymerRanges(unit), unit.elements);
-        this.value = createPolymerBackbonePair(structure, unit);
-        this.hasNext = this.polymerIt.hasNext;
-    }
-}

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

@@ -0,0 +1,126 @@
+/**
+ * Copyright (c) 2021 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author Alexander Rose <alexander.rose@weirdbyte.de>
+ */
+import { Segmentation } from '../../../../../mol-data/int';
+import { SortedRanges } from '../../../../../mol-data/int/sorted-ranges';
+import { ElementIndex, ResidueIndex, Unit } from '../../../../../mol-model/structure';
+import { MoleculeType } from '../../../../../mol-model/structure/model/types';
+import { getPolymerRanges } from '../polymer';
+
+export type PolymerBackboneLinkCallback = (indexA: ElementIndex, indexB: ElementIndex, groupA: number, groupB: number, moleculeType: MoleculeType) => void
+
+export function eachPolymerBackboneLink(unit: Unit, callback: PolymerBackboneLinkCallback) {
+    switch (unit.kind) {
+        case Unit.Kind.Atomic: return eachAtomicPolymerBackboneLink(unit, callback);
+        case Unit.Kind.Spheres:
+        case Unit.Kind.Gaussians:
+            return eachCoarsePolymerBackboneLink(unit, callback);
+    }
+}
+
+function eachAtomicPolymerBackboneLink(unit: Unit.Atomic, callback: PolymerBackboneLinkCallback) {
+    const cyclicPolymerMap = unit.model.atomicRanges.cyclicPolymerMap;
+    const polymerIt = SortedRanges.transientSegments(getPolymerRanges(unit), unit.elements);
+    const residueIt = Segmentation.transientSegments(unit.model.atomicHierarchy.residueAtomSegments, unit.elements);
+    const traceElementIndex = unit.model.atomicHierarchy.derived.residue.traceElementIndex as ArrayLike<ElementIndex>; // can assume it won't be -1 for polymer residues
+    const { moleculeType } = unit.model.atomicHierarchy.derived.residue;
+
+    let indexA = -1 as ResidueIndex;
+    let indexB = -1 as ResidueIndex;
+    let isFirst = true;
+    let firstGroup = -1;
+    let i = 0;
+    while (polymerIt.hasNext) {
+        isFirst = true;
+        firstGroup = i;
+        residueIt.setSegment(polymerIt.move());
+        while (residueIt.hasNext) {
+            if (isFirst) {
+                const index_1 = residueIt.move().index;
+                ++i;
+                if (!residueIt.hasNext)
+                    continue;
+                isFirst = false;
+                indexB = index_1;
+            }
+            const index = residueIt.move().index;
+            indexA = indexB;
+            indexB = index;
+            callback(traceElementIndex[indexA], traceElementIndex[indexB], i - 1, i, moleculeType[indexA]);
+            ++i;
+        }
+        if (cyclicPolymerMap.has(indexB)) {
+            indexA = indexB;
+            indexB = cyclicPolymerMap.get(indexA)!;
+            callback(traceElementIndex[indexA], traceElementIndex[indexB], i - 1, firstGroup, moleculeType[indexA]);
+        }
+    }
+}
+
+function eachCoarsePolymerBackboneLink(unit: Unit.Spheres | Unit.Gaussians, callback: PolymerBackboneLinkCallback) {
+    const polymerIt = SortedRanges.transientSegments(getPolymerRanges(unit), unit.elements);
+    const { elements } = unit;
+
+    let isFirst = true;
+    let i = 0;
+    while (polymerIt.hasNext) {
+        isFirst = true;
+        const _a = polymerIt.move(), start = _a.start, end = _a.end;
+        for (let j = start, jl = end; j < jl; ++j) {
+            if (isFirst) {
+                ++j;
+                ++i;
+                if (j > jl)
+                    continue;
+                isFirst = false;
+            }
+            callback(elements[j - 1], elements[j], i - 1, i, 0 /* Unknown */);
+            ++i;
+        }
+    }
+}
+
+//
+
+export type PolymerBackboneElementCallback = (index: ElementIndex, group: number) => void
+
+export function eachPolymerBackboneElement(unit: Unit, callback: PolymerBackboneElementCallback) {
+    switch (unit.kind) {
+        case Unit.Kind.Atomic: return eachAtomicPolymerBackboneElement(unit, callback);
+        case Unit.Kind.Spheres:
+        case Unit.Kind.Gaussians:
+            return eachCoarsePolymerBackboneElement(unit, callback);
+    }
+}
+
+export function eachAtomicPolymerBackboneElement(unit: Unit.Atomic, callback: PolymerBackboneElementCallback) {
+    const polymerIt = SortedRanges.transientSegments(getPolymerRanges(unit), unit.elements);
+    const residueIt = Segmentation.transientSegments(unit.model.atomicHierarchy.residueAtomSegments, unit.elements);
+    const traceElementIndex = unit.model.atomicHierarchy.derived.residue.traceElementIndex as ArrayLike<ElementIndex>; // can assume it won't be -1 for polymer residues
+
+    let i = 0;
+    while (polymerIt.hasNext) {
+        residueIt.setSegment(polymerIt.move());
+        while (residueIt.hasNext) {
+            const index = residueIt.move().index;
+            callback(traceElementIndex[index], i);
+            ++i;
+        }
+    }
+}
+
+function eachCoarsePolymerBackboneElement(unit: Unit.Spheres | Unit.Gaussians, callback: PolymerBackboneElementCallback) {
+    const polymerIt = SortedRanges.transientSegments(getPolymerRanges(unit), unit.elements);
+    const { elements } = unit;
+
+    let i = 0;
+    while (polymerIt.hasNext) {
+        const _a = polymerIt.move(), start = _a.start, end = _a.end;
+        for (let j = start, jl = end; j < jl; ++j) {
+            callback(elements[j], i);
+            ++i;
+        }
+    }
+}