Updated ElementSet API

src/apps/render-test/state.ts

@@ -120,11 +120,11 @@ export default class State {
         async function createSpacefills (structure: Structure) {
             const spacefills: RenderObject[] = []
             const { elements, units } = structure;
-            const unitIds = ElementSet.unitIds(elements);
+            const unitIds = ElementSet.unitIndices(elements);
             for (let i = 0, _i = unitIds.length; i < _i; i++) {
                 const unitId = unitIds[i];
                 const unit = units[unitId];
-                const atomGroup = ElementSet.unitGetByIndex(elements, i);
+                const atomGroup = ElementSet.groupAt(elements, i);
 
                 const spacefill = Spacefill()
                 spacefills.push(...await Run(spacefill.create(unit, atomGroup), log, 1))

src/mol-geo/representation/structure/index.ts

@@ -36,8 +36,8 @@ export class StructureRepresentation {
                 ElementGroup.hashCode,
                 (a, b) => units[a.key].model.id === units[b.key].model.id && OrderedSet.areEqual(a.elements, b.elements));
 
-            for (let i = 0, _i = ElementSet.unitCount(elements); i < _i; i++) {
-                const group = ElementSet.unitGetByIndex(elements, i);
+            for (let i = 0, _i = ElementSet.groupCount(elements); i < _i; i++) {
+                const group = ElementSet.groupAt(elements, i);
                 uniqueGroups.add(group.key, group);
 
             }

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

@@ -26,7 +26,7 @@ describe('atom set', () => {
         expect(setToPairs(set)).toEqual([p(10, 11)]);
         expect(ElementSet.elementHas(set, p(10, 11))).toBe(true);
         expect(ElementSet.elementHas(set, p(11, 11))).toBe(false);
-        expect(ElementSet.elementGetAt(set, 0)).toBe(p(10, 11));
+        expect(ElementSet.elementAt(set, 0)).toBe(p(10, 11));
         expect(ElementSet.elementCount(set)).toBe(1);
     });
 
@@ -35,7 +35,7 @@ describe('atom set', () => {
         expect(setToPairs(set)).toEqual([p(10, 11)]);
         expect(ElementSet.elementHas(set, p(10, 11))).toBe(true);
         expect(ElementSet.elementHas(set, p(11, 11))).toBe(false);
-        expect(ElementSet.elementGetAt(set, 0)).toBe(p(10, 11));
+        expect(ElementSet.elementAt(set, 0)).toBe(p(10, 11));
         expect(ElementSet.elementCount(set)).toBe(1);
     });
 
@@ -51,7 +51,7 @@ describe('atom set', () => {
         expect(ElementSet.elementHas(set, p(3, 0))).toBe(true);
         expect(ElementSet.elementHas(set, p(1, 7))).toBe(true);
         for (let i = 0; i < ElementSet.elementCount(set); i++) {
-            expect(Element.areEqual(ElementSet.elementGetAt(set, i), ret[i])).toBe(true);
+            expect(Element.areEqual(ElementSet.elementAt(set, i), ret[i])).toBe(true);
         }
     });
 
@@ -62,8 +62,8 @@ describe('atom set', () => {
         gen.add(1, OrderedSet.ofSingleton(3));
         const set = gen.getSet();
 
-        expect(ElementSet.unitGetById(set, 0)).toBe(ElementSet.unitGetById(template, 0));
-        expect(ElementSet.unitGetById(set, 1)).toBe(ElementSet.unitGetById(template, 1));
+        expect(ElementSet.groupFromUnitIndex(set, 0)).toBe(ElementSet.groupFromUnitIndex(template, 0));
+        expect(ElementSet.groupFromUnitIndex(set, 1)).toBe(ElementSet.groupFromUnitIndex(template, 1));
         expect(set).toBe(template);
     });
 
@@ -74,8 +74,8 @@ describe('atom set', () => {
         gen.add(1, OrderedSet.ofSingleton(4));
         const set = gen.getSet();
 
-        expect(ElementSet.unitGetById(set, 0)).toBe(ElementSet.unitGetById(template, 0));
-        expect(ElementSet.unitGetById(set, 1) === ElementSet.unitGetById(template, 1)).toBe(false);
+        expect(ElementSet.groupFromUnitIndex(set, 0)).toBe(ElementSet.groupFromUnitIndex(template, 0));
+        expect(ElementSet.groupFromUnitIndex(set, 1) === ElementSet.groupFromUnitIndex(template, 1)).toBe(false);
         expect(set === template).toBe(false);
     });
 
@@ -89,9 +89,9 @@ describe('atom set', () => {
 
         const u0 = ElementSet.union([p01, p02, p06], template);
         const u1 = ElementSet.union([p01, p02, p06, p13], template);
-        expect(ElementSet.unitGetById(u0, 0)).toBe(ElementSet.unitGetById(template, 0));
-        expect(ElementSet.unitGetById(u1, 0)).toBe(ElementSet.unitGetById(template, 0));
-        expect(ElementSet.unitGetById(u1, 1)).toBe(ElementSet.unitGetById(template, 1));
+        expect(ElementSet.groupFromUnitIndex(u0, 0)).toBe(ElementSet.groupFromUnitIndex(template, 0));
+        expect(ElementSet.groupFromUnitIndex(u1, 0)).toBe(ElementSet.groupFromUnitIndex(template, 0));
+        expect(ElementSet.groupFromUnitIndex(u1, 1)).toBe(ElementSet.groupFromUnitIndex(template, 1));
         expect(u1).toBe(template);
     });
 
@@ -108,7 +108,7 @@ describe('atom set', () => {
         }
         const ms = gen.getSet();
         for (let i = 0; i < control.length; i++) {
-            expect(Element.areEqual(ElementSet.elementGetAt(ms, i), control[i])).toBe(true);
+            expect(Element.areEqual(ElementSet.elementAt(ms, i), control[i])).toBe(true);
         }
 
         for (let i = 0; i < control.length; i++) {

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

@@ -46,14 +46,14 @@ export function atoms(params?: Partial<AtomGroupsQueryParams>): Query.Provider {
 function atomGroupsLinear(atomTest: Element.Predicate): Query.Provider {
     return async (structure, ctx) => {
         const { elements, units } = structure;
-        const unitIds = ElementSet.unitIds(elements);
+        const unitIds = ElementSet.unitIndices(elements);
         const l = Element.Location();
         const builder = ElementSet.LinearBuilder(elements);
 
         for (let i = 0, _i = unitIds.length; i < _i; i++) {
             const unitId = unitIds[i];
             l.unit = units[unitId];
-            const set = ElementSet.unitGetByIndex(elements, i).elements;
+            const set = ElementSet.groupAt(elements, i).elements;
 
             builder.beginUnit();
             for (let j = 0, _j = OrderedSet.size(set); j < _j; j++) {
@@ -72,7 +72,7 @@ function atomGroupsLinear(atomTest: Element.Predicate): Query.Provider {
 function atomGroupsSegmented({ entityTest, chainTest, residueTest, atomTest }: AtomGroupsQueryParams): Query.Provider {
     return async (structure, ctx) => {
         const { elements, units } = structure;
-        const unitIds = ElementSet.unitIds(elements);
+        const unitIds = ElementSet.unitIndices(elements);
         const l = Element.Location();
         const builder = ElementSet.LinearBuilder(elements);
 
@@ -80,7 +80,7 @@ function atomGroupsSegmented({ entityTest, chainTest, residueTest, atomTest }: A
             const unitId = unitIds[i];
             const unit = units[unitId];
             l.unit = unit;
-            const set = ElementSet.unitGetByIndex(elements, i).elements;
+            const set = ElementSet.groupAt(elements, i).elements;
 
             builder.beginUnit();
             const chainsIt = Segmentation.transientSegments(unit.hierarchy.chainSegments, set);
@@ -164,7 +164,7 @@ class LinearGroupingBuilder {
 function atomGroupsGrouped({ entityTest, chainTest, residueTest, atomTest, groupBy }: AtomGroupsQueryParams): Query.Provider {
     return async (structure, ctx) => {
         const { elements, units } = structure;
-        const unitIds = ElementSet.unitIds(elements);
+        const unitIds = ElementSet.unitIndices(elements);
         const l = Element.Location();
         const builder = new LinearGroupingBuilder(structure);
 
@@ -172,7 +172,7 @@ function atomGroupsGrouped({ entityTest, chainTest, residueTest, atomTest, group
             const unitId = unitIds[i];
             const unit = units[unitId];
             l.unit = unit;
-            const set = ElementSet.unitGetByIndex(elements, i).elements;
+            const set = ElementSet.groupAt(elements, i).elements;
 
             const chainsIt = Segmentation.transientSegments(unit.hierarchy.chainSegments, set);
             const residuesIt = Segmentation.transientSegments(unit.hierarchy.residueSegments, set);

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

@@ -35,7 +35,7 @@ namespace Selection {
 
     export function getAt(sel: Selection, i: number): Structure {
         if (isSingleton(sel)) {
-            const atom = ElementSet.elementGetAt(sel.set, i);
+            const atom = ElementSet.elementAt(sel.set, i);
             return Structure.create(sel.structure.units, ElementSet.singleton(atom, sel.structure.elements));
         }
         return Structure.create(sel.structure.units, sel.sets[i]);

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

@@ -9,10 +9,19 @@ import { sortArray } from 'mol-data/util/sort'
 import { hash1 } from 'mol-data/util/hash-functions'
 import Element from '../../element'
 import ElementGroup from '../group'
+import { ElementSetLookup3D } from '../../util/lookup3d'
+import Structure from '../../structure';
 
 /** Long and painful implementation starts here */
 
-export type ElementSetImpl = { groups: IntMap<ElementGroup>, offsets: Int32Array, hashCode: number, keys: SortedArray }
+export type ElementSetImpl = {
+    groups: IntMap<ElementGroup>,
+    offsets: Int32Array,
+    hashCode: number,
+    keys: SortedArray,
+
+    __lookup3d__?: ElementSetLookup3D
+}
 
 export const Empty: ElementSetImpl = { groups: IntMap.Empty, offsets: new Int32Array(1), hashCode: 0, keys: SortedArray.Empty };
 
@@ -270,6 +279,13 @@ export function Generator() {
     return new AtomSetGenerator();
 }
 
+export function getLookup3d(s: Structure) {
+    const set = s.elements as any as ElementSetImpl;
+    if (set.__lookup3d__) return set.__lookup3d__;
+    set.__lookup3d__ = ElementSetLookup3D.create(s);
+    return set.__lookup3d__;
+}
+
 /** When adding groups, compare them to existing ones. If they all match, return the whole original set. */
 class ChildGenerator {
     private keys: number[] = [];

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

@@ -9,26 +9,32 @@ import Element from '../element'
 import ElementGroup from './group'
 import * as Impl from './impl/set'
 import * as Builders from './impl/set-builder'
+import { ElementSetLookup3D } from '../util/lookup3d';
+import Structure from '../structure';
 
-/** A map-like representation of grouped atom set */
+/**
+ * A map-like representation of grouped atom set
+ *
+ * Essentially corresponds to the type { [unitId: number]: ElementGroup }.
+ */
 namespace ElementSet {
     export const Empty: ElementSet = Impl.Empty as any;
 
     export const ofAtoms: (elements: ArrayLike<Element>, template: ElementSet) => ElementSet = Impl.ofElements as any;
     export const singleton: (element: Element, template: ElementSet) => ElementSet = Impl.singleton as any;
 
-    export const unitCount: (set: ElementSet) => number = Impl.keyCount as any;
-    export const unitIds: (set: ElementSet) => SortedArray = Impl.getKeys as any;
-    export const unitHas: (set: ElementSet, id: number) => boolean = Impl.hasKey as any;
-    export const unitGetId: (set: ElementSet, i: number) => number = Impl.getKey as any;
+    export const unitIndices: (set: ElementSet) => SortedArray = Impl.getKeys as any;
+    export const unitHas: (set: ElementSet, index: number) => boolean = Impl.hasKey as any;
 
-    export const unitGetById: (set: ElementSet, key: number) => ElementGroup = Impl.getByKey as any;
-    export const unitGetByIndex: (set: ElementSet, i: number) => ElementGroup = Impl.getByIndex as any;
+    export const groupCount: (set: ElementSet) => number = Impl.keyCount as any;
+    export const groupUnitIndex: (set: ElementSet, index: number) => number = Impl.getKey as any;
+    export const groupFromUnitIndex: (set: ElementSet, unitId: number) => ElementGroup = Impl.getByKey as any;
+    export const groupAt: (set: ElementSet, index: number) => ElementGroup = Impl.getByIndex as any;
 
     export const elementCount: (set: ElementSet) => number = Impl.size as any;
     export const elementHas: (set: ElementSet, x: Element) => boolean = Impl.hasAtom as any;
     export const elementIndexOf: (set: ElementSet, x: Element) => number = Impl.indexOf as any;
-    export const elementGetAt: (set: ElementSet, i: number) => Element = Impl.getAt as any;
+    export const elementAt: (set: ElementSet, i: number) => Element = Impl.getAt as any;
     export const elements: (set: ElementSet) => Iterator<Element> = Impl.values as any;
 
     export const hashCode: (set: ElementSet) => number = Impl.hashCode as any;
@@ -49,12 +55,13 @@ namespace ElementSet {
     export interface TemplateGenerator { add(unit: number, set: OrderedSet): void, getSet(): ElementSet }
     export const TemplateGenerator: (template: ElementSet) => TemplateGenerator = Impl.TemplateGenerator as any
 
-    // TODO: bounding sphere
+    export const getLookup3d: (s: Structure) => ElementSetLookup3D = Impl.getLookup3d;
+
     // TODO: distance, areWithIn?
     // TODO: check connected
     // TODO: add "parent" property? how to avoid using too much memory? Transitive parents? Parent unlinking?
 }
 
-interface ElementSet { '@type': 'element-set' | Element['@type'] }
+interface ElementSet { '@type': 'element-set' }
 
 export default ElementSet

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

@@ -12,15 +12,12 @@ import Unit from './unit'
 import ElementSet from './element/set'
 import ElementGroup from './element/group'
 import Element from './element'
-import { StructureLookup3D } from './util/lookup3d';
 
 // A structure is a pair of "units" and an element set.
 // Each unit contains the data and transformation of its corresponding elements.
 interface Structure {
     readonly units: ReadonlyArray<Unit>,
-    readonly elements: ElementSet,
-
-    __lookup3d__?: StructureLookup3D
+    readonly elements: ElementSet
 }
 
 namespace Structure {
@@ -77,7 +74,7 @@ namespace Structure {
     export function getModels(s: Structure) {
         const { units, elements } = s;
         const arr = UniqueArray.create<Model['id'], Model>();
-        const ids = ElementSet.unitIds(elements);
+        const ids = ElementSet.unitIndices(elements);
         for (let i = 0; i < ids.length; i++) {
             const u = units[ids[i]];
             UniqueArray.add(arr, u.model.id, u.model);
@@ -86,9 +83,7 @@ namespace Structure {
     }
 
     export function getLookup3d(s: Structure) {
-        if (s.__lookup3d__) return s.__lookup3d__;
-        s.__lookup3d__ = StructureLookup3D.create(s);
-        return s.__lookup3d__;
+        return ElementSet.getLookup3d(s);
     }
 
     export function getBoundary(s: Structure) {

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

@@ -32,12 +32,12 @@ function buildAssemblyImpl(structure: Structure, name: string) {
             if (Selection.structureCount(selection) === 0) continue;
             const { units, elements } = Selection.unionStructure(selection);
 
-            const unitIds = ElementSet.unitIds(elements);
+            const unitIds = ElementSet.unitIndices(elements);
 
             for (const oper of g.operators) {
                 for (let uI = 0, _uI = unitIds.length; uI < _uI; uI++) {
                     const unit = units[unitIds[uI]];
-                    assembler.add(Unit.withOperator(unit, oper), ElementSet.unitGetByIndex(elements, uI));
+                    assembler.add(Unit.withOperator(unit, oper), ElementSet.groupAt(elements, uI));
                 }
             }
         }

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

@@ -20,9 +20,9 @@ function computeStructureBoundary(s: Structure): { box: Box3D, sphere: Sphere3D
     let radiusSq = 0;
     let size = 0;
 
-    for (let i = 0, _i = ElementSet.unitCount(elements); i < _i; i++) {
-        const group = ElementSet.unitGetByIndex(elements, i);
-        const { x, y, z } = units[ElementSet.unitGetId(elements, i)];
+    for (let i = 0, _i = ElementSet.groupCount(elements); i < _i; i++) {
+        const group = ElementSet.groupAt(elements, i);
+        const { x, y, z } = units[ElementSet.groupUnitIndex(elements, i)];
 
         size += ElementGroup.size(group);
         for (let j = 0, _j = ElementGroup.size(group); j < _j; j++) {
@@ -48,9 +48,9 @@ function computeStructureBoundary(s: Structure): { box: Box3D, sphere: Sphere3D
         cz /= size;
     }
 
-    for (let i = 0, _i = ElementSet.unitCount(elements); i < _i; i++) {
-        const group = ElementSet.unitGetByIndex(elements, i);
-        const { x, y, z } = units[ElementSet.unitGetId(elements, i)];
+    for (let i = 0, _i = ElementSet.groupCount(elements); i < _i; i++) {
+        const group = ElementSet.groupAt(elements, i);
+        const { x, y, z } = units[ElementSet.groupUnitIndex(elements, i)];
 
         size += ElementGroup.size(group);
         for (let j = 0, _j = ElementGroup.size(group); j < _j; j++) {

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

@@ -12,10 +12,10 @@ import { Vec3 } from 'mol-math/linear-algebra';
 import { OrderedSet } from 'mol-data/int';
 import { computeStructureBoundary } from './boundary';
 
-interface StructureLookup3D extends Lookup3D<Element> {}
+interface ElementSetLookup3D extends Lookup3D<Element> {}
 
-namespace StructureLookup3D {
-    class Impl implements StructureLookup3D {
+namespace ElementSetLookup3D {
+    class Impl implements ElementSetLookup3D {
         private unitLookup: Lookup3D;
         private result = Result.create<Element>();
         private pivot = Vec3.zero();
@@ -28,8 +28,8 @@ namespace StructureLookup3D {
 
             for (let t = 0, _t = closeUnits.count; t < _t; t++) {
                 const i = closeUnits.indices[t];
-                const unitId = ElementSet.unitGetId(elements, i);
-                const group = ElementSet.unitGetByIndex(elements, i);
+                const unitId = ElementSet.groupUnitIndex(elements, i);
+                const group = ElementSet.groupAt(elements, i);
                 const unit = units[unitId];
                 Vec3.set(this.pivot, x, y, z);
                 if (!unit.operator.isIdentity) {
@@ -52,8 +52,8 @@ namespace StructureLookup3D {
 
             for (let t = 0, _t = closeUnits.count; t < _t; t++) {
                 const i = closeUnits.indices[t];
-                const unitId = ElementSet.unitGetId(elements, i);
-                const group = ElementSet.unitGetByIndex(elements, i);
+                const unitId = ElementSet.groupUnitIndex(elements, i);
+                const group = ElementSet.groupAt(elements, i);
                 const unit = units[unitId];
                 Vec3.set(this.pivot, x, y, z);
                 if (!unit.operator.isIdentity) {
@@ -70,7 +70,7 @@ namespace StructureLookup3D {
 
         constructor(private structure: Structure) {
             const { units, elements } = structure;
-            const unitCount = ElementSet.unitCount(elements);
+            const unitCount = ElementSet.groupCount(elements);
             const xs = new Float32Array(unitCount);
             const ys = new Float32Array(unitCount);
             const zs = new Float32Array(unitCount);
@@ -78,8 +78,8 @@ namespace StructureLookup3D {
 
             const center = Vec3.zero();
             for (let i = 0; i < unitCount; i++) {
-                const group = ElementSet.unitGetByIndex(elements, i);
-                const unit = units[ElementSet.unitGetId(elements, i)];
+                const group = ElementSet.groupAt(elements, i);
+                const unit = units[ElementSet.groupUnitIndex(elements, i)];
                 const lookup = Unit.getLookup3d(unit, group);
                 const s = lookup.boundary.sphere;
 
@@ -96,9 +96,9 @@ namespace StructureLookup3D {
         }
     }
 
-    export function create(s: Structure): StructureLookup3D {
+    export function create(s: Structure): ElementSetLookup3D {
         return new Impl(s);
     }
 }
 
-export { StructureLookup3D }
+export { ElementSetLookup3D }

src/perf-tests/structure.ts

@@ -120,14 +120,14 @@ export namespace PropertyAccess {
 
     function sumProperty(structure: Structure, p: Element.Property<number>) {
         const { elements, units } = structure;
-        const unitIds = ElementSet.unitIds(elements);
+        const unitIds = ElementSet.unitIndices(elements);
         const l = Element.Location();
 
         let s = 0;
 
         for (let i = 0, _i = unitIds.length; i < _i; i++) {
             l.unit = units[unitIds[i]];
-            const set = ElementSet.unitGetByIndex(elements, i);
+            const set = ElementSet.groupAt(elements, i);
 
 
             for (let j = 0, _j = ElementGroup.size(set); j < _j; j++) {
@@ -141,7 +141,7 @@ export namespace PropertyAccess {
 
     function sumPropertySegmented(structure: Structure, p: Element.Property<number>) {
         const { elements, units } = structure;
-        const unitIds = ElementSet.unitIds(elements);
+        const unitIds = ElementSet.unitIndices(elements);
         const l = Element.Location();
 
         let s = 0;
@@ -150,7 +150,7 @@ export namespace PropertyAccess {
         for (let i = 0, _i = unitIds.length; i < _i; i++) {
             const unit = units[unitIds[i]];
             l.unit = unit;
-            const set = ElementSet.unitGetByIndex(elements, i);
+            const set = ElementSet.groupAt(elements, i);
 
             const chainsIt = Segmentation.transientSegments(unit.hierarchy.chainSegments, set.elements);
             const residues = unit.hierarchy.residueSegments;
@@ -312,9 +312,9 @@ export namespace PropertyAccess {
             (a, b) => a.unit.model.id === b.unit.model.id && (a.group.key === b.group.key && OrderedSet.areEqual(a.group.elements, b.group.elements))
         );
 
-        for (let i = 0, _i = ElementSet.unitCount(elements); i < _i; i++) {
-            const group = ElementSet.unitGetByIndex(elements, i);
-            const unitId = ElementSet.unitGetId(elements, i);
+        for (let i = 0, _i = ElementSet.groupCount(elements); i < _i; i++) {
+            const group = ElementSet.groupAt(elements, i);
+            const unitId = ElementSet.groupUnitIndex(elements, i);
             uniqueGroups.add(unitId, { unit: units[unitId], group });
         }