Wrap query in a task

src/mol-model/structure/model/formats/mmcif/assembly.ts

@@ -8,7 +8,7 @@ import { Mat4, Tensor } from 'mol-math/linear-algebra'
 import SymmetryOperator from 'mol-math/geometry/symmetry-operator'
 import Format from '../../format'
 import { Assembly, OperatorGroup, OperatorGroups } from '../../properties/symmetry'
-import { Queries as Q } from '../../../query'
+import { Queries as Q, Query } from '../../../query'
 
 import mmCIF_Format = Format.mmCIF
 
@@ -57,10 +57,10 @@ function operatorGroupsProvider(generators: Generator[], matrices: Matrices): ()
             const operatorList = parseOperatorList(gen.expression);
             const operatorNames = expandOperators(operatorList);
             const operators = getAssemblyOperators(matrices, operatorNames, operatorOffset);
-            const selector = Q.generators.atoms({ chainTest: Q.pred.and(
+            const selector = Query(Q.generators.atoms({ chainTest: Q.pred.and(
                 Q.pred.eq(Q.props.unit.operator_name, SymmetryOperator.DefaultName),
                 Q.pred.inSet(Q.props.chain.label_asym_id, gen.asymIds)
-            )});
+            )}));
             groups[groups.length] = { selector, operators };
             operatorOffset += operators.length;
         }

src/mol-model/structure/query/generators.ts

@@ -10,7 +10,7 @@ import P from './properties'
 import { Structure, ElementSet, Element } from '../structure'
 import { OrderedSet, Segmentation } from 'mol-data/int'
 
