<commit_before>/**
 * Copyright (c) 2021-2022 mol* contributors, licensed under MIT, See LICENSE file for more info.
 *
 * @author <NAME> <<EMAIL>>
 * @author <NAME> <<EMAIL>>
 */

import { Segmentation } from '../../../../mol-data/int';
import { MinimizeRmsd } from '../../../../mol-math/linear-algebra/3d/minimize-rmsd';
import { SIFTSMapping } from '../../../../mol-model-props/sequence/sifts-mapping';
import { ElementIndex } from '../../model/indexing';
import { Structure } from '../structure';
import { Unit } from '../unit';

export interface AlignmentResultEntry {
    transform: MinimizeRmsd.Result,
    pivot: number,
    other: number
}

export interface AlignmentResult {
    entries: AlignmentResultEntry[],
    zeroOverlapPairs: [number, number][],
    failedPairs: [number, number][]
}

export function alignAndSuperposeWithSIFTSMapping(structures: Structure[], options?: { traceOnly?: boolean }): AlignmentResult {
    const indexMap = new Map<string, IndexEntry>();

    for (let i = 0; i < structures.length; i++) {
        buildIndex(structures[i], indexMap, i);
    }

    const index = Array.from(indexMap.values());

    // TODO: support non-first structure pivots
    const pairs = findPairs(structures.length, index);

    const zeroOverlapPairs: AlignmentResult['zeroOverlapPairs'] = [];
    const failedPairs: AlignmentResult['failedPairs'] = [];


    const entries: AlignmentResultEntry[] = [];
    for (const p of pairs) {
        if (p.count === 0) {
            zeroOverlapPairs.push([p.i, p.j]);
        } else {
            const [a, b] = getPositionTables(index, p.i, p.j, p.count, options?.traceOnly ?? true);
            const transform = MinimizeRmsd.compute({ a, b });
            if (Number.isNaN(transform.rmsd)) {
                failedPairs.push([p.i, p.j]);
            } else {
                entries.push({ transform, pivot: p.i, other: p.j });
            }
        }
    }

    return { entries, zeroOverlapPairs, failedPairs };
}

function getPositionTables(index: IndexEntry[], pivot: number, other: number, N: number, traceOnly: boolean) {
    const xs = MinimizeRmsd.Positions.empty(N);
    const ys = MinimizeRmsd.Positions.empty(N);

    let o = 0;
    for (const { pivots } of index) {
        const a = pivots[pivot];
        const b = pivots[other];
        if (!a || !b) continue;

        const l = Math.min(a[2] - a[1], b[2] - b[1]);

        // TODO: check if residue types match?
        for (let i = 0; i < l; i++) {
            const eI = (a[1] + i) as ElementIndex;
            const eJ = (b[1] + i) as ElementIndex;
            if (traceOnly && (!traceAtom(a[0], eI) || !traceAtom(b[0], eJ))) continue;

            xs.x[o] = a[0].conformation.x(eI);
            xs.y[o] = a[0].conformation.y(eI);
            xs.z[o] = a[0].conformation.z(eI);

            ys.x[o] = b[0].conformation.x(eJ);
            ys.y[o] = b[0].conformation.y(eJ);
            ys.z[o] = b[0].conformation.z(eJ);
            o++;
        }
    }

    if (traceOnly) {
        xs.x = xs.x.slice(0, o);
        xs.y = xs.y.slice(0, o);
        xs.z = xs.z.slice(0, o);
        ys.x = ys.x.slice(0, o);
        ys.y = ys.y.slice(0, o);
        ys.z = ys.z.slice(0, o);
    }
    return [xs, ys];
}

function traceAtom(unit: Unit.Atomic, eI: ElementIndex): boolean {
    const l = unit.model.atomicHierarchy.atoms.label_atom_id.value(eI);
    return l === 'CA' || l === 'BB' || l === `C4'`;
}

