rustytorch/demos/rtx-alphafold-demo/src/sample_data.rs

//! Sample protein data for the AlphaFold-Lite demo.

use alphafold_shared::{SampleProtein, get_sample_proteins};

/// Get all sample proteins for the demo.
#[must_use]
pub fn list_samples() -> Vec<SampleProtein> {
    get_sample_proteins()
}

/// Get a sample protein by name.
#[must_use]
pub fn get_sample_by_name(name: &str) -> Option<SampleProtein> {
    get_sample_proteins()
        .into_iter()
        .find(|p| p.name.to_lowercase().contains(&name.to_lowercase()))
}

/// Get a sample protein by UniProt/PDB ID.
#[must_use]
pub fn get_sample_by_id(id: &str) -> Option<SampleProtein> {
    get_sample_proteins()
        .into_iter()
        .find(|p| p.id.to_uppercase() == id.to_uppercase())
}

/// Generate a random peptide sequence.
#[must_use]
pub fn generate_random_peptide(length: usize, seed: u64) -> String {
    use rand::SeedableRng;
    use rand::prelude::IndexedRandom;

    let amino_acids = [
        'A', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'K', 'L', 'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'V',
        'W', 'Y',
    ];

    let mut rng = rand::rngs::StdRng::seed_from_u64(seed);
    (0..length)
        .map(|_| *amino_acids.choose(&mut rng).unwrap())
        .collect()
}

/// Generate a helical peptide (helix-favoring residues).
#[must_use]
pub fn generate_helical_peptide(length: usize) -> String {
    // Helix-favoring: A, E, L, M, K, R
    let helix_residues = ['A', 'E', 'L', 'M', 'K', 'R'];
    use rand::SeedableRng;
    use rand::prelude::IndexedRandom;

    let mut rng = rand::rngs::StdRng::seed_from_u64(42);
    (0..length)
        .map(|_| *helix_residues.choose(&mut rng).unwrap())
        .collect()
}

/// Generate a beta-sheet peptide (strand-favoring residues).
#[must_use]
pub fn generate_strand_peptide(length: usize) -> String {
    // Strand-favoring: V, I, Y, F, W, T
    let strand_residues = ['V', 'I', 'Y', 'F', 'W', 'T'];
    use rand::SeedableRng;
    use rand::prelude::IndexedRandom;

    let mut rng = rand::rngs::StdRng::seed_from_u64(43);
    (0..length)
        .map(|_| *strand_residues.choose(&mut rng).unwrap())
        .collect()
}

/// Well-known protein structures for reference.
#[derive(Debug, Clone)]
pub struct ReferenceProtein {
    /// Protein name
    pub name: &'static str,
    /// PDB ID
    pub pdb_id: &'static str,
    /// Sequence
    pub sequence: &'static str,
    /// Description
    pub description: &'static str,
    /// Resolution (Angstroms)
    pub resolution: f32,
    /// Number of chains
    pub num_chains: usize,
}

/// Get well-known reference proteins from PDB.
#[must_use]
pub fn get_reference_proteins() -> Vec<ReferenceProtein> {
    vec![
        ReferenceProtein {
            name: "Crambin",
            pdb_id: "1CRN",
            sequence: "TTCCPSIVARSNFNVCRLPGTPEAICATYTGCIIIPGATCPGDYAN",
            description: "Small, well-characterized plant protein",
            resolution: 0.54,
            num_chains: 1,
        },
        ReferenceProtein {
            name: "Rubredoxin",
            pdb_id: "1IRO",
            sequence: "MKKYVCTVCGYEYDPAEGDPDNGVKPGTSFDDLPADWVCPVCGAPKSEFERVED",
            description: "Iron-sulfur protein",
            resolution: 0.95,
            num_chains: 1,
        },
        ReferenceProtein {
            name: "Trp-cage",
            pdb_id: "1L2Y",
            sequence: "NLYIQWLKDGGPSSGRPPPS",
            description: "Designed miniprotein, one of smallest folded proteins",
            resolution: 0.0, // NMR
            num_chains: 1,
        },
        ReferenceProtein {
            name: "Villin headpiece",
            pdb_id: "1VII",
            sequence: "LSDEDFKAVFGMTRSAFANLPLWKQQNLKKEKGLF",
            description: "Fast-folding protein domain",
            resolution: 0.0, // NMR
            num_chains: 1,
        },
        ReferenceProtein {
            name: "Chignolin",
            pdb_id: "1UAO",
            sequence: "GYDPETGTWG",
            description: "Designed 10-residue beta-hairpin",
            resolution: 0.0, // NMR
            num_chains: 1,
        },
    ]
}

/// Get motif sequences for testing.
#[must_use]
pub fn get_motif_sequences() -> Vec<(&'static str, &'static str)> {
    vec![
        ("Helix-Turn-Helix", "AAAAAEELLLLLLKKKPPPPGGGGAAAAAELLLLLKKK"),
        ("Beta-Hairpin", "VVVVVYYYYYGGGGPPPVVVVVYYYYYGGGG"),
        ("Alpha-Beta", "AAAAAAELLLLMMMVVVVVIIIIYYYYAAAAELL"),
        ("Coiled-Coil", "LEELKKKLEELKKKLEELKKKLEELKKK"),
        ("Zinc Finger", "CPVCGKAFRSQHLGIHQRSH"),
    ]
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_list_samples() {
        let samples = list_samples();
        assert!(!samples.is_empty());
    }

    #[test]
    fn test_get_sample_by_name() {
        let insulin = get_sample_by_name("Insulin");
        assert!(insulin.is_some());
        assert_eq!(insulin.unwrap().name, "Insulin");
    }

    #[test]
    fn test_get_sample_by_id() {
        let gfp = get_sample_by_id("1EMA");
        assert!(gfp.is_some());
        assert!(gfp.unwrap().name.contains("GFP"));
    }

    #[test]
    fn test_generate_random_peptide() {
        let peptide = generate_random_peptide(20, 42);
        assert_eq!(peptide.len(), 20);
        assert!(peptide.chars().all(|c| c.is_ascii_uppercase()));
    }

    #[test]
    fn test_generate_helical_peptide() {
        let peptide = generate_helical_peptide(30);
        assert_eq!(peptide.len(), 30);
        // Should only contain helix-favoring residues
        assert!(peptide.chars().all(|c| "AELMKR".contains(c)));
    }

    #[test]
    fn test_generate_strand_peptide() {
        let peptide = generate_strand_peptide(30);
        assert_eq!(peptide.len(), 30);
        // Should only contain strand-favoring residues
        assert!(peptide.chars().all(|c| "VIYFWT".contains(c)));
    }

    #[test]
    fn test_reference_proteins() {
        let refs = get_reference_proteins();
        assert!(!refs.is_empty());

        for protein in &refs {
            assert!(!protein.sequence.is_empty());
            assert!(!protein.pdb_id.is_empty());
        }
    }

    #[test]
    fn test_motif_sequences() {
        let motifs = get_motif_sequences();
        assert!(!motifs.is_empty());

        for (name, seq) in &motifs {
            assert!(!name.is_empty());
            assert!(!seq.is_empty());
        }
    }
}