basic partial charge and pdbt support

- AtomPartialCharge format property provider
- add from PDBQT and MOL2
- color theme

src/mol-model-formats/structure/mol2.ts

@@ -15,12 +15,13 @@ import { EntityBuilder } from './common/entity';
 import { ModelFormat } from '../format';
 import { IndexPairBonds } from './property/bonds/index-pair';
 import { Mol2File } from '../../mol-io/reader/mol2/schema';
+import { AtomPartialCharge } from './property/partial-charge';
 
 async function getModels(mol2: Mol2File, ctx: RuntimeContext): Promise<Model[]> {
     const models: Model[] = [];
 
     for (let i = 0, il = mol2.structures.length; i < il; ++i) {
-        const { atoms, bonds } = mol2.structures[i];
+        const { atoms, bonds, molecule } = mol2.structures[i];
 
         const A = Column.ofConst('A', atoms.count, Column.Schema.str);
 
@@ -70,6 +71,11 @@ async function getModels(mol2: Mol2File, ctx: RuntimeContext): Promise<Model[]>
             const pairBonds = IndexPairBonds.fromData({ pairs: { indexA, indexB, order }, count: bonds.count });
             IndexPairBonds.Provider.set(_models[0], pairBonds);
 
+            AtomPartialCharge.Provider.set(_models[0], {
+                data: atoms.charge,
+                type: molecule.charge_type
+            });
+
             models.push(_models[0]);
         }
     }

src/mol-model-formats/structure/pdb.ts

@@ -11,12 +11,35 @@ import { Model } from '../../mol-model/structure/model';
 import { Task } from '../../mol-task';
 import { MmcifFormat } from './mmcif';
 import { createModels } from './basic/parser';
+import { Column } from '../../mol-data/db';
+import { AtomPartialCharge } from './property/partial-charge';
 
 export function trajectoryFromPDB(pdb: PdbFile): Task<Model.Trajectory> {
     return Task.create('Parse PDB', async ctx => {
         await ctx.update('Converting to mmCIF');
         const cif = await pdbToMmCif(pdb);
         const format = MmcifFormat.fromFrame(cif);
-        return createModels(format.data.db, format, ctx);
+        const models = await createModels(format.data.db, format, ctx);
+        const partial_charge = cif.categories['atom_site']?.getField('partial_charge');
+        if (partial_charge) {
+            // TODO works only for single, unsorted model, to work generally
+            //      would need to do model splitting again
+            if (models.length === 1) {
+                const srcIndex = models[0].atomicHierarchy.atoms.sourceIndex;
+                const isIdentity = Column.isIdentity(srcIndex);
+                const srcIndexArray = isIdentity ? void 0 : srcIndex.toArray({ array: Int32Array });
+
+                const q = partial_charge.toFloatArray();
+                const partialCharge = srcIndexArray
+                    ? Column.ofFloatArray(Column.mapToArray(srcIndex, i => q[i], Float32Array))
+                    : Column.ofFloatArray(q);
+
+                AtomPartialCharge.Provider.set(models[0], {
+                    data: partialCharge,
+                    type: 'GASTEIGER' // from PDBQT
+                });
+            }
+        }
+        return models;
     });
 }

src/mol-model-formats/structure/pdb/atom-site.ts