-export const all: Query = s => Selection.Singletons(s, s.elements);
+export const all: Query.Provider = async (s, ctx) => Selection.Singletons(s, s.elements);
 
 export interface AtomQueryParams {
     entityTest: Element.Predicate,
@@ -27,7 +27,7 @@ export interface AtomGroupsQueryParams extends AtomQueryParams {
 export function residues(params?: Partial<AtomQueryParams>) { return atoms({ ...params, groupBy: P.residue.key }); }
 export function chains(params?: Partial<AtomQueryParams>) { return atoms({ ...params, groupBy: P.chain.key }); }
 
-export function atoms(params?: Partial<AtomGroupsQueryParams>): Query {
+export function atoms(params?: Partial<AtomGroupsQueryParams>): Query.Provider {
     if (!params || (!params.atomTest && !params.residueTest && !params.chainTest && !params.entityTest && !params.groupBy)) return all;
     if (!!params.atomTest && !params.residueTest && !params.chainTest && !params.entityTest && !params.groupBy) return atomGroupsLinear(params.atomTest);
 
@@ -43,8 +43,8 @@ export function atoms(params?: Partial<AtomGroupsQueryParams>): Query {
     return atomGroupsGrouped(normalized);
 }
 
-function atomGroupsLinear(atomTest: Element.Predicate): Query {
-    return structure => {
+function atomGroupsLinear(atomTest: Element.Predicate): Query.Provider {
+    return async (structure, ctx) => {
         const { elements, units } = structure;
         const unitIds = ElementSet.unitIds(elements);
         const l = Element.Location();
@@ -61,14 +61,16 @@ function atomGroupsLinear(atomTest: Element.Predicate): Query {
                 if (atomTest(l)) builder.addToUnit(l.atom);
             }
             builder.commitUnit(unitId);
+
+            if (ctx.shouldUpdate) await ctx.update({ message: 'Atom Groups', current: 0, max: unitIds.length });
         }
 
         return Selection.Singletons(structure, builder.getSet());
     };
 }
 
-function atomGroupsSegmented({ entityTest, chainTest, residueTest, atomTest }: AtomGroupsQueryParams): Query {
-    return structure => {
+function atomGroupsSegmented({ entityTest, chainTest, residueTest, atomTest }: AtomGroupsQueryParams): Query.Provider {
+    return async (structure, ctx) => {
         const { elements, units } = structure;
         const unitIds = ElementSet.unitIds(elements);
         const l = Element.Location();
@@ -104,6 +106,8 @@ function atomGroupsSegmented({ entityTest, chainTest, residueTest, atomTest }: A
                 }
             }
             builder.commitUnit(unitId);
+
+            if (ctx.shouldUpdate) await ctx.update({ message: 'Atom Groups', current: 0, max: unitIds.length });
         }
 
         return Selection.Singletons(structure, builder.getSet());
@@ -157,8 +161,8 @@ class LinearGroupingBuilder {
     constructor(private structure: Structure) { }
 }
 
-function atomGroupsGrouped({ entityTest, chainTest, residueTest, atomTest, groupBy }: AtomGroupsQueryParams): Query {
-    return structure => {
+function atomGroupsGrouped({ entityTest, chainTest, residueTest, atomTest, groupBy }: AtomGroupsQueryParams): Query.Provider {
+    return async (structure, ctx) => {
         const { elements, units } = structure;
         const unitIds = ElementSet.unitIds(elements);
         const l = Element.Location();
@@ -192,6 +196,8 @@ function atomGroupsGrouped({ entityTest, chainTest, residueTest, atomTest, group
                     }
                 }
             }
+
+            if (ctx.shouldUpdate) await ctx.update({ message: 'Atom Groups', current: 0, max: unitIds.length });
         }
 
         return builder.getSelection();

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

@@ -4,10 +4,21 @@
  * @author David Sehnal <david.sehnal@gmail.com>
  */
 
+import { Task, RuntimeContext } from 'mol-task'
 import { Structure } from '../structure'
 import Selection from './selection'
 
 // TODO: Query { (s: Structure): Computation<Selection> }
 
-interface Query { (s: Structure): Selection }
+interface Query { (s: Structure): Task<Selection>, provider: Query.Provider }
+function Query(q: Query.Provider): Query {
+    const ret = (s => Task.create('Query', ctx => q(s, ctx))) as Query;
+    ret.provider = q;
+    return ret;
+}
+
+namespace Query {
+    export interface Provider { (s: Structure, ctx: RuntimeContext): Promise<Selection> }
+}
+
 export default Query
\ No newline at end of file

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

@@ -9,6 +9,7 @@ import ElementSet from './element/set'
 import Unit from './unit'
 import { Selection } from '../query'
 import { ModelSymmetry } from '../model'
+import { Task } from 'mol-task';
 
 namespace Symmetry {
     export const buildAssembly = buildAssemblyImpl;
@@ -17,6 +18,7 @@ namespace Symmetry {
 export default Symmetry;
 
 function buildAssemblyImpl(structure: Structure, name: string) {
+    return Task.create('Build Symmetry', async ctx => {
         const models = Structure.getModels(structure);
         if (models.length !== 1) throw new Error('Can only build assemblies from structures based on 1 model.');
 
@@ -26,7 +28,7 @@ function buildAssemblyImpl(structure: Structure, name: string) {
         const assembler = Structure.Builder();
 
         for (const g of assembly.operatorGroups) {
-        const selection = g.selector(structure);
+            const selection = await ctx.runChild(g.selector(structure));
             if (Selection.structureCount(selection) === 0) continue;
             const { units, elements } = Selection.unionStructure(selection);
 
@@ -41,4 +43,5 @@ function buildAssemblyImpl(structure: Structure, name: string) {
         }
 
         return assembler.getStructure();
+    });
 }
\ No newline at end of file

src/perf-tests/structure.ts

@@ -11,7 +11,7 @@ import * as fs from 'fs'
 import fetch from 'node-fetch'
 import CIF from 'mol-io/reader/cif'
 
-import { Structure, Model, Queries as Q, Element, ElementGroup, ElementSet, Selection, Symmetry } from 'mol-model/structure'
+import { Structure, Model, Queries as Q, Element, ElementGroup, ElementSet, Selection, Symmetry, Query } from 'mol-model/structure'
 import { Segmentation } from 'mol-data/int'
 
 import to_mmCIF from 'mol-model/structure/export/mmcif'
@@ -293,14 +293,18 @@ export namespace PropertyAccess {
         console.log(to_mmCIF('test', s));
     }
 
-    export function testAssembly(id: string, s: Structure) {
+    export async function testAssembly(id: string, s: Structure) {
         console.time('assembly')
-        const a = Symmetry.buildAssembly(s, '1');
+        const a = await Run(Symmetry.buildAssembly(s, '1'));
         console.timeEnd('assembly')
         fs.writeFileSync(`${DATA_DIR}/${id}_assembly.bcif`, to_mmCIF(id, a, true));
         console.log('exported');
     }
 
+    function query(q: Query, s: Structure) {
+        return Run((q(s)));
+    }
+
     export async function run() {
         const { structures, models/*, mmcif*/ } = await getBcif('1cbs');
         // const { structures, models } = await getBcif('3j3q');
@@ -350,26 +354,26 @@ export namespace PropertyAccess {
         //const auth_asym_id = Q.props.chain.auth_asym_id;
         //const set =  new Set(['A', 'B', 'C', 'D']);
         //const q = Q.generators.atomGroups({ atomTest: l => auth_seq_id(l) < 3 });
-        const q = Q.generators.atoms({ atomTest: Q.pred.eq(Q.props.residue.auth_comp_id, 'ALA') });
+        const q = Query(Q.generators.atoms({ atomTest: Q.pred.eq(Q.props.residue.auth_comp_id, 'ALA') }));
         const P = Q.props
         //const q0 = Q.generators.atoms({ atomTest: l => auth_comp_id(l) === 'ALA' });
-        const q1 = Q.generators.atoms({ residueTest: l => auth_comp_id(l) === 'ALA' });
-        const q2 = Q.generators.atoms({ residueTest: l => auth_comp_id(l) === 'ALA', groupBy: Q.props.residue.key });
-        const q3 = Q.generators.atoms({
+        const q1 = Query(Q.generators.atoms({ residueTest: l => auth_comp_id(l) === 'ALA' }));
+        const q2 = Query(Q.generators.atoms({ residueTest: l => auth_comp_id(l) === 'ALA', groupBy: Q.props.residue.key }));
+        const q3 = Query(Q.generators.atoms({
             chainTest: Q.pred.inSet(P.chain.auth_asym_id, ['A', 'B', 'C', 'D']),
             residueTest: Q.pred.eq(P.residue.auth_comp_id, 'ALA')
-        });
-        q(structures[0]);
+        }));
+        await query(q, structures[0]);
         //console.log(to_mmCIF('test', Selection.union(q0r)));
 
         console.time('q1')
-        q1(structures[0]);
+        await query(q1, structures[0]);
         console.timeEnd('q1')
         console.time('q1')
-        q1(structures[0]);
+        await query(q1, structures[0]);
         console.timeEnd('q1')
         console.time('q2')
-        const q2r = q2(structures[0]);
+        const q2r = await query(q2, structures[0]);
         console.timeEnd('q2')
         console.log(Selection.structureCount(q2r));
         //console.log(q1(structures[0]));
@@ -379,8 +383,8 @@ export namespace PropertyAccess {
         suite
             //.add('test q', () => q1(structures[0]))
             //.add('test q', () => q(structures[0]))
-            .add('test q1', () => q1(structures[0]))
-            .add('test q3', () => q3(structures[0]))
+            .add('test q1', async () => await query(q1, structures[0]))
+            .add('test q3', async () => await query(q3, structures[0]))
             //.add('test int', () => sumProperty(structures[0], l => col(l.atom)))
             // .add('sum residue', () => sumPropertyResidue(structures[0], l => l.unit.hierarchy.residues.auth_seq_id.value(l.unit.residueIndex[l.atom])))