structure visuals refactoring

src/mol-geo/representation/structure/complex-visual.ts

@@ -6,47 +6,57 @@
 
 import { Structure } from 'mol-model/structure';
 import { Visual } from '..';
-import { MeshRenderObject } from 'mol-gl/render-object';
+import { MeshRenderObject, LinesRenderObject, PointsRenderObject, DirectVolumeRenderObject } from 'mol-gl/render-object';
 import { Mesh } from '../../geometry/mesh/mesh';
 import { RuntimeContext } from 'mol-task';
 import { LocationIterator } from '../../util/location-iterator';
-import { createComplexMeshRenderObject } from './visual/util/common';
-import { StructureProps, VisualUpdateState, StructureMeshParams } from '.';
+import { createComplexMeshRenderObject, sizeChanged, colorChanged, UnitKind, UnitKindOptions } from './visual/util/common';
+import { StructureProps, VisualUpdateState, StructureMeshParams, StructureParams } from '.';
 import { deepEqual, ValueCell } from 'mol-util';
 import { PickingId } from '../../geometry/picking';
 import { Loci, isEveryLoci, EmptyLoci } from 'mol-model/loci';
 import { MarkerAction, applyMarkerAction } from '../../geometry/marker-data';
 import { Interval } from 'mol-data/int';
-import { updateRenderableState } from '../../geometry/geometry';
+import { updateRenderableState, Geometry } from '../../geometry/geometry';
 import { createColors } from '../../geometry/color-data';
-import { UnitKindOptions, UnitKind } from './units-visual';
 import { MultiSelectParam, paramDefaultValues } from 'mol-view/parameter';
+import { RenderableValues } from 'mol-gl/renderable/schema';
+import { createSizes } from 'mol-geo/geometry/size-data';
 
 export interface  ComplexVisual<P extends StructureProps> extends Visual<Structure, P> { }
 
