/** * Copyright (c) 2019-2022 mol* contributors, licensed under MIT, See LICENSE file for more info. * * @author Alexander Rose * @author Ludovic Autin */ import { StateAction, StateBuilder, StateTransformer, State } 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, CompartmentPrimitives } from './data'; import { getFromPdb, getFromCellPackDB, IngredientFiles, parseCif, parsePDBfile, getStructureMean, getFromOPM } from './util'; import { Model, Structure, StructureSymmetry, StructureSelection, QueryContext, Unit, Trajectory } from '../../mol-model/structure'; import { trajectoryFromMmCIF } from '../../mol-model-formats/structure/mmcif'; import { trajectoryFromPDB } from '../../mol-model-formats/structure/pdb'; import { Mat4, Vec3, Quat } from '../../mol-math/linear-algebra'; import { SymmetryOperator } from '../../mol-math/geometry'; import { Task, RuntimeContext } from '../../mol-task'; import { StateTransforms } from '../../mol-plugin-state/transforms'; import { ParseCellPack, StructureFromCellpack, DefaultCellPackBaseUrl, StructureFromAssemblies, CreateCompartmentSphere } from './state'; import { MolScriptBuilder as MS } from '../../mol-script/language/builder'; import { getMatFromResamplePoints } from './curve'; import { compile } from '../../mol-script/runtime/query/compiler'; import { CellpackPackingPreset, CellpackMembranePreset } from './preset'; import { Asset } from '../../mol-util/assets'; import { Color } from '../../mol-util/color'; import { objectForEach } from '../../mol-util/object'; import { readFromFile } from '../../mol-util/data-source'; import { ColorNames } from '../../mol-util/color/names'; function getCellPackModelUrl(fileName: string, baseUrl: string) { return `${baseUrl}/results/${fileName}`; } class TrajectoryCache { private map = new Map(); set(id: string, trajectory: Trajectory) { this.map.set(id, trajectory); } get(id: string) { return this.map.get(id); } } async function getModel(plugin: PluginContext, id: string, ingredient: Ingredient, baseUrl: string, trajCache: TrajectoryCache, location: string, file?: Asset.File ) { const assetManager = plugin.managers.asset; const modelIndex = (ingredient.source.model) ? parseInt(ingredient.source.model) : 0; let surface = (ingredient.ingtype) ? (ingredient.ingtype === 'transmembrane') : false; if (location === 'surface') surface = true; let trajectory = trajCache.get(id); const assets: Asset.Wrapper[] = []; if (!trajectory) { if (file) { if (file.name.endsWith('.cif')) { const text = await plugin.runTask(assetManager.resolve(file, 'string')); assets.push(text); const cif = (await parseCif(plugin, text.data)).blocks[0]; trajectory = await plugin.runTask(trajectoryFromMmCIF(cif)); } else if (file.name.endsWith('.bcif')) { const binary = await plugin.runTask(assetManager.resolve(file, 'binary')); assets.push(binary); const cif = (await parseCif(plugin, binary.data)).blocks[0]; trajectory = await plugin.runTask(trajectoryFromMmCIF(cif)); } else if (file.name.endsWith('.pdb')) { const text = await plugin.runTask(assetManager.resolve(file, 'string')); assets.push(text); const pdb = await parsePDBfile(plugin, text.data, id); trajectory = await plugin.runTask(trajectoryFromPDB(pdb)); } else { throw new Error(`unsupported file type '${file.name}'`); } } else if (id.match(/^[1-9][a-zA-Z0-9]{3,3}$/i)) { if (surface) { try { const data = await getFromOPM(plugin, id, assetManager); assets.push(data.asset); data.pdb.id! = id.toUpperCase(); trajectory = await plugin.runTask(trajectoryFromPDB(data.pdb)); } catch (e) { // fallback to getFromPdb // console.error(e); const { mmcif, asset } = await getFromPdb(plugin, id, assetManager); assets.push(asset); trajectory = await plugin.runTask(trajectoryFromMmCIF(mmcif)); } } else { const { mmcif, asset } = await getFromPdb(plugin, id, assetManager); assets.push(asset); trajectory = await plugin.runTask(trajectoryFromMmCIF(mmcif)); } } else { const data = await getFromCellPackDB(plugin, id, baseUrl, assetManager); assets.push(data.asset); if ('pdb' in data) { trajectory = await plugin.runTask(trajectoryFromPDB(data.pdb)); } else { trajectory = await plugin.runTask(trajectoryFromMmCIF(data.mmcif)); } } trajCache.set(id, trajectory!); } const model = await plugin.resolveTask(trajectory?.getFrameAtIndex(modelIndex)!); return { model, assets }; } async function getStructure(plugin: PluginContext, model: Model, source: Ingredient, props: { assembly?: string } = {}) { let structure = Structure.ofModel(model); const { assembly } = props; if (assembly) { structure = await plugin.runTask(StructureSymmetry.buildAssembly(structure, assembly)); } let query; if (source.source.selection) { const sel = source.source.selection; // selection can have the model ID as well. remove it const asymIds: string[] = sel.replace(/ /g, '').replace(/:/g, '').split('or').slice(1); query = MS.struct.modifier.union([ MS.struct.generator.atomGroups({ 'chain-test': MS.core.set.has([MS.set(...asymIds), MS.ammp('auth_asym_id')]) }) ]); } else { query = MS.struct.modifier.union([ MS.struct.generator.atomGroups({ 'entity-test': MS.core.rel.eq([MS.ammp('entityType'), 'polymer']) }) ]); } const compiled = compile(query); const result = compiled(new QueryContext(structure)); structure = StructureSelection.unionStructure(result); // change here if possible the label ? // structure.label = source.name; return structure; } function getTransformLegacy(trans: Vec3, rot: Quat) { const q: Quat = Quat.create(-rot[3], rot[0], rot[1], rot[2]); const m: Mat4 = Mat4.fromQuat(Mat4(), q); Mat4.transpose(m, m); Mat4.scale(m, m, Vec3.create(-1.0, 1.0, -1.0)); Mat4.setTranslation(m, trans); return m; } function getTransform(trans: Vec3, rot: Quat) { const q: Quat = Quat.create(-rot[0], rot[1], rot[2], -rot[3]); const m: Mat4 = Mat4.fromQuat(Mat4(), q); const p: Vec3 = Vec3.create(-trans[0], trans[1], trans[2]); Mat4.setTranslation(m, p); return m; } function getResultTransforms(results: Ingredient['results'], legacy: boolean) { if (legacy) return results.map((r: Ingredient['results'][0]) => getTransformLegacy(r[0], r[1])); else return results.map((r: Ingredient['results'][0]) => getTransform(r[0], r[1])); } function getCurveTransforms(ingredient: Ingredient) { const n = ingredient.nbCurve || 0; const instances: Mat4[] = []; let segmentLength = 3.4; if (ingredient.uLength) { segmentLength = ingredient.uLength; } else if (ingredient.radii) { segmentLength = ingredient.radii[0].radii ? ingredient.radii[0].radii[0] * 2.0 : 3.4; } let resampling: boolean = false; for (let i = 0; i < n; ++i) { const cname = `curve${i}`; if (!(cname in ingredient)) { console.warn(`Expected '${cname}' in ingredient`); continue; } const _points = ingredient[cname] as Vec3[]; if (_points.length <= 2) { // TODO handle curve with 2 or less points continue; } // test for resampling const distance: number = Vec3.distance(_points[0], _points[1]); if (distance >= segmentLength + 2.0) { // console.info(distance); resampling = true; } const points = new Float32Array(_points.length * 3); for (let i = 0, il = _points.length; i < il; ++i) Vec3.toArray(_points[i], points, i * 3); const newInstances = getMatFromResamplePoints(points, segmentLength, resampling); instances.push(...newInstances); } return instances; } function getAssembly(name: string, transforms: Mat4[], structure: Structure) { const builder = Structure.Builder({ label: name }); const { units } = structure; for (let i = 0, il = transforms.length; i < il; ++i) { const id = `${i + 1}`; const op = SymmetryOperator.create(id, transforms[i], { assembly: { id, operId: i, operList: [id] } }); for (const unit of units) { builder.addWithOperator(unit, op); } } return builder.getStructure(); } async function getCurve(name: string, transforms: Mat4[], model: Model) { const structure = Structure.ofModel(model); const assembly = getAssembly(name, transforms, structure); return assembly; } async function getIngredientStructure(plugin: PluginContext, ingredient: Ingredient, baseUrl: string, ingredientFiles: IngredientFiles, trajCache: TrajectoryCache, location: 'surface' | 'interior' | 'cytoplasme') { const { name, source, results, nbCurve } = ingredient; if (source.pdb === 'None') return; const file = ingredientFiles[source.pdb]; if (!file) { // TODO can these be added to the library? if (name === 'HIV1_CAhex_0_1_0') return; // 1VU4CtoH_hex.pdb if (name === 'HIV1_CAhexCyclophilA_0_1_0') return; // 1AK4fitTo1VU4hex.pdb if (name === 'iLDL') return; // EMD-5239 if (name === 'peptides') return; // peptide.pdb if (name === 'lypoglycane') return; } // model id in case structure is NMR const { model, assets } = await getModel(plugin, source.pdb || name, ingredient, baseUrl, trajCache, location, file); if (!model) return; let structure: Structure; if (nbCurve) { structure = await getCurve(name, getCurveTransforms(ingredient), model); } else { if ((!results || results.length === 0)) return; let bu: string|undefined = source.bu ? source.bu : undefined; if (bu) { if (bu === 'AU') { bu = undefined; } else { bu = bu.slice(2); } } structure = await getStructure(plugin, model, ingredient, { assembly: bu }); // transform with offset and pcp let legacy: boolean = true; const pcp = ingredient.principalVector ? ingredient.principalVector : ingredient.principalAxis; if (pcp) { legacy = false; const structureMean = getStructureMean(structure); Vec3.negate(structureMean, structureMean); const m1: Mat4 = Mat4.identity(); Mat4.setTranslation(m1, structureMean); structure = Structure.transform(structure, m1); if (ingredient.offset) { const o: Vec3 = Vec3.create(ingredient.offset[0], ingredient.offset[1], ingredient.offset[2]); if (!Vec3.exactEquals(o, Vec3())) { // -1, 1, 4e-16 ?? if (location !== 'surface') { Vec3.negate(o, o); } const m: Mat4 = Mat4.identity(); Mat4.setTranslation(m, o); structure = Structure.transform(structure, m); } } if (pcp) { const p: Vec3 = Vec3.create(pcp[0], pcp[1], pcp[2]); if (!Vec3.exactEquals(p, Vec3.unitZ)) { const q: Quat = Quat.identity(); Quat.rotationTo(q, p, Vec3.unitZ); const m: Mat4 = Mat4.fromQuat(Mat4(), q); structure = Structure.transform(structure, m); } } } structure = getAssembly(name, getResultTransforms(results, legacy), structure); } return { structure, assets }; } export function createStructureFromCellPack(plugin: PluginContext, packing: CellPacking, baseUrl: string, ingredientFiles: IngredientFiles) { return Task.create('Create Packing Structure', async ctx => { const { ingredients, location, name } = packing; const assets: Asset.Wrapper[] = []; const trajCache = new TrajectoryCache(); const structures: Structure[] = []; const colors: Color[] = []; for (const iName in ingredients) { if (ctx.shouldUpdate) await ctx.update(iName); const ingredientStructure = await getIngredientStructure(plugin, ingredients[iName], baseUrl, ingredientFiles, trajCache, location); if (ingredientStructure) { structures.push(ingredientStructure.structure); assets.push(...ingredientStructure.assets); const c = ingredients[iName].color; if (c) { colors.push(Color.fromNormalizedRgb(c[0], c[1], c[2])); } else { colors.push(Color.fromNormalizedRgb(1, 0, 0)); } } } if (ctx.shouldUpdate) await ctx.update(`${name} - units`); const units: Unit[] = []; let offsetInvariantId = 0; let offsetChainGroupId = 0; for (const s of structures) { if (ctx.shouldUpdate) await ctx.update(`${s.label}`); let maxInvariantId = 0; const maxChainGroupId = 0; for (const u of s.units) { const invariantId = u.invariantId + offsetInvariantId; const chainGroupId = u.chainGroupId + offsetChainGroupId; if (u.invariantId > maxInvariantId) maxInvariantId = u.invariantId; units.push(Unit.create(units.length, invariantId, chainGroupId, u.traits, u.kind, u.model, u.conformation.operator, u.elements, u.props)); } offsetInvariantId += maxInvariantId + 1; offsetChainGroupId += maxChainGroupId + 1; } if (ctx.shouldUpdate) await ctx.update(`${name} - structure`); const structure = Structure.create(units, { label: name + '.' + location }); for (let i = 0, il = structure.models.length; i < il; ++i) { Model.TrajectoryInfo.set(structure.models[i], { size: il, index: i }); } return { structure, assets, colors: colors }; }); } async function handleHivRna(plugin: PluginContext, packings: CellPacking[], baseUrl: string) { for (let i = 0, il = packings.length; i < il; ++i) { if (packings[i].name === 'HIV1_capsid_3j3q_PackInner_0_1_0' || packings[i].name === 'HIV_capsid') { const url = Asset.getUrlAsset(plugin.managers.asset, `${baseUrl}/extras/rna_allpoints.json`); const json = await plugin.runTask(plugin.managers.asset.resolve(url, 'json', false)); const points = json.data.points as number[]; const curve0: Vec3[] = []; for (let j = 0, jl = points.length; j < jl; j += 3) { curve0.push(Vec3.fromArray(Vec3(), points, j)); } packings[i].ingredients['RNA'] = { source: { pdb: 'RNA_U_Base.pdb', transform: { center: false } }, results: [], name: 'RNA', nbCurve: 1, curve0 }; } } } async function loadMembrane(plugin: PluginContext, name: string, state: State, params: LoadCellPackModelParams) { let file: Asset.File | undefined = undefined; if (params.ingredients !== null) { const fileName = `${name}.bcif`; for (const f of params.ingredients) { if (fileName === f.name) { file = f; break; } } if (!file) { // check for cif directly const cifileName = `${name}.cif`; for (const f of params.ingredients) { if (cifileName === f.name) { file = f; break; } } } } let legacy_membrane: boolean = false; // temporary variable until all membrane are converted to the new correct cif format let geometry_membrane: boolean = false; // membrane can be a mesh geometry let b = state.build().toRoot(); if (file) { if (file.name.endsWith('.cif')) { b = b.apply(StateTransforms.Data.ReadFile, { file, isBinary: false, label: file.name }, { state: { isGhost: true } }); } else if (file.name.endsWith('.bcif')) { b = b.apply(StateTransforms.Data.ReadFile, { file, isBinary: true, label: file.name }, { state: { isGhost: true } }); } } else { if (name.toLowerCase().endsWith('.bcif')) { const url = Asset.getUrlAsset(plugin.managers.asset, `${params.baseUrl}/membranes/${name}`); b = b.apply(StateTransforms.Data.Download, { url, isBinary: true, label: name }, { state: { isGhost: true } }); } else if (name.toLowerCase().endsWith('.cif')) { const url = Asset.getUrlAsset(plugin.managers.asset, `${params.baseUrl}/membranes/${name}`); b = b.apply(StateTransforms.Data.Download, { url, isBinary: false, label: name }, { state: { isGhost: true } }); } else if (name.toLowerCase().endsWith('.ply')) { const url = Asset.getUrlAsset(plugin.managers.asset, `${params.baseUrl}/geometries/${name}`); b = b.apply(StateTransforms.Data.Download, { url, isBinary: false, label: name }, { state: { isGhost: true } }); geometry_membrane = true; } else { const url = Asset.getUrlAsset(plugin.managers.asset, `${params.baseUrl}/membranes/${name}.bcif`); b = b.apply(StateTransforms.Data.Download, { url, isBinary: true, label: name }, { state: { isGhost: true } }); legacy_membrane = true; } } const props = { type: { name: 'assembly' as const, params: { id: '1' } } }; if (legacy_membrane) { // old membrane const membrane = await b.apply(StateTransforms.Data.ParseCif, undefined, { state: { isGhost: true } }) .apply(StateTransforms.Model.TrajectoryFromMmCif, undefined, { state: { isGhost: true } }) .apply(StateTransforms.Model.ModelFromTrajectory, undefined, { state: { isGhost: true } }) .apply(StructureFromAssemblies, undefined, { state: { isGhost: true } }) .commit({ revertOnError: true }); const membraneParams = { ignoreLight: params.preset.adjustStyle, representation: params.preset.representation, }; await CellpackMembranePreset.apply(membrane, membraneParams, plugin); } else if (geometry_membrane) { await b.apply(StateTransforms.Data.ParsePly, undefined, { state: { isGhost: true } }) .apply(StateTransforms.Model.ShapeFromPly) .apply(StateTransforms.Representation.ShapeRepresentation3D, { xrayShaded: true, doubleSided: true, coloring: { name: 'uniform', params: { color: ColorNames.orange } } }) .commit({ revertOnError: true }); } else { const membrane = await b.apply(StateTransforms.Data.ParseCif, undefined, { state: { isGhost: true } }) .apply(StateTransforms.Model.TrajectoryFromMmCif, undefined, { state: { isGhost: true } }) .apply(StateTransforms.Model.ModelFromTrajectory, undefined, { state: { isGhost: true } }) .apply(StateTransforms.Model.StructureFromModel, props, { state: { isGhost: true } }) .commit({ revertOnError: true }); const membraneParams = { ignoreLight: params.preset.adjustStyle, representation: params.preset.representation, }; await CellpackMembranePreset.apply(membrane, membraneParams, plugin); } } async function handleMembraneSpheres(state: State, primitives: CompartmentPrimitives) { const nSpheres = primitives.positions!.length / 3; // console.log('ok mb ', nSpheres); // TODO : take in account the type of the primitives. for (let j = 0; j < nSpheres; j++) { await state.build() .toRoot() .apply(CreateCompartmentSphere, { center: Vec3.create( primitives.positions![j * 3 + 0], primitives.positions![j * 3 + 1], primitives.positions![j * 3 + 2] ), radius: primitives!.radii![j] }) .commit(); } } async function loadPackings(plugin: PluginContext, runtime: RuntimeContext, state: State, params: LoadCellPackModelParams) { const ingredientFiles = params.ingredients || []; let cellPackJson: StateBuilder.To>; let resultsFile: Asset.File | null = params.results; if (params.source.name === 'id') { const url = Asset.getUrlAsset(plugin.managers.asset, getCellPackModelUrl(params.source.params, params.baseUrl)); cellPackJson = state.build().toRoot() .apply(StateTransforms.Data.Download, { url, isBinary: false, label: params.source.params }, { state: { isGhost: true } }); } else { const file = params.source.params; if (!file?.file) { plugin.log.error('No file selected'); return; } let modelFile: Asset.File; if (file.name.toLowerCase().endsWith('.zip')) { const data = await readFromFile(file.file, 'zip').runInContext(runtime); if (data['model.json']) { modelFile = Asset.File(new File([data['model.json']], 'model.json')); } else { throw new Error('model.json missing from zip file'); } if (data['results.bin']) { resultsFile = Asset.File(new File([data['results.bin']], 'results.bin')); } objectForEach(data, (v, k) => { if (k === 'model.json') return; if (k === 'results.bin') return; ingredientFiles.push(Asset.File(new File([v], k))); }); } else { modelFile = file; } cellPackJson = state.build().toRoot() .apply(StateTransforms.Data.ReadFile, { file: modelFile, isBinary: false, label: modelFile.name }, { state: { isGhost: true } }); } const cellPackBuilder = cellPackJson .apply(StateTransforms.Data.ParseJson, undefined, { state: { isGhost: true } }) .apply(ParseCellPack, { resultsFile, baseUrl: params.baseUrl }); const cellPackObject = await state.updateTree(cellPackBuilder).runInContext(runtime); const { packings } = cellPackObject.obj!.data; await handleHivRna(plugin, packings, params.baseUrl); for (let i = 0, il = packings.length; i < il; ++i) { const p = { packing: i, baseUrl: params.baseUrl, ingredientFiles }; const packing = await state.build() .to(cellPackBuilder.ref) .apply(StructureFromCellpack, p) .commit({ revertOnError: true }); const packingParams = { traceOnly: params.preset.traceOnly, ignoreLight: params.preset.adjustStyle, representation: params.preset.representation, }; await CellpackPackingPreset.apply(packing, packingParams, plugin); if (packings[i].compartment) { if (params.membrane === 'lipids') { if (packings[i].compartment!.geom_type) { if (packings[i].compartment!.geom_type === 'file') { // TODO: load mesh files or vertex,faces data await loadMembrane(plugin, packings[i].compartment!.filename!, state, params); } else if (packings[i].compartment!.compartment_primitives) { await handleMembraneSpheres(state, packings[i].compartment!.compartment_primitives!); } } else { // try loading membrane from repo as a bcif file or from the given list of files. if (params.membrane === 'lipids') { await loadMembrane(plugin, packings[i].name, state, params); } } } else if (params.membrane === 'geometry') { if (packings[i].compartment!.compartment_primitives) { await handleMembraneSpheres(state, packings[i].compartment!.compartment_primitives!); } else if (packings[i].compartment!.geom_type === 'file') { if (packings[i].compartment!.filename!.toLowerCase().endsWith('.ply')) { await loadMembrane(plugin, packings[i].compartment!.filename!, state, params); } } } } } } const LoadCellPackModelParams = { source: PD.MappedStatic('id', { 'id': PD.Select('InfluenzaModel2.json', [ ['blood_hiv_immature_inside.json', 'Blood HIV immature'], ['HIV_immature_model.json', 'HIV immature'], ['Blood_HIV.json', 'Blood HIV'], ['HIV-1_0.1.6-8_mixed_radii_pdb.json', 'HIV'], ['influenza_model1.json', 'Influenza envelope'], ['InfluenzaModel2.json', 'Influenza complete'], ['ExosomeModel.json', 'Exosome Model'], ['MycoplasmaGenitalium.json', 'Mycoplasma Genitalium curated model'], ] as const, { description: 'Download the model definition with `id` from the server at `baseUrl.`' }), 'file': PD.File({ accept: '.json,.cpr,.zip', description: 'Open model definition from .json/.cpr file or open .zip file containing model definition plus ingredients.', label: 'Recipe file' }), }, { options: [['id', 'Id'], ['file', 'File']] }), baseUrl: PD.Text(DefaultCellPackBaseUrl), results: PD.File({ accept: '.bin', description: 'open results file in binary format from cellpackgpu for the specified recipe', label: 'Results file' }), membrane: PD.Select('lipids', PD.arrayToOptions(['lipids', 'geometry', 'none'])), ingredients: PD.FileList({ accept: '.cif,.bcif,.pdb', label: 'Ingredient files' }), preset: PD.Group({ traceOnly: PD.Boolean(false), adjustStyle: PD.Boolean(true), representation: PD.Select('gaussian-surface', PD.arrayToOptions(['spacefill', 'gaussian-surface', 'point', 'orientation'] as const)) }, { isExpanded: true }) }; type LoadCellPackModelParams = PD.Values export const LoadCellPackModel = StateAction.build({ display: { name: 'Load CellPack', description: 'Open or download a model' }, params: LoadCellPackModelParams, from: PSO.Root })(({ state, params }, ctx: PluginContext) => Task.create('CellPack Loader', async taskCtx => { if (params.preset.adjustStyle) { ctx.managers.interactivity.setProps({ granularity: 'chain' }); ctx.managers.structure.component.setOptions({ ... ctx.managers.structure.component.state.options, visualQuality: 'custom', ignoreLight: true, showHydrogens: false, }); ctx.canvas3d?.setProps({ multiSample: { mode: 'off' }, cameraClipping: { far: false }, renderer: { colorMarker: false }, marking: { enabled: true, ghostEdgeStrength: 1, }, postprocessing: { occlusion: { name: 'on', params: { samples: 32, radius: 8, bias: 1, blurKernelSize: 15, resolutionScale: 1, } }, shadow: { name: 'on', params: { bias: 0.6, maxDistance: 80, steps: 3, tolerance: 1.0, } }, outline: { name: 'on', params: { scale: 1, threshold: 0.33, color: ColorNames.black, includeTransparent: true, } } } }); } await loadPackings(ctx, taskCtx, state, params); }));