Refactored Selection

src/mol-math/geometry/symmetry-operator.ts

@@ -38,6 +38,7 @@ namespace SymmetryOperator {
 
     export interface CoordinateMapper { (index: number, slot: Vec3): Vec3 }
     export interface ArrayMapping {
+        readonly operator: SymmetryOperator,
         readonly invariantPosition: CoordinateMapper,
         readonly position: CoordinateMapper,
         x(index: number): number,

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

@@ -10,7 +10,7 @@ import P from './properties'
 import { Structure, AtomSet, Atom } from '../structure'
 import { OrderedSet, Segmentation } from 'mol-data/int'
 
-export const all: Query = s => s;
+export const all: Query = s => Selection.Singletons(s, s.atoms);
 
 export interface AtomGroupsParams {
     entityTest: Atom.Predicate,
@@ -56,7 +56,7 @@ function atomGroupsLinear(atomTest: Atom.Predicate): Query {
             builder.commitUnit(unitId);
         }
 
-        return Structure.create(units, builder.getSet());
+        return Selection.Singletons(structure, builder.getSet());
     };
 }
 
@@ -99,7 +99,7 @@ function atomGroupsSegmented({ entityTest, chainTest, residueTest, atomTest }: A
             builder.commitUnit(unitId);
         }
 
-        return Structure.create(units, builder.getSet());
+        return Selection.Singletons(structure, builder.getSet());
     };
 }
 
@@ -124,27 +124,27 @@ class LinearGroupingBuilder {
         return true;
     }
 
-    private singletonStructure(): Structure {
+    private singletonSelection(): Selection {
         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.ofAtoms(atoms, this.structure.atoms));
+        return Selection.Singletons(this.structure, AtomSet.ofAtoms(atoms, this.structure.atoms));
     }
 
     private fullSelection() {
-        const ret: Structure[] = [];
+        const sets: AtomSet[] = [];
         for (let i = 0, _i = this.builders.length; i < _i; i++) {
-            ret[i] = Structure.create(this.structure.units, this.builders[i].getSet());
+            sets[i] = this.builders[i].getSet();
         }
-        return ret;
+        return Selection.Seq(this.structure, sets);
     }
 
     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();
+        if (len === 0) return Selection.Empty(this.structure);
+        if (len === 1) return Selection.Singletons(this.structure, this.builders[0].getSet());
+        if (this.allSingletons()) return this.singletonSelection();
         return this.fullSelection();
     }
 

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

@@ -4,131 +4,94 @@
  * @author David Sehnal <david.sehnal@gmail.com>
  */
 
-import Iterator from 'mol-data/iterator'
 import { HashSet } from 'mol-data/util'
-import { Structure, Atom, AtomSet } from '../structure'
+import { Structure, AtomSet } from '../structure'
 
