mol-plugin: split structure into seq/wtr/lig before visualization

src/mol-model/structure/query.ts

@@ -10,13 +10,15 @@ export * from './query/context'
 import * as generators from './query/queries/generators'
 import * as modifiers from './query/queries/modifiers'
 import * as combinators from './query/queries/combinators'
+import * as internal from './query/queries/internal'
 import pred from './query/predicates'
 
 export const Queries = {
     generators,
     modifiers,
     combinators,
-    pred
+    pred,
+    internal
 }
 
 export { StructureSelection, StructureQuery }
\ No newline at end of file

src/mol-model/structure/query/queries/internal.ts

+/**
+ * Copyright (c) 2017 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author David Sehnal <david.sehnal@gmail.com>
+ */
+
+import { Segmentation } from 'mol-data/int';
+import StructureElement from 'mol-model/structure/structure/element';
+import { StructureProperties as P, Unit } from '../../structure';
+import Structure from '../../structure/structure';
+import { StructureQuery } from '../query';
+import { StructureSelection } from '../selection';
+
+export function sequence(): StructureQuery {
+    return ctx => {
+        const { inputStructure } = ctx;
+        const l = StructureElement.create();
+
+        const units: Unit[] = [];
+        for (const unit of inputStructure.units) {
+            l.unit = unit;
+            const elements = unit.elements;
+            l.element = elements[0];
+            if (P.entity.type(l) !== 'polymer') continue;
+
+            const residuesIt = Segmentation.transientSegments(unit.model.atomicHierarchy.residueAtomSegments, elements);
+            let residueCount = 0;
+            while (residuesIt.hasNext) {
+                residueCount++;
+                residuesIt.move();
+            }
+
+            if (residueCount < 8) continue;
+
+            units.push(unit);
+        }
+        return StructureSelection.Singletons(inputStructure, new Structure(units));
+    };
+}
+
+export function water(): StructureQuery {
+    return ctx => {
+        const { inputStructure } = ctx;
+        const l = StructureElement.create();
+
+        const units: Unit[] = [];
+        for (const unit of inputStructure.units) {
+            l.unit = unit;
+            const elements = unit.elements;
+            l.element = elements[0];
+            if (P.entity.type(l) !== 'water') continue;
+            units.push(unit);
+        }
+        return StructureSelection.Singletons(inputStructure, new Structure(units));
+    };
+}
+
+export function lidangs(): StructureQuery {
+    return ctx => {
+        const { inputStructure } = ctx;
+        const l = StructureElement.create();
+
+        const units: Unit[] = [];
+        for (const unit of inputStructure.units) {
+            l.unit = unit;
+            const elements = unit.elements;
+            l.element = elements[0];
+            if (P.entity.type(l) === 'water') continue;
+            if (P.entity.type(l) === 'polymer') {
+                const residuesIt = Segmentation.transientSegments(unit.model.atomicHierarchy.residueAtomSegments, elements);
+                let residueCount = 0;
+                while (residuesIt.hasNext) {
+                    residueCount++;
+                    residuesIt.move();
+                }
+
+                if (residueCount >= 8) continue;
+            }
+
+            units.push(unit);
+        }
+        return StructureSelection.Singletons(inputStructure, new Structure(units));
+    };
+}

src/mol-plugin/state/actions/basic.ts

@@ -10,6 +10,7 @@ import { StateTransforms } from '../transforms';
 import { ParamDefinition as PD } from 'mol-util/param-definition';
 import { StateSelection } from 'mol-state/state/selection';
 import { CartoonParams } from 'mol-repr/structure/representation/cartoon';
+import { BallAndStickParams } from 'mol-repr/structure/representation/ball-and-stick';
 
 export const CreateStructureFromPDBe = StateAction.create<PluginStateObject.Root, void, { id: string }>({
     from: [PluginStateObject.Root],
@@ -33,22 +34,21 @@ export const CreateStructureFromPDBe = StateAction.create<PluginStateObject.Root
         //     ])
         // });
 
-        const newTree = b.toRoot()
+        const root = b.toRoot()
             .apply(StateTransforms.Data.Download, { url })
             .apply(StateTransforms.Data.ParseCif)
             .apply(StateTransforms.Model.TrajectoryFromMmCif, {})
             .apply(StateTransforms.Model.ModelFromTrajectory, { modelIndex: 0 })
-            .apply(StateTransforms.Model.StructureAssemblyFromModel)
-            // .apply(StateTransforms.Model.CreateStructureSelection, { query, label: 'ALA residues' })
-            .apply(StateTransforms.Representation.StructureRepresentation3D, {
-                type: {
-                    name: 'cartoon',
-                    params: PD.getDefaultValues(CartoonParams)
-                }
-            })
-            .getTree();
+            .apply(StateTransforms.Model.StructureAssemblyFromModel);
+
+        root.apply(StateTransforms.Model.StructureComplexElement, { type: 'sequence' })
+            .apply(StateTransforms.Representation.StructureRepresentation3D, { type: { name: 'cartoon', params: PD.getDefaultValues(CartoonParams) } });
+        root.apply(StateTransforms.Model.StructureComplexElement, { type: 'ligands' })
+            .apply(StateTransforms.Representation.StructureRepresentation3D, { type: { name: 'ball-and-stick', params: PD.getDefaultValues(BallAndStickParams) } });
+        root.apply(StateTransforms.Model.StructureComplexElement, { type: 'water' })
+            .apply(StateTransforms.Representation.StructureRepresentation3D, { type: { name: 'ball-and-stick', params: { ...PD.getDefaultValues(BallAndStickParams), alpha: 0.51 } } });
 
-        return state.update(newTree);
+        return state.update(root.getTree());
     }
 });
 

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