function findPairs(N: number, index: IndexEntry[]) {
    const pairwiseCounts: number[][] = [];
    for (let i = 0; i < N; i++) {
        pairwiseCounts[i] = [];
        for (let j = 0; j < N; j++) pairwiseCounts[i][j] = 0;
    }

    for (const { pivots } of index) {
        for (let i = 0; i < N; i++) {
            if (!pivots[i]) continue;

            const lI = pivots[i]![2] - pivots[i]![1];

            for (let j = i + 1; j < N; j++) {
                if (!pivots[j]) continue;

                const lJ = pivots[j]![2] - pivots[j]![1];<commit_msg>support shuffled atoms when aligning SIFTS trace<commit_after>/**
 * Copyright (c) 2021-2022 mol* contributors, licensed under MIT, See LICENSE file for more info.
 *
 * @author <NAME> <<EMAIL>>
 * @author <NAME> <<EMAIL>>
 */

import { Segmentation } from '../../../../mol-data/int';
import { MinimizeRmsd } from '../../../../mol-math/linear-algebra/3d/minimize-rmsd';
import { SIFTSMapping } from '../../../../mol-model-props/sequence/sifts-mapping';
import { ElementIndex } from '../../model/indexing';
import { Structure } from '../structure';
import { Unit } from '../unit';

export interface AlignmentResultEntry {
    transform: MinimizeRmsd.Result,
    pivot: number,
    other: number
}

export interface AlignmentResult {
    entries: AlignmentResultEntry[],
    zeroOverlapPairs: [number, number][],
    failedPairs: [number, number][]
}

export function alignAndSuperposeWithSIFTSMapping(structures: Structure[], options?: { traceOnly?: boolean }): AlignmentResult {
    const indexMap = new Map<string, IndexEntry>();

    for (let i = 0; i < structures.length; i++) {
        buildIndex(structures[i], indexMap, i);
    }

    const index = Array.from(indexMap.values());

    // TODO: support non-first structure pivots
    const pairs = findPairs(structures.length, index);

    const zeroOverlapPairs: AlignmentResult['zeroOverlapPairs'] = [];
    const failedPairs: AlignmentResult['failedPairs'] = [];


    const entries: AlignmentResultEntry[] = [];
    for (const p of pairs) {
        if (p.count === 0) {
            zeroOverlapPairs.push([p.i, p.j]);
        } else {
            const [a, b] = getPositionTables(index, p.i, p.j, p.count, options?.traceOnly ?? true);
            const transform = MinimizeRmsd.compute({ a, b });
            if (Number.isNaN(transform.rmsd)) {
                failedPairs.push([p.i, p.j]);
            } else {
                entries.push({ transform, pivot: p.i, other: p.j });
            }
        }
    }

    return { entries, zeroOverlapPairs, failedPairs };
}

function getPositionTables(index: IndexEntry[], pivot: number, other: number, N: number, traceOnly: boolean) {
    const xs = MinimizeRmsd.Positions.empty(N);
    const ys = MinimizeRmsd.Positions.empty(N);

    let o = 0;
    for (const { pivots } of index) {
        let a = pivots[pivot];
        let b = pivots[other];
        if (!a || !b) continue;

        if (traceOnly) {
            a = findTraceAtom(...a);
            b = findTraceAtom(...b);
            if (!a || !b) continue;
        }

        const l = Math.min(a[2] - a[1], b[2] - b[1]);

        // TODO: check if residue types match?
        for (let i = 0; i < l; i++) {
            let eI = (a[1] + i) as ElementIndex;
            xs.x[o] = a[0].conformation.x(eI);
            xs.y[o] = a[0].conformation.y(eI);
            xs.z[o] = a[0].conformation.z(eI);

            eI = (b[1] + i) as ElementIndex;
            ys.x[o] = b[0].conformation.x(eI);
            ys.y[o] = b[0].conformation.y(eI);
            ys.z[o] = b[0].conformation.z(eI);
            o++;
        }
    }

    if (traceOnly) {
        xs.x = xs.x.slice(0, o);
        xs.y = xs.y.slice(0, o);
        xs.z = xs.z.slice(0, o);
        ys.x = ys.x.slice(0, o);
        ys.y = ys.y.slice(0, o);
        ys.z = ys.z.slice(0, o);
    }
    return [xs, ys];
}

function findTraceAtom(unit: Unit.Atomic, start: ElementIndex, end: ElementIndex): [Unit.Atomic, ElementIndex, ElementIndex] | undefined {
    for (let i = start; i < end; i++) {
        const l = unit.model.atomicHierarchy.atoms.label_atom_id.value(i);
        if (l === 'CA' || l === 'BB') {
            return [unit, i, i + 1 as ElementIndex];
        }
    }
}

function findPairs(N: number, index: IndexEntry[]) {
    const pairwiseCounts: number[][] = [];
    for (let i = 0; i < N; i++) {
        pairwiseCounts[i] = [];
        for (let j = 0; j < N; j++) pairwiseCounts[i][j] = 0;
    }

    for (const { pivots } of index) {
        for (let i = 0; i < N; i++) {
            if (!pivots[i]) continue;

            const lI = pivots[i]![2] - pivots[i]![1];

            for (let j = i + 1; j < N; j++) {
                if (!pivots[j]) continue;

                const lJ = pivots[j]![2] - pivots[j]![1];
{'fraction_non_alphanumeric': 0.1747844042268918, 'fraction_numerical': 0.010202842220332808, 'mean_word_length': 2.2290196078431372, 'pattern_counts': {'":': 0, '<': 30, '<?xml version=': 0, '>': 17, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '830299', 'n_tokens_mistral': 2970, 'n_tokens_neox': 2818, 'n_words': 690}