-type Selection =
-    | Structure // each atom is interpreted as a singleton structure
-    | Structure[]
-
-// TODO: Do not allow to change unit set in the middle of a query, create a differnt language to control assemblies etc.
-/*
-type Selection =
-  | { kind: 'sequence', structure: Structure, sets: AtomSet[] }
-  | { kind: 'atom-set', structure: Structure, set: AtomSet }
-
-  structure allows for good unions.
-*/
+// A selection is a pair of a Structure and a sequence of unique AtomSets
+type Selection = Selection.Singletons | Selection.Seq
 
 namespace Selection {
-    export const Empty: Selection = [];
+    // If each element of the selection is a singleton, we can use a more efficient representation.
+    export interface Singletons { readonly kind: 'singletons', readonly structure: Structure, readonly set: AtomSet }
+    export interface Seq { readonly kind: 'seq', readonly structure: Structure, readonly sets: ReadonlyArray<AtomSet> }
 
-    function isStructure(x: Selection): x is Structure { return !!(x as Structure).units && !!(x as Structure).atoms; }
+    export function Singletons(structure: Structure, set: AtomSet): Singletons { return { kind: 'singletons', structure, set } }
+    export function Seq(structure: Structure, sets: AtomSet[]): Seq { return { kind: 'seq', structure, sets } }
 
-    export const isOfSingletons = isStructure
+    export function Empty(structure: Structure): Selection { return Seq(structure, []); };
 
-    export function structureCount(sel: Selection) {
-        if (isStructure(sel)) return AtomSet.atomCount(sel.atoms);
-        return sel.length;
-    }
+    export function isSingleton(s: Selection): s is Singletons { return s.kind === 'singletons'; }
+    export function isEmpty(s: Selection) { return isSingleton(s) ? AtomSet.atomCount(s.set) === 0 : s.sets.length === 0; }
 
-    export function union(sel: Selection): Structure {
-        if (isStructure(sel)) return sel;
-        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 Structure.create(unionUnits(sel), AtomSet.union(sets, AtomSet.Empty));
+    export function structureCount(sel: Selection) {
+        if (isSingleton(sel)) return AtomSet.atomCount(sel.set);
+        return sel.sets.length;
     }
 
-    export function structures(sel: Selection): Iterator<Structure> {
-        if (isStructure(sel)) {
-            const units = sel.units;
-            return Iterator.map<Atom, Structure>(AtomSet.atoms(sel.atoms), atoms => Structure.create(units, atoms));
-        }
-        return Iterator.Array(sel);
+    export function unionStructure(sel: Selection): Structure {
+        if (isEmpty(sel)) return Structure.Empty(sel.structure.units);
+        if (isSingleton(sel)) return Structure.create(sel.structure.units, sel.set);
+        return Structure.create(sel.structure.units, AtomSet.union(sel.sets, sel.structure.atoms));
     }
 
     export function getAt(sel: Selection, i: number): Structure {
-        if (isStructure(sel)) {
-            return Structure.create(sel.units, AtomSet.atomGetAt(sel.atoms, i));
+        if (isSingleton(sel)) {
+            const atom = AtomSet.atomGetAt(sel.set, i);
+            return Structure.create(sel.structure.units, AtomSet.ofAtoms([atom], sel.structure.atoms));
         }
-        return sel[i];
+        return Structure.create(sel.structure.units, sel.sets[i]);
     }
 
     export interface Builder {
-        add(s: Structure): void,
+        add(set: AtomSet): void,
         getSelection(): Selection
     }
 
+    function getSelection(structure: Structure, sets: AtomSet[], allSingletons: boolean) {
+        const len = sets.length;
+        if (len === 0) return Empty(structure);
+        if (len === 1) return Singletons(structure, sets[0]);
+        if (allSingletons) return Singletons(structure, AtomSet.union(sets, structure.atoms));
+        return Seq(structure, sets);
+    }
+
     class LinearBuilderImpl implements Builder {
-        private structures: Structure[] = [];
+        private sets: AtomSet[] = [];
         private allSingletons = true;
 
-        add(s: Structure) {
-            const atomCount = AtomSet.atomCount(s.atoms);
+        add(atoms: AtomSet) {
+            const atomCount = AtomSet.atomCount(atoms);
             if (atomCount === 0) return;
-            this.structures[this.structures.length] = s;
+            this.sets[this.sets.length] = atoms;
             if (atomCount !== 1) this.allSingletons = false;
         }
 
-        getSelection() {
-            const len = this.structures.length;
-            if (len === 0) return Empty;
-            if (len === 1) return this.structures[0];
-            if (this.allSingletons) return union(this.structures);
-            return this.structures;
-        }
+        getSelection() { return getSelection(this.structure, this.sets, this.allSingletons); }
 
-        constructor() { }
+        constructor(private structure: Structure) { }
     }
 
     class HashBuilderImpl implements Builder {
-        private structures: Structure[] = [];
+        private sets: AtomSet[] = [];
         private allSingletons = true;
-        private sets = HashSet(AtomSet.hashCode, AtomSet.areEqual);
+        private uniqueSets = HashSet(AtomSet.hashCode, AtomSet.areEqual);
 
-        add(s: Structure) {
-            const atomCount = AtomSet.atomCount(s.atoms);
-            if (atomCount === 0 || !this.sets.add(s.atoms)) return;
-            this.structures[this.structures.length] = s;
+        add(atoms: AtomSet) {
+            const atomCount = AtomSet.atomCount(atoms);
+            if (atomCount === 0 || !this.uniqueSets.add(atoms)) return;
+            this.sets[this.sets.length] = atoms;
             if (atomCount !== 1) this.allSingletons = false;
         }
 
-        getSelection() {
-            const len = this.structures.length;
-            if (len === 0) return Empty;
-            if (len === 1) return this.structures[0];
-            if (this.allSingletons) return union(this.structures);
-            return this.structures;
-        }
+        getSelection() { return getSelection(this.structure, this.sets, this.allSingletons); }
 
-        constructor() { }
+        constructor(private structure: Structure) { }
     }
 
-    export function LinearBuilder(): Builder { return new LinearBuilderImpl(); }
-    export function UniqueBuilder(): Builder { return new HashBuilderImpl(); }
+    export function LinearBuilder(structure: Structure): Builder { return new LinearBuilderImpl(structure); }
+    export function UniqueBuilder(structure: Structure): Builder { return new HashBuilderImpl(structure); }
 
     // TODO: spatial lookup
 }
 
 export default Selection
\ No newline at end of file
-
-function unionUnits(xs: Structure[]): Structure['units'] {
-    return xs[0].units;
-//    let prev = xs[0].units;
-    // let sameUnits = true;
-    // for (let i = 1, _i = xs.length; i < _i; i++) {
-    //     if (xs[i].units !== prev) sameUnits = false;
-    // }
-    // if (sameUnits) return prev;
-
-    // const ret = [...prev];
-    // for (let i = 1, _i = xs.length; i < _i; i++) {
-    //     const units = xs[i].units;
-    //     if (units !== prev) IntMap.addFrom(ret, units);
-    //     prev = units;
-    // }
-
-    //return ret;
-}

src/mol-model/structure/structure/atom/group.ts

@@ -24,7 +24,7 @@ namespace AtomGroup {
     }
 
     export function create(unit: Unit, atoms: OrderedSet): AtomGroup {
-        if (OrderedSet.areEqual(atoms, unit.naturalGroup.atoms)) return unit.naturalGroup;
+        if (OrderedSet.areEqual(atoms, unit.fullGroup.atoms)) return unit.fullGroup;
         return createNew(atoms);
     }
 

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

@@ -34,7 +34,7 @@ namespace AtomSet {
     export const areEqual: (a: AtomSet, b: AtomSet) => boolean = Impl.areEqual as any;
     export const areIntersecting: (a: AtomSet, b: AtomSet) => boolean = Impl.areIntersecting as any;
 
-    export const union: (sets: AtomSet[], template: AtomSet) => AtomSet = Impl.unionMany as any;
+    export const union: (sets: ArrayLike<AtomSet>, template: AtomSet) => AtomSet = Impl.unionMany as any;
     export const intersect: (a: AtomSet, b: AtomSet) => AtomSet = Impl.intersect as any;
     export const subtract: (a: AtomSet, b: AtomSet) => AtomSet = Impl.subtract as any;
 

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

@@ -13,18 +13,16 @@ import AtomSet from './atom/set'
 import AtomGroup from './atom/group'
 import Atom from './atom'
 
-
-interface Structure extends Readonly<{
-    units: Unit[],
-    atoms: AtomSet
-}> { }
+// A structure is a pair of "units" and an atom set.
+// Each unit contains the data and transformation of its corresponding atoms.
+interface Structure {
+    readonly units: ReadonlyArray<Unit>,
+    readonly atoms: AtomSet
+}
 
 namespace Structure {
-    export const Empty: Structure = { units: [], atoms: AtomSet.Empty };
-
-    export function create(units: Unit[], atoms: AtomSet): Structure {
-        return { units, atoms };
-    }
+    export function create(units: ReadonlyArray<Unit>, atoms: AtomSet): Structure { return { units, atoms }; }
+    export function Empty(units: ReadonlyArray<Unit>): Structure { return create(units, AtomSet.Empty); };
 
     export function ofData(format: Format) {
         const models = Model.create(format);
@@ -38,7 +36,7 @@ namespace Structure {
         for (let c = 0; c < chains.count; c++) {
             const group = AtomGroup.createNew(OrderedSet.ofBounds(chains.segments[c], chains.segments[c + 1]));
             const unit = Unit.create(model, SymmetryOperator.Default, group);
-            builder.add(unit, unit.naturalGroup);
+            builder.add(unit, unit.fullGroup);
         }
 
         return builder.getStructure();
@@ -61,7 +59,7 @@ namespace Structure {
         add(unit: Unit, atoms: AtomGroup) { const id = this.addUnit(unit); this.setAtoms(id, atoms); }
         addUnit(unit: Unit) { const id = this._unitId++; this.units[id] = unit; return id; }
         setAtoms(unitId: number, atoms: AtomGroup) { this.atoms.add(unitId, atoms); this.atomCount += AtomGroup.size(atoms); }
-        getStructure(): Structure { return this.atomCount > 0 ? Structure.create(this.units, this.atoms.getSet()) : Empty; }
+        getStructure(): Structure { return this.atomCount > 0 ? Structure.create(this.units, this.atoms.getSet()) : Empty(this.units); }
     }
 
     export function Builder(): Builder { return new BuilderImpl(); }

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

@@ -28,7 +28,7 @@ function buildAssemblyImpl(structure: Structure, name: string) {
     for (const g of assembly.operatorGroups) {
         const selection = g.selector(structure);
         if (Selection.structureCount(selection) === 0) continue;
-        const { units, atoms } = Selection.union(selection);
+        const { units, atoms } = Selection.unionStructure(selection);
 
         const unitIds = AtomSet.unitIds(atoms);
 

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

@@ -8,16 +8,23 @@ import SymmetryOperator from 'mol-math/geometry/symmetry-operator'
 import AtomGroup from './atom/group'
 import { Model } from '../model'
 
+// A bulding block of a structure that corresponds
+// to a "natural group of atoms" (most often a "chain")
+// together with a tranformation (rotation and translation)
+// that is dynamically applied to the underlying atom set.
+//
+// An atom set can be referenced by multiple diffrent units which
+// makes construction of assemblies and spacegroups very efficient.
 interface Unit extends SymmetryOperator.ArrayMapping {
     // Provides access to the underlying data.
     readonly model: Model,
 
-    // Determines the operation applied to this unit.
-    // The transform and and inverse are baked into the "getPosition" function
-    readonly operator: SymmetryOperator,
-
     // The "full" atom group corresponding to this unit.
-    readonly naturalGroup: AtomGroup,
+    // Every selection is a subset of this atoms group.
+    // Things like inter-unit bonds or spatial lookups
+    // can be be implemented efficiently as "views" of the
+    // full group.
+    readonly fullGroup: AtomGroup,
 
     // Reference some commonly accessed things for faster access.
     readonly residueIndex: ArrayLike<number>,
@@ -29,14 +36,14 @@ interface Unit extends SymmetryOperator.ArrayMapping {
 }
 
 namespace Unit {
-    export function create(model: Model, operator: SymmetryOperator, naturalGroup: AtomGroup): Unit {
+    export function create(model: Model, operator: SymmetryOperator, fullGroup: AtomGroup): Unit {
         const h = model.hierarchy;
         const { invariantPosition, position, x, y, z } = SymmetryOperator.createMapping(operator, model.conformation);
 
         return {
             model,
             operator,
-            naturalGroup,
+            fullGroup,
             residueIndex: h.residueSegments.segmentMap,
             chainIndex: h.chainSegments.segmentMap,
             hierarchy: model.hierarchy,
@@ -48,7 +55,7 @@ namespace Unit {
     }
 
     export function withOperator(unit: Unit, operator: SymmetryOperator) {
-        return create(unit.model, SymmetryOperator.compose(unit.operator, operator), unit.naturalGroup);
+        return create(unit.model, SymmetryOperator.compose(unit.operator, operator), unit.fullGroup);
     }
 }
 

src/perf-tests/structure.ts

@@ -36,11 +36,9 @@ function *test() {
 
 async function runIt<T>(itP: () => IterableIterator<T>) {
     const it = itP();
-    let lastValue: T | undefined;
     while (true) {
         const { value, done } = it.next();
         if (done) return value;
-        lastValue = value;
     }
 }