Structure data model tweaks

src/mol-data/util/equivalence-classes.ts

@@ -17,6 +17,7 @@ class EquivalenceClassesImpl<K, V> {
         return { id, keys, value };
     }
 
+    // Return the group representative.
     add(key: K, a: V) {
         const hash = this.getHash(a);
         if (this.byHash.has(hash)) {
@@ -25,16 +26,16 @@ class EquivalenceClassesImpl<K, V> {
                 const group = groups[i];
                 if (this.areEqual(a, group.value)) {
                     group.keys[group.keys.length] = key;
-                    return group.id;
+                    return group.value;
                 }
             }
             const group = this.createGroup(key, a);
             groups[groups.length] = group;
-            return group.id;
+            return group.value;
         } else {
             const group = this.createGroup(key, a);
             this.byHash.set(hash, [group]);
-            return group.id;
+            return group.value;
         }
     }
 

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

@@ -5,7 +5,7 @@
  * @author David Sehnal <david.sehnal@gmail.com>
  */
 
-import { Structure, StructureSymmetry } from 'mol-model/structure';
+import { Structure, StructureSymmetry, Unit } from 'mol-model/structure';
 import { Task } from 'mol-task'
 import { RenderObject } from 'mol-gl/scene';
 import { Representation, RepresentationProps } from '..';
@@ -14,7 +14,7 @@ import { Representation, RepresentationProps } from '..';
 
 export interface UnitsRepresentation<P> {
     renderObjects: ReadonlyArray<RenderObject>
-    create: (group: StructureSymmetry.UnitGroup, props: P) => Task<void>
+    create: (group: Unit.SymmetryGroup, props: P) => Task<void>
     update: (props: P) => Task<boolean>
 }
 
