TmMolStar/src/extensions/tmdet/tmdet-color-theme.ts

import { StructureElement, Unit } from '../../mol-model/structure';
import { ColorTheme } from '../../mol-theme/color';
import { ThemeDataContext } from '../../mol-theme/theme';
import { Color } from '../../mol-util/color';
import { ColorNames } from '../../mol-util/color/names';
import { ParamDefinition as PD } from '../../mol-util/param-definition';
import { Location } from '../../mol-model/location';
import { PDBTMChain } from './types';


export type TmDetColorThemeParams = {
    pdbtmDescriptor: any
};

export function TmDetColorTheme(
    ctx: ThemeDataContext,
    props: PD.Values<TmDetColorThemeParams>
): ColorTheme<TmDetColorThemeParams> {
    console.log('ColorTheme CTX', ctx);
    const descriptorChains = props.pdbtmDescriptor.chains as PDBTMChain[];
    console.log('ColorTheme PROPS', descriptorChains);

    return {
        factory: TmDetColorTheme,
        granularity: 'group', //'residue' as any,
        color: (location: Location) => getColor(location, descriptorChains),
        props: props,
        description: 'TMDet...',
    };
}

const DefaultResidueColor = ColorNames.black;

//TODO: colors of not curated sites
const siteColors = [
    Color.fromArray([255,100,100], 0), // Side1
    Color.fromArray([100,100,255], 0), // Side2
    Color.fromArray([255,255,  0], 0), // TM alpha
    Color.fromArray([255,255,  0], 0), // TM beta
    Color.fromArray([255,127,  0], 0), // TM re-entrant loop
    Color.fromArray([0,255,    0], 0), // Interfacial Helix
    Color.fromArray([196,196,196], 0), // Unknow localization
    Color.fromArray([0,255,    0], 0)  // Membrane Inside
];


function getColor(location: Location, chains: PDBTMChain[]): Color {
    let color = DefaultResidueColor;

    // TODO: How to solve cases when chain operation occurs?
    const chainList = createResidueListsPerChain(chains);
    if (StructureElement.Location.is(location) && Unit.isAtomic(location.unit)) {
        const atomicHierarchy = location.unit.model.atomicHierarchy;
        const residueIdx = StructureElement.residueIndex(location);
        const chainIdx = StructureElement.Location.chainIndex(location);

        const authId = atomicHierarchy.residues.auth_seq_id.value(residueIdx).toString();
        const chainId = atomicHierarchy.chains.auth_asym_id.value(chainIdx).toString();

        color = residueColor(chainList, chainId, authId);
        console.log(`chain: ${chainId} res: ${authId}`);
    }
    return color;
}

type ChainList = {
    chainId: string,
    residues: {
        authId: string|undefined,
        siteId: number
    }[]
}[];

function createResidueListsPerChain(chains: PDBTMChain[]) {
    const chainList: ChainList = [];

    chains.forEach((chain: PDBTMChain) => {
        chainList.push({ chainId: chain.chain_label, residues: [] });
        chain.residues.forEach((residue) => {
            chainList[chainList.length - 1].residues.push({
                authId: residue.pdb_res_label,
                siteId: residue.site_data![0].site_id_ref - 1
            });
        });
    });

    return chainList;
}

function residueColor(chains: ChainList, chainId: string, residueId: string): Color {
    let color = DefaultResidueColor;

    const chain = chains.filter((chain) => chain.chainId === chainId)[0];
    if (chain) {
        const residue = chain.residues.filter((res) => res.authId === residueId)[0];
        if (residue) {
            color = siteColors[residue.siteId];
        }
    }

    return color;
}

// Provider

export const TmDetColorThemeProvider: ColorTheme.Provider<TmDetColorThemeParams, 'tmdet-custom-color-theme'> = {
    name: 'tmdet-custom-color-theme',
    label: 'TMDet Topology Theme',
    category: 'TMDet',
    factory: TmDetColorTheme,
    getParams: () => ({ pdbtmDescriptor: { isOptional: true } }),
    defaultValues: { pdbtmDescriptor: undefined },
    isApplicable: () => true,
};