@@ -7,11 +7,12 @@
 import { PluginStateTransform } from '../objects';
 import { PluginStateObject as SO } from '../objects';
 import { Task } from 'mol-task';
-import { Model, Format, Structure, ModelSymmetry, StructureSymmetry, QueryContext, StructureSelection as Sel } from 'mol-model/structure';
+import { Model, Format, Structure, ModelSymmetry, StructureSymmetry, QueryContext, StructureSelection as Sel, StructureQuery, Queries } from 'mol-model/structure';
 import { ParamDefinition as PD } from 'mol-util/param-definition';
 import Expression from 'mol-script/language/expression';
 import { compile } from 'mol-script/runtime/query/compiler';
 import { MolScriptBuilder } from 'mol-script/language/builder';
+import { StateObject } from 'mol-state';
 
 export { TrajectoryFromMmCif }
 namespace TrajectoryFromMmCif { export interface Params { blockHeader?: string } }
@@ -140,3 +141,34 @@ const StructureSelection = PluginStateTransform.Create<SO.Molecule.Structure, SO
         return new SO.Molecule.Structure(s, label);
     }
 });
+
+export { StructureComplexElement }
+namespace StructureComplexElement { export interface Params { type: 'sequence' | 'water' | 'ligands' } }
+const StructureComplexElement = PluginStateTransform.Create<SO.Molecule.Structure, SO.Molecule.Structure, StructureComplexElement.Params>({
+    name: 'structure-complex-element',
+    display: {
+        name: 'Complex Element',
+        description: 'Create a molecular structure from the specified model.'
+    },
+    from: [SO.Molecule.Structure],
+    to: [SO.Molecule.Structure],
+    params: () => ({ type: PD.Text('sequence', { isHidden: true }) }),
+    apply({ a, params }) {
+        // TODO: update function.
+
+        let query: StructureQuery, label: string;
+        switch (params.type) {
+            case 'sequence': query = Queries.internal.sequence(); label = 'Sequence'; break;
+            case 'water': query = Queries.internal.water(); label = 'Water'; break;
+            case 'ligands': query = Queries.internal.lidangs(); label = 'Ligands'; break;
+            default: throw new Error(`${params.type} is a valid complex element.`);
+        }
+
+        const result = query(new QueryContext(a.data));
+        const s = Sel.unionStructure(result);
+
+        if (s.elementCount === 0) return StateObject.Null;
+        return new SO.Molecule.Structure(s, { label, description: structureDesc(s) });
+    }
+});
+

src/mol-repr/structure/representation/ball-and-stick.ts

@@ -32,7 +32,7 @@ export const BallAndStickParams = {
     unitKinds: PD.MultiSelect<UnitKind>(['atomic'], UnitKindOptions),
     sizeFactor: PD.Numeric(0.3, { min: 0.01, max: 10, step: 0.01 }),
     sizeAspectRatio: PD.Numeric(2/3, { min: 0.01, max: 3, step: 0.01 }),
-    colorTheme: PD.Mapped('polymer-index', BuiltInColorThemeOptions, name => PD.Group((BuiltInColorThemes as { [k: string]: ColorTheme.Provider<any> })[name].getParams({}))),
+    colorTheme: PD.Mapped('element-symbol', BuiltInColorThemeOptions, name => PD.Group((BuiltInColorThemes as { [k: string]: ColorTheme.Provider<any> })[name].getParams({}))),
     visuals: PD.MultiSelect<BallAndStickVisualName>(['element-sphere', 'intra-link', 'inter-link'], BallAndStickVisualOptions),
 }
 export type BallAndStickParams = typeof BallAndStickParams