Bladeren bron

Updated mol-geo to new Structure data model

David Sehnal 7 jaren geleden
bovenliggende
commit
0427d52d97

+ 2 - 2
src/apps/render-test/state.ts

@@ -17,7 +17,7 @@ import Spacefill, { SpacefillProps } from 'mol-geo/representation/structure/spac
 import Point, { PointProps } from 'mol-geo/representation/structure/point'
 
 import { Run } from 'mol-task'
-import { Symmetry, Structure, Model } from 'mol-model/structure'
+import { StructureSymmetry, Structure, Model } from 'mol-model/structure'
 
 // import mcubes from './utils/mcubes'
 import { getModelFromPdbId, getModelFromFile, log, Volume, getVolumeFromEmdId } from './utils'
@@ -110,7 +110,7 @@ export default class State {
         let structure: Structure
         const assemblies = model.symmetry.assemblies
         if (assemblies.length) {
-            structure = await Run(Symmetry.buildAssembly(Structure.ofModel(model), assembly || '1'), log, 500)
+            structure = await Run(StructureSymmetry.buildAssembly(Structure.ofModel(model), assembly || '1'), log, 500)
         } else {
             structure = Structure.ofModel(model)
         }

+ 26 - 31
src/apps/structure-info/model.ts

@@ -10,7 +10,7 @@ require('util.promisify').shim();
 
 // import { Table } from 'mol-data/db'
 import CIF from 'mol-io/reader/cif'
-import { Model, Structure, Element, ElementSet, Unit, ElementGroup, Queries } from 'mol-model/structure'
+import { Model, Structure, Element, Unit, Queries } from 'mol-model/structure'
 // import { Run, Progress } from 'mol-task'
 import { OrderedSet } from 'mol-data/int';
 import { Table } from 'mol-data/db';
@@ -41,31 +41,29 @@ export function atomLabel(model: Model, aI: number) {
 
 
 export function printBonds(structure: Structure) {
-    const { units, elements } = structure;
-    const unitIds = ElementSet.unitIndices(elements);
+    // TODO
+    // for (const unit of structure.units) {
+    //     const unit = units[OrderedSet.getAt(unitIds, i)];
+    //     const group = ElementSet.groupFromUnitIndex(elements, OrderedSet.getAt(unitIds, i));
 
-    for (let i = 0, _i = OrderedSet.size(unitIds); i < _i; i++) {
-        const unit = units[OrderedSet.getAt(unitIds, i)];
-        const group = ElementSet.groupFromUnitIndex(elements, OrderedSet.getAt(unitIds, i));
+    //     const { count, offset, neighbor } = Unit.getGroupBonds(unit, group);
+    //     const { model }  = unit;
 
-        const { count, offset, neighbor } = Unit.getGroupBonds(unit, group);
-        const { model }  = unit;
+    //     if (!count) continue;
 
-        if (!count) continue;
+    //     for (let j = 0; j < offset.length - 1; ++j) {
+    //         const start = offset[j];
+    //         const end = offset[j + 1];
 
-        for (let j = 0; j < offset.length - 1; ++j) {
-            const start = offset[j];
-            const end = offset[j + 1];
+    //         if (end <= start) continue;
 
-            if (end <= start) continue;
-
-            const aI = ElementGroup.getAt(group, j);
-            for (let _bI = start; _bI < end; _bI++) {
-                const bI = ElementGroup.getAt(group, neighbor[_bI])
-                console.log(`${atomLabel(model, aI)} -- ${atomLabel(model, bI)}`);
-            }
-        }
-    }
+    //         const aI = ElementGroup.getAt(group, j);
+    //         for (let _bI = start; _bI < end; _bI++) {
+    //             const bI = ElementGroup.getAt(group, neighbor[_bI])
+    //             console.log(`${atomLabel(model, aI)} -- ${atomLabel(model, bI)}`);
+    //         }
+    //     }
+    // }
 }
 
 export function printSequence(model: Model) {
@@ -83,26 +81,23 @@ export function printSequence(model: Model) {
 
 export function printUnits(structure: Structure) {
     console.log('Units\n=============');
-    const { elements, units } = structure;
-    const unitIds = ElementSet.unitIndices(elements);
     const l = Element.Location();
 
-    for (let i = 0, _i = unitIds.length; i < _i; i++) {
-        const unitId = unitIds[i];
-        l.unit = units[unitId];
-        const set = ElementSet.groupAt(elements, i).elements;
-        const size = OrderedSet.size(set);
+    for (const unit of structure.units) {
+        l.unit = unit;
+        const elements = unit.elements;
+        const size = OrderedSet.size(elements);
 
         if (Unit.isAtomic(l.unit)) {
-            console.log(`Atomic unit ${unitId}: ${size} elements`);
+            console.log(`Atomic unit ${unit.id}: ${size} elements`);
         } else if (Unit.isCoarse(l.unit)) {
-            console.log(`Coarse unit ${unitId} (${Unit.isSpheres(l.unit) ? 'spheres' : 'gaussians'}): ${size} elements.`);
+            console.log(`Coarse unit ${unit.id} (${Unit.isSpheres(l.unit) ? 'spheres' : 'gaussians'}): ${size} elements.`);
 
             const props = Queries.props.coarse_grained;
             const seq = l.unit.model.sequence;
 
             for (let j = 0, _j = Math.min(size, 10); j < _j; j++) {
-                l.element = OrderedSet.getAt(set, j);
+                l.element = OrderedSet.getAt(elements, j);
 
                 const residues: string[] = [];
                 const start = props.seq_id_begin(l), end = props.seq_id_end(l);

+ 13 - 43
src/mol-geo/representation/structure/index.ts

@@ -2,11 +2,10 @@
  * Copyright (c) 2018 mol* contributors, licensed under MIT, See LICENSE file for more info.
  *
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
+ * @author David Sehnal <david.sehnal@gmail.com>
  */
 
-import { ElementGroup, ElementSet, Structure, Unit } from 'mol-model/structure';
-import { EquivalenceClasses } from 'mol-data/util';
-import { OrderedSet } from 'mol-data/int'
+import { Structure, StructureSymmetry } from 'mol-model/structure';
 import { Task } from 'mol-task'
 import { RenderObject } from 'mol-gl/scene';
 import { Representation, RepresentationProps } from '..';
@@ -15,7 +14,7 @@ import { Representation, RepresentationProps } from '..';
 
 export interface UnitsRepresentation<P> {
     renderObjects: ReadonlyArray<RenderObject>
-    create: (units: ReadonlyArray<Unit>, elementGroup: ElementGroup, props: P) => Task<void>
+    create: (group: StructureSymmetry.UnitGroup, props: P) => Task<void>
     update: (props: P) => Task<boolean>
 }
 
@@ -27,8 +26,7 @@ export interface StructureRepresentation<P extends RepresentationProps = {}> ext
 
 interface GroupRepresentation<T> {
     repr: UnitsRepresentation<T>
-    units: Unit[]
-    elementGroup: ElementGroup
+    group: StructureSymmetry.UnitGroup
 }
 
 export function StructureRepresentation<P>(reprCtor: () => UnitsRepresentation<P>): StructureRepresentation<P> {
@@ -39,40 +37,12 @@ export function StructureRepresentation<P>(reprCtor: () => UnitsRepresentation<P
         renderObjects,
         create(structure: Structure, props: P = {} as P) {
             return Task.create('StructureRepresentation.create', async ctx => {
-                const { elements, units } = structure;
-                const uniqueGroups = EquivalenceClasses<number, { unit: Unit, group: ElementGroup }>(
-                    ({ unit, group }) => ElementGroup.hashCode(group),
-                    (a, b) => a.unit.model.id === b.unit.model.id && OrderedSet.areEqual(a.group.elements, b.group.elements)
-                );
-
-                // const uniqueTransformations = EquivalenceClasses<number, { unit: Unit, group: ElementGroup }>(
-                //     ({ unit, group }) => unit.operator.matrix.join(','),
-                //     (a, b) => Mat4.areEqual(a.unit.operator.matrix, b.unit.operator.matrix, EPSILON.Value)
-                // );
-
-                const unitIndices = ElementSet.unitIndices(elements);
-                for (let i = 0, _i = unitIndices.length; i < _i; i++) {
-                    const unitIndex = unitIndices[i];
-                    const group = ElementSet.groupFromUnitIndex(elements, unitIndex);
-                    const unit = units[unitIndex]
-                    uniqueGroups.add(unitIndex, { unit, group });
-                    // uniqueTransformations.add(unitIndex, { unit, group });
-                }
-
-                // console.log({ uniqueGroups, uniqueTransformations })
-
-                for (let i = 0, il = uniqueGroups.groups.length; i < il; i++) {
-                    const groupUnits: Unit[] = []
-                    const group = uniqueGroups.groups[i]
-                    // console.log('group', i)
-                    for (let j = 0, jl = group.length; j < jl; j++) {
-                        groupUnits.push(units[group[j]])
-                    }
-                    const elementGroup = ElementSet.groupFromUnitIndex(elements, group[0])
+                const groups = StructureSymmetry.getTransformGroups(structure);
+                for (let i = 0; i < groups.length; i++) {
+                    const group = groups[i];
                     const repr = reprCtor()
-                    groupReprs.push({ repr, units: groupUnits, elementGroup })
-                    await ctx.update({ message: 'Building structure unit representations...', current: i, max: il });
-                    await ctx.runChild(repr.create(groupUnits, elementGroup, props));
+                    groupReprs.push({ repr, group })
+                    await ctx.runChild(repr.create(group, props), { message: 'Building structure unit representations...', current: i, max: groups.length });
                     renderObjects.push(...repr.renderObjects)
                 }
             });
@@ -83,10 +53,10 @@ export function StructureRepresentation<P>(reprCtor: () => UnitsRepresentation<P
                 renderObjects.length = 0 // clear
                 for (let i = 0, il = groupReprs.length; i < il; ++i) {
                     const groupRepr = groupReprs[i]
-                    const { repr, units, elementGroup } = groupRepr
-                    await ctx.update({ message: 'Updating structure unit representations...', current: i, max: il });
-                    if (!await ctx.runChild(repr.update(props))) {
-                        await ctx.runChild(repr.create(units, elementGroup, props))
+                    const { repr, group } = groupRepr
+                    const state = { message: 'Updating structure unit representations...', current: i, max: il };
+                    if (!await ctx.runChild(repr.update(props), state)) {
+                        await ctx.runChild(repr.create(group, props), state)
                     }
                     renderObjects.push(...repr.renderObjects)
                 }

+ 19 - 17
src/mol-geo/representation/structure/point.ts

@@ -2,12 +2,12 @@
  * Copyright (c) 2018 mol* contributors, licensed under MIT, See LICENSE file for more info.
  *
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
+ * @author David Sehnal <david.sehnal@gmail.com>
  */
 
 import { ValueCell } from 'mol-util/value-cell'
 import { createPointRenderObject, RenderObject, PointRenderObject } from 'mol-gl/scene'
-import { OrderedSet } from 'mol-data/int'
-import { Unit, ElementGroup, Element } from 'mol-model/structure';
+import { Unit, Element, StructureSymmetry } from 'mol-model/structure';
 import { Task } from 'mol-task'
 import { fillSerial } from 'mol-gl/renderable/util';
 
@@ -16,6 +16,7 @@ import VertexMap from '../../shape/vertex-map';
 import { ColorTheme, SizeTheme } from '../../theme';
 import { createTransforms, createColors, createSizes } from './utils';
 import { deepEqual } from 'mol-util';
+import { SortedArray } from 'mol-data/int';
 
 export const DefaultPointProps = {
     colorTheme: { name: 'instance-index' } as ColorTheme,
@@ -25,16 +26,17 @@ export const DefaultPointProps = {
 }
 export type PointProps = Partial<typeof DefaultPointProps>
 
-export function createPointVertices(unit: Unit, elementGroup: ElementGroup) {
-    const elementCount = OrderedSet.size(elementGroup.elements)
+export function createPointVertices(unit: Unit) {
+    const elements = unit.elements
+    const elementCount = elements.length
     const vertices = new Float32Array(elementCount * 3)
 
-    const { x, y, z } = unit
+    const { x, y, z } = unit.conformation
     const l = Element.Location()
     l.unit = unit
 
     for (let i = 0; i < elementCount; i++) {
-        l.element = ElementGroup.getAt(elementGroup, i)
+        l.element = elements[i];
         const i3 = i * 3
         vertices[i3] = x(l.element)
         vertices[i3 + 1] = y(l.element)
@@ -49,21 +51,21 @@ export default function Point(): UnitsRepresentation<PointProps> {
     let curProps = DefaultPointProps
 
     let _units: ReadonlyArray<Unit>
-    let _elementGroup: ElementGroup
+    let _elements: SortedArray
 
     return {
         renderObjects,
-        create(units: ReadonlyArray<Unit>, elementGroup: ElementGroup, props: PointProps = {}) {
+        create(group: StructureSymmetry.UnitGroup, props: PointProps = {}) {
             return Task.create('Point.create', async ctx => {
                 renderObjects.length = 0 // clear
                 curProps = { ...DefaultPointProps, ...props }
 
-                _units = units
-                _elementGroup = elementGroup
+                _units = group.units
+                _elements = group.elements;
 
                 const { colorTheme, sizeTheme, alpha, visible } = curProps
-                const elementCount = OrderedSet.size(elementGroup.elements)
-                const unitCount = units.length
+                const elementCount = _elements.length
+                const unitCount = _units.length
 
                 const vertexMap = VertexMap.create(
                     elementCount,
@@ -73,16 +75,16 @@ export default function Point(): UnitsRepresentation<PointProps> {
                 )
 
                 await ctx.update('Computing point vertices');
-                const vertices = createPointVertices(units[0], elementGroup)
+                const vertices = createPointVertices(_units[0])
 
                 await ctx.update('Computing point transforms');
-                const transforms = createTransforms(units)
+                const transforms = createTransforms(group)
 
                 await ctx.update('Computing point colors');
-                const color = createColors(units, elementGroup, vertexMap, colorTheme)
+                const color = createColors(group, vertexMap, colorTheme)
 
                 await ctx.update('Computing point sizes');
-                const size = createSizes(units, elementGroup, vertexMap, sizeTheme)
+                const size = createSizes(group, vertexMap, sizeTheme)
 
                 points = createPointRenderObject({
                     objectId: 0,
@@ -106,7 +108,7 @@ export default function Point(): UnitsRepresentation<PointProps> {
         },
         update(props: PointProps) {
             return Task.create('Point.update', async ctx => {
-                if (!points || !_units || !_elementGroup) return false
+                if (!points || !_units || !_elements) return false
 
                 const newProps = { ...curProps, ...props }
                 if (deepEqual(curProps, newProps)) {

+ 13 - 12
src/mol-geo/representation/structure/spacefill.ts

@@ -2,6 +2,7 @@
  * Copyright (c) 2018 mol* contributors, licensed under MIT, See LICENSE file for more info.
  *
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
+ * @author David Sehnal <david.sehnal@gmail.com>
  */
 
 import { ValueCell } from 'mol-util/value-cell'
@@ -9,8 +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 { OrderedSet } from 'mol-data/int'
-import { Unit, ElementGroup, Element, Queries } from 'mol-model/structure';
+import { Unit, Element, Queries, StructureSymmetry } from 'mol-model/structure';
 import { UnitsRepresentation } from './index';
 import { Task } from 'mol-task'
 import { MeshBuilder } from '../../shape/mesh-builder';
@@ -28,9 +28,10 @@ export const DefaultSpacefillProps = {
 }
 export type SpacefillProps = Partial<typeof DefaultSpacefillProps>
 
-function createSpacefillMesh(unit: Unit, elementGroup: ElementGroup, detail: number) {
+function createSpacefillMesh(unit: Unit, detail: number) {
     return Task.create('Sphere mesh', async ctx => {
-        const elementCount = OrderedSet.size(elementGroup.elements)
+        const { elements } = unit;
+        const elementCount = elements.length;
         const vertexCount = elementCount * icosahedronVertexCount(detail)
         const meshBuilder = MeshBuilder.create(vertexCount)
 
@@ -47,12 +48,12 @@ function createSpacefillMesh(unit: Unit, elementGroup: ElementGroup, detail: num
         const v = Vec3.zero()
         const m = Mat4.identity()
 
-        const { x, y, z } = unit
+        const { x, y, z } = unit.conformation
         const l = Element.Location()
         l.unit = unit
 
         for (let i = 0; i < elementCount; i++) {
-            l.element = ElementGroup.getAt(elementGroup, i)
+            l.element = elements[i]
             v[0] = x(l.element)
             v[1] = y(l.element)
             v[2] = z(l.element)
@@ -76,23 +77,23 @@ export default function Spacefill(): UnitsRepresentation<SpacefillProps> {
 
     return {
         renderObjects,
-        create(units: ReadonlyArray<Unit>, elementGroup: ElementGroup, props: SpacefillProps = {}) {
+        create(group: StructureSymmetry.UnitGroup, props: SpacefillProps = {}) {
             return Task.create('Spacefill.create', async ctx => {
                 renderObjects.length = 0 // clear
 
                 const { detail, colorTheme, alpha, visible, doubleSided } = { ...DefaultSpacefillProps, ...props }
 
                 await ctx.update('Computing spacefill mesh');
-                const mesh = await ctx.runChild(createSpacefillMesh(units[0], elementGroup, detail))
+                const mesh = await ctx.runChild(createSpacefillMesh(group.units[0], detail))
                 // console.log(mesh)
 
                 const vertexMap = VertexMap.fromMesh(mesh)
 
                 await ctx.update('Computing spacefill transforms');
-                const transforms = createTransforms(units)
+                const transforms = createTransforms(group)
 
                 await ctx.update('Computing spacefill colors');
-                const color = createColors(units, elementGroup, vertexMap, colorTheme)
+                const color = createColors(group, vertexMap, colorTheme)
 
                 spheres = createMeshRenderObject({
                     objectId: 0,
@@ -107,9 +108,9 @@ export default function Spacefill(): UnitsRepresentation<SpacefillProps> {
                     transform: ValueCell.create(transforms),
                     index: mesh.indexBuffer,
 
-                    instanceCount: units.length,
+                    instanceCount: group.units.length,
                     indexCount: mesh.triangleCount,
-                    elementCount: OrderedSet.size(elementGroup.elements),
+                    elementCount: group.elements.length,
                     positionCount: mesh.vertexCount
                 })
                 renderObjects.push(spheres)

+ 11 - 10
src/mol-geo/representation/structure/utils.ts

@@ -2,9 +2,10 @@
  * Copyright (c) 2018 mol* contributors, licensed under MIT, See LICENSE file for more info.
  *
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
+ * @author David Sehnal <david.sehnal@gmail.com>
  */
 
-import { Unit, ElementGroup } from 'mol-model/structure';
+import { StructureSymmetry } from 'mol-model/structure';
 import { Mat4 } from 'mol-math/linear-algebra'
 
 import { createUniformColor } from '../../util/color-data';
@@ -14,35 +15,35 @@ import VertexMap from '../../shape/vertex-map';
 import { ColorTheme, SizeTheme } from '../../theme';
 import { elementIndexColorData, elementSymbolColorData, instanceIndexColorData, chainIdColorData } from '../../theme/structure/color';
 
-export function createTransforms(units: ReadonlyArray<Unit>) {
+export function createTransforms({ units }: StructureSymmetry.UnitGroup) {
     const unitCount = units.length
     const transforms = new Float32Array(unitCount * 16)
     for (let i = 0; i < unitCount; i++) {
-        Mat4.toArray(units[i].operator.matrix, transforms, i * 16)
+        Mat4.toArray(units[i].conformation.operator.matrix, transforms, i * 16)
     }
     return transforms
 }
 
-export function createColors(units: ReadonlyArray<Unit>, elementGroup: ElementGroup, vertexMap: VertexMap, props: ColorTheme) {
+export function createColors(group: StructureSymmetry.UnitGroup, vertexMap: VertexMap, props: ColorTheme) {
     switch (props.name) {
         case 'atom-index':
-            return elementIndexColorData({ units, elementGroup, vertexMap })
+            return elementIndexColorData({ group, vertexMap })
         case 'chain-id':
-            return chainIdColorData({ units, elementGroup, vertexMap })
+            return chainIdColorData({ group, vertexMap })
         case 'element-symbol':
-            return elementSymbolColorData({ units, elementGroup, vertexMap })
+            return elementSymbolColorData({ group, vertexMap })
         case 'instance-index':
-            return instanceIndexColorData({ units, elementGroup, vertexMap })
+            return instanceIndexColorData({ group, vertexMap })
         case 'uniform':
             return createUniformColor(props)
     }
 }
 
-export function createSizes(units: ReadonlyArray<Unit>, elementGroup: ElementGroup, vertexMap: VertexMap, props: SizeTheme) {
+export function createSizes(group: StructureSymmetry.UnitGroup, vertexMap: VertexMap, props: SizeTheme) {
     switch (props.name) {
         case 'uniform':
             return createUniformSize(props)
         case 'vdw':
-            return elementSizeData({ units, elementGroup, vertexMap })
+            return elementSizeData({ group, vertexMap })
     }
 }

+ 5 - 5
src/mol-geo/theme/structure/color/chain-id.ts

@@ -4,7 +4,7 @@
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
  */
 
-import { ElementGroup, Unit, Queries, Element } from 'mol-model/structure';
+import { Unit, Queries, Element } from 'mol-model/structure';
 
 import { StructureColorDataProps } from '.';
 import { createAttributeOrElementColor } from '../../../util/color-data';
@@ -18,8 +18,8 @@ function createChainIdMap(unit: Unit) {
     let count: number
     let asym_id: Column<string>
     if (Unit.isAtomic(unit)) {
-        asym_id = unit.hierarchy.chains.label_asym_id
-        count = unit.hierarchy.chains._rowCount
+        asym_id = unit.model.hierarchy.chains.label_asym_id
+        count = unit.model.hierarchy.chains._rowCount
     } else if (Unit.isCoarse(unit)) {
         asym_id = unit.sites.asym_id
         count = unit.sites.count
@@ -39,7 +39,7 @@ function createChainIdMap(unit: Unit) {
 }
 
 export function chainIdColorData(props: StructureColorDataProps) {
-    const { units, elementGroup, vertexMap } = props
+    const { group: { units, elements }, vertexMap } = props
     const unit = units[0]
 
     const { map, count } = createChainIdMap(unit)
@@ -59,7 +59,7 @@ export function chainIdColorData(props: StructureColorDataProps) {
 
     return createAttributeOrElementColor(vertexMap, {
         colorFn: (elementIdx: number) => {
-            l.element = ElementGroup.getAt(elementGroup, elementIdx)
+            l.element = elements[elementIdx]
             return scale.color(map.get(asym_id(l)) || 0)
         },
         vertexMap

+ 2 - 3
src/mol-geo/theme/structure/color/element-index.ts

@@ -6,13 +6,12 @@
 
 import { ColorScale } from 'mol-util/color';
 import { StructureColorDataProps } from '.';
-import { OrderedSet } from 'mol-data/int';
 import { createElementInstanceColor } from '../../../util/color-data';
 
 export function elementIndexColorData(props: StructureColorDataProps) {
-    const { units, elementGroup, vertexMap } = props
+    const { group: { units, elements }, vertexMap } = props
     const instanceCount = units.length
-    const elementCount = OrderedSet.size(elementGroup.elements)
+    const elementCount = elements.length
 
     const domain = [ 0, instanceCount * elementCount - 1 ]
     const scale = ColorScale.create({ domain })

+ 2 - 3
src/mol-geo/theme/structure/color/element-symbol.ts

@@ -7,7 +7,6 @@
 import { ElementSymbol } from 'mol-model/structure/model/types';
 import { Color } from 'mol-util/color';
 import { StructureColorDataProps } from '.';
-import { OrderedSet } from 'mol-data/int';
 import { createAttributeOrElementColor } from '../../../util/color-data';
 
 // from Jmol http://jmol.sourceforge.net/jscolors/ (or 0xFFFFFF)
@@ -23,11 +22,11 @@ export function elementSymbolColor(element: ElementSymbol): Color {
 }
 
 export function elementSymbolColorData(props: StructureColorDataProps) {
-    const { units, elementGroup, vertexMap } = props
+    const { group: { units, elements }, vertexMap } = props
     const { type_symbol } = units[0].model.hierarchy.atoms
     return createAttributeOrElementColor(vertexMap, {
         colorFn: (elementIdx: number) => {
-            const e = OrderedSet.getAt(elementGroup.elements, elementIdx)
+            const e = elements[elementIdx]
             return elementSymbolColor(type_symbol.value(e))
         },
         vertexMap

+ 2 - 3
src/mol-geo/theme/structure/color/index.ts

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

+ 1 - 1
src/mol-geo/theme/structure/color/instance-index.ts

@@ -9,7 +9,7 @@ import { StructureColorDataProps } from '.';
 import { createInstanceColor } from '../../../util/color-data';
 
 export function instanceIndexColorData(props: StructureColorDataProps) {
-    const { units } = props
+    const { group: { units } } = props
     const instanceCount = units.length
 
     const domain = [ 0, instanceCount - 1 ]

+ 5 - 4
src/mol-geo/theme/structure/size/element.ts

@@ -4,14 +4,15 @@
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
  */
 
-import { ElementGroup, Element, Unit, Queries } from 'mol-model/structure';
+import { Element, Unit, Queries } from 'mol-model/structure';
 import { StructureSizeDataProps } from '.';
 import { createAttributeSize } from '../../../util/size-data';
 
 /** Create attribute data with the size of an element, i.e. vdw for atoms and radius for coarse spheres */
 export function elementSizeData(props: StructureSizeDataProps) {
-    const { units, elementGroup, vertexMap } = props
-    const unit = units[0]
+    const { group, vertexMap } = props
+    const unit = group.units[0]
+    const elements = group.elements;
     let radius: Element.Property<number>
     if (Unit.isAtomic(unit)) {
         radius = Queries.props.atom.vdw_radius
@@ -22,7 +23,7 @@ export function elementSizeData(props: StructureSizeDataProps) {
     l.unit = unit
     return createAttributeSize({
         sizeFn: (elementIdx: number) => {
-            l.element = ElementGroup.getAt(elementGroup, elementIdx)
+            l.element = elements[elementIdx]
             return radius(l)
         },
         vertexMap

+ 2 - 3
src/mol-geo/theme/structure/size/index.ts

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

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

@@ -7,6 +7,6 @@
 import Element from './structure/element'
 import Structure from './structure/structure'
 import Unit from './structure/unit'
-import Symmetry from './structure/symmetry'
+import StructureSymmetry from './structure/symmetry'
 
-export { Element, Structure, Unit, Symmetry }
+export { Element, Structure, Unit, StructureSymmetry }

+ 52 - 23
src/mol-model/structure/structure/symmetry.ts

@@ -8,37 +8,66 @@ import Structure from './structure'
 import { Selection } from '../query'
 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';
 
-namespace Symmetry {
-    export const buildAssembly = buildAssemblyImpl;
-}
+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 default Symmetry;
+    export function buildAssembly(structure: Structure, name: string) {
+        return Task.create('Build Symmetry', async ctx => {
+            const models = Structure.getModels(structure);
+            if (models.length !== 1) throw new Error('Can only build assemblies from structures based on 1 model.');
 
-function buildAssemblyImpl(structure: Structure, name: string) {
-    return Task.create('Build Symmetry', async ctx => {
-        const models = Structure.getModels(structure);
-        if (models.length !== 1) throw new Error('Can only build assemblies from structures based on 1 model.');
+            const assembly = ModelSymmetry.findAssembly(models[0], name);
+            if (!assembly) throw new Error(`Assembly '${name}' is not defined.`);
 
-        const assembly = ModelSymmetry.findAssembly(models[0], name);
-        if (!assembly) throw new Error(`Assembly '${name}' is not defined.`);
+            const assembler = Structure.Builder();
 
-        const assembler = Structure.Builder();
+            for (const g of assembly.operatorGroups) {
+                const selection = await ctx.runChild(g.selector(structure));
+                if (Selection.structureCount(selection) === 0) {
+                    continue;
+                }
+                const { units } = Selection.unionStructure(selection);
 
-        for (const g of assembly.operatorGroups) {
-            const selection = await ctx.runChild(g.selector(structure));
-            if (Selection.structureCount(selection) === 0) {
-                continue;
+                for (const oper of g.operators) {
+                    for (const unit of units) {
+                        assembler.addWithOperator(unit, oper);
+                    }
+                }
             }
-            const { units } = Selection.unionStructure(selection);
 
-            for (const oper of g.operators) {
-                for (const unit of units) {
-                    assembler.addWithOperator(unit, oper);
-                }
+            return assembler.getStructure();
+        });
+    }
+
+    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);
+
+        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));
+
+            for (const id of group) {
+                const unit = s.unitMap.get(id);
+                setGrouping.add(unit.id, unit);
+            }
+
+            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)) });
             }
         }
 
-        return assembler.getStructure();
-    });
-}
+        return ret;
+    }
+}
+
+export default StructureSymmetry;

+ 2 - 2
src/perf-tests/structure.ts

@@ -11,7 +11,7 @@ import * as fs from 'fs'
 import fetch from 'node-fetch'
 import CIF from 'mol-io/reader/cif'
 
-import { Structure, Model, Queries as Q, Element, Selection, Symmetry, Query } from 'mol-model/structure'
+import { Structure, Model, Queries as Q, Element, Selection, StructureSymmetry, Query } from 'mol-model/structure'
 //import { Segmentation, OrderedSet } from 'mol-data/int'
 
 import to_mmCIF from 'mol-model/structure/export/mmcif'
@@ -291,7 +291,7 @@ export namespace PropertyAccess {
 
     export async function testAssembly(id: string, s: Structure) {
         console.time('assembly')
-        const a = await Run(Symmetry.buildAssembly(s, '1'));
+        const a = await Run(StructureSymmetry.buildAssembly(s, '1'));
         console.timeEnd('assembly')
         fs.writeFileSync(`${DATA_DIR}/${id}_assembly.bcif`, to_mmCIF(id, a, true));
         console.log('exported');