From 5cc28c9471b9dd5a11602303bbab5d2faa9ec8dc Mon Sep 17 00:00:00 2001 From: David Sehnal Date: Sun, 8 Dec 2024 18:21:13 +0100 Subject: [PATCH] Add ModelWithCoordinates transform (#1378) * Add ModelWithCoordinates transform * PR feedback --- CHANGELOG.md | 1 + .../plugin/transforms/custom-conformation.md | 45 +++++++++++++++++++ docs/mkdocs.yml | 1 + src/mol-model/structure/model/model.ts | 32 +++++++++---- src/mol-plugin-state/transforms/model.ts | 38 +++++++++++++++- 5 files changed, 107 insertions(+), 10 deletions(-) create mode 100644 docs/docs/plugin/transforms/custom-conformation.md diff --git a/CHANGELOG.md b/CHANGELOG.md index 77ba18da2..94e70c006 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,7 @@ Note that since we don't clearly distinguish between a public and private interf ## [Unreleased] +- Add `ModelWithCoordinates` decorator transform. - Fix outlines on transparent background using illumination mode (#1364) - Fix transparent depth texture artifacts using illumination mode diff --git a/docs/docs/plugin/transforms/custom-conformation.md b/docs/docs/plugin/transforms/custom-conformation.md new file mode 100644 index 000000000..72dd73d75 --- /dev/null +++ b/docs/docs/plugin/transforms/custom-conformation.md @@ -0,0 +1,45 @@ +# Assign custom conformation to a Model + +This document shows how to update model conformation dynamically using the `ModelWithCoordinates` transforms. If this does not work well with your particular use case, it is suggested to write a custom version of `ModelWithCoordinates` with similar usage as outlined in this document. + +```ts +async function animateFirstXCoordinateExample(plugin: PluginContext, url: string, format: BuiltInTrajectoryFormat) { + // Load data + const _data = await plugin.builders.data.download({ url }); + const trajectory = await plugin.builders.structure.parseTrajectory(_data, format); + const hierarchy = await this.plugin.builders.structure.hierarchy.applyPreset(trajectory, 'default'); + if (!hierarchy) return; + + // Insert ModelWithCoordinates cell to be updated in the loop bellow + const coordinatesNode = await plugin.build().to(hierarchy!.model).insert(ModelWithCoordinates).commit(); + + const x0 = hierarchy!.model.data!.atomicConformation.x[0]; + let xOffset = 0; + async function animateFirstXCoord() { + // Normally, the whole conformation would come from an API/library call, but here we fake it: + const { x, y, z } = hierarchy!.model.data!.atomicConformation; + const nextX = [...(x as number[])]; + nextX[0] = x0 + xOffset; + xOffset += 0.05; + if (xOffset > 1) xOffset = 0; + + // Construct new coodinate frame from the data and commit the update. + // Rest of the state tree will reconcile automatically. + await plugin.build().to(coordinatesNode).update({ + atomicCoordinateFrame: { + elementCount: x.length, + time: { value: 0, unit: 'step' }, + xyzOrdering: { isIdentity: true }, + x: nextX, + y, + z, + } + }).commit(); + + requestAnimationFrame(animateFirstXCoord); + } + animateFirstXCoord(); +} + +// animateFirstXCoordinateExample('https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/CID/2244/record/SDF/?record_type=3d', 'sdf'); +``` \ No newline at end of file diff --git a/docs/mkdocs.yml b/docs/mkdocs.yml index 943273f98..00cf25208 100644 --- a/docs/mkdocs.yml +++ b/docs/mkdocs.yml @@ -40,6 +40,7 @@ nav: - CIF Schemas: 'plugin/cif-schemas.md' - State Transforms: - Custom Trajectory: 'plugin/transforms/custom-trajectory.md' + - Custom Conformation: 'plugin/transforms/custom-conformation.md' - Data Access Tools: - 'data-access-tools/model-server.md' - Volume Server: diff --git a/src/mol-model/structure/model/model.ts b/src/mol-model/structure/model/model.ts index c312c3788..5cf19527d 100644 --- a/src/mol-model/structure/model/model.ts +++ b/src/mol-model/structure/model/model.ts @@ -15,7 +15,7 @@ import { SaccharideComponentMap } from '../structure/carbohydrates/constants'; import { ModelFormat } from '../../../mol-model-formats/format'; import { calcModelCenter, getAsymIdCount } from './util'; import { Vec3 } from '../../../mol-math/linear-algebra'; -import { Coordinates } from '../coordinates'; +import { Coordinates, Frame } from '../coordinates'; import { Topology } from '../topology'; import { Task } from '../../../mol-task'; import { IndexPairBonds } from '../../../mol-model-formats/structure/property/bonds/index-pair'; @@ -96,9 +96,7 @@ export namespace Model { const trajectory: Model[] = []; const { frames } = coordinates; - const srcIndex = model.atomicHierarchy.atomSourceIndex; - const isIdentity = Column.isIdentity(srcIndex); - const srcIndexArray = isIdentity ? void 0 : srcIndex.toArray({ array: Int32Array }); + const srcIndexArray = getSourceIndexArray(model); const coarseGrained = isCoarseGrained(model); const elementCount = model.atomicHierarchy.atoms._rowCount; @@ -112,11 +110,7 @@ export namespace Model { ...model, id: UUID.create22(), modelNum: i, - atomicConformation: Coordinates.getAtomicConformation(f, { - atomId: model.atomicConformation.atomId, - occupancy: model.atomicConformation.occupancy, - B_iso_or_equiv: model.atomicConformation.B_iso_or_equiv - }, srcIndexArray), + atomicConformation: getAtomicConformationFromFrame(model, f), // TODO: add support for supplying sphere and gaussian coordinates in addition to atomic coordinates? // coarseConformation: coarse.conformation, customProperties: new CustomProperties(), @@ -137,6 +131,18 @@ export namespace Model { return { trajectory, srcIndexArray }; } + function getSourceIndexArray(model: Model): ArrayLike | undefined { + const srcIndex = model.atomicHierarchy.atomSourceIndex; + let srcIndexArray: ArrayLike | undefined = undefined; + if ('__srcIndexArray__' in model._staticPropertyData) { + srcIndexArray = model._dynamicPropertyData.__srcIndexArray__; + } else { + srcIndexArray = Column.isIdentity(srcIndex) ? void 0 : srcIndex.toArray({ array: Int32Array }); + model._dynamicPropertyData.__srcIndexArray__ = srcIndexArray; + } + return srcIndexArray; + } + export function trajectoryFromModelAndCoordinates(model: Model, coordinates: Coordinates): Trajectory { return new ArrayTrajectory(_trajectoryFromModelAndCoordinates(model, coordinates).trajectory); } @@ -162,6 +168,14 @@ export namespace Model { }); } + export function getAtomicConformationFromFrame(model: Model, frame: Frame) { + return Coordinates.getAtomicConformation(frame, { + atomId: model.atomicConformation.atomId, + occupancy: model.atomicConformation.occupancy, + B_iso_or_equiv: model.atomicConformation.B_iso_or_equiv + }, getSourceIndexArray(model)); + } + const CenterProp = '__Center__'; export function getCenter(model: Model): Vec3 { if (model._dynamicPropertyData[CenterProp]) return model._dynamicPropertyData[CenterProp]; diff --git a/src/mol-plugin-state/transforms/model.ts b/src/mol-plugin-state/transforms/model.ts index 7729158e0..0e0758254 100644 --- a/src/mol-plugin-state/transforms/model.ts +++ b/src/mol-plugin-state/transforms/model.ts @@ -17,7 +17,7 @@ import { trajectoryFromGRO } from '../../mol-model-formats/structure/gro'; import { trajectoryFromCCD, trajectoryFromMmCIF } from '../../mol-model-formats/structure/mmcif'; import { trajectoryFromPDB } from '../../mol-model-formats/structure/pdb'; import { topologyFromPsf } from '../../mol-model-formats/structure/psf'; -import { Coordinates, Model, Queries, QueryContext, Structure, StructureElement, StructureQuery, StructureSelection as Sel, Topology, ArrayTrajectory, Trajectory } from '../../mol-model/structure'; +import { Coordinates, Model, Queries, QueryContext, Structure, StructureElement, StructureQuery, StructureSelection as Sel, Topology, ArrayTrajectory, Trajectory, Frame } from '../../mol-model/structure'; import { PluginContext } from '../../mol-plugin/context'; import { MolScriptBuilder } from '../../mol-script/language/builder'; import { Expression } from '../../mol-script/language/expression'; @@ -78,6 +78,7 @@ export { TrajectoryFromMOL2 }; export { TrajectoryFromCube }; export { TrajectoryFromCifCore }; export { ModelFromTrajectory }; +export { ModelWithCoordinates }; export { StructureFromTrajectory }; export { StructureFromModel }; export { TransformStructureConformation }; @@ -705,6 +706,41 @@ const TransformStructureConformation = PluginStateTransform.BuiltIn({ // } }); +type ModelWithCoordinates = typeof ModelWithCoordinates +const ModelWithCoordinates = PluginStateTransform.BuiltIn({ + name: 'model-with-coordinates', + display: { name: 'Model With Coordinates', description: 'Updates the current model with provided coordinate frame' }, + from: SO.Molecule.Model, + to: SO.Molecule.Model, + params: { + frameIndex: PD.Optional(PD.Numeric(0, undefined, { isHidden: true })), + frameCount: PD.Optional(PD.Numeric(1, undefined, { isHidden: true })), + atomicCoordinateFrame: PD.Optional(PD.Value(undefined, { isHidden: true })), + }, + isDecorator: true, +})({ + apply({ a, params }) { + if (!params.atomicCoordinateFrame) { + return a; + } + const model: Model = { ...a.data, atomicConformation: Model.getAtomicConformationFromFrame(a.data, params.atomicCoordinateFrame) }; + Model.TrajectoryInfo.set(model, { index: params.frameIndex ?? 0, size: params.frameCount ?? 1 }); + return new SO.Molecule.Model(model, { label: a.label, description: a.description }); + }, + update: ({ a, b, oldParams, newParams }) => { + if (oldParams.atomicCoordinateFrame === newParams.atomicCoordinateFrame) { + return StateTransformer.UpdateResult.Unchanged; + } + if (!newParams.atomicCoordinateFrame) { + b.data = a.data; + } else { + b.data = { ...b.data, atomicConformation: Model.getAtomicConformationFromFrame(b.data, newParams.atomicCoordinateFrame) }; + } + Model.TrajectoryInfo.set(b.data, { index: newParams.frameIndex ?? 0, size: newParams.frameCount ?? 1 }); + return StateTransformer.UpdateResult.Updated; + }, +}); + type StructureSelectionFromExpression = typeof StructureSelectionFromExpression const StructureSelectionFromExpression = PluginStateTransform.BuiltIn({ name: 'structure-selection-from-expression',