predicates

src/mol-data/structure/query.ts

@@ -7,11 +7,13 @@
 import Selection from './query/selection'
 import Query from './query/query'
 import * as generators from './query/generators'
-import * as props from './query/properties'
+import props from './query/properties'
+import pred from './query/predicates'
 
 export const Queries = {
     generators,
-    props
+    props,
+    pred
 }
 
 export { Selection, Query }

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

@@ -6,11 +6,13 @@
 
 import Query from './query'
 import Selection from './selection'
-import * as P from './properties'
+import P from './properties'
 import { Structure, AtomSet, Atom } from '../structure'
 import { OrderedSet, Segmentation } from 'mol-base/collections/integer'
 
-export interface AtomGroupsSpec {
+export const all: Query = s => s;
+
+export interface AtomGroupsParams {
     entityTest: Atom.Predicate,
     chainTest: Atom.Predicate,
     residueTest: Atom.Predicate,
@@ -18,21 +20,19 @@ export interface AtomGroupsSpec {
     groupBy: Atom.Property<any>
 }
 
-export const all: Query = s => s;
+export function atoms(params?: Partial<AtomGroupsParams>): Query {
+    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);
 
-export function atomGroups(spec?: Partial<AtomGroupsSpec>): Query {
-    if (!spec || (!spec.atomTest && !spec.residueTest && !spec.chainTest && !spec.entityTest && !spec.groupBy)) return all;
-    if (!!spec.atomTest && !spec.residueTest && !spec.chainTest && !spec.entityTest && !spec.groupBy) return atomGroupsLinear(spec.atomTest);
-
-    const normalized: AtomGroupsSpec = {
-        entityTest: spec.entityTest || P.constant.true,
-        chainTest: spec.entityTest || P.constant.true,
-        residueTest: spec.residueTest || P.constant.true,
-        atomTest: spec.atomTest || P.constant.true,
-        groupBy: spec.entityTest || P.constant.zero,
+    const normalized: AtomGroupsParams = {
+        entityTest: params.entityTest || P.constant.true,
+        chainTest: params.chainTest || P.constant.true,
+        residueTest: params.residueTest || P.constant.true,
+        atomTest: params.atomTest || P.constant.true,
+        groupBy: params.groupBy || P.constant.zero,
     };
 
-    if (!spec.groupBy) return atomGroupsSegmented(normalized)
+    if (!params.groupBy) return atomGroupsSegmented(normalized)
     return atomGroupsGrouped(normalized);
 }
 
@@ -60,7 +60,7 @@ function atomGroupsLinear(atomTest: Atom.Predicate): Query {
     };
 }
 
-function atomGroupsSegmented({ entityTest, chainTest, residueTest, atomTest }: AtomGroupsSpec): Query {
+function atomGroupsSegmented({ entityTest, chainTest, residueTest, atomTest }: AtomGroupsParams): Query {
     return structure => {
         const { atoms, units } = structure;
         const unitIds = AtomSet.unitIds(atoms);
@@ -151,8 +151,43 @@ class LinearGroupingBuilder {
     constructor(private structure: Structure) { }
 }
 
-function atomGroupsGrouped({ entityTest, chainTest, residueTest, atomTest, groupBy }: AtomGroupsSpec): Query {
+function atomGroupsGrouped({ entityTest, chainTest, residueTest, atomTest, groupBy }: AtomGroupsParams): Query {
     return structure => {
-        throw 'nyi'
+        const { atoms, units } = structure;
+        const unitIds = AtomSet.unitIds(atoms);
+        const l = Atom.Location();
+        const builder = new LinearGroupingBuilder(structure);
+
+        for (let i = 0, _i = unitIds.length; i < _i; i++) {
+            const unitId = unitIds[i];
+            const unit = units[unitId];
+            l.unit = unit;
+            const set = AtomSet.unitGetByIndex(atoms, i);
+
+            const chainsIt = Segmentation.transientSegments(unit.hierarchy.chainSegments, set);
+            const residuesIt = Segmentation.transientSegments(unit.hierarchy.residueSegments, set);
+            while (chainsIt.hasNext) {
+                const chainSegment = chainsIt.move();
+                l.atom = OrderedSet.getAt(set, chainSegment.start);
+                // test entity and chain
+                if (!entityTest(l) || !chainTest(l)) continue;
+
+                residuesIt.setSegment(chainSegment);
+                while (residuesIt.hasNext) {
+                    const residueSegment = residuesIt.move();
+                    l.atom = OrderedSet.getAt(set, residueSegment.start);
+
+                    // test residue
+                    if (!residueTest(l)) continue;
+
+                    for (let j = residueSegment.start, _j = residueSegment.end; j < _j; j++) {
+                        l.atom = OrderedSet.getAt(set, j);
+                        if (atomTest(l)) builder.add(groupBy(l), unitId, l.atom);
+                    }
+                }
+            }
+        }
+
+        return builder.getSelection();
     };
 }

src/mol-data/structure/query/predicates.ts

@@ -0,0 +1,99 @@
+/**
+ * Copyright (c) 2017 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author David Sehnal <david.sehnal@gmail.com>
+ */
+
+import { Atom } from '../structure'
+import P from './properties'
+
+namespace Predicates {
+    interface SetLike<A> { has(v: A): boolean }
+    function isSetLike<A>(x: any): x is SetLike<A> { return !!x && !!x.has }
+
+    export function eq<A>(p: Atom.Property<A>, value: A): Atom.Predicate { return l => p(l) === value; }
+    export function lt<A>(p: Atom.Property<A>, value: A): Atom.Predicate { return l => p(l) < value; }
+    export function lte<A>(p: Atom.Property<A>, value: A): Atom.Predicate { return l => p(l) <= value; }
+    export function gt<A>(p: Atom.Property<A>, value: A): Atom.Predicate { return l => p(l) > value; }
+    export function gte<A>(p: Atom.Property<A>, value: A): Atom.Predicate { return l => p(l) >= value; }
+
+    export function inSet<A>(p: Atom.Property<A>, values: SetLike<A> | ArrayLike<A>): Atom.Predicate {
+        if (isSetLike(values)) {
+            return l => values.has(p(l));
+        } else {
+            const set = new Set<A>();
+            for (let i = 0; i < values.length; i++) set.add(values[i]);
+            return l => set.has(p(l));
+        }
+    }
+
+    export function and(...ps: Atom.Predicate[]): Atom.Predicate {
+        switch (ps.length) {
+            case 0: return P.constant.true;
+            case 1: return ps[0];
+            case 2: {
+                const a = ps[0], b = ps[1];
+                return l => a(l) && b(l);
+            }
+            case 3: {
+                const a = ps[0], b = ps[1], c = ps[2];
+                return l => a(l) && b(l) && c(l);
+            }
+            case 4: {
+                const a = ps[0], b = ps[1], c = ps[2], d = ps[3];
+                return l => a(l) && b(l) && c(l) && d(l);
+            }
+            case 5: {
+                const a = ps[0], b = ps[1], c = ps[2], d = ps[3], e = ps[4];
+                return l => a(l) && b(l) && c(l) && d(l) && e(l);
+            }
+            case 6: {
+                const a = ps[0], b = ps[1], c = ps[2], d = ps[3], e = ps[4], f = ps[5];
+                return l => a(l) && b(l) && c(l) && d(l) && e(l) && f(l);
+            }
+            default: {
+                const count = ps.length;
+                return l => {
+                    for (let i = 0; i < count; i++) if (!ps[i]) return false;
+                    return true;
+                }
+            }
+        }
+    }
+
+    export function or(...ps: Atom.Predicate[]): Atom.Predicate {
+        switch (ps.length) {
+            case 0: return P.constant.true;
+            case 1: return ps[0];
+            case 2: {
+                const a = ps[0], b = ps[1];
+                return l => a(l) || b(l);
+            }
+            case 3: {
+                const a = ps[0], b = ps[1], c = ps[2];
+                return l => a(l) || b(l) || c(l);
+            }
+            case 4: {
+                const a = ps[0], b = ps[1], c = ps[2], d = ps[3];
+                return l => a(l) || b(l) || c(l) || d(l);
+            }
+            case 5: {
+                const a = ps[0], b = ps[1], c = ps[2], d = ps[3], e = ps[4];
+                return l => a(l) || b(l) || c(l) || d(l) || e(l);
+            }
+            case 6: {
+                const a = ps[0], b = ps[1], c = ps[2], d = ps[3], e = ps[4], f = ps[5];
+                return l => a(l) || b(l) || c(l) || d(l) || e(l) || f(l);
+            }
+            default: {
+                const count = ps.length;
+                return l => {
+                    for (let i = 0; i < count; i++) if (ps[i]) return true;
+                    return false;
+                }
+            }
+        }
+    }
+}
+
+export default Predicates

src/mol-data/structure/query/properties.ts

@@ -6,21 +6,33 @@
 
 import { Atom } from '../structure'
 
-export const constant = {
+const constant = {
     true: Atom.property(l => true),
     false: Atom.property(l => false),
     zero: Atom.property(l => 0)
 }
 
-export const atom = {
+const atom = {
     type_symbol: Atom.property(l => l.unit.hierarchy.atoms.type_symbol.value(l.atom))
 }
 
-export const residue = {
+const residue = {
+    key: Atom.property(l => l.unit.hierarchy.residueKey.value(l.unit.residueIndex[l.atom])),
+
     auth_seq_id: Atom.property(l => l.unit.hierarchy.residues.auth_seq_id.value(l.unit.residueIndex[l.atom])),
     auth_comp_id: Atom.property(l => l.unit.hierarchy.residues.auth_comp_id.value(l.unit.residueIndex[l.atom]))
 }
 
-export const chain = {
+const chain = {
     auth_asym_id: Atom.property(l => l.unit.hierarchy.chains.auth_asym_id.value(l.unit.chainIndex[l.atom]))
-}
+}
+
+const Properties = {
+    constant,
+    atom,
+    residue,
+    chain
+}
+
+type Properties = typeof Properties
+export default Properties

src/perf-tests/structure.ts

@@ -10,7 +10,7 @@ import * as util from 'util'
 import * as fs from 'fs'
 import CIF from 'mol-io/reader/cif'
 
-import { Structure, Model, Queries as Q, Atom, AtomSet } from 'mol-data/structure'
+import { Structure, Model, Queries as Q, Atom, AtomSet, Selection } from 'mol-data/structure'
 import { OrderedSet as OrdSet, Segmentation } from 'mol-base/collections/integer'
 
 require('util.promisify').shim();
@@ -235,8 +235,8 @@ export namespace PropertyAccess {
     // }
 
     export async function run() {
-        const { structures, models } = await readCIF('./examples/1cbs_full.bcif');
-        //const { structures, models } = await readCIF('e:/test/quick/1jj2_full.bcif');
+        //const { structures, models } = await readCIF('./examples/1cbs_full.bcif');
+        const { structures, models } = await readCIF('e:/test/quick/1jj2_full.bcif');
         //const { structures, models } = await readCIF('e:/test/quick/3j3q_updated.cif');
 
         console.log('parsed');
@@ -258,9 +258,14 @@ 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.atomGroups({ atomTest: l => auth_comp_id(l) === 'ALA' });
-        const q1 = Q.generators.atomGroups({ residueTest: l => auth_comp_id(l) === 'ALA' });
-        //const q2 = Q.generators.atomGroups({ chainTest: l => set.has(auth_asym_id(l)),  residueTest: l => auth_comp_id(l) === 'ALA' });
+        const q = Q.generators.atoms({ atomTest: Q.pred.eq(Q.props.residue.auth_comp_id, 'ALA') });
+        //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({
+            chainTest: Q.pred.inSet(Q.props.chain.auth_asym_id, ['A', 'B', 'C', 'D']),
+            residueTest: Q.pred.eq(Q.props.residue.auth_comp_id, 'ALA')
+        });
         q(structures[0]);
         console.time('q1')
         q1(structures[0]);
@@ -268,12 +273,19 @@ export namespace PropertyAccess {
         console.time('q1')
         q1(structures[0]);
         console.timeEnd('q1')
+        console.time('q2')
+        const q2r = q2(structures[0]);
+        console.timeEnd('q2')
+        console.log(Selection.structureCount(q2r))
         //console.log(q1(structures[0]));
 
         //const col = models[0].conformation.atomId.value;
         const suite = new B.Suite();
         suite
-            .add('test q', () => q1(structures[0]))
+            //.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 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])))