<commit_before>import { ResidueIndex } from 'mol-model/structure';
import { MoleculeType } from 'mol-model/structure/model/types';
import { GridLookup3D } from 'mol-math/geometry';
import { Vec3 } from 'mol-math/linear-algebra';
import { SortedArray } from 'mol-data/int';
import { ElementIndex } from 'mol-model/structure/model/indexing';
import { AtomicHierarchy, AtomicConformation } from '../atomic';

export function computeModelASA(hierarchy: AtomicHierarchy, conformation: AtomicConformation) {
    const { lookup3d, proteinResidues } = calcAtomicTraceLookup3D(hierarchy, conformation)
    const backboneIndices = calcBackboneAtomIndices(hierarchy, proteinResidues)

    const residueCount = proteinResidues.length
    const numberOfSpherePoints = 960

    const ctx: ASAContext = {
        probeSize: 1.4,
        spherePoints: generateSpherePoints(numberOfSpherePoints),
        cons: 4.0 * Math.PI / numberOfSpherePoints,

        hierarchy,
        proteinResidues
   }

    calculateASA(ctx)
}

function calculateASA(ctx: ASAContext) {

}

function generateSpherePoints(numberOfSpherePoints: number): Vec3[] {
    const points: Vec3[] = []
    const inc = Math.PI

    return points
}

interface ASAContext {
    probeSize: number,
    spherePoints: Vec3[],
    cons: number,

    hierarchy: AtomicHierarchy,
    proteinResidues: SortedArray<ResidueIndex>
}

interface BackboneAtomIndices {
    cIndices: ArrayLike<ElementIndex | -1>
    hIndices: ArrayLike<ElementIndex | -1>
    oIndices: ArrayLike<ElementIndex | -1>
    nIndices: ArrayLike<ElementIndex | -1>
}

function calcAtomicTraceLookup3D(hierarchy: AtomicHierarchy, conformation: AtomicConformation) {
    const { x, y, z } = conformation;
    const { moleculeType, traceElementIndex } = hierarchy.derived.residue
    const indices: number[] = []
    const _proteinResidues: number[] = []
    for (let i = 0, il = moleculeType.length; i < il; ++i) {
        if (moleculeType[i] === MoleculeType.protein) {
            indices[indices.length] = traceElementIndex[i]
            _proteinResidues[_proteinResidues.length] = i
        }
    }
    const lookup3d = GridLookup3D({ x, y, z, indices: SortedArray.ofSortedArray(indices) }, 4);
    const proteinResidues = SortedArray.ofSortedArray<ResidueIndex>(_proteinResidues)
    return { lookup3d, proteinResidues }
}


function calcBackboneAtomIndices(hierarchy: AtomicHierarchy, proteinResidues: SortedArray<ResidueIndex>): BackboneAtomIndices {
    const residueCount = proteinResidues.length
    const { index } = hierarchy

    const c = new Int32Array(residueCount)
    const h = new Int32Array(residueCount)
    const o = new Int32Array(residueCount)
    const n = new Int32Array(residueCount)

    for (let i = 0, il = residueCount; i < il; ++i) {
        const rI = proteinResidues[i]
        c[i] = index.findAtomOnResidue(rI, 'C')
        h[i] = index.findAtomOnResidue(rI, 'H')
        o[i] = index.findAtomOnResidue(rI, 'O')
        n[i] = index.findAtomOnResidue(rI, 'N')
    }

    return {
        cIndices: c as unknown as ArrayLike<ElementIndex | -1>,
        hIndices: h as unknown as ArrayLike<ElementIndex | -1>,
        oIndices: o as unknown as ArrayLike<ElementIndex | -1>,
        nIndices: n as unknown as ArrayLike<ElementIndex | -1>,
    }
}<commit_msg>computation of atom radii<commit_after>import { Vec3 } from 'mol-math/linear-algebra';
import { AtomicHierarchy, AtomicConformation } from '../atomic';
import { ElementSymbol, VdwRadii } from '../../types';
import { skipUntil } from 'rxjs/operators';
import { atomicHet } from 'mol-model/structure/query/queries/internal';
import { compile } from 'mol-script/runtime/query/compiler';

const defaultNumberOfPoints = 960;
const defaultProbeSize = 1.4;
const trigonalCarbonVdw = 1.76;
const tetrahedralCarbonVdw = 1.87;
const trigonalNitrogenVdw = 1.65;
const tetrahedralNitrogenVdw = 1.50;
/** deviating radii from the definition in types.ts */
const oxygenVdw = 1.4;
const sulfurVdw = 1.85;

