fix volume controls disapearing after assembly change

src/structure-viewer/helpers.ts

@@ -4,6 +4,9 @@
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
  */
 
+import { StructureControlsHelper } from './ui/structure';
+import { StructureViewer } from '.';
+
 export type SupportedFormats = 'cif' | 'bcif' | 'pdb'
 export interface LoadParams {
     /** URL pointing to a structure file  */
@@ -27,4 +30,9 @@ export enum StateElements {
     Assembly = 'assembly',
 
     VolumeStreaming = 'volume-streaming',
+}
+
+export interface StructureViewerState {
+    structureControlsHelper: StructureControlsHelper
+    experimentalData: StructureViewer['experimentalData']
 }

src/structure-viewer/index.ts

@@ -14,7 +14,7 @@ import { StateTransforms } from 'molstar/lib/mol-plugin/state/transforms';
 import { PluginStateObject as PSO } from 'molstar/lib/mol-plugin/state/objects';
 import { AnimateModelIndex } from 'molstar/lib/mol-plugin/state/animation/built-in';
 import { StateBuilder, StateSelection } from 'molstar/lib/mol-state';
-import { LoadParams, SupportedFormats, StateElements } from './helpers';
+import { LoadParams, SupportedFormats, StateElements, StructureViewerState } from './helpers';
 import { ControlsWrapper, ViewportWrapper } from './ui/controls';
 import { Scheduler } from 'molstar/lib/mol-task';
 import { InitVolumeStreaming, CreateVolumeStreamingInfo } from 'molstar/lib/mol-plugin/behavior/dynamic/volume-streaming/transformers';
@@ -43,7 +43,6 @@ export type StructureViewerProps = typeof DefaultStructureViewerProps
 
 export class StructureViewer {
     private readonly plugin: PluginContext;
-    private readonly structureControlsHelper: StructureControlsHelper;
     private readonly props: Readonly<StructureViewerProps>
 
     constructor(target: string | HTMLElement, props: Partial<StructureViewerProps> = {}) {
@@ -81,11 +80,15 @@ export class StructureViewer {
             }
         });
 
+        (this.plugin.customState as StructureViewerState) = {
+            structureControlsHelper: new StructureControlsHelper(this.plugin),
+            experimentalData: this.experimentalData
+        };
+
         this.props = { ...DefaultStructureViewerProps, ...props }
 
         const renderer = this.plugin.canvas3d.props.renderer;
         PluginCommands.Canvas3D.SetSettings.dispatch(this.plugin, { settings: { renderer: { ...renderer, backgroundColor: ColorNames.white } } });
-        this.structureControlsHelper = new StructureControlsHelper(this.plugin)
     }
 
     get state() {
@@ -120,11 +123,9 @@ export class StructureViewer {
         const modelTree = this.model(this.download(state.build().toRoot(), url, isBinary), format);
         await this.applyState(modelTree);
 
-        await this.structureControlsHelper.setAssembly(assemblyId)
+        await (this.plugin.customState as StructureViewerState).structureControlsHelper.setAssembly(assemblyId)
 
         Scheduler.setImmediate(() => PluginCommands.Camera.Reset.dispatch(this.plugin, { }));
-
-        this.experimentalData.init()
     }
 
     async loadPdbId(pdbId: string, assemblyId = 'deposited') {

src/structure-viewer/ui/structure.tsx

@@ -6,7 +6,7 @@
 
 import * as React from 'react';
 import { CollapsableControls, CollapsableState } from 'molstar/lib/mol-plugin/ui/base';
-import { StateElements } from '../helpers';
+import { StateElements, StructureViewerState } from '../helpers';
 import { ParameterControls } from 'molstar/lib/mol-plugin/ui/controls/parameters';
 import { ParamDefinition as PD } from 'molstar/lib/mol-util/param-definition';
 import { PluginCommands } from 'molstar/lib/mol-plugin/command';
@@ -26,6 +26,10 @@ export class StructureControlsHelper {
         return PluginCommands.State.Update.dispatch(this.plugin, { state: this.plugin.state.dataState, tree });
     }
 
+    get experimentalData () {
+        return (this.plugin.customState as StructureViewerState).experimentalData
+    }
+
     async preset() {
         await this.plugin.helpers.structureRepresentation.preset()
     }
@@ -56,6 +60,7 @@ export class StructureControlsHelper {
         }
         await this.applyState(tree)
         await this.preset()
+        await this.experimentalData.init()
     }
 
     async setModel(modelIndex: number) {
@@ -98,11 +103,12 @@ export class StructureControlsHelper {
 }
 
 export class StructureControls<P, S extends StructureControlsState> extends CollapsableControls<P, S> {
-    structureControlsHelper: StructureControlsHelper
-
     constructor(props: P, context?: any) {
         super(props, context);
-        this.structureControlsHelper = new StructureControlsHelper(this.plugin)
+    }
+
+    get structureControlsHelper () {
+        return (this.plugin.customState as StructureViewerState).structureControlsHelper
     }
 
     async setColorTheme(theme: { [k: string]: string }) {