TmMolStar/src/mol-model-props/integrative/cross-link-restraint/format.ts

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

import { Model } from '../../../mol-model/structure/model/model'
import { Table } from '../../../mol-data/db'
import { mmCIF_Schema } from '../../../mol-io/reader/cif/schema/mmcif';
import { Unit, CustomPropertyDescriptor } from '../../../mol-model/structure';
import { ElementIndex } from '../../../mol-model/structure/model/indexing';
import { FormatPropertyProvider } from '../../../mol-model-formats/structure/common/property';

export { ModelCrossLinkRestraint }

interface ModelCrossLinkRestraint {
    getIndicesByElement: (element: ElementIndex, kind: Unit.Kind) => number[]
    data: Table<mmCIF_Schema['ihm_cross_link_restraint']>
}

namespace ModelCrossLinkRestraint {
    export const Descriptor: CustomPropertyDescriptor = {
        name: 'ihm_cross_link_restraint',
        // TODO cifExport
    }

    export const Provider = FormatPropertyProvider.create<ModelCrossLinkRestraint>(Descriptor)

    export function fromTable(table: Table<mmCIF_Schema['ihm_cross_link_restraint']>, model: Model): ModelCrossLinkRestraint {

        const p1 = {
            entity_id: table.entity_id_1,
            asym_id: table.asym_id_1,
            seq_id: table.seq_id_1,
            atom_id: table.atom_id_1,
        }

        const p2: typeof p1 = {
            entity_id: table.entity_id_2,
            asym_id: table.asym_id_2,
            seq_id: table.seq_id_2,
            atom_id: table.atom_id_2,
        }

        function _add(map: Map<ElementIndex, number[]>, element: ElementIndex, row: number) {
            const indices = map.get(element)
            if (indices) indices.push(row)
            else map.set(element, [ row ])
        }

        function add(row: number, ps: typeof p1) {
            const entityId = ps.entity_id.value(row)
            const asymId = ps.asym_id.value(row)
            const seqId = ps.seq_id.value(row)

            if (table.model_granularity.value(row) === 'by-atom') {
                const atomicElement = model.atomicHierarchy.index.findAtom({
                    auth_seq_id: seqId,
                    label_asym_id: asymId,
                    label_atom_id: ps.atom_id.value(row),
                    label_entity_id: entityId,
                })
                if (atomicElement >= 0) _add(atomicElementMap, atomicElement as ElementIndex, row)
            } else if (model.coarseHierarchy.isDefined) {
                const sphereElement = model.coarseHierarchy.spheres.findSequenceKey(entityId, asymId, seqId)
                if (sphereElement >= 0) {
                    _add(sphereElementMap, sphereElement, row)
                } else {
                    const gaussianElement = model.coarseHierarchy.gaussians.findSequenceKey(entityId, asymId, seqId)
                    if (gaussianElement >= 0) _add(gaussianElementMap, gaussianElement, row)
                }
            }
        }

        function getMapByKind(kind: Unit.Kind) {
            switch (kind) {
                case Unit.Kind.Atomic: return atomicElementMap;
                case Unit.Kind.Spheres: return sphereElementMap;
                case Unit.Kind.Gaussians: return gaussianElementMap;
            }
        }

        /** map from atomic element to cross link indices */
        const atomicElementMap: Map<ElementIndex, number[]> = new Map()
        /** map from sphere element to cross link indices */
        const sphereElementMap: Map<ElementIndex, number[]> = new Map()
        /** map from gaussian element to cross link indices */
        const gaussianElementMap: Map<ElementIndex, number[]> = new Map()

        const emptyIndexArray: number[] = [];

        for (let i = 0; i < table._rowCount; ++i) {
            add(i, p1)
            add(i, p2)
        }

        return {
            getIndicesByElement: (element: ElementIndex, kind: Unit.Kind) => {
                const map = getMapByKind(kind)
                const idx = map.get(element)
                return idx !== undefined ? idx : emptyIndexArray
            },
            data: table
        }
    }
}