first query working

src/mol-data/query.ts

@@ -4,9 +4,9 @@
  * @author David Sehnal <david.sehnal@gmail.com>
  */
 
-import { Structure } from './structure'
 import Selection from './query/selection'
+import Query from './query/query'
+import * as generators from './query/generators'
+import * as props from './query/properties'
 
-interface Query { (s: Structure): Selection }
-
-export default Query
+export { Selection, Query, generators, props }

src/mol-data/query/generators.ts

@@ -4,9 +4,77 @@
  * @author David Sehnal <david.sehnal@gmail.com>
  */
 
-// import Selection from './selection'
-// import AtomSet from '../structure/atom-set'
-// import Atom from '../structure/atom'
+import Query from './query'
+//import Selection from './selection'
+import * as P from './properties'
+import { AtomSet, Atom } from '../structure'
+import { OrderedSet } from '../../mol-base/collections/integer'
+
+export interface AtomGroupsSpec {
+    entityTest: Atom.Predicate,
+    chainTest: Atom.Predicate,
+    residueTest: Atom.Predicate,
+    atomTest: Atom.Predicate,
+    groupBy: Atom.Property<any>
+}
+
+export const all: Query = s => s;
+
+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,
+    };
+
+    if (!spec.groupBy) return atomGroupsSegmented(normalized)
+    return atomGroupsGrouped(normalized);
+}
+
+function atomGroupsLinear(atomTest: Atom.Predicate): Query {
+    return structure => {
+        const { atoms, units } = structure;
+        const unitIds = AtomSet.unitIds(atoms);
+        const l = Atom.Location();
+        const builder = AtomSet.Builder(atoms);
+
+        for (let i = 0, _i = unitIds.length; i < _i; i++) {
+            const unitId = unitIds[i];
+            l.unit = units[unitId];
+            const set = AtomSet.unitGetByIndex(atoms, i);
+
+            builder.beginUnit();
+            for (let j = 0, _j = OrderedSet.size(set); j < _j; j++) {
+                l.atom = OrderedSet.getAt(set, j);
+                if (atomTest(l)) builder.addToUnit(l.atom);
+            }
+            builder.commitUnit(unitId);
+        }
+
+        return { units, atoms: builder.getSet() };
+    };
+}
+
+function atomGroupsSegmented({ entityTest, chainTest, residueTest, atomTest }: AtomGroupsSpec): Query {
+    return structure => {
+
+
+        throw 'nyi'
+    };
+}
+
+function atomGroupsGrouped({ entityTest, chainTest, residueTest, atomTest, groupBy }: AtomGroupsSpec): Query {
+    return structure => {
+
+
+        throw 'nyi'
+    };
+}
 
 // class LinearGroupingBuilder {
 

src/mol-data/query/properties.ts

@@ -0,0 +1,25 @@
+/**
+ * Copyright (c) 2017 molio contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author David Sehnal <david.sehnal@gmail.com>
+ */
+
+import { Atom } from '../structure'
+
+export const constant = {
+    true: Atom.property(l => true),
+    false: Atom.property(l => false),
+    zero: Atom.property(l => 0)
+}
+
+export const atom = {
+    type_symbol: Atom.property(l => l.unit.hierarchy.atoms.type_symbol.value(l.atom))
+}
+
+export const residue = {
+    auth_seq_id: Atom.property(l => l.unit.hierarchy.residues.auth_seq_id.value(l.unit.residueIndex[l.atom]))
+}
+
+export const chain = {
+    auth_asym_id: Atom.property(l => l.unit.hierarchy.chains.auth_asym_id.value(l.unit.chainIndex[l.atom]))
+}

src/mol-data/structure/atom.ts

@@ -30,6 +30,8 @@ namespace Atom {
         l.unit = structure.units[unit(atom)];
         l.atom = index(atom);
     }
+
+    export function property<T>(p: Property<T>) { return p; }
 }
 
 export default Atom

src/mol-data/structure/atom/set/builder.ts

@@ -33,20 +33,24 @@ class Builder {
     getSet(): AtomSet {
         const sets: { [key: number]: OrderedSet } = Object.create(null);
 
+        let allEqual = this.keys.length === AtomSet.unitCount(this.parent);
+
         for (let i = 0, _i = this.keys.length; i < _i; i++) {
             const k = this.keys[i];
             const unit = this.units[k];
             const l = unit.length;
             if (!this.sorted && l > 1) sortArray(unit);
-            if (l === 1) {
-                sets[k] = OrderedSet.ofSingleton(unit[0]);
+
+            const set = l === 1 ? OrderedSet.ofSingleton(unit[0]) : OrderedSet.ofSortedArray(unit);
+            const parentSet = AtomSet.unitGetById(this.parent, k);
+            if (OrderedSet.areEqual(set, parentSet)) {
+                sets[k] = parentSet;
             } else {
-                const set = OrderedSet.ofSortedArray(unit);
-                const parentSet = AtomSet.unitGetById(this.parent, k);
-                sets[k] = OrderedSet.areEqual(set, parentSet) ? parentSet : set;
+                sets[k] = set;
+                allEqual = false;
             }
         }
-        return AtomSet.create(sets);
+        return allEqual ? this.parent : AtomSet.create(sets);
     }
 
     constructor(private parent: AtomSet, private sorted: boolean) { }

src/perf-tests/structure.ts

@@ -12,6 +12,7 @@ import CIF from '../mol-io/reader/cif'
 
 import Model from '../mol-data/Model'
 import { Structure, Atom, AtomSet } from '../mol-data/structure'
+import * as Q from '../mol-data/query'
 import { OrderedSet as OrdSet, Segmentation } from '../mol-base/collections/integer'
 
 require('util.promisify').shim();
@@ -237,8 +238,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');
@@ -253,6 +254,13 @@ export namespace PropertyAccess {
         console.log(sumDirect(structures[0]));
         console.log('r', sumPropertyResidue(structures[0], l => l.unit.hierarchy.residues.auth_seq_id.value(l.unit.residueIndex[l.atom])));
 
+        //const authSeqId = Atom.property(l => l.unit.hierarchy.residues.auth_seq_id.value(l.unit.residueIndex[l.atom]));
+
+        const auth_seq_id = Q.props.residue.auth_seq_id;
+        const q = Q.generators.atomGroups({ atomTest: l => auth_seq_id(l) < 3 });
+        const qr = q(structures[0]);
+        console.log(qr);
+
         //const col = models[0].conformation.atomId.value;
         const suite = new B.Suite();
         suite