@@ -1,5 +1,5 @@
 /**
- * Copyright (c) 2019 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ * Copyright (c) 2019-2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
  *
  * @author David Sehnal <david.sehnal@gmail.com>
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
@@ -32,11 +32,13 @@ export function getAtomSiteTemplate(data: string, count: number) {
         B_iso_or_equiv: ts(),
         type_symbol: ts(),
         pdbx_PDB_model_num: str(),
-        label_entity_id: str()
+        label_entity_id: str(),
+
+        partial_charge: ts(),
     };
 }
 
-export function getAtomSite(sites: AtomSiteTemplate): { [K in keyof mmCIF_Schema['atom_site']]?: CifField } {
+export function getAtomSite(sites: AtomSiteTemplate): { [K in keyof mmCIF_Schema['atom_site'] | 'partial_charge']?: CifField } {
     const auth_asym_id = CifField.ofTokens(sites.auth_asym_id);
     const auth_atom_id = CifField.ofTokens(sites.auth_atom_id);
     const auth_comp_id = CifField.ofTokens(sites.auth_comp_id);
@@ -66,11 +68,13 @@ export function getAtomSite(sites: AtomSiteTemplate): { [K in keyof mmCIF_Schema
         type_symbol: CifField.ofTokens(sites.type_symbol),
 
         pdbx_PDB_ins_code: CifField.ofTokens(sites.pdbx_PDB_ins_code),
-        pdbx_PDB_model_num: CifField.ofStrings(sites.pdbx_PDB_model_num)
+        pdbx_PDB_model_num: CifField.ofStrings(sites.pdbx_PDB_model_num),
+
+        partial_charge: CifField.ofTokens(sites.partial_charge)
     };
 }
 
-export function addAtom(sites: AtomSiteTemplate, model: string, data: Tokenizer, s: number, e: number) {
+export function addAtom(sites: AtomSiteTemplate, model: string, data: Tokenizer, s: number, e: number, isPdbqt: boolean) {
     const { data: str } = data;
     const length = e - s;
 
@@ -133,10 +137,14 @@ export function addAtom(sites: AtomSiteTemplate, model: string, data: Tokenizer,
     }
 
     // 73 - 76        LString(4)      Segment identifier, left-justified.
-    // ignored
+    if (isPdbqt) {
+        TokenBuilder.addToken(sites.partial_charge, Tokenizer.trim(data, s + 70, s + 76));
+    } else {
+        // ignored
+    }
 
     // 77 - 78        LString(2)      Element symbol, right-justified.
-    if (length >= 78) {
+    if (length >= 78 && !isPdbqt) {
         Tokenizer.trim(data, s + 76, s + 78);
 
         if (data.tokenStart < data.tokenEnd) {

src/mol-model-formats/structure/pdb/to-cif.ts

@@ -1,5 +1,5 @@
 /**
- * Copyright (c) 2019 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ * Copyright (c) 2019-2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
  *
  * @author David Sehnal <david.sehnal@gmail.com>
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
@@ -47,6 +47,7 @@ export async function pdbToMmCif(pdb: PdbFile): Promise<CifFrame> {
     const heteroNames: [string, string][] = [];
 
     let modelNum = 0, modelStr = '';
+    let isPdbqt = false;
 
     for (let i = 0, _i = lines.count; i < _i; i++) {
         let s = indices[2 * i], e = indices[2 * i + 1];
@@ -54,11 +55,13 @@ export async function pdbToMmCif(pdb: PdbFile): Promise<CifFrame> {
             case 'A':
                 if (substringStartsWith(data, s, e, 'ATOM  ')) {
                     if (!modelNum) { modelNum++; modelStr = '' + modelNum; }
-                    addAtom(atomSite, modelStr, tokenizer, s, e);
+                    addAtom(atomSite, modelStr, tokenizer, s, e, isPdbqt);
                 } else if (substringStartsWith(data, s, e, 'ANISOU')) {
                     addAnisotropic(anisotropic, modelStr, tokenizer, s, e);
                 }
                 break;
+            case 'B':
+                if (substringStartsWith(data, s, e, 'BRANCH')) isPdbqt = true;
             case 'C':
                 if (substringStartsWith(data, s, e, 'CRYST1')) {
                     helperCategories.push(...parseCryst1(pdb.id || '?', data.substring(s, e)));
@@ -75,10 +78,12 @@ export async function pdbToMmCif(pdb: PdbFile): Promise<CifFrame> {
                     i = j - 1;
                 }
                 break;
+            case 'E':
+                if (substringStartsWith(data, s, e, 'ENDROOT')) isPdbqt = true;
             case 'H':
                 if (substringStartsWith(data, s, e, 'HETATM')) {
                     if (!modelNum) { modelNum++; modelStr = '' + modelNum; }
-                    addAtom(atomSite, modelStr, tokenizer, s, e);
+                    addAtom(atomSite, modelStr, tokenizer, s, e, isPdbqt);
                 } else if (substringStartsWith(data, s, e, 'HELIX')) {
                     let j = i + 1;
                     while (true) {
@@ -129,6 +134,8 @@ export async function pdbToMmCif(pdb: PdbFile): Promise<CifFrame> {
                     }
                     helperCategories.push(...parseRemark350(lines, i, j));
                     i = j - 1;
+                } else if (substringStartsWith(data, s, e, 'ROOT')) {
+                    isPdbqt = true;
                 }
                 break;
             case 'S':
@@ -144,6 +151,8 @@ export async function pdbToMmCif(pdb: PdbFile): Promise<CifFrame> {
                 }
                 // TODO: SCALE record => cif.atom_sites.fract_transf_matrix, cif.atom_sites.fract_transf_vector
                 break;
+            case 'T':
+                if (substringStartsWith(data, s, e, 'TORSDOF')) isPdbqt = true;
         }
     }
 
@@ -161,10 +170,13 @@ export async function pdbToMmCif(pdb: PdbFile): Promise<CifFrame> {
         atomSite.label_entity_id[i] = entityBuilder.getEntityId(compId, moleculeType, asymIds.value(i));
     }
 
+    const atom_site = getAtomSite(atomSite);
+    if (!isPdbqt) delete atom_site.partial_charge;
+
     const categories = {
         entity: CifCategory.ofTable('entity', entityBuilder.getEntityTable()),
         chem_comp: CifCategory.ofTable('chem_comp', componentBuilder.getChemCompTable()),
-        atom_site: CifCategory.ofFields('atom_site', getAtomSite(atomSite)),
+        atom_site: CifCategory.ofFields('atom_site', atom_site),
         atom_site_anisotrop: CifCategory.ofFields('atom_site_anisotrop', getAnisotropic(anisotropic))
     } as any;
 

src/mol-model-formats/structure/property/partial-charge.ts

@@ -0,0 +1,24 @@
+/**
+ * Copyright (c) 2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author Alexander Rose <alexander.rose@weirdbyte.de>
+ */
+
+import { Column } from '../../../mol-data/db';
+import { CustomPropertyDescriptor } from '../../../mol-model/custom-property';
+import { FormatPropertyProvider } from '../common/property';
+
+export { AtomPartialCharge };
+
+interface AtomPartialCharge {
+    data: Column<number>
+    type?: string
+}
+
+namespace AtomPartialCharge {
+    export const Descriptor: CustomPropertyDescriptor = {
+        name: 'atom_partial_charge',
+    };
+
+    export const Provider = FormatPropertyProvider.create<AtomPartialCharge>(Descriptor);
+}

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