-export const ComplexMeshParams = {
-    ...StructureMeshParams,
-    unitKinds: MultiSelectParam<UnitKind>('Unit Kind', '', [ 'atomic', 'spheres' ], UnitKindOptions),
+const ComplexParams = {
+    ...StructureParams,
+    unitKinds: MultiSelectParam<UnitKind>('Unit Kind', '', ['atomic', 'spheres'], UnitKindOptions),
 }
-export const DefaultComplexMeshProps = paramDefaultValues(ComplexMeshParams)
-export type ComplexMeshProps = typeof DefaultComplexMeshProps
+const DefaultComplexProps = paramDefaultValues(ComplexParams)
+type ComplexProps = typeof DefaultComplexProps
+
+type ComplexRenderObject = MeshRenderObject | LinesRenderObject | PointsRenderObject | DirectVolumeRenderObject
 
-export interface ComplexMeshVisualBuilder<P extends ComplexMeshProps> {
+interface ComplexVisualBuilder<P extends ComplexProps, G extends Geometry> {
     defaultProps: P
-    createMesh(ctx: RuntimeContext, structure: Structure, props: P, mesh?: Mesh): Promise<Mesh>
+    createGeometry(ctx: RuntimeContext, structure: Structure, props: P, geometry?: G): Promise<G>
     createLocationIterator(structure: Structure): LocationIterator
     getLoci(pickingId: PickingId, structure: Structure, id: number): Loci
     mark(loci: Loci, structure: Structure, apply: (interval: Interval) => boolean): boolean,
     setUpdateState(state: VisualUpdateState, newProps: P, currentProps: P): void
 }
 
-export function ComplexMeshVisual<P extends ComplexMeshProps>(builder: ComplexMeshVisualBuilder<P>): ComplexVisual<P> {
-    const { defaultProps, createMesh, createLocationIterator, getLoci, mark, setUpdateState } = builder
+interface ComplexVisualGeometryBuilder<P extends ComplexProps, G extends Geometry> extends ComplexVisualBuilder<P, G> {
+    createEmptyGeometry(geometry?: G): G
+    createRenderObject(ctx: RuntimeContext, structure: Structure, geometry: Geometry, locationIt: LocationIterator, currentProps: P): Promise<ComplexRenderObject>
+    updateValues(values: RenderableValues, newProps: P): void
+}
+
+export function ComplexVisual<P extends ComplexMeshProps>(builder: ComplexVisualGeometryBuilder<P, Geometry>): ComplexVisual<P> {
+    const { defaultProps, createGeometry, createLocationIterator, getLoci, mark, setUpdateState } = builder
+    const { createRenderObject, updateValues } = builder
     const updateState = VisualUpdateState.create()
 
-    let renderObject: MeshRenderObject | undefined
+    let renderObject: ComplexRenderObject | undefined
     let currentProps: P
-    let mesh: Mesh
+    let geometry: Geometry
     let currentStructure: Structure
     let locationIt: LocationIterator
     let conformationHash: number
@@ -56,10 +66,10 @@ export function ComplexMeshVisual<P extends ComplexMeshProps>(builder: ComplexMe
         currentStructure = structure
 
         conformationHash = Structure.conformationHash(currentStructure)
-        mesh = await createMesh(ctx, currentStructure, currentProps, mesh)
+        geometry = await createGeometry(ctx, currentStructure, currentProps, geometry)
 
         locationIt = createLocationIterator(structure)
-        renderObject = await createComplexMeshRenderObject(ctx, structure, mesh, locationIt, currentProps)
+        renderObject = await createRenderObject(ctx, structure, geometry, locationIt, currentProps)
     }
 
     async function update(ctx: RuntimeContext, props: Partial<P>) {
@@ -77,25 +87,30 @@ export function ComplexMeshVisual<P extends ComplexMeshProps>(builder: ComplexMe
             updateState.createGeometry = true
         }
 
-        if (!deepEqual(newProps.sizeTheme, currentProps.sizeTheme)) updateState.createGeometry = true
-        if (!deepEqual(newProps.colorTheme, currentProps.colorTheme)) updateState.updateColor = true
-        // if (!deepEqual(newProps.unitKinds, currentProps.unitKinds)) updateState.createMesh = true // TODO
+        if (colorChanged(currentProps, newProps)) updateState.updateColor = true
+        if (!deepEqual(newProps.unitKinds, currentProps.unitKinds)) updateState.createGeometry = true
 
         //
 
         if (updateState.createGeometry) {
-            mesh = await createMesh(ctx, currentStructure, newProps, mesh)
-            ValueCell.update(renderObject.values.drawCount, mesh.triangleCount * 3)
+            geometry = await createGeometry(ctx, currentStructure, newProps, geometry)
+            ValueCell.update(renderObject.values.drawCount, Geometry.getDrawCount(geometry))
             updateState.updateColor = true
         }
 
+        if (updateState.updateSize) {
+            // not all geometries have size data, so check here
+            if ('uSize' in renderObject.values) {
+                await createSizes(ctx, locationIt, newProps, renderObject.values)
+            }
+        }
+
         if (updateState.updateColor) {
             await createColors(ctx, locationIt, newProps, renderObject.values)
         }
 
-        // TODO why do I need to cast here?
-        Mesh.updateValues(renderObject.values, newProps as ComplexMeshProps)
-        updateRenderableState(renderObject.state, newProps as ComplexMeshProps)
+        updateValues(renderObject.values, newProps)
+        updateRenderableState(renderObject.state, newProps)
 
         currentProps = newProps
         return true
@@ -147,4 +162,28 @@ export function ComplexMeshVisual<P extends ComplexMeshProps>(builder: ComplexMe
             renderObject = undefined
         }
     }
+}
+
+// mesh
+
+export const ComplexMeshParams = {
+    ...StructureMeshParams,
+    unitKinds: MultiSelectParam<UnitKind>('Unit Kind', '', [ 'atomic', 'spheres' ], UnitKindOptions),
+}
+export const DefaultComplexMeshProps = paramDefaultValues(ComplexMeshParams)
+export type ComplexMeshProps = typeof DefaultComplexMeshProps
+
+export interface ComplexMeshVisualBuilder<P extends ComplexMeshProps> extends ComplexVisualBuilder<P, Mesh> { }
+
+export function ComplexMeshVisual<P extends ComplexMeshProps>(builder: ComplexMeshVisualBuilder<P>): ComplexVisual<P> {
+    return ComplexVisual({
+        ...builder,
+        setUpdateState: (state: VisualUpdateState, newProps: P, currentProps: P) => {
+            builder.setUpdateState(state, newProps, currentProps)
+            if (sizeChanged(currentProps, newProps)) state.createGeometry = true
+        },
+        createEmptyGeometry: Mesh.createEmpty,
+        createRenderObject: createComplexMeshRenderObject,
+        updateValues: Mesh.updateValues
+    })
 }

src/mol-geo/representation/structure/representation/ball-and-stick.ts

@@ -16,7 +16,7 @@ import { InterUnitLinkVisual } from '../visual/inter-unit-link-cylinder';
 import { SizeThemeName, SizeThemeOptions } from 'mol-view/theme/size';
 import { getQualityProps } from '../../util';
 import { paramDefaultValues, SelectParam, NumberParam, MultiSelectParam } from 'mol-view/parameter';
-import { UnitKind, UnitKindOptions } from '../units-visual';
+import { UnitKind, UnitKindOptions } from '../visual/util/common';
 
 export const BallAndStickParams = {
     ...ElementSphereParams,

src/mol-geo/representation/structure/units-visual.ts

@@ -6,7 +6,7 @@
 
 import { Unit, Structure } from 'mol-model/structure';
 import { RepresentationProps, Visual } from '../';
-import { VisualUpdateState, StructureMeshParams, StructurePointsParams, StructureLinesParams, StructureDirectVolumeParams, StructureProps, StructureParams } from '.';
+import { VisualUpdateState, StructureMeshParams, StructurePointsParams, StructureLinesParams, StructureDirectVolumeParams, StructureParams } from '.';
 import { RuntimeContext } from 'mol-task';
 import { PickingId } from '../../geometry/picking';
 import { LocationIterator } from '../../util/location-iterator';
@@ -14,27 +14,18 @@ import { Mesh } from '../../geometry/mesh/mesh';
 import { MarkerAction, applyMarkerAction, createMarkers } from '../../geometry/marker-data';
 import { Loci, isEveryLoci, EmptyLoci } from 'mol-model/loci';
 import { MeshRenderObject, PointsRenderObject, LinesRenderObject, DirectVolumeRenderObject } from 'mol-gl/render-object';
-import { createUnitsMeshRenderObject, createUnitsPointsRenderObject, createUnitsTransform, createUnitsLinesRenderObject, createUnitsDirectVolumeRenderObject } from './visual/util/common';
+import { createUnitsMeshRenderObject, createUnitsPointsRenderObject, createUnitsTransform, createUnitsLinesRenderObject, createUnitsDirectVolumeRenderObject, UnitKind, UnitKindOptions, includesUnitKind, colorChanged, sizeChanged } from './visual/util/common';
 import { deepEqual, ValueCell, UUID } from 'mol-util';
 import { Interval } from 'mol-data/int';
 import { Points } from '../../geometry/points/points';
 import { updateRenderableState, Geometry } from '../../geometry/geometry';
-import { createColors, ColorProps } from '../../geometry/color-data';
-import { createSizes, SizeProps } from '../../geometry/size-data';
+import { createColors } from '../../geometry/color-data';
+import { createSizes } from '../../geometry/size-data';
 import { Lines } from '../../geometry/lines/lines';
 import { MultiSelectParam, paramDefaultValues } from 'mol-view/parameter';
 import { DirectVolume } from '../../geometry/direct-volume/direct-volume';
 import { RenderableValues } from 'mol-gl/renderable/schema';
 
-export const UnitKindInfo = {
-    'atomic': {},
-    'spheres': {},
-    'gaussians': {},
-}
-export type UnitKind = keyof typeof UnitKindInfo
-export const UnitKindNames = Object.keys(UnitKindInfo)
-export const UnitKindOptions = UnitKindNames.map(n => [n, n] as [UnitKind, string])
-
 export type StructureGroup = { structure: Structure, group: Unit.SymmetryGroup }
 
 export interface UnitsVisual<P extends RepresentationProps = {}> extends Visual<StructureGroup, P> { }
@@ -46,30 +37,6 @@ function sameGroupConformation(groupA: Unit.SymmetryGroup, groupB: Unit.Symmetry
     )
 }
 
-function includesUnitKind(unitKinds: UnitKind[], unit: Unit) {
-    for (let i = 0, il = unitKinds.length; i < il; ++i) {
-        if (Unit.isAtomic(unit) && unitKinds[i] === 'atomic') return true
-        if (Unit.isSpheres(unit) && unitKinds[i] === 'spheres') return true
-        if (Unit.isGaussians(unit) && unitKinds[i] === 'gaussians') return true
-    }
-    return false
-}
-
-function sizeChanged(oldProps: SizeProps, newProps: SizeProps) {
-    return (
-        oldProps.sizeTheme !== newProps.sizeTheme ||
-        oldProps.sizeValue !== newProps.sizeValue ||
-        oldProps.sizeFactor !== newProps.sizeFactor
-    )
-}
-
-function colorChanged(oldProps: ColorProps, newProps: ColorProps) {
-    return (
-        oldProps.colorTheme !== newProps.colorTheme ||
-        oldProps.colorValue !== newProps.colorValue
-    )
-}
-
 const UnitsParams = {
     ...StructureParams,
     unitKinds: MultiSelectParam<UnitKind>('Unit Kind', '', ['atomic', 'spheres'], UnitKindOptions),
@@ -79,7 +46,7 @@ type UnitsProps = typeof DefaultUnitsProps
 
 type UnitsRenderObject = MeshRenderObject | LinesRenderObject | PointsRenderObject | DirectVolumeRenderObject
 
-interface UnitsVisualBuilder<P extends StructureProps, G extends Geometry> {
+interface UnitsVisualBuilder<P extends UnitsProps, G extends Geometry> {
     defaultProps: P
     createGeometry(ctx: RuntimeContext, unit: Unit, structure: Structure, props: P, geometry?: G): Promise<G>
     createLocationIterator(group: Unit.SymmetryGroup): LocationIterator
@@ -88,7 +55,7 @@ interface UnitsVisualBuilder<P extends StructureProps, G extends Geometry> {
     setUpdateState(state: VisualUpdateState, newProps: P, currentProps: P): void
 }
 
-interface UnitsVisualGeometryBuilder<P extends StructureProps, G extends Geometry> extends UnitsVisualBuilder<P, G> {
+interface UnitsVisualGeometryBuilder<P extends UnitsProps, G extends Geometry> extends UnitsVisualBuilder<P, G> {
     createEmptyGeometry(geometry?: G): G
     createRenderObject(ctx: RuntimeContext, group: Unit.SymmetryGroup, geometry: Geometry, locationIt: LocationIterator, currentProps: P): Promise<UnitsRenderObject>
     updateValues(values: RenderableValues, newProps: P): void

src/mol-geo/representation/structure/visual/carbohydrate-link-cylinder.ts

@@ -74,7 +74,7 @@ export type CarbohydrateLinkProps = typeof DefaultCarbohydrateLinkProps
 export function CarbohydrateLinkVisual(): ComplexVisual<CarbohydrateLinkProps> {
     return ComplexMeshVisual<CarbohydrateLinkProps>({
         defaultProps: DefaultCarbohydrateLinkProps,
-        createMesh: createCarbohydrateLinkCylinderMesh,
+        createGeometry: createCarbohydrateLinkCylinderMesh,
         createLocationIterator: CarbohydrateLinkIterator,
         getLoci: getLinkLoci,
         mark: markLink,

src/mol-geo/representation/structure/visual/carbohydrate-symbol-mesh.ts

@@ -157,7 +157,7 @@ export type CarbohydrateSymbolProps = typeof DefaultCarbohydrateSymbolProps
 export function CarbohydrateSymbolVisual(): ComplexVisual<CarbohydrateSymbolProps> {
     return ComplexMeshVisual<CarbohydrateSymbolProps>({
         defaultProps: DefaultCarbohydrateSymbolProps,
-        createMesh: createCarbohydrateSymbolMesh,
+        createGeometry: createCarbohydrateSymbolMesh,
         createLocationIterator: CarbohydrateElementIterator,
         getLoci: getCarbohydrateLoci,
         mark: markCarbohydrate,

src/mol-geo/representation/structure/visual/cross-link-restraint-cylinder.ts

@@ -62,7 +62,7 @@ export type CrossLinkRestraintProps = typeof DefaultCrossLinkRestraintProps
 export function CrossLinkRestraintVisual(): ComplexVisual<CrossLinkRestraintProps> {
     return ComplexMeshVisual<CrossLinkRestraintProps>({
         defaultProps: DefaultCrossLinkRestraintProps,
-        createMesh: createCrossLinkRestraintCylinderMesh,
+        createGeometry: createCrossLinkRestraintCylinderMesh,
         createLocationIterator: CrossLinkRestraintIterator,
         getLoci: getLinkLoci,
         mark: markLink,

src/mol-geo/representation/structure/visual/inter-unit-link-cylinder.ts

@@ -62,7 +62,7 @@ export type InterUnitLinkProps = typeof DefaultInterUnitLinkProps
 export function InterUnitLinkVisual(): ComplexVisual<InterUnitLinkProps> {
     return ComplexMeshVisual<InterUnitLinkProps>({
         defaultProps: DefaultInterUnitLinkProps,
-        createMesh: createInterUnitLinkCylinderMesh,
+        createGeometry: createInterUnitLinkCylinderMesh,
         createLocationIterator: LinkIterator.fromStructure,
         getLoci: getLinkLoci,
         mark: markLink,

src/mol-geo/representation/structure/visual/util/common.ts

@@ -16,6 +16,8 @@ import { createRenderableState } from '../../../../geometry/geometry';
 import { Mat4 } from 'mol-math/linear-algebra';
 import { Lines } from '../../../../geometry/lines/lines';
 import { DirectVolume } from '../../../../geometry/direct-volume/direct-volume';
+import { SizeProps } from 'mol-geo/geometry/size-data';
+import { ColorProps } from 'mol-geo/geometry/color-data';
 
 export function createUnitsTransform({ units }: Unit.SymmetryGroup, transformData?: TransformData) {
     const unitCount = units.length
@@ -27,6 +29,39 @@ export function createUnitsTransform({ units }: Unit.SymmetryGroup, transformDat
     return createTransform(array, unitCount, transformData)
 }
 
+export const UnitKindInfo = {
+    'atomic': {},
+    'spheres': {},
+    'gaussians': {},
+}
+export type UnitKind = keyof typeof UnitKindInfo
+export const UnitKindNames = Object.keys(UnitKindInfo)
+export const UnitKindOptions = UnitKindNames.map(n => [n, n] as [UnitKind, string])
+
+export function includesUnitKind(unitKinds: UnitKind[], unit: Unit) {
+    for (let i = 0, il = unitKinds.length; i < il; ++i) {
+        if (Unit.isAtomic(unit) && unitKinds[i] === 'atomic') return true
+        if (Unit.isSpheres(unit) && unitKinds[i] === 'spheres') return true
+        if (Unit.isGaussians(unit) && unitKinds[i] === 'gaussians') return true
+    }
+    return false
+}
+
+export function sizeChanged(oldProps: SizeProps, newProps: SizeProps) {
+    return (
+        oldProps.sizeTheme !== newProps.sizeTheme ||
+        oldProps.sizeValue !== newProps.sizeValue ||
+        oldProps.sizeFactor !== newProps.sizeFactor
+    )
+}
+
+export function colorChanged(oldProps: ColorProps, newProps: ColorProps) {
+    return (
+        oldProps.colorTheme !== newProps.colorTheme ||
+        oldProps.colorValue !== newProps.colorValue
+    )
+}
+
 // mesh
 
 type StructureMeshProps = Mesh.Props & StructureProps