@@ -26,7 +26,7 @@ export interface StructureRepresentation<P extends RepresentationProps = {}> ext
 
 interface GroupRepresentation<T> {
     repr: UnitsRepresentation<T>
-    group: StructureSymmetry.UnitGroup
+    group: Unit.SymmetryGroup
 }
 
 export function StructureRepresentation<P>(reprCtor: () => UnitsRepresentation<P>): StructureRepresentation<P> {

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

@@ -7,7 +7,7 @@
 
 import { ValueCell } from 'mol-util/value-cell'
 import { createPointRenderObject, RenderObject, PointRenderObject } from 'mol-gl/scene'
-import { Unit, Element, StructureSymmetry } from 'mol-model/structure';
+import { Unit, Element } from 'mol-model/structure';
 import { Task } from 'mol-task'
 import { fillSerial } from 'mol-gl/renderable/util';
 
@@ -55,7 +55,7 @@ export default function Point(): UnitsRepresentation<PointProps> {
 
     return {
         renderObjects,
-        create(group: StructureSymmetry.UnitGroup, props: PointProps = {}) {
+        create(group: Unit.SymmetryGroup, props: PointProps = {}) {
             return Task.create('Point.create', async ctx => {
                 renderObjects.length = 0 // clear
                 curProps = { ...DefaultPointProps, ...props }

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

@@ -10,7 +10,7 @@ import { ValueCell } from 'mol-util/value-cell'
 import { RenderObject, createMeshRenderObject, MeshRenderObject } from 'mol-gl/scene'
 // import { createColorTexture } from 'mol-gl/util';
 import { Vec3, Mat4 } from 'mol-math/linear-algebra'
-import { Unit, Element, Queries, StructureSymmetry } from 'mol-model/structure';
+import { Unit, Element, Queries } from 'mol-model/structure';
 import { UnitsRepresentation } from './index';
 import { Task } from 'mol-task'
 import { MeshBuilder } from '../../shape/mesh-builder';
@@ -77,7 +77,7 @@ export default function Spacefill(): UnitsRepresentation<SpacefillProps> {
 
     return {
         renderObjects,
-        create(group: StructureSymmetry.UnitGroup, props: SpacefillProps = {}) {
+        create(group: Unit.SymmetryGroup, props: SpacefillProps = {}) {
             return Task.create('Spacefill.create', async ctx => {
                 renderObjects.length = 0 // clear
 

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

@@ -5,7 +5,7 @@
  * @author David Sehnal <david.sehnal@gmail.com>
  */
 
-import { StructureSymmetry } from 'mol-model/structure';
+import { Unit } from 'mol-model/structure';
 import { Mat4 } from 'mol-math/linear-algebra'
 
 import { createUniformColor } from '../../util/color-data';
@@ -15,7 +15,7 @@ import VertexMap from '../../shape/vertex-map';
 import { ColorTheme, SizeTheme } from '../../theme';
 import { elementIndexColorData, elementSymbolColorData, instanceIndexColorData, chainIdColorData } from '../../theme/structure/color';
 
-export function createTransforms({ units }: StructureSymmetry.UnitGroup) {
+export function createTransforms({ units }: Unit.SymmetryGroup) {
     const unitCount = units.length
     const transforms = new Float32Array(unitCount * 16)
     for (let i = 0; i < unitCount; i++) {
@@ -24,7 +24,7 @@ export function createTransforms({ units }: StructureSymmetry.UnitGroup) {
     return transforms
 }
 
-export function createColors(group: StructureSymmetry.UnitGroup, vertexMap: VertexMap, props: ColorTheme) {
+export function createColors(group: Unit.SymmetryGroup, vertexMap: VertexMap, props: ColorTheme) {
     switch (props.name) {
         case 'atom-index':
             return elementIndexColorData({ group, vertexMap })
@@ -39,7 +39,7 @@ export function createColors(group: StructureSymmetry.UnitGroup, vertexMap: Vert
     }
 }
 
-export function createSizes(group: StructureSymmetry.UnitGroup, vertexMap: VertexMap, props: SizeTheme) {
+export function createSizes(group: Unit.SymmetryGroup, vertexMap: VertexMap, props: SizeTheme) {
     switch (props.name) {
         case 'uniform':
             return createUniformSize(props)

src/mol-geo/theme/structure/color/index.ts

@@ -4,11 +4,11 @@
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
  */
 
-import { StructureSymmetry } from 'mol-model/structure';
+import { Unit } from 'mol-model/structure';
 import VertexMap from '../../../shape/vertex-map';
 
 export interface StructureColorDataProps {
-    group: StructureSymmetry.UnitGroup,
+    group: Unit.SymmetryGroup,
     vertexMap: VertexMap
 }
 

src/mol-geo/theme/structure/size/index.ts

@@ -4,11 +4,11 @@
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
  */
 
-import { StructureSymmetry } from 'mol-model/structure';
+import { Unit } from 'mol-model/structure';
 import VertexMap from '../../../shape/vertex-map';
 
 export interface StructureSizeDataProps {
-    group: StructureSymmetry.UnitGroup,
+    group: Unit.SymmetryGroup,
     vertexMap: VertexMap
 }
 

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

@@ -12,6 +12,7 @@ import { sortArray, sort, arraySwap, hash1 } from 'mol-data/util';
 import Element from './element'
 import Unit from './unit'
 import { StructureLookup3D } from './util/lookup3d';
+import StructureSymmetry from './symmetry';
 
 class Structure {
     readonly unitMap: IntMap<Unit>;
@@ -184,6 +185,8 @@ namespace Structure {
             const newUnits: Unit[] = [];
             sortArray(this.ids);
 
+            const symmGroups = StructureSymmetry.UnitEquivalenceBuilder();
+
             for (let i = 0, _i = this.ids.length; i < _i; i++) {
                 const id = this.ids[i];
                 const parent = this.parent.unitMap.get(id);
@@ -194,13 +197,16 @@ namespace Structure {
                 // if the length is the same, just copy the old unit.
                 if (unit.length === parent.elements.length) {
                     newUnits[newUnits.length] = parent;
+                    symmGroups.add(parent.id, parent);
                     continue;
                 }
 
                 if (!this.isSorted && l > 1) sortArray(unit);
 
-                // TODO: checking if two units are equal in the same symmetry group to reuse computed properties.
-                newUnits[newUnits.length] = parent.getChild(SortedArray.ofSortedArray(unit));
+                let child = parent.getChild(SortedArray.ofSortedArray(unit));
+                const pivot = symmGroups.add(child.id, child);
+                if (child !== pivot) child = pivot.applyOperator(child.id, child.conformation.operator, true);
+                newUnits[newUnits.length] = child;
             }
 
             return create(newUnits);

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

@@ -10,13 +10,9 @@ import { ModelSymmetry } from '../model'
 import { Task } from 'mol-task';
 import { SortedArray } from 'mol-data/int';
 import Unit from './unit';
-import { EquivalenceClasses } from 'mol-data/util';
+import { EquivalenceClasses, hash2 } from 'mol-data/util';
 
 namespace StructureSymmetry {
-    // Units that have the same elements but differ with operator only.
-    export type UnitGroup = { readonly elements: SortedArray, readonly units: ReadonlyArray<Unit> }
-    export type TransformGroups = ReadonlyArray<UnitGroup>
-
     export function buildAssembly(structure: Structure, name: string) {
         return Task.create('Build Symmetry', async ctx => {
             const models = Structure.getModels(structure);
@@ -45,25 +41,26 @@ namespace StructureSymmetry {
         });
     }
 
-    export function getTransformGroups(s: Structure): StructureSymmetry.TransformGroups {
-        // group everything by the "invariantId"
-        const invariantGroups = EquivalenceClasses<number, Unit>(u => u.invariantId, (a, b) => a.invariantId === b.invariantId && a.model.id === b.model.id);
-        for (const u of s.units) invariantGroups.add(u.id, u);
+    function hashUnit(u: Unit) {
+        return hash2(u.invariantId, SortedArray.hashCode(u.elements));
+    }
 
-        const ret: UnitGroup[] = [];
-        // group everything by the "element array"
-        for (const group of invariantGroups.groups) {
-            const setGrouping = EquivalenceClasses<number, Unit>(u => SortedArray.hashCode(u.elements), (a, b) => SortedArray.areEqual(a.elements, b.elements));
+    function areUnitsEquivalent(a: Unit, b: Unit) {
+        return a.invariantId === b.invariantId && a.model.id === b.model.id && SortedArray.areEqual(a.elements, b.elements);
+    }
 
-            for (const id of group) {
-                const unit = s.unitMap.get(id);
-                setGrouping.add(unit.id, unit);
-            }
+    export function UnitEquivalenceBuilder() {
+        return EquivalenceClasses<number, Unit>(hashUnit, areUnitsEquivalent);
+    }
 
-            for (const eqUnits of setGrouping.groups) {
-                const first = s.unitMap.get(eqUnits[0]);
-                ret.push({ elements: first.elements, units: eqUnits.map(id => s.unitMap.get(id)) });
-            }
+    export function getTransformGroups(s: Structure): ReadonlyArray<Unit.SymmetryGroup> {
+        const groups = UnitEquivalenceBuilder();
+        for (const u of s.units) groups.add(u.id, u);
+
+        const ret: Unit.SymmetryGroup[] = [];
+        for (const eqUnits of groups.groups) {
+            const first = s.unitMap.get(eqUnits[0]);
+            ret.push({ elements: first.elements, units: eqUnits.map(id => s.unitMap.get(id)) });
         }
 
         return ret;

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

@@ -32,9 +32,8 @@ namespace Unit {
         }
     }
 
-    export function applyOperator(id: number, unit: Unit, operator: SymmetryOperator): Unit {
-        return create(id, unit.kind, unit.model, SymmetryOperator.compose(unit.conformation.operator, operator), unit.elements);
-    }
+    // A group of units that differ only by symmetry operators.
+    export type SymmetryGroup = { readonly elements: SortedArray, readonly units: ReadonlyArray<Unit> }
 
     export interface Base {
         readonly id: number,
@@ -45,7 +44,7 @@ namespace Unit {
         readonly conformation: SymmetryOperator.ArrayMapping,
 
         getChild(elements: SortedArray): Unit,
-        applyOperator(id: number, operator: SymmetryOperator): Unit,
+        applyOperator(id: number, operator: SymmetryOperator, dontCompose?: boolean /* = false */): Unit,
 
         readonly lookup3d: Lookup3D
     }
@@ -77,8 +76,8 @@ namespace Unit {
             return new Atomic(this.id, this.invariantId, this.model, elements, this.conformation);
         }
 
-        applyOperator(id: number, operator: SymmetryOperator): Unit {
-            const op = SymmetryOperator.compose(this.conformation.operator, operator);
+        applyOperator(id: number, operator: SymmetryOperator, dontCompose = false): Unit {
+            const op = dontCompose ? operator : SymmetryOperator.compose(this.conformation.operator, operator);
             return new Atomic(id, this.invariantId, this.model, this.elements, SymmetryOperator.createMapping(op, this.model.atomSiteConformation));
         }
 
@@ -129,8 +128,8 @@ namespace Unit {
             return createCoarse(this.id, this.invariantId, this.model, this.kind, this.sites, elements, this.conformation);
         }
 
-        applyOperator(id: number, operator: SymmetryOperator): Unit {
-            const op = SymmetryOperator.compose(this.conformation.operator, operator);
+        applyOperator(id: number, operator: SymmetryOperator, dontCompose = false): Unit {
+            const op = dontCompose ? operator : SymmetryOperator.compose(this.conformation.operator, operator);
             return createCoarse(id, this.invariantId, this.model, this.kind, this.sites, this.elements, SymmetryOperator.createMapping(op, this.sites));
         }
 

src/perf-tests/structure.ts

@@ -292,8 +292,13 @@ export namespace PropertyAccess {
     export async function testAssembly(id: string, s: Structure) {
         console.time('assembly')
         const a = await Run(StructureSymmetry.buildAssembly(s, '1'));
+        //const auth_comp_id = Q.props.residue.auth_comp_id;
+        //const q1 = Query(Q.generators.atoms({ residueTest: l => auth_comp_id(l) === 'ALA' }));
+        //const alas = await query(q1, a);
+
         console.timeEnd('assembly')
         fs.writeFileSync(`${DATA_DIR}/${id}_assembly.bcif`, to_mmCIF(id, a, true));
+        //fs.writeFileSync(`${DATA_DIR}/${id}_assembly.bcif`, to_mmCIF(id, Selection.unionStructure(alas), true));
         console.log('exported');
     }
 
@@ -404,8 +409,8 @@ export namespace PropertyAccess {
             //.add('test q', () => q1(structures[0]))
             //.add('test q', () => q(structures[0]))
             .add('test int', () => sumProperty(structures[0], l => col(l.element)))
-            .add('test q1', async () => await q1(structures[0]))
-            .add('test q3', async () => await q3(structures[0]))
+            .add('test q1', async () => await query(q1, structures[0]))
+            .add('test q3', async () => await query(q3, structures[0]))
             // .add('sum residue', () => sumPropertyResidue(structures[0], l => l.unit.hierarchy.residues.auth_seq_id.value(l.unit.residueIndex[l.atom])))
 
             // .add('baseline', () =>  baseline(models[0]))