mol-plugin: basic structure labels support

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

@@ -201,6 +201,7 @@ namespace StructureElement {
                 let i = 0;
                 while (i < len) {
                     const rI = residueIndex[unitElements[OrderedSet.getAt(indices, i)]];
+                    i++;
                     while (i < len && residueIndex[unitElements[OrderedSet.getAt(indices, i)]] === rI) {
                         i++;
                     }

src/mol-plugin/index.ts

@@ -46,6 +46,7 @@ export const DefaultPluginSpec: PluginSpec = {
         PluginSpec.Action(StateTransforms.Model.ModelFromTrajectory),
         PluginSpec.Action(StateTransforms.Volume.VolumeFromCcp4),
         PluginSpec.Action(StateTransforms.Representation.StructureRepresentation3D),
+        PluginSpec.Action(StateTransforms.Representation.StructureLabels3D),
         PluginSpec.Action(StateTransforms.Representation.ExplodeStructureRepresentation3D),
         PluginSpec.Action(StateTransforms.Representation.VolumeRepresentation3D),
 

src/mol-plugin/state/objects.ts

@@ -15,6 +15,7 @@ import { VolumeRepresentation } from 'mol-repr/volume/representation';
 import { StateObject, StateTransformer } from 'mol-state';
 import { Ccp4File } from 'mol-io/reader/ccp4/schema';
 import { Dsn6File } from 'mol-io/reader/dsn6/schema';
+import { ShapeRepresentation } from 'mol-repr/shape/representation';
 
 export type TypeClass = 'root' | 'data' | 'prop'
 
@@ -67,7 +68,7 @@ export namespace PluginStateObject {
         export class Trajectory extends Create<ReadonlyArray<_Model>>({ name: 'Trajectory', typeClass: 'Object' }) { }
         export class Model extends Create<_Model>({ name: 'Model', typeClass: 'Object' }) { }
         export class Structure extends Create<_Structure>({ name: 'Structure', typeClass: 'Object' }) { }
-        export class Representation3D extends CreateRepresentation3D<StructureRepresentation<any>>({ name: 'Structure 3D' }) { }
+        export class Representation3D extends CreateRepresentation3D<StructureRepresentation<any> | ShapeRepresentation<any, any, any>>({ name: 'Structure 3D' }) { }
     }
 
     export namespace Volume {

src/mol-plugin/state/transforms/representation.ts

@@ -21,8 +21,18 @@ import { BuiltInSizeThemeName, SizeTheme } from 'mol-theme/size';
 import { createTheme, ThemeRegistryContext } from 'mol-theme/theme';
 import { ParamDefinition as PD } from 'mol-util/param-definition';
 import { PluginStateObject as SO, PluginStateTransform } from '../objects';
+import { Text } from 'mol-geo/geometry/text/text';
+import { ColorNames } from 'mol-util/color/tables';
+import { getLabelRepresentation } from 'mol-plugin/util/structure-labels';
+import { ShapeRepresentation } from 'mol-repr/shape/representation';
 
-export namespace StructureRepresentation3DHelpers {
+export { StructureRepresentation3D }
+export { StructureRepresentation3DHelpers }
+export { StructureLabels3D}
+export { ExplodeStructureRepresentation3D }
+export { VolumeRepresentation3D }
+
+namespace StructureRepresentation3DHelpers {
     export function getDefaultParams(ctx: PluginContext, name: BuiltInStructureRepresentationsName, structure: Structure, structureParams?: Partial<PD.Values<StructureParams>>): StateTransformer.Params<StructureRepresentation3D> {
         const type = ctx.structureRepresentation.registry.get(name);
 
@@ -100,9 +110,7 @@ export namespace StructureRepresentation3DHelpers {
         })
     }
 }
-export { StructureRepresentation3D };
-export { ExplodeStructureRepresentation3D };
-export { VolumeRepresentation3D };
+
 type StructureRepresentation3D = typeof StructureRepresentation3D
 const StructureRepresentation3D = PluginStateTransform.BuiltIn({
     name: 'structure-representation-3d',
@@ -151,9 +159,9 @@ const StructureRepresentation3D = PluginStateTransform.BuiltIn({
         })
     }
 })({
-    canAutoUpdate({ oldParams, newParams }) {
-        // TODO: allow for small molecules
-        return oldParams.type.name === newParams.type.name;
+    canAutoUpdate({ a, oldParams, newParams }) {
+        // TODO: other criteria as well?
+        return a.data.elementCount < 10000 && oldParams.type.name === newParams.type.name;
     },
     apply({ a, params }, plugin: PluginContext) {
         return Task.create('Structure Representation', async ctx => {
@@ -177,6 +185,45 @@ const StructureRepresentation3D = PluginStateTransform.BuiltIn({
     }
 });
 
+
+type StructureLabels3D = typeof StructureLabels3D
+const StructureLabels3D = PluginStateTransform.BuiltIn({
+    name: 'structure-labels-3d',
+    display: '3D Labels',
+    from: SO.Molecule.Structure,
+    to: SO.Molecule.Representation3D,
+    params: {
+        // TODO: other targets
+        target: PD.Select<'elements' | 'residues'>('residues', [['residues', 'Residues'], ['elements', 'Elements']]),
+        options: PD.Group({
+            ...Text.Params,
+
+            background: PD.Boolean(true),
+            backgroundMargin: PD.Numeric(0.2, { min: 0, max: 1, step: 0.01 }),
+            backgroundColor: PD.Color(ColorNames.snow),
+            backgroundOpacity: PD.Numeric(0.9, { min: 0, max: 1, step: 0.01 }),
+        })
+    }
+})({
+    canAutoUpdate({ a, oldParams, newParams }) {
+        // TODO: find good criteria
+        return false;
+    },
+    apply({ a, params }) {
+        return Task.create('Structure Labels', async ctx => {
+            const repr = await getLabelRepresentation(ctx, a.data, params);
+            return new SO.Molecule.Representation3D(repr, { label: `Labels`, description: params.target });
+        });
+    },
+    update({ a, b, newParams }) {
+        return Task.create('Structure Labels', async ctx => {
+            await getLabelRepresentation(ctx, a.data, newParams, b.data as ShapeRepresentation<any, any, any>);
+            return StateTransformer.UpdateResult.Updated;
+        });
+    }
+});
+
+
 type ExplodeStructureRepresentation3D = typeof ExplodeStructureRepresentation3D
 const ExplodeStructureRepresentation3D = PluginStateTransform.BuiltIn({
     name: 'explode-structure-representation-3d',

src/mol-plugin/util/structure-labels.ts

@@ -0,0 +1,111 @@
+/**
+ * Copyright (c) 2019 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author David Sehnal <david.sehnal@gmail.com>
+ */
+
+import { Structure, StructureElement, StructureProperties, Unit } from 'mol-model/structure';
+import { StateTransformer } from 'mol-state';
+import { StructureLabels3D } from '../state/transforms/representation';
+import { ShapeRepresentation } from 'mol-repr/shape/representation';
+import { Vec3 } from 'mol-math/linear-algebra';
+import { Text } from 'mol-geo/geometry/text/text';
+import { TextBuilder } from 'mol-geo/geometry/text/text-builder';
+import { Shape } from 'mol-model/shape';
+import { ColorNames } from 'mol-util/color/tables';
+import { RuntimeContext } from 'mol-task';
+import { ParamDefinition as PD } from 'mol-util/param-definition';
+import { BoundaryHelper } from 'mol-math/geometry/boundary-helper';
+
+interface LabelsData {
+    texts: string[],
+    positions: Vec3[],
+    sizes: number[],
+    depths: number[]
+}
+
+function getLabelsText(data: LabelsData, props: PD.Values<Text.Params>, text?: Text) {
+    const { texts, positions, depths } = data
+    const textBuilder = TextBuilder.create(props, texts.length * 10, texts.length * 10 / 2, text)
+    for (let i = 0, il = texts.length; i < il; ++i) {
+        const p = positions[i]
+        textBuilder.add(texts[i], p[0], p[1], p[2], depths[i], i)
+    }
+    return textBuilder.getText()
+}
+
+export async function getLabelRepresentation(ctx: RuntimeContext, structure: Structure, params: StateTransformer.Params<StructureLabels3D>, prev?: ShapeRepresentation<LabelsData, Text, Text.Params>) {
+    const repr = prev || ShapeRepresentation(getLabelsShape, Text.Utils);
+    const data = getLabelData(structure, params.target);
+    await repr.createOrUpdate(params.options, data).runInContext(ctx);
+    return repr;
+}
+
+function getLabelsShape(ctx: RuntimeContext, data: LabelsData, props: PD.Values<Text.Params>, shape?: Shape<Text>) {
+    const geo = getLabelsText(data, props, shape && shape.geometry);
+    return Shape.create('Scene Labels', geo, () => ColorNames.dimgrey, g => data.sizes[g], () => '')
+}
+
+const boundaryHelper = new BoundaryHelper();
+function getLabelData(structure: Structure, level: 'elements' | 'residues'): LabelsData {
+    const data: LabelsData = { texts: [], positions: [], sizes: [], depths: [] };
+
+    const l = StructureElement.create();
+    const { units } = structure;
+
+    const { auth_atom_id } = StructureProperties.atom;
+    const { auth_seq_id, auth_comp_id } = StructureProperties.residue;
+    const { auth_asym_id } = StructureProperties.chain;
+    const p = Vec3.zero();
+
+    for (const unit of units) {
+        // TODO: support coarse models
+
+        if (unit.kind !== Unit.Kind.Atomic) continue;
+        l.unit = unit;
+        const elements = unit.elements;
+
+        const pos = unit.conformation.position;
+
+        if (level === 'elements') {
+            for (let j = 0, _j = elements.length; j < _j; j++) {
+                l.element = elements[j];
+
+                pos(l.element, p);
+                data.texts.push(auth_atom_id(l));
+                data.positions.push(Vec3.clone(p));
+                data.sizes.push(1);
+                data.depths.push(2);
+            }
+        } else {
+            const residueIndex = unit.model.atomicHierarchy.residueAtomSegments.index;
+
+            let i = 0, len = elements.length;
+            while (i < len) {
+                const start = i, rI = residueIndex[elements[i]];
+                i++;
+                while (i < len && residueIndex[elements[i]] === rI) i++;
+
+                boundaryHelper.reset(0);
+                for (let eI = start; eI < i; eI++) {
+                    pos(elements[eI], p);
+                    boundaryHelper.boundaryStep(p, 0);
+                }
+                boundaryHelper.finishBoundaryStep();
+                for (let eI = start; eI < i; eI++) {
+                    pos(elements[eI], p);
+                    boundaryHelper.extendStep(p, 0);
+                }
+
+                l.element = elements[start];
+
+                data.texts.push(`${auth_comp_id(l)} ${auth_seq_id(l)}:${auth_asym_id(l)}`);
+                data.positions.push(Vec3.clone(boundaryHelper.center));
+                data.sizes.push(Math.max(1, boundaryHelper.radius / 5));
+                data.depths.push(boundaryHelper.radius);
+            }
+        }
+    }
+
+    return data;
+}