TmMolStar/src/extensions/tmdet/labeling.ts

/**
 * 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>
 */

import { DebugUtil } from '../../apps/tm-viewer';
import { CustomElementProperty } from '../../mol-model-props/common/custom-element-property';
import { Model, ElementIndex } from '../../mol-model/structure';
import { Loci } from '../../mol-model/loci';
import { LociLabel, LociLabelProvider } from '../../mol-plugin-state/manager/loci-label';
import { TmDetChainListCache, TmDetDescriptorCache } from './prop';
import { createResidueListsPerChain } from './tmdet-color-theme';
import { ChainList, getResidue } from './types';
import { OrderedSet } from '../../mol-data/int';

const siteLabels = [
    "Side1",
    "Side2",
    "TM alpha",
    "TM beta",
    "TM re-entrant loop",
    "Interfacial Helix",
    "Unknown localization",
    "Membrane Inside"
];
const DefaultResidueLabel = 6; // Unknown localization


export const LabeledResidues = CustomElementProperty.create<number>({
    label: 'TMDet Topology Site Labels',
    name: 'tmdet-topology-site-labels',
    getData(model: Model) {
        const pdbId = model.entryId;
        DebugUtil.log('start getData', model.label, pdbId);

        const map = new Map<ElementIndex, number>();
        const descriptor = TmDetDescriptorCache.get(pdbId);
        if (!descriptor) {
            return { value: map };
        }

        // TODO: integrate with the coloring version
        const getSiteId = function(chains: ChainList, chainId: string, residueId: string): number {
            let siteId = DefaultResidueLabel;

            const residue = getResidue(chains, chainId, residueId);
            if (residue) {
                siteId = residue.siteId;
            }

            return siteId;
        };

        let chainList =  TmDetChainListCache.get(pdbId);
        if (!chainList) {
            chainList = createResidueListsPerChain(descriptor.chains);
            TmDetChainListCache.set(pdbId, chainList);
        }
        DebugUtil.log('getData', model.atomicHierarchy.chainAtomSegments);
        // DebugUtil.log('getData: auth_comp_id', model.atomicHierarchy.atoms.auth_comp_id.toArray());
        const residueIndex = model.atomicHierarchy.residueAtomSegments.index;
        const chainIndex = model.atomicHierarchy.chainAtomSegments.index;
        for (let i = 0, _i = model.atomicHierarchy.atoms._rowCount; i < _i; i++) {
            const residueIdx = residueIndex[i];
            const chainIdx = chainIndex[i];
            const residueAuthId = model.atomicHierarchy.residues.auth_seq_id.value(residueIdx).toString();
            const residueName = model.atomicHierarchy.atoms.auth_comp_id.value(i).toString();
            const chainAuthId = model.atomicHierarchy.chains.auth_asym_id.value(chainIdx).toString();
            const chainLabel = model.atomicHierarchy.chains.label_asym_id.value(chainIdx).toString();
            //DebugUtil.log(`${i} ${residueIdx} ${residueName} ${residueAuthId} ${chainIdx} ${chainAuthId} ${chainLabel}`);
    
            map.set(i as ElementIndex, getSiteId(chainList, chainAuthId, residueAuthId));
        }
        DebugUtil.log('end of getData');
        return { value: map };
    },
    getLabel(e) {
        return siteLabels[e];
    }
});

// function residueLabel(chains: ChainList, chainId: string, residueId: string): string {
//     let label = DefaultResidueLabel;

//     const residue = getResidue(chains, chainId, residueId);
//     if (residue) {
//         label = siteLabels[residue.siteId];
//     }

//     return color;
// }


export const TmLabelProvider: LociLabelProvider = {
    label: (loci: Loci): LociLabel => {
        if (loci.kind == 'element-loci') {
            DebugUtil.log('LabelProvider:', loci.elements[0].unit.model.atomicHierarchy.chains.auth_asym_id.toArray());
            DebugUtil.log('LabelProvider:', loci.elements[0].unit.model.atomicHierarchy.chains.label_asym_id.toArray());
            DebugUtil.log('LabelProvider:', loci.elements[0].unit.model.atomicHierarchy.chains.label_entity_id.toArray());
            DebugUtil.log('IsLocation:', loci.structure.elementLocations());
            DebugUtil.log('start:', OrderedSet.start(loci.elements[0].indices));
        }
        return "hello";
    }
}