tweaking data model

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

@@ -30,6 +30,15 @@ describe('atom set', () => {
         expect(AtomSet.atomCount(set)).toBe(1);
     });
 
+    it('singleton atom', () => {
+        const set = AtomSet.singleton(p(10, 11), AtomSet.Empty);
+        expect(setToPairs(set)).toEqual([p(10, 11)]);
+        expect(AtomSet.atomHas(set, p(10, 11))).toBe(true);
+        expect(AtomSet.atomHas(set, p(11, 11))).toBe(false);
+        expect(AtomSet.atomGetAt(set, 0)).toBe(p(10, 11));
+        expect(AtomSet.atomCount(set)).toBe(1);
+    });
+
     it('multi', () => {
         const gen = AtomSet.Generator();
         gen.add(1, AtomGroup.createNew(OrderedSet.ofSortedArray([4, 6, 7])));

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

@@ -12,7 +12,7 @@ import { OrderedSet, Segmentation } from 'mol-data/int'
 
 export const all: Query = s => Selection.Singletons(s, s.atoms);
 
-export interface AtomGroupsParams {
+export interface AtomQueryParams {
     entityTest: Atom.Predicate,
     chainTest: Atom.Predicate,
     residueTest: Atom.Predicate,
@@ -20,11 +20,18 @@ export interface AtomGroupsParams {
     groupBy: Atom.Property<any>
 }
 
-export function atoms(params?: Partial<AtomGroupsParams>): Query {
+export interface AtomGroupsQueryParams extends AtomQueryParams {
+    groupBy: Atom.Property<any>
+}
+
+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 {
     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);
 
-    const normalized: AtomGroupsParams = {
+    const normalized: AtomGroupsQueryParams = {
         entityTest: params.entityTest || P.constant.true,
         chainTest: params.chainTest || P.constant.true,
         residueTest: params.residueTest || P.constant.true,
@@ -60,7 +67,7 @@ function atomGroupsLinear(atomTest: Atom.Predicate): Query {
     };
 }
 
-function atomGroupsSegmented({ entityTest, chainTest, residueTest, atomTest }: AtomGroupsParams): Query {
+function atomGroupsSegmented({ entityTest, chainTest, residueTest, atomTest }: AtomGroupsQueryParams): Query {
     return structure => {
         const { atoms, units } = structure;
         const unitIds = AtomSet.unitIds(atoms);
@@ -133,17 +140,16 @@ class LinearGroupingBuilder {
     }
 
     private fullSelection() {
-        const sets: AtomSet[] = [];
+        const sets: AtomSet[] = new Array(this.builders.length);
         for (let i = 0, _i = this.builders.length; i < _i; i++) {
             sets[i] = this.builders[i].getSet();
         }
-        return Selection.Seq(this.structure, sets);
+        return Selection.Sequence(this.structure, sets);
     }
 
     getSelection(): Selection {
         const len = this.builders.length;
         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();
     }
@@ -151,7 +157,7 @@ class LinearGroupingBuilder {
     constructor(private structure: Structure) { }
 }
 
-function atomGroupsGrouped({ entityTest, chainTest, residueTest, atomTest, groupBy }: AtomGroupsParams): Query {
+function atomGroupsGrouped({ entityTest, chainTest, residueTest, atomTest, groupBy }: AtomGroupsQueryParams): Query {
     return structure => {
         const { atoms, units } = structure;
         const unitIds = AtomSet.unitIds(atoms);

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

@@ -8,17 +8,16 @@ import { HashSet } from 'mol-data/util'
 import { Structure, AtomSet } from '../structure'
 
 // A selection is a pair of a Structure and a sequence of unique AtomSets
-type Selection = Selection.Singletons | Selection.Seq
+type Selection = Selection.Singletons | Selection.Sequence
 
 namespace 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> }
+    export interface Sequence { readonly kind: 'sequence', readonly structure: Structure, readonly sets: ReadonlyArray<AtomSet> }
 
     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 function Empty(structure: Structure): Selection { return Seq(structure, []); };
+    export function Sequence(structure: Structure, sets: AtomSet[]): Sequence { return { kind: 'sequence', structure, sets } }
+    export function Empty(structure: Structure): Selection { return Sequence(structure, []); };
 
     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; }
@@ -37,11 +36,30 @@ namespace Selection {
     export function getAt(sel: Selection, i: number): Structure {
         if (isSingleton(sel)) {
             const atom = AtomSet.atomGetAt(sel.set, i);
-            return Structure.create(sel.structure.units, AtomSet.ofAtoms([atom], sel.structure.atoms));
+            return Structure.create(sel.structure.units, AtomSet.singleton(atom, sel.structure.atoms));
         }
         return Structure.create(sel.structure.units, sel.sets[i]);
     }
 
+    export function toStructures(sel: Selection): Structure[] {
+        const { units } = sel.structure;
+        if (isSingleton(sel)) {
+            const ret: Structure[] = new Array(AtomSet.atomCount(sel.set));
+            const atoms = AtomSet.atoms(sel.set);
+            let offset = 0;
+            while (atoms.hasNext) {
+                const atom = atoms.move();
+                ret[offset++] = Structure.create(units, AtomSet.singleton(atom, sel.structure.atoms))
+            }
+            return ret;
+        } else {
+            const { sets } = sel;
+            const ret: Structure[] = new Array(sets.length);
+            for (let i = 0, _i = sets.length; i < _i; i++) ret[i] = Structure.create(units, sets[i]);
+            return ret;
+        }
+    }
+
     export interface Builder {
         add(set: AtomSet): void,
         getSelection(): Selection
@@ -50,9 +68,8 @@ namespace 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);
+        return Sequence(structure, sets);
     }
 
     class LinearBuilderImpl implements Builder {

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

@@ -20,6 +20,10 @@ export function ofAtoms(atoms: ArrayLike<Atom>, template: AtomSetImpl): AtomSetI
     return ofAtomsImpl(atoms, template);
 }
 
+export function singleton(atom: Atom, template: AtomSetImpl) {
+    return singletonImpl(atom, template);
+}
+
 export function getKeys(set: AtomSetImpl): SortedArray {
     return set.keys;
 }
@@ -358,6 +362,21 @@ function ofAtomsImpl(xs: ArrayLike<Atom>, template: AtomSetImpl) {
     return generator.getSet();
 }
 
+function singletonImpl(atom: Atom, template: AtomSetImpl) {
+    const k = Atom.unit(atom), i = Atom.index(atom);
+    const { groups } = template;
+    const gs = IntMap.Mutable<AtomGroup>();
+    if (groups.has(k)) {
+        const g = groups.get(k);
+        if (AtomGroup.size(g) === 1 && AtomGroup.getAt(g, 0) === i) {
+            gs.set(k, g);
+            return create([k], gs);
+        }
+    }
+    gs.set(k, AtomGroup.createNew(OS.ofSingleton(i)));
+    return create([k], gs);
+}
+
 function getOffsetIndex(xs: ArrayLike<number>, value: number) {
     let min = 0, max = xs.length - 1;
     while (min < max) {

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

@@ -15,6 +15,7 @@ namespace AtomSet {
     export const Empty: AtomSet = Impl.Empty as any;
 
     export const ofAtoms: (atoms: ArrayLike<Atom>, template: AtomSet) => AtomSet = Impl.ofAtoms as any;
+    export const singleton: (atom: Atom, template: AtomSet) => AtomSet = Impl.singleton as any;
 
     export const unitCount: (set: AtomSet) => number = Impl.keyCount as any;
     export const unitIds: (set: AtomSet) => SortedArray = Impl.getKeys as any;