@@ -120,7 +120,7 @@ export namespace Model {
         return _trajectoryFromModelAndCoordinates(model, coordinates).trajectory;
     }
 
-    function invertIndex(xs: ArrayLike<number>) {
+    export function invertIndex(xs: ArrayLike<number>) {
         const ret = new Int32Array(xs.length);
         for (let i = 0, _i = xs.length; i < _i; i++) {
             ret[xs[i]] = i;

src/mol-plugin-state/formats/trajectory.ts

@@ -89,7 +89,7 @@ export const PdbProvider: TrajectoryFormatProvider = {
     label: 'PDB',
     description: 'PDB',
     category: Category,
-    stringExtensions: ['pdb', 'ent'],
+    stringExtensions: ['pdb', 'ent', 'pdbqt'],
     parse: directTrajectory(StateTransforms.Model.TrajectoryFromPDB),
     visuals: defaultVisuals
 };

src/mol-theme/color.ts

@@ -1,5 +1,5 @@
 /**
- * Copyright (c) 2018-2019 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ * Copyright (c) 2018-2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
  *
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
  */
@@ -32,6 +32,7 @@ import { ModelIndexColorThemeProvider } from './color/model-index';
 import { OccupancyColorThemeProvider } from './color/occupancy';
 import { OperatorNameColorThemeProvider } from './color/operator-name';
 import { OperatorHklColorThemeProvider } from './color/operator-hkl';
+import { PartialChargeColorThemeProvider } from './color/partial-charge';
 
 export type LocationColor = (location: Location, isSecondary: boolean) => Color
 
@@ -91,6 +92,7 @@ namespace ColorTheme {
         'occupancy': OccupancyColorThemeProvider,
         'operator-hkl': OperatorHklColorThemeProvider,
         'operator-name': OperatorNameColorThemeProvider,
+        'partial-charge': PartialChargeColorThemeProvider,
         'polymer-id': PolymerIdColorThemeProvider,
         'polymer-index': PolymerIndexColorThemeProvider,
         'residue-name': ResidueNameColorThemeProvider,

src/mol-theme/color/partial-charge.ts

@@ -0,0 +1,66 @@
+/**
+ * Copyright (c) 2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author Alexander Rose <alexander.rose@weirdbyte.de>
+ */
+
+import { Color, ColorScale } from '../../mol-util/color';
+import { StructureElement, Unit, Bond, ElementIndex } from '../../mol-model/structure';
+import { Location } from '../../mol-model/location';
+import { ColorTheme } from '../color';
+import { ParamDefinition as PD } from '../../mol-util/param-definition';
+import { ThemeDataContext } from '../theme';
+import { AtomPartialCharge } from '../../mol-model-formats/structure/property/partial-charge';
+
+const DefaultPartialChargeColor = Color(0xffff99);
+const Description = `Assigns a color based on the partial charge of an atom.`;
+
+export const PartialChargeColorThemeParams = {
+    domain: PD.Interval([-1, 1]),
+    list: PD.ColorList('red-white-blue', { presetKind: 'scale' }),
+};
+export type PartialChargeColorThemeParams = typeof PartialChargeColorThemeParams
+export function getPartialChargeColorThemeParams(ctx: ThemeDataContext) {
+    return PartialChargeColorThemeParams; // TODO return copy
+}
+
+function getPartialCharge(unit: Unit, element: ElementIndex) {
+    return AtomPartialCharge.Provider.get(unit.model)?.data.value(element);
+}
+
+export function PartialChargeColorTheme(ctx: ThemeDataContext, props: PD.Values<PartialChargeColorThemeParams>): ColorTheme<PartialChargeColorThemeParams> {
+    const scale = ColorScale.create({
+        domain: props.domain,
+        listOrName: props.list.colors,
+    });
+
+    function color(location: Location): Color {
+        if (StructureElement.Location.is(location)) {
+            const q = getPartialCharge(location.unit, location.element);
+            return q !== undefined ? scale.color(q) : DefaultPartialChargeColor;
+        } else if (Bond.isLocation(location)) {
+            const q = getPartialCharge(location.aUnit, location.aUnit.elements[location.aIndex]);
+            return q !== undefined ? scale.color(q) : DefaultPartialChargeColor;
+        }
+        return DefaultPartialChargeColor;
+    }
+
+    return {
+        factory: PartialChargeColorTheme,
+        granularity: 'group',
+        color,
+        props,
+        description: Description,
+        legend: scale ? scale.legend : undefined
+    };
+}
+
+export const PartialChargeColorThemeProvider: ColorTheme.Provider<PartialChargeColorThemeParams, 'partial-charge'> = {
+    name: 'partial-charge',
+    label: 'Partial Charge',
+    category: ColorTheme.Category.Atom,
+    factory: PartialChargeColorTheme,
+    getParams: getPartialChargeColorThemeParams,
+    defaultValues: PD.getDefaultValues(PartialChargeColorThemeParams),
+    isApplicable: (ctx: ThemeDataContext) => !!ctx.structure && ctx.structure.models.some(m => AtomPartialCharge.Provider.get(m) !== undefined)
+};