improved assembly symmetry init and cage size

src/mol-model-props/rcsb/assembly-symmetry.ts

@@ -8,7 +8,7 @@ import { AssemblySymmetryQuery, AssemblySymmetryQueryVariables } from './graphql
 import query from './graphql/symmetry.gql';
 
 import { ParamDefinition as PD } from '../../mol-util/param-definition'
-import { CustomPropertyDescriptor, Structure, Model } from '../../mol-model/structure';
+import { CustomPropertyDescriptor, Structure, Model, StructureSelection, QueryContext } from '../../mol-model/structure';
 import { Database as _Database, Column } from '../../mol-data/db'
 import { GraphQLClient } from '../../mol-util/graphql-client';
 import { CustomProperty } from '../common/custom-property';
@@ -16,6 +16,9 @@ import { NonNullableArray } from '../../mol-util/type-helpers';
 import { CustomStructureProperty } from '../common/custom-structure-property';
 import { MmcifFormat } from '../../mol-model-formats/structure/mmcif';
 import { ReadonlyVec3 } from '../../mol-math/linear-algebra/3d/vec3';
+import { SetUtils } from '../../mol-util/set';
+import { MolScriptBuilder as MS } from '../../mol-script/language/builder';
+import { compile } from '../../mol-script/runtime/query/compiler';
 
 const BiologicalAssemblyNames = new Set([
     'author_and_software_defined_assembly',
@@ -83,6 +86,34 @@ export namespace AssemblySymmetry {
     export function isRotationAxes(x: AssemblySymmetryValue['rotation_axes']): x is RotationAxes {
         return !!x && x.length > 0
     }
+
+    export function getAsymIds(assemblySymmetry: AssemblySymmetryValue) {
+        const asymIds = new Set<string>()
+        for (const c of assemblySymmetry.clusters) {
+            if (!c?.members) continue
+            for (const m of c.members) {
+                if (m?.asym_id) asymIds.add(m.asym_id)
+            }
+        }
+        return SetUtils.toArray(asymIds)
+    }
+
+    function getAsymIdsStructure(structure: Structure, asymIds: string[]) {
+        const query = MS.struct.modifier.union([
+            MS.struct.generator.atomGroups({
+                'chain-test': MS.core.set.has([MS.set(...asymIds), MS.ammp('label_asym_id')])
+            })
+        ])
+        const compiled = compile<StructureSelection>(query)
+        const result = compiled(new QueryContext(structure))
+        return StructureSelection.unionStructure(result)
+    }
+
+    /** Returns structure limited to all cluster member chains */
+    export function getStructure(structure: Structure, assemblySymmetry: AssemblySymmetryValue) {
+        const asymIds = AssemblySymmetry.getAsymIds(assemblySymmetry)
+        return asymIds.length > 0 ? getAsymIdsStructure(structure, asymIds) : structure
+    }
 }
 
 export function getSymmetrySelectParam(structure?: Structure) {

src/mol-model-props/rcsb/representations/assembly-symmetry.ts

@@ -7,7 +7,7 @@
 import { ParamDefinition as PD } from '../../../mol-util/param-definition';
 import { AssemblySymmetryValue, AssemblySymmetryProvider, AssemblySymmetry } from '../assembly-symmetry';
 import { MeshBuilder } from '../../../mol-geo/geometry/mesh/mesh-builder';
-import { Vec3, Mat4 } from '../../../mol-math/linear-algebra';
+import { Vec3, Mat4, Mat3 } from '../../../mol-math/linear-algebra';
 import { addCylinder } from '../../../mol-geo/geometry/mesh/builder/cylinder';
 import { Mesh } from '../../../mol-geo/geometry/mesh/mesh';
 import { RuntimeContext } from '../../../mol-task';
@@ -32,6 +32,7 @@ import { Mutable } from '../../../mol-util/type-helpers';
 import { equalEps } from '../../../mol-math/linear-algebra/3d/common';
 import { Structure } from '../../../mol-model/structure';
 import { isInteger } from '../../../mol-util/number';
+import { Sphere3D } from '../../../mol-math/geometry';
 
 const OrderColors = ColorMap({
     '2': ColorNames.deepskyblue,
@@ -71,8 +72,9 @@ const CageParams = {
 type CageParams = typeof CageParams
 
 const AssemblySymmetryVisuals = {
-    'axes': (ctx: RepresentationContext, getParams: RepresentationParamsGetter<Structure, AxesParams>) => ShapeRepresentation(getAxesShape, Mesh.Utils, { modifyState: s => ({ ...s, markerActions: MarkerActions.Highlighting }) }),
+    // cage should come before 'axes' so that the representative loci uses the cage shape
     'cage': (ctx: RepresentationContext, getParams: RepresentationParamsGetter<Structure, CageParams>) => ShapeRepresentation(getCageShape, Mesh.Utils, { modifyState: s => ({ ...s, markerActions: MarkerActions.Highlighting }) }),
+    'axes': (ctx: RepresentationContext, getParams: RepresentationParamsGetter<Structure, AxesParams>) => ShapeRepresentation(getAxesShape, Mesh.Utils, { modifyState: s => ({ ...s, markerActions: MarkerActions.Highlighting }) }),
 }
 
 export const AssemblySymmetryParams = {
@@ -271,13 +273,43 @@ function setSymbolTransform(t: Mat4, symbol: string, axes: AssemblySymmetry.Rota
             } else {
                 Vec3.copy(up, Vec3.unitX)
             }
-            const sizeXY = (structure.lookup3d.boundary.sphere.radius * 2) * 0.8
             Mat4.targetTo(t, eye, target, up)
-            Mat4.scale(t, t, Vec3.create(sizeXY, sizeXY, size))
+
+            const { sphere } = structure.lookup3d.boundary
+            let sizeXY = (sphere.radius * 2) * 0.8 // fallback for missing extrema
+            if (Sphere3D.hasExtrema(sphere)) {
+                const n = Mat3.directionTransform(Mat3(), t)
+                const dirs = unitCircleDirections.map(d => Vec3.transformMat3(Vec3(), d, n))
+                sizeXY = getMaxProjectedDistance(sphere.extrema, dirs, sphere.center)
+            }
+
+            Mat4.scale(t, t, Vec3.create(sizeXY, sizeXY, size * 0.9))
         }
     }
 }
 
+const unitCircleDirections = (function() {
+    const dirs: Vec3[] = []
+    const circle = polygon(12, false, 1)
+    for (let i = 0, il = circle.length; i < il; i += 3) {
+        dirs.push(Vec3.fromArray(Vec3(), circle, i))
+    }
+    return dirs
+})()
+const tmpProj = Vec3()
+
+function getMaxProjectedDistance(points: Vec3[], directions: Vec3[], center: Vec3) {
+    let maxDist = 0
+    for (const p of points) {
+        for (const d of directions) {
+            Vec3.projectPointOnVector(tmpProj, p, d, center)
+            const dist = Vec3.distance(tmpProj, center)
+            if (dist > maxDist) maxDist = dist
+        }
+    }
+    return maxDist
+}
+
 function getCageMesh(data: Structure, props: PD.Values<CageParams>, mesh?: Mesh) {
     const assemblySymmetry = AssemblySymmetryProvider.get(data).value!
     const { scale } = props
@@ -285,6 +317,8 @@ function getCageMesh(data: Structure, props: PD.Values<CageParams>, mesh?: Mesh)
     const { rotation_axes, symbol } = assemblySymmetry
     if (!AssemblySymmetry.isRotationAxes(rotation_axes)) return Mesh.createEmpty(mesh)
 
+    const structure = AssemblySymmetry.getStructure(data, assemblySymmetry)
+
     const cage = getSymbolCage(symbol)
     if (!cage) return Mesh.createEmpty(mesh)
 
@@ -294,7 +328,7 @@ function getCageMesh(data: Structure, props: PD.Values<CageParams>, mesh?: Mesh)
 
     const builderState = MeshBuilder.createState(256, 128, mesh)
     builderState.currentGroup = 0
-    setSymbolTransform(t, symbol, rotation_axes, size, data)
+    setSymbolTransform(t, symbol, rotation_axes, size, structure)
     Vec3.scale(tmpCenter, Vec3.add(tmpCenter, start, end), 0.5)
     Mat4.setTranslation(t, tmpCenter)
     MeshBuilder.addCage(builderState, t, cage, radius, 1, 8)

src/mol-plugin/behavior/dynamic/custom-props/rcsb/assembly-symmetry.ts

@@ -148,9 +148,8 @@ const AssemblySymmetry3D = PluginStateTransform.BuiltIn({
 
 //
 
-const AssemblySymmetryPresetParams = {
+export const AssemblySymmetryPresetParams = {
     ...StructureRepresentationPresetProvider.CommonParams,
-    symmetryIndex: PD.Numeric(0)
 }
 
 export const AssemblySymmetryPreset = StructureRepresentationPresetProvider({

src/mol-plugin/behavior/dynamic/custom-props/rcsb/ui/assembly-symmetry.tsx

@@ -7,11 +7,16 @@
 import * as React from 'react';
 import { CollapsableState, CollapsableControls } from '../../../../../../mol-plugin-ui/base';
 import { ApplyActionControl } from '../../../../../../mol-plugin-ui/state/apply-action';
-import { InitAssemblySymmetry3D, AssemblySymmetry3D } from '../assembly-symmetry';
+import { InitAssemblySymmetry3D, AssemblySymmetry3D, AssemblySymmetryPreset } from '../assembly-symmetry';
 import { AssemblySymmetryProvider,  AssemblySymmetryProps, AssemblySymmetryDataProvider } from '../../../../../../mol-model-props/rcsb/assembly-symmetry';
 import { ParameterControls } from '../../../../../../mol-plugin-ui/controls/parameters';
 import { ParamDefinition as PD } from '../../../../../../mol-util/param-definition';
 import { StructureHierarchyManager } from '../../../../../../mol-plugin-state/manager/structure/hierarchy';
+import { StateAction } from '../../../../../../mol-state';
+import { PluginStateObject } from '../../../../../../mol-plugin-state/objects';
+import { PluginContext } from '../../../../../context';
+import { Task } from '../../../../../../mol-task';
+import { PluginCommands } from '../../../../../commands';
 
 interface AssemblySymmetryControlState extends CollapsableState {
     isBusy: boolean
@@ -56,7 +61,7 @@ export class AssemblySymmetryControls extends CollapsableControls<{}, AssemblySy
     renderEnable() {
         const pivot = this.pivot;
         if (!pivot.cell.parent) return null;
-        return <ApplyActionControl state={pivot.cell.parent} action={InitAssemblySymmetry3D} initiallyCollapsed={true} nodeRef={pivot.cell.transform.ref} simpleApply={{ header: 'Enable', icon: 'check' }} />;
+        return <ApplyActionControl state={pivot.cell.parent} action={EnableAssemblySymmetry3D} initiallyCollapsed={true} nodeRef={pivot.cell.transform.ref} simpleApply={{ header: 'Enable', icon: 'check' }} />;
     }
 
     renderNoSymmetries() {
@@ -132,4 +137,12 @@ export class AssemblySymmetryControls extends CollapsableControls<{}, AssemblySy
         if (this.enable) return this.renderEnable();
         return this.renderParams();
     }
-}
+}
+
+const EnableAssemblySymmetry3D = StateAction.build({
+    from: PluginStateObject.Molecule.Structure,
+})(({ a, ref, state }, plugin: PluginContext) => Task.create('Enable Assembly Symmetry', async ctx => {
+    const action = InitAssemblySymmetry3D.create({})
+    await PluginCommands.State.ApplyAction(plugin, { state, action, ref })
+    await AssemblySymmetryPreset.apply(ref, Object.create(null), plugin)
+}));