ui integration

src/extensions/membrane-orientation/behavior.ts

@@ -8,11 +8,13 @@
 import { ParamDefinition as PD } from '../../mol-util/param-definition';
 import { StructureRepresentationPresetProvider, PresetStructureRepresentations } from '../../mol-plugin-state/builder/structure/representation-preset';
 import { MembraneOrientationProvider, MembraneOrientation } from './membrane-orientation';
-import { StateObjectRef } from '../../mol-state';
+import { StateObjectRef, StateAction, StateTransformer, StateTransform } from '../../mol-state';
 import { Task } from '../../mol-task';
 import { PluginBehavior } from '../../mol-plugin/behavior';
-import { MembraneOrientationRepresentationProvider } from './representation';
+import { MembraneOrientationRepresentationProvider, MembraneOrientationParams, MembraneOrientationRepresentation } from './representation';
 import { HydrophobicityColorThemeProvider } from '../../mol-theme/color/hydrophobicity';
+import { PluginStateObject, PluginStateTransform } from '../../mol-plugin-state/objects';
+import { PluginContext } from '../../mol-plugin/context';
 
 export const MembraneOrientationData = PluginBehavior.create<{ autoAttach: boolean }>({
     name: 'membrane-orientation-prop',
@@ -25,10 +27,10 @@ export const MembraneOrientationData = PluginBehavior.create<{ autoAttach: boole
         private provider = MembraneOrientationProvider
 
         register(): void {
+            console.log('beh @ register');
+            this.ctx.state.data.actions.add(InitMembraneOrientation3D);
             this.ctx.customStructureProperties.register(this.provider, this.params.autoAttach);
 
-            this.ctx.representation.structure.themes.colorThemeRegistry.add(HydrophobicityColorThemeProvider);
-
             this.ctx.representation.structure.registry.add(MembraneOrientationRepresentationProvider);
             this.ctx.query.structure.registry.add(MembraneOrientation.isTransmembrane);
 
@@ -43,10 +45,9 @@ export const MembraneOrientationData = PluginBehavior.create<{ autoAttach: boole
         }
 
         unregister() {
+            this.ctx.state.data.actions.remove(InitMembraneOrientation3D);
             this.ctx.customStructureProperties.unregister(this.provider.descriptor.name);
 
-            this.ctx.representation.structure.themes.colorThemeRegistry.remove(HydrophobicityColorThemeProvider);
-
             this.ctx.representation.structure.registry.remove(MembraneOrientationRepresentationProvider);
             this.ctx.query.structure.registry.remove(MembraneOrientation.isTransmembrane);
 
@@ -58,13 +59,75 @@ export const MembraneOrientationData = PluginBehavior.create<{ autoAttach: boole
     })
 });
 
+export const InitMembraneOrientation3D = StateAction.build({
+    display: {
+        name: 'Membrane Orientation',
+        description: 'Initialize Membrane Orientation planes and rims. Data calculated with ANVIL algorithm.'
+    },
+    from: PluginStateObject.Molecule.Structure,
+    isApplicable: (a) => MembraneOrientationProvider.isApplicable(a.data)
+})(({ a, ref, state }, plugin: PluginContext) => Task.create('Init Membrane Orientation', async ctx => {
+    try {
+        const propCtx = { runtime: ctx, assetManager: plugin.managers.asset };
+        await MembraneOrientationProvider.attach(propCtx, a.data);
+    } catch(e) {
+        plugin.log.error(`Membrane Orientation: ${e}`);
+        return;
+    }
+    const tree = state.build().to(ref)
+        .applyOrUpdateTagged('membrane-orientation-3d', MembraneOrientation3D);
+    await state.updateTree(tree).runInContext(ctx);
+}));
+
+export { MembraneOrientation3D };
+
+type MembraneOrientation3D = typeof MembraneOrientation3D
+const MembraneOrientation3D = PluginStateTransform.BuiltIn({
+    name: 'membrane-orientation-3d',
+    display: {
+        name: 'Membrane Orientation',
+        description: 'Membrane Orientation planes and rims. Data calculated with ANVIL algorithm.'
+    },
+    from: PluginStateObject.Molecule.Structure,
+    to: PluginStateObject.Shape.Representation3D,
+    params: (a) => {
+        return {
+            ...MembraneOrientationParams,
+        };
+    }
+})({
+    canAutoUpdate({ oldParams, newParams }) {
+        return true;
+    },
+    apply({ a, params }, plugin: PluginContext) {
+        return Task.create('Membrane Orientation', async ctx => {
+            await MembraneOrientationProvider.attach({ runtime: ctx, assetManager: plugin.managers.asset }, a.data);
+            const repr = MembraneOrientationRepresentation({ webgl: plugin.canvas3d?.webgl, ...plugin.representation.structure.themes }, () => MembraneOrientationParams);
+            await repr.createOrUpdate(params, a.data).runInContext(ctx);
+            return new PluginStateObject.Shape.Representation3D({ repr, source: a });
+        });
+    },
+    update({ a, b, newParams }, plugin: PluginContext) {
+        return Task.create('Membrane Orientation', async ctx => {
+            await MembraneOrientationProvider.attach({ runtime: ctx, assetManager: plugin.managers.asset }, a.data);
+            const props = { ...b.data.repr.props, ...newParams };
+            await b.data.repr.createOrUpdate(props, a.data).runInContext(ctx);
+            return StateTransformer.UpdateResult.Updated;
+        });
+    },
+    isApplicable(a) {
+        return MembraneOrientationProvider.isApplicable(a.data);
+    }
+});
+
 export const MembraneOrientationPreset = StructureRepresentationPresetProvider({
     id: 'preset-membrane-orientation',
     display: {
         name: 'Membrane Orientation', group: 'Annotation',
-        description: 'Initialize orientation of membrane layers. Data calculated with ANVIL algorithm.' // TODO add ' or obtained via RCSB PDB'
+        description: 'Shows orientation of membrane layers. Data calculated with ANVIL algorithm.' // TODO add ' or obtained via RCSB PDB'
     },
     isApplicable(a) {
+        console.log(`present applicable: ${MembraneOrientationProvider.isApplicable(a.data)}`);
         return MembraneOrientationProvider.isApplicable(a.data);
     },
     params: () => StructureRepresentationPresetProvider.CommonParams,
@@ -73,11 +136,24 @@ export const MembraneOrientationPreset = StructureRepresentationPresetProvider({
         const structure  = structureCell?.obj?.data;
         if (!structureCell || !structure) return {};
 
+        if (!MembraneOrientationProvider.get(structure).value) {
             await plugin.runTask(Task.create('Membrane Orientation', async runtime => {
+                console.log('present: attaching');
                 await MembraneOrientationProvider.attach({ runtime, assetManager: plugin.managers.asset }, structure);
             }));
+        }
 
+        const membraneOrientation = await tryCreateMembraneOrientation(plugin, structureCell);
         const colorTheme = HydrophobicityColorThemeProvider.name as any;
-        return await PresetStructureRepresentations.auto.apply(ref, { ...params, theme: { globalName: colorTheme, focus: { name: colorTheme } } }, plugin);
+        const preset = await PresetStructureRepresentations.auto.apply(ref, { ...params, theme: { globalName: colorTheme, focus: { name: colorTheme } } }, plugin);
+
+        return { components: preset.components, representations: { ...preset.representations, membraneOrientation } };
     }
 });
+
+export function tryCreateMembraneOrientation(plugin: PluginContext, structure: StateObjectRef<PluginStateObject.Molecule.Structure>, params?: StateTransformer.Params<MembraneOrientation3D>, initialState?: Partial<StateTransform.State>) {
+    const state = plugin.state.data;
+    const membraneOrientation = state.build().to(structure)
+        .applyOrUpdateTagged('membrane-orientation-3d', MembraneOrientation3D, params, { state: initialState });
+    return membraneOrientation.commit({ revertOnError: true });
+}

src/extensions/membrane-orientation/representation.ts

@@ -74,11 +74,13 @@ export type MembraneOrientationParams = typeof MembraneOrientationParams
 export type MembraneOrientationProps = PD.Values<MembraneOrientationParams>
 
 export function getMembraneOrientationParams(ctx: ThemeRegistryContext, structure: Structure) {
+    console.log('rep @ getParams');
     return PD.clone(MembraneOrientationParams);
 }
 
 export type MembraneOrientationRepresentation = StructureRepresentation<MembraneOrientationParams>
 export function MembraneOrientationRepresentation(ctx: RepresentationContext, getParams: RepresentationParamsGetter<Structure, MembraneOrientationParams>): MembraneOrientationRepresentation {
+    console.log('rep @ create');
     return Representation.createMulti('Membrane Orientation', ctx, getParams, StructureRepresentationStateBuilder, MembraneOrientationVisuals as unknown as Representation.Def<Structure, MembraneOrientationParams>);
 }
 
@@ -95,6 +97,7 @@ export const MembraneOrientationRepresentationProvider = StructureRepresentation
 });
 
 function getBilayerRims(ctx: RuntimeContext, data: Structure, props: BilayerRimsProps): Shape<Lines> {
+    console.log('rims');
     const { p1, p2, centroid, normal, radius } = MembraneOrientationProvider.get(data).value!;
     const scaledRadius = props.radiusFactor * radius;
     const builder = LinesBuilder.create(128, 64);
@@ -131,6 +134,7 @@ function getCircle(p: Vec3, centroid: Vec3, normal: Vec3, radius: number) {
 }
 
 function getBilayerPlanes(ctx: RuntimeContext, data: Structure, props: BilayerPlanesProps, shape?: Shape<Mesh>): Shape<Mesh> {
+    console.log('planes');
     const { p1, p2, centroid, normal, radius } = MembraneOrientationProvider.get(data).value!;
     const state = MeshBuilder.createState(128, 64, shape && shape.geometry);
     const scaledRadius = props.radiusFactor * radius;