wip, cellpack loader

src/apps/viewer/extensions/cellpack/model.ts

@@ -8,7 +8,7 @@ import { StateAction } from '../../../../mol-state';
 import { PluginContext } from '../../../../mol-plugin/context';
 import { PluginStateObject as PSO } from '../../../../mol-plugin/state/objects';
 import { ParamDefinition as PD } from '../../../../mol-util/param-definition';
-import { Ingredient, CellPacking } from './data';
+import { Ingredient, CellPacking, Cell } from './data';
 import { getFromPdb, getFromCellPackDB } from './util';
 import { Model, Structure, StructureSymmetry, StructureSelection, QueryContext } from '../../../../mol-model/structure';
 import { trajectoryFromMmCIF } from '../../../../mol-model-formats/structure/mmcif';
@@ -28,6 +28,11 @@ import { compile } from '../../../../mol-script/runtime/query/compiler';
 import { UniformColorThemeProvider } from '../../../../mol-theme/color/uniform';
 import { ThemeRegistryContext } from '../../../../mol-theme/theme';
 import { ColorTheme } from '../../../../mol-theme/color';
+import { _parse_mmCif } from '../../../../mol-model-formats/structure/mmcif/parser';
+import { ModelFormat } from '../../../../mol-model-formats/structure/format';
+import { CifCategory, CifField } from '../../../../mol-io/reader/cif';
+import { mmCIF_Schema } from '../../../../mol-io/reader/cif/schema/mmcif';
+import { Column } from '../../../../mol-data/db';
 
 function getCellPackModelUrl(fileName: string, baseUrl: string) {
     return `${baseUrl}/results/${fileName}`
@@ -118,6 +123,109 @@ function getAssembly(transforms: Mat4[], structure: Structure) {
     return builder.getStructure();
 }
 
+function getCifCurve(name: string, transforms: Mat4[], model: Model) {
+    const d = model.sourceData.data.atom_site
+    const n = d._rowCount
+    const rowCount = n * transforms.length
+
+    const { offsets, count } = model.atomicHierarchy.chainAtomSegments
+
+    const x = d.Cartn_x.toArray()
+    const y = d.Cartn_y.toArray()
+    const z = d.Cartn_z.toArray()
+
+    const Cartn_x = new Float32Array(rowCount)
+    const Cartn_y = new Float32Array(rowCount)
+    const Cartn_z = new Float32Array(rowCount)
+    const map = new Uint32Array(rowCount)
+    const seq = new Int32Array(rowCount)
+    let offset = 0
+    for (let c = 0; c < count; ++c) {
+        const cStart = offsets[c]
+        const cEnd = offsets[c + 1]
+        const cLength = cEnd - cStart
+        for (let t = 0, tl = transforms.length; t < tl; ++t) {
+            const m = transforms[t]
+            for (let j = cStart; j < cEnd; ++j) {
+                const i = offset + j - cStart
+                const xj = x[j], yj = y[j], zj = z[j]
+                Cartn_x[i] = m[0] * xj + m[4] * yj + m[8] * zj + m[12]
+                Cartn_y[i] = m[1] * xj + m[5] * yj + m[9] * zj + m[13]
+                Cartn_z[i] = m[2] * xj + m[6] * yj + m[10] * zj + m[14]
+                map[i] = j
+                seq[i] = t + 1
+            }
+            offset += cLength
+        }
+    }
+
+    function multColumn<T>(column: Column<T>) {
+        const array = column.toArray()
+        return Column.ofLambda({
+            value: row => array[map[row]],
+            areValuesEqual: (rowA, rowB) => map[rowA] === map[rowB] || array[map[rowA]] === array[map[rowB]],
+            rowCount, schema: column.schema
+        })
+    }
+
+    const _atom_site: CifCategory.SomeFields<mmCIF_Schema['atom_site']> = {
+        auth_asym_id: CifField.ofColumn(multColumn(d.auth_asym_id)),
+        auth_atom_id: CifField.ofColumn(multColumn(d.auth_atom_id)),
+        auth_comp_id: CifField.ofColumn(multColumn(d.auth_comp_id)),
+        auth_seq_id: CifField.ofNumbers(seq),
+
+        B_iso_or_equiv: CifField.ofColumn(Column.ofConst(0, rowCount, Column.Schema.float)),
+        Cartn_x: CifField.ofNumbers(Cartn_x),
+        Cartn_y: CifField.ofNumbers(Cartn_y),
+        Cartn_z: CifField.ofNumbers(Cartn_z),
+        group_PDB: CifField.ofColumn(Column.ofConst('ATOM', rowCount, Column.Schema.str)),
+        id: CifField.ofColumn(Column.ofLambda({
+            value: row => row,
+            areValuesEqual: (rowA, rowB) => rowA === rowB,
+            rowCount, schema: d.id.schema,
+        })),
+
+        label_alt_id: CifField.ofColumn(multColumn(d.label_alt_id)),
+
+        label_asym_id: CifField.ofColumn(multColumn(d.label_asym_id)),
+        label_atom_id: CifField.ofColumn(multColumn(d.label_atom_id)),
+        label_comp_id: CifField.ofColumn(multColumn(d.label_comp_id)),
+        label_seq_id: CifField.ofNumbers(seq),
+        label_entity_id: CifField.ofColumn(Column.ofConst('1', rowCount, Column.Schema.str)),
+
+        occupancy: CifField.ofColumn(Column.ofConst(1, rowCount, Column.Schema.float)),
+        type_symbol: CifField.ofColumn(multColumn(d.type_symbol)),
+
+        pdbx_PDB_ins_code: CifField.ofColumn(Column.ofConst('', rowCount, Column.Schema.str)),
+        pdbx_PDB_model_num: CifField.ofColumn(Column.ofConst(1, rowCount, Column.Schema.int)),
+    }
+
+    const categories = {
+        entity: CifCategory.ofTable('entity', model.sourceData.data.entity),
+        chem_comp: CifCategory.ofTable('chem_comp', model.sourceData.data.chem_comp),
+        atom_site: CifCategory.ofFields('atom_site', _atom_site)
+    }
+
+    return {
+        header: name,
+        categoryNames: Object.keys(categories),
+        categories
+    };
+}
+
+async function getCurve(name: string, transforms: Mat4[], model: Model) {
+    const cif = getCifCurve(name, transforms, model)
+
+    const curveModelTask = Task.create('Curve Model', async ctx => {
+        const format = ModelFormat.mmCIF(cif)
+        const models = await _parse_mmCif(format, ctx)
+        return models[0]
+    })
+
+    const curveModel = await curveModelTask.run()
+    return getStructure(curveModel)
+}
+
 async function getIngredientStructure(ingredient: Ingredient, baseUrl: string) {
     const { name, source, results, nbCurve } = ingredient
 
@@ -133,10 +241,12 @@ async function getIngredientStructure(ingredient: Ingredient, baseUrl: string) {
     const model = await getModel(source.pdb || name, baseUrl)
     if (!model) return
 
+    if (nbCurve) {
+        return getCurve(name, getCurveTransforms(ingredient), model)
+    } else {
         const structure = await getStructure(model, { assembly: source.biomt ? '1' : undefined })
-    const transforms = nbCurve ? getCurveTransforms(ingredient) : getResultTransforms(results)
-    const assembly = getAssembly(transforms, structure)
-    return assembly
+        return getAssembly(getResultTransforms(results), structure)
+    }
 }
 
 export function createStructureFromCellPack(packing: CellPacking, baseUrl: string) {
@@ -144,14 +254,16 @@ export function createStructureFromCellPack(packing: CellPacking, baseUrl: strin
         const { ingredients, name } = packing
         const structures: Structure[] = []
         for (const iName in ingredients) {
-            if (ctx.shouldUpdate) ctx.update(iName)
+            if (ctx.shouldUpdate) await ctx.update(iName)
             const s = await getIngredientStructure(ingredients[iName], baseUrl)
             if (s) structures.push(s)
         }
 
+        if (ctx.shouldUpdate) await ctx.update(`${name} - units`)
         const builder = Structure.Builder({ label: name })
         let offsetInvariantId = 0
         for (const s of structures) {
+            if (ctx.shouldUpdate) await ctx.update(`${s.label}`)
             let maxInvariantId = 0
             for (const u of s.units) {
                 const invariantId = u.invariantId + offsetInvariantId
@@ -161,6 +273,7 @@ export function createStructureFromCellPack(packing: CellPacking, baseUrl: strin
             offsetInvariantId += maxInvariantId
         }
 
+        if (ctx.shouldUpdate) await ctx.update(`${name} - structure`)
         const s = builder.getStructure()
         return s
     })
@@ -175,6 +288,7 @@ export const LoadCellPackModel = StateAction.build({
             ['HIV-1_0.1.6-8_mixed_radii_pdb.cpr', 'HIV-1_0.1.6-8_mixed_radii_pdb'],
             ['influenza_model1.json', 'influenza_model1'],
             ['Mycoplasma1.5_mixed_pdb_fixed.cpr', 'Mycoplasma1.5_mixed_pdb_fixed'],
+            ['curveTest', 'Curve Test'],
         ]),
         baseUrl: PD.Text(DefaultCellPackBaseUrl),
         preset: PD.Group({
@@ -192,14 +306,50 @@ export const LoadCellPackModel = StateAction.build({
 
     const root = state.build().toRoot();
 
-    const cellPackBuilder = root
+    let cellPackBuilder: any
+
+    if (params.id === 'curveTest') {
+        const url = `${params.baseUrl}/extras/rna_allpoints.json`
+        const data = await ctx.fetch({ url, type: 'string' }).runInContext(taskCtx);
+        const { points } = await (new Response(data)).json() as { points: number[] }
+        const curve0: Vec3[] = []
+        for (let j = 0, jl = Math.min(points.length, 3 * 100); j < jl; j += 3) {
+            curve0.push(Vec3.fromArray(Vec3(), points, j))
+        }
+        const cell: Cell = {
+            recipe: { setupfile: '', paths: [], version: '', name: 'Curve Test' },
+            compartments: {
+                'CurveCompartment': {
+                    interior: {
+                        ingredients: {
+                            'CurveIngredient': {
+                                source: { pdb: 'RNA_U_Base.pdb', transform: { center: false } },
+                                results: [],
+                                name: 'RNA',
+                                nbCurve: 1,
+                                curve0
+                            }
+                        }
+                    }
+                }
+            }
+        }
+
+        cellPackBuilder = root
+            .apply(StateTransforms.Data.ImportJson, { data: cell }, { state: { isGhost: true } })
+            .apply(ParseCellPack)
+    } else {
+        cellPackBuilder = root
             .apply(StateTransforms.Data.Download, { url, isBinary: false, label: params.id }, { state: { isGhost: true } })
             .apply(StateTransforms.Data.ParseJson, undefined, { state: { isGhost: true } })
             .apply(ParseCellPack)
 
+
+    }
+
     const cellPackObject = await state.updateTree(cellPackBuilder).runInContext(taskCtx)
     const { packings } = cellPackObject.data
-    let tree = state.build().to(cellPackBuilder.ref);
+    const tree = state.build().to(cellPackBuilder.ref);
 
     const isHiv = (
         params.id === 'BloodHIV1.0_mixed_fixed_nc1.cpr' ||
@@ -270,7 +420,9 @@ export const LoadCellPackModel = StateAction.build({
             )
     }
 
+    console.time('cellpack')
     await state.updateTree(tree).runInContext(taskCtx);
+    console.timeEnd('cellpack')
 }));
 
 function getReprParams(ctx: PluginContext, params: { representation: 'spacefill' | 'gaussian-surface' | 'point', traceOnly: boolean }) {

src/apps/viewer/extensions/cellpack/util.ts

@@ -14,8 +14,8 @@ async function parseCif(data: string|Uint8Array) {
     return parsed.result;
 }
 
-async function parsePDBfile(data: string) {
-    const comp = parsePDB(data);
+async function parsePDBfile(data: string, id: string) {
+    const comp = parsePDB(data, id);
     const parsed = await comp.run();
     if (parsed.isError) throw parsed;
     return parsed.result;
@@ -26,9 +26,9 @@ async function downloadCif(url: string, isBinary: boolean) {
     return parseCif(isBinary ? new Uint8Array(await data.arrayBuffer()) : await data.text());
 }
 
-async function downloadPDB(url: string) {
+async function downloadPDB(url: string, id: string) {
     const data = await fetch(url);
-    return parsePDBfile(await data.text());
+    return parsePDBfile(await data.text(), id);
 }
 
 export async function getFromPdb(id: string) {
@@ -42,6 +42,7 @@ function getCellPackDataUrl(id: string, baseUrl: string) {
 }
 
 export async function getFromCellPackDB(id: string, baseUrl: string) {
-    const parsed = await downloadPDB(getCellPackDataUrl(id, baseUrl));
+    const name = id.endsWith('.pdb') ? id.substring(0, id.length - 4) : id
+    const parsed = await downloadPDB(getCellPackDataUrl(id, baseUrl), name);
     return parsed;
 }
\ No newline at end of file