TmMolStar/src/mol-theme/color/sequence-id.ts

/**
 * Copyright (c) 2018 mol* contributors, licensed under MIT, See LICENSE file for more info.
 *
 * @author Alexander Rose <alexander.rose@weirdbyte.de>
 */

import { Unit, StructureElement, Bond, ElementIndex } from '../../mol-model/structure';

import { ColorScale, Color } from '../../mol-util/color';
import { Location } from '../../mol-model/location';
import { ColorTheme } from '../color';
import { ParamDefinition as PD } from '../../mol-util/param-definition';
import { ThemeDataContext } from '../../mol-theme/theme';

const DefaultColor = Color(0xCCCCCC);
const Description = 'Gives every polymer residue a color based on its `seq_id` value.';

export const SequenceIdColorThemeParams = {
    list: PD.ColorList('turbo', { presetKind: 'scale' }),
};
export type SequenceIdColorThemeParams = typeof SequenceIdColorThemeParams
export function getSequenceIdColorThemeParams(ctx: ThemeDataContext) {
    return SequenceIdColorThemeParams; // TODO return copy
}

function getSeqId(unit: Unit, element: ElementIndex): number {
    const { model } = unit;
    switch (unit.kind) {
        case Unit.Kind.Atomic:
            const residueIndex = model.atomicHierarchy.residueAtomSegments.index[element];
            return model.atomicHierarchy.residues.label_seq_id.value(residueIndex);
        case Unit.Kind.Spheres:
            return Math.round(
                (model.coarseHierarchy.spheres.seq_id_begin.value(element) +
                    model.coarseHierarchy.spheres.seq_id_end.value(element)) / 2
            );
        case Unit.Kind.Gaussians:
            return Math.round(
                (model.coarseHierarchy.gaussians.seq_id_begin.value(element) +
                    model.coarseHierarchy.gaussians.seq_id_end.value(element)) / 2
            );
    }
}

function getSequenceLength(unit: Unit, element: ElementIndex) {
    const { model } = unit;
    let entityId = '';
    switch (unit.kind) {
        case Unit.Kind.Atomic:
            const chainIndex = model.atomicHierarchy.chainAtomSegments.index[element];
            entityId = model.atomicHierarchy.chains.label_entity_id.value(chainIndex);
            break;
        case Unit.Kind.Spheres:
            entityId = model.coarseHierarchy.spheres.entity_id.value(element);
            break;
        case Unit.Kind.Gaussians:
            entityId = model.coarseHierarchy.gaussians.entity_id.value(element);
            break;
    }
    if (entityId === '') return 0;

    const entityIndex = model.entities.getEntityIndex(entityId);
    if (entityIndex === -1) return 0;

    const entity = model.sequence.byEntityKey[entityIndex];
    if (entity === undefined) return 0;

    return entity.sequence.length;
}

export function SequenceIdColorTheme(ctx: ThemeDataContext, props: PD.Values<SequenceIdColorThemeParams>): ColorTheme<SequenceIdColorThemeParams> {
    const scale = ColorScale.create({
        listOrName: props.list.colors,
        minLabel: 'Start',
        maxLabel: 'End',
    });
    const color = (location: Location): Color => {
        if (StructureElement.Location.is(location)) {
            const { unit, element } = location;
            const seq_id = getSeqId(unit, element);
            if (seq_id > 0) {
                const seqLen = getSequenceLength(unit, element);
                if (seqLen) {
                    scale.setDomain(0, seqLen - 1);
                    return scale.color(seq_id);
                }
            }
        } else if (Bond.isLocation(location)) {
            const { aUnit, aIndex } = location;
            const seq_id = getSeqId(aUnit, aUnit.elements[aIndex]);
            if (seq_id > 0) {
                const seqLen = getSequenceLength(aUnit, aUnit.elements[aIndex]);
                if (seqLen) {
                    scale.setDomain(0, seqLen - 1);
                    return scale.color(seq_id);
                }
            }
        }
        return DefaultColor;
    };

    return {
        factory: SequenceIdColorTheme,
        granularity: 'group',
        color,
        props,
        description: Description,
        legend: scale ? scale.legend : undefined
    };
}

export const SequenceIdColorThemeProvider: ColorTheme.Provider<SequenceIdColorThemeParams, 'sequence-id'> = {
    name: 'sequence-id',
    label: 'Sequence Id',
    category: ColorTheme.Category.Residue,
    factory: SequenceIdColorTheme,
    getParams: getSequenceIdColorThemeParams,
    defaultValues: PD.getDefaultValues(SequenceIdColorThemeParams),
    isApplicable: (ctx: ThemeDataContext) => !!ctx.structure
};