model export extension

src/apps/viewer/index.ts

@@ -14,6 +14,7 @@ import { MAQualityAssessment } from '../../extensions/model-archive/quality-asse
 import { QualityAssessmentPLDDTPreset, QualityAssessmentQmeanPreset } from '../../extensions/model-archive/quality-assessment/behavior';
 import { QualityAssessment } from '../../extensions/model-archive/quality-assessment/prop';
 import { Mp4Export } from '../../extensions/mp4-export';
+import { ModelExport } from '../../extensions/model-export';
 import { PDBeStructureQualityReport } from '../../extensions/pdbe';
 import { RCSBAssemblySymmetry, RCSBValidationReport } from '../../extensions/rcsb';
 import { DownloadStructure, PdbDownloadProvider } from '../../mol-plugin-state/actions/structure';
@@ -62,6 +63,7 @@ const Extensions = {
     'rcsb-validation-report': PluginSpec.Behavior(RCSBValidationReport),
     'anvil-membrane-orientation': PluginSpec.Behavior(ANVILMembraneOrientation),
     'g3d': PluginSpec.Behavior(G3DFormat),
+    'model-export': PluginSpec.Behavior(ModelExport),
     'mp4-export': PluginSpec.Behavior(Mp4Export),
     'geo-export': PluginSpec.Behavior(GeometryExport),
     'ma-quality-assessment': PluginSpec.Behavior(MAQualityAssessment),

src/extensions/model-export/export.ts

@@ -0,0 +1,63 @@
+/**
+ * Copyright (c) 2021 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author David Sehnal <david.sehnal@gmail.com>
+ */
+
+import { utf8ByteCount, utf8Write } from '../../mol-io/common/utf8';
+import { to_mmCIF } from '../../mol-model/structure';
+import { PluginContext } from '../../mol-plugin/context';
+import { Task } from '../../mol-task';
+import { getFormattedTime } from '../../mol-util/date';
+import { download } from '../../mol-util/download';
+import { zip } from '../../mol-util/zip/zip';
+
+export async function exportHierarchy(plugin: PluginContext, options?: { format?: 'cif' | 'bcif' }) {
+    try {
+        await _exportHierarchy(plugin, options);
+    } catch (e) {
+        plugin.log.error(`Export failed: ${e}`);
+    }
+}
+
+async function _exportHierarchy(plugin: PluginContext, options?: { format?: 'cif' | 'bcif' }) {
+    const format = options?.format ?? 'cif';
+    const { structures } = plugin.managers.structure.hierarchy.current;
+
+    const files: [name: string, data: string | Uint8Array][] = [];
+    const entryMap = new Map<string, number>();
+
+    for (const _s of structures) {
+        const s = _s.cell.obj?.data;
+        if (!s) continue;
+        if (s.models.length > 1) {
+            plugin.log.warn(`[Export] Skipping ${_s.cell.obj?.label}: Multimodel exports not supported.`);
+        }
+
+        const name = entryMap.has(s.model.entryId)
+            ? `${s.model.entryId}_${entryMap.get(s.model.entryId)! + 1}.${format}`
+            : `${s.model.entryId}.${format}`;
+        entryMap.set(s.model.entryId, (entryMap.get(s.model.entryId) ?? 0) + 1);
+        files.push([name, to_mmCIF(s.model.entryId, s, format === 'bcif', { copyAllCategories: true })]);
+    }
+
+    if (files.length === 1) {
+        download(new Blob([files[0][1]]), files[0][0]);
+    } else if (files.length > 1) {
+        const zipData: { [key: string]: Uint8Array } = {};
+        for (const [fn, data] of files) {
+            if (data instanceof Uint8Array) {
+                zipData[fn] = data;
+            } else {
+                const bytes = new Uint8Array(utf8ByteCount(data));
+                utf8Write(bytes, 0, data);
+                zipData[fn] = bytes;
+            }
+        }
+        const task = Task.create('Export Models', async ctx => {
+            return zip(ctx, zipData);
+        });
+        const buffer = await plugin.runTask(task);
+        download(new Blob([new Uint8Array(buffer, 0, buffer.byteLength)]), `structures_${getFormattedTime()}.zip`);
+    }
+}

src/extensions/model-export/index.ts

@@ -0,0 +1,30 @@
+/**
+ * Copyright (c) 2021 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author David Sehnal <david.sehnal@gmail.com>
+ */
+
+import { PluginBehavior } from '../../mol-plugin/behavior/behavior';
+import { ModelExportUI } from './ui';
+
+export const ModelExport = PluginBehavior.create<{}>({
+    name: 'extension-model-export',
+    category: 'misc',
+    display: {
+        name: 'Model Export'
+    },
+    ctor: class extends PluginBehavior.Handler<{}> {
+        register(): void {
+            this.ctx.customStructureControls.set('model-export', ModelExportUI as any);
+        }
+
+        update() {
+            return false;
+        }
+
+        unregister() {
+            this.ctx.customStructureControls.delete('model-export');
+        }
+    },
+    params: () => ({})
+});

src/extensions/model-export/ui.tsx

@@ -0,0 +1,70 @@
+/**
+ * Copyright (c) 2021 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author David Sehnal <david.sehnal@gmail.com>
+ */
+
+import { useState } from 'react';
+import { CollapsableControls, CollapsableState } from '../../mol-plugin-ui/base';
+import { Button } from '../../mol-plugin-ui/controls/common';
+import { GetAppSvg } from '../../mol-plugin-ui/controls/icons';
+import { ParameterControls } from '../../mol-plugin-ui/controls/parameters';
+import { useBehavior } from '../../mol-plugin-ui/hooks/use-behavior';
+import { PluginContext } from '../../mol-plugin/context';
+import { ParamDefinition as PD } from '../../mol-util/param-definition';
+import { exportHierarchy } from './export';
+
+export class ModelExportUI extends CollapsableControls<{}, {}> {
+    protected defaultState(): CollapsableState {
+        return {
+            header: 'Export Models',
+            isCollapsed: true,
+            brand: { accent: 'cyan', svg: GetAppSvg }
+        };
+    }
+    protected renderControls(): JSX.Element | null {
+        return <ExportControls plugin={this.plugin} />;
+    }
+}
+
+const Params = {
+    format: PD.Select<'cif' | 'bcif'>('cif', [['cif', 'mmCIF'], ['bcif', 'Binary mmCIF']])
+};
+// type ParamValue = PD.Values<typeof Params>;
+const DefaultParams = PD.getDefaultValues(Params);
+
+function ExportControls({ plugin }: { plugin: PluginContext }) {
+    const [params, setParams] = useState(DefaultParams);
+    const [exporting, setExporting] = useState(false);
+    useBehavior(plugin.managers.structure.hierarchy.behaviors.selection); // triggers UI update
+    const isBusy = useBehavior(plugin.behaviors.state.isBusy);
+    const hierarchy = plugin.managers.structure.hierarchy.current;
+
+    let label: string = 'Nothing to Export';
+    if (hierarchy.structures.length === 1) {
+        label = 'Export';
+    } if (hierarchy.structures.length > 1) {
+        label = 'Export (as ZIP)';
+    }
+
+    const onExport = async () => {
+        setExporting(true);
+        try {
+            await exportHierarchy(plugin, { format: params.format });
+        } finally {
+            setExporting(false);
+        }
+    };
+
+    return <>
+        <ParameterControls params={Params} values={params} onChangeValues={setParams} isDisabled={isBusy || exporting} />
+        <Button
+            onClick={onExport}
+            style={{ marginTop: 1 }}
+            disabled={isBusy || hierarchy.structures.length === 0 || exporting}
+            commit={hierarchy.structures.length ? 'on' : 'off'}
+        >
+            {label}
+        </Button>
+    </>;
+}

src/mol-model/structure/export/mmcif.ts

@@ -250,10 +250,10 @@ function encode_mmCIF_categories_copyAll(encoder: CifWriter.Encoder, ctx: CifExp
 }
 
 
-function to_mmCIF(name: string, structure: Structure, asBinary = false) {
+function to_mmCIF(name: string, structure: Structure, asBinary = false, params?: encode_mmCIF_categories_Params) {
     const enc = CifWriter.createEncoder({ binary: asBinary });
     enc.startDataBlock(name);
-    encode_mmCIF_categories(enc, structure);
+    encode_mmCIF_categories(enc, structure, params);
     return enc.getData();
 }