sturcuture factoring, queries

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

@@ -5,8 +5,9 @@
  */
 
 import Query from './query'
+import Selection from './selection'
 import * as P from './properties'
-import { AtomSet, Atom } from '../structure'
+import { Structure, AtomSet, Atom } from '../structure'
 import { OrderedSet, Segmentation } from 'mol-base/collections/integer'
 
 export interface AtomGroupsSpec {
@@ -40,7 +41,7 @@ function atomGroupsLinear(atomTest: Atom.Predicate): Query {
         const { atoms, units } = structure;
         const unitIds = AtomSet.unitIds(atoms);
         const l = Atom.Location();
-        const builder = AtomSet.Builder(atoms);
+        const builder = AtomSet.LinearBuilder(atoms);
 
         for (let i = 0, _i = unitIds.length; i < _i; i++) {
             const unitId = unitIds[i];
@@ -55,7 +56,7 @@ function atomGroupsLinear(atomTest: Atom.Predicate): Query {
             builder.commitUnit(unitId);
         }
 
-        return { units, atoms: builder.getSet() };
+        return Structure.create(units, builder.getSet());
     };
 }
 
@@ -64,7 +65,7 @@ function atomGroupsSegmented({ entityTest, chainTest, residueTest, atomTest }: A
         const { atoms, units } = structure;
         const unitIds = AtomSet.unitIds(atoms);
         const l = Atom.Location();
-        const builder = AtomSet.Builder(atoms);
+        const builder = AtomSet.LinearBuilder(atoms);
 
         for (let i = 0, _i = unitIds.length; i < _i; i++) {
             const unitId = unitIds[i];
@@ -98,16 +99,60 @@ function atomGroupsSegmented({ entityTest, chainTest, residueTest, atomTest }: A
             builder.commitUnit(unitId);
         }
 
-        return { units, atoms: builder.getSet() };
+        return Structure.create(units, builder.getSet());
     };
 }
 
+class LinearGroupingBuilder {
+    private builders: AtomSet.Builder[] = [];
+    private builderMap: { [key: string]: AtomSet.Builder } = Object.create(null);
+
+    add(key: any, unit: number, atom: number) {
+        let b = this.builderMap[key];
+        if (!b) {
+            b = AtomSet.LinearBuilder(this.structure.atoms);
+            this.builders[this.builders.length] = b;
+            this.builderMap[key] = b;
+        }
+        b.add(unit, atom);
+    }
+
+    private allSingletons() {
+        for (let i = 0, _i = this.builders.length; i < _i; i++) {
+            if (this.builders[i].atomCount > 1) return false;
+        }
+        return true;
+    }
+
+    private singletonStructure(): Structure {
+        const atoms: Atom[] = Atom.createEmptyArray(this.builders.length);
+        for (let i = 0, _i = this.builders.length; i < _i; i++) {
+            atoms[i] = this.builders[i].singleton();
+        }
+        return Structure.create(this.structure.units, AtomSet.create(atoms));
+    }
+
+    private fullSelection() {
+        const ret: Structure[] = [];
+        for (let i = 0, _i = this.builders.length; i < _i; i++) {
+            ret[i] = Structure.create(this.structure.units, this.builders[i].getSet());
+        }
+        return ret;
+    }
+
+    getSelection(): Selection {
+        const len = this.builders.length;
+        if (len === 0) return Selection.Empty;
+        if (len === 1) return Structure.create(this.structure.units, this.builders[0].getSet());
+        if (this.allSingletons()) return this.singletonStructure();
+        return this.fullSelection();
+    }
+
+    constructor(private structure: Structure) { }
+}
+
 function atomGroupsGrouped({ entityTest, chainTest, residueTest, atomTest, groupBy }: AtomGroupsSpec): Query {
     return structure => {
         throw 'nyi'
     };
-}
-
-// class LinearGroupingBuilder {
-
-// }
+}

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

@@ -27,20 +27,20 @@ namespace Selection {
         if (!sel.length) return Structure.Empty;
         const sets = [];
         for (let i = 0, _i = sel.length; i < _i; i++) sets[sets.length] = sel[i].atoms;
-        return { units: unionUnits(sel), atoms: AtomSet.unionMany(sets) };
+        return Structure.create(unionUnits(sel), AtomSet.unionMany(sets));
     }
 
     export function structures(sel: Selection): Iterator<Structure> {
         if (isStructure(sel)) {
             const units = sel.units;
-            return Iterator.map<Atom, Structure>(AtomSet.atoms(sel.atoms), atoms => ({ units, atoms }));
+            return Iterator.map<Atom, Structure>(AtomSet.atoms(sel.atoms), atoms => Structure.create(units, atoms));
         }
         return Iterator.Array(sel);
     }
 
     export function getAt(sel: Selection, i: number): Structure {
         if (isStructure(sel)) {
-            return { units: sel.units, atoms: AtomSet.atomGetAt(sel.atoms, i) };
+            return Structure.create(sel.units, AtomSet.atomGetAt(sel.atoms, i));
         }
         return sel[i];
     }

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

@@ -20,6 +20,8 @@ namespace Atom {
     export const areEqual: (a: Atom, b: Atom) => boolean = Tuple.areEqual;
     export const hashCode: (a: Atom) => number = Tuple.hashCode;
 
+    export function createEmptyArray(n: number): Atom[] { return new Float64Array(n) as any; }
+
     /** All the information required to access atom properties */
     export interface Location { unit: Unit, atom: number }
     export function Location(): Location { return { unit: {} as any, atom: 0 }; }

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

@@ -7,7 +7,7 @@
 import { OrderedSet, SortedArray, Iterator } from 'mol-base/collections/integer'
 import Atom from '../atom'
 import * as Impl from './set/impl'
-import createBuilder from './set/builder'
+import createBuilder, { Builder as AtomSetBuilder } from './set/builder'
 
 /** A map-like representation of grouped atom set */
 namespace AtomSet {
@@ -38,8 +38,9 @@ namespace AtomSet {
     export const intersect: (a: AtomSet, b: AtomSet) => AtomSet = Impl.intersect as any;
     export const subtract: (a: AtomSet, b: AtomSet) => AtomSet = Impl.subtract as any;
 
-    export function SortedBuilder(parent: AtomSet) { return createBuilder(parent, true); }
-    export function Builder(parent: AtomSet) { return createBuilder(parent, false); }
+    export type Builder = AtomSetBuilder
+    export function LinearBuilder(parent: AtomSet): Builder { return createBuilder(parent, true); }
+    export function UnsortedBuilder(parent: AtomSet): Builder { return createBuilder(parent, false); }
 
     // TODO: bounding sphere
     // TODO: distance, areWithIn?

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

@@ -5,14 +5,17 @@
  */
 
 import AtomSet from '../set'
+import Atom from '../../atom'
 import { OrderedSet } from 'mol-base/collections/integer'
 import { sortArray } from 'mol-base/collections/sort'
 
-class Builder {
+export class Builder {
     private keys: number[] = [];
     private units: number[][] = Object.create(null);
     private currentUnit: number[] = [];
 
+    atomCount = 0;
+
     add(u: number, a: number) {
         const unit = this.units[u];
         if (!!unit) { unit[unit.length] = a; }
@@ -20,10 +23,11 @@ class Builder {
             this.units[u] = [a];
             this.keys[this.keys.length] = u;
         }
+        this.atomCount++;
     }
 
     beginUnit() { this.currentUnit = this.currentUnit.length > 0 ? [] : this.currentUnit; }
-    addToUnit(a: number) { this.currentUnit[this.currentUnit.length] = a; }
+    addToUnit(a: number) { this.currentUnit[this.currentUnit.length] = a; this.atomCount++; }
     commitUnit(u: number) {
         if (this.currentUnit.length === 0) return;
         this.keys[this.keys.length] = u;
@@ -53,6 +57,11 @@ class Builder {
         return allEqual ? this.parent : AtomSet.create(sets);
     }
 
+    singleton(): Atom {
+        const u = this.keys[0];
+        return Atom.create(u, this.units[u][0]);
+    }
+
     constructor(private parent: AtomSet, private sorted: boolean) { }
 }
 

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

@@ -467,7 +467,7 @@ function unionN(sets: ArrayLike<AtomSetImpl>, eCount: { count: number }) {
     eCount.count = countE;
     if (!countN) return Empty;
     if (countN === sets.length) return ofAtoms(sets as ArrayLike<Atom>);
-    const packed = new Float64Array(countN);
+    const packed = Atom.createEmptyArray(countN);
     let offset = 0;
     for (let i = 0, _i = sets.length; i < _i; i++) {
         const s = sets[i];

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

@@ -18,6 +18,10 @@ interface Structure extends Readonly<{
 namespace Structure {
     export const Empty = { units: {}, atoms: AtomSet.Empty };
 
+    export function create(units: Structure['units'], atoms: AtomSet): Structure {
+        return { units, atoms };
+    }
+
     export function ofData(format: Format) {
         const models = Model.create(format);
         return models.map(ofModel);
@@ -50,7 +54,7 @@ namespace Structure {
 
         addUnit(unit: Unit) { this.units[unit.id] = unit; }
         addAtoms(unitId: number, atoms: OrderedSet) { this.atoms[unitId] = atoms; this.atomCount += OrderedSet.size(atoms); }
-        getStructure(): Structure { return this.atomCount > 0 ? { units: this.units, atoms: AtomSet.create(this.atoms) } : Empty; }
+        getStructure(): Structure { return this.atomCount > 0 ? Structure.create(this.units, AtomSet.create(this.atoms)) : Empty; }
     }
 
     export function Builder(): Builder { return new BuilderImpl(); }

src/perf-tests/sets.ts

@@ -200,7 +200,7 @@ export namespace Union {
 
 export namespace Build {
     function createSorted() {
-        const b = AtomSet.SortedBuilder(AtomSet.Empty);
+        const b = AtomSet.LinearBuilder(AtomSet.Empty);
         for (let i = 0; i < 10; i++) {
             for (let j = 0; j < 1000; j++) {
                 b.add(i, j);
@@ -210,7 +210,7 @@ export namespace Build {
     }
 
     function createByUnit() {
-        const b = AtomSet.SortedBuilder(AtomSet.Empty);
+        const b = AtomSet.LinearBuilder(AtomSet.Empty);
         for (let i = 0; i < 10; i++) {
             b.beginUnit();
             for (let j = 0; j < 1000; j++) {

src/perf-tests/structure.ts

@@ -236,7 +236,7 @@ export namespace PropertyAccess {
 
     export async function run() {
         const { structures, models } = await readCIF('./examples/1cbs_full.bcif');
-        //const { structures, models } = await readCIF('e:/test/quick/3j3q_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');