export function computeModelASA(hierarchy: AtomicHierarchy, conformation: AtomicConformation) {
    const numberOfSpherePoints = defaultNumberOfPoints;

    const ctx: ASAContext = {
        probeSize: defaultProbeSize,
        spherePoints: generateSpherePoints(numberOfSpherePoints),
        cons: 4.0 * Math.PI / numberOfSpherePoints,

        hierarchy,
        conformation
    }

    calculateASA(ctx);
}

function calculateASA(ctx: ASAContext) {
    console.log(ctx.hierarchy)
    // 1. extract all heavy atoms
    // 2. assign radius to all heavy atoms - depends on element
    // 3. for each residue
    // 3a. calculate the individual ASA of each atom
    // 3b. sum up
    // 3c. (optionally) normalize by maximum value expected for a particular amino acid - needs lookup of max values

    // const { type_symbol } = ctx.hierarchy.atoms;
    // console.log(type_symbol.value(0));

    const { atoms, residues, residueAtomSegments, derived } = ctx.hierarchy;
    const { type_symbol } = atoms;
    const radii: number[] = [];

    // TODO can be more elegant by obtaining residue info once and then using offset to navigate to next residue
    for (let i = 0; i < type_symbol.rowCount; ++i) {
        // skip hydrogen atoms
        if (type_symbol.value(i) === 'H' || type_symbol.value(i) === 'D' || type_symbol.value(i) === 'T')
            continue;

        // determine group this atom belongs to
        const groupIndex = residueAtomSegments.index[i];

        // skip entities not part of a peptide chain
        if (derived.residue.moleculeType[groupIndex] !== 4)
            continue;

        const atomId = atoms.label_atom_id.value(i);
        const compId = residues.label_comp_id.value(groupIndex);
        const element = type_symbol.value(i);
        radii[radii.length] = determineRadius(atomId, element, compId);
    }
}

/**
 * Gets the van der Waals radius of the given atom following the values defined by Chothia (1976)
 * J.Mol.Biol.105,1-14. NOTE: the vdw values defined by the paper assume no Hydrogens and thus "inflates" slightly
 * the heavy atoms to account for Hydrogens. Thus this method cannot be used in a structure that contains Hydrogens!
 */
function determineRadius(atomId: string, element: ElementSymbol, compId: string): number {
    switch (element) {
        case 'O':
        return oxygenVdw;
        case 'S':
        return sulfurVdw;
        case 'N':
        return atomId === 'NZ' ? tetrahedralNitrogenVdw : trigonalNitrogenVdw;
        case 'C':
        if (atomId === 'C' || atomId === 'CE1' || atomId === 'CE2' || atomId === 'CE3' ||
            atomId === 'CH2' || atomId === 'CZ' || atomId === 'CZ2' || atomId === 'CZ3') {
            return trigonalCarbonVdw;
        } else if (atomId === 'CA' || atomId === 'CB' || atomId === 'CE' || atomId === 'CG1' ||
            atomId === 'CG2') {
            return tetrahedralCarbonVdw;
        } else if (compId === 'PHE' || compId === 'TRP' || compId === 'TYR' || compId === 'HIS' ||
            compId === 'ASP' || compId === 'ASN') {
            return trigonalCarbonVdw;
        } else if (compId === 'PRO' || compId === 'LYS' || compId === 'ARG' || compId === 'MET' ||
            compId === 'ILE' || compId === 'LEU') {
            return tetrahedralCarbonVdw;
        } else if (compId === 'GLU' || compId === 'GLN') {
            return atomId === 'CD' ? trigonalCarbonVdw : tetrahedralCarbonVdw;
        }
        default:
        return (VdwRadii as any)[atomId];
    }
}

/** Creates a collection of points on a sphere by the Golden Section Spiral algorithm. */
function generateSpherePoints(numberOfSpherePoints: number): Vec3[] {
    const points: Vec3[] = [];
    const inc = Math.PI * (3.0 - Math.sqrt(5.0));
    const offset = 2.0 / numberOfSpherePoints;
    for (let k = 0; k < numberOfSpherePoints; ++k) {
        const y = k * offset - 1.0 + (offset / 2.0);
        const r = Math.sqrt(1.0 - y * y);
        const phi = k * inc;
        points[points.length] = [Math.cos(phi), y, Math.sin(phi) * r] as Vec3;
    }
    return points;
}

interface ASAContext {
    probeSize: number,
    spherePoints: Vec3[],
    cons: number,

    hierarchy: AtomicHierarchy,
    conformation: AtomicConformation
}
{'fraction_non_alphanumeric': 0.12099817407181984, 'fraction_numerical': 0.012903225806451613, 'mean_word_length': 3.1162324649298596, 'pattern_counts': {'":': 0, '<': 18, '<?xml version=': 0, '>': 14, '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': '7322988', 'n_tokens_mistral': 2718, 'n_tokens_neox': 2577, 'n_words': 693}