gecco 0.4.0

Gene Cluster prediction with Conditional Random Fields
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117

//! Supervised classifier to predict the type of a gene cluster.

pub mod backend;

use std::collections::BTreeSet;

use anyhow::Result;

use crate::model::{Cluster, ClusterType};

/// A binarizer that converts between ClusterType and binary vectors.
pub struct TypeBinarizer {
    pub classes: Vec<String>,
}

impl TypeBinarizer {
    pub fn new(classes: Vec<String>) -> Self {
        Self { classes }
    }

    /// Convert a list of ClusterType into a binary matrix.
    pub fn transform(&self, types: &[ClusterType]) -> Vec<Vec<f64>> {
        types
            .iter()
            .map(|ct| {
                self.classes
                    .iter()
                    .map(|cls| if ct.names.contains(cls) { 1.0 } else { 0.0 })
                    .collect()
            })
            .collect()
    }

    /// Convert a binary matrix back to ClusterType instances.
    pub fn inverse_transform(&self, matrix: &[Vec<bool>]) -> Vec<ClusterType> {
        matrix
            .iter()
            .map(|row| {
                let names: BTreeSet<String> = row
                    .iter()
                    .zip(self.classes.iter())
                    .filter(|(&flag, _)| flag)
                    .map(|(_, cls)| cls.clone())
                    .collect();
                ClusterType { names }
            })
            .collect()
    }
}

/// Trait for a random forest backend used for type classification.
pub trait RandomForestModel: Send + Sync {
    /// Fit the model on training data.
    fn fit(&mut self, x: &[Vec<f64>], y: &[Vec<f64>]) -> Result<()>;
    /// Predict class probabilities for each sample.
    /// Returns shape [n_samples][n_classes] probabilities.
    fn predict_proba(&self, x: &[Vec<f64>]) -> Result<Vec<Vec<f64>>>;
}

/// Classifier that predicts cluster biosynthetic types.
pub struct TypeClassifier {
    pub binarizer: TypeBinarizer,
    pub domains: Vec<String>,
    model: Option<Box<dyn RandomForestModel>>,
}

impl TypeClassifier {
    pub fn new(classes: Vec<String>) -> Self {
        Self {
            binarizer: TypeBinarizer::new(classes),
            domains: Vec::new(),
            model: None,
        }
    }

    pub fn set_model(&mut self, model: Box<dyn RandomForestModel>) {
        self.model = Some(model);
    }

    pub fn set_domains(&mut self, domains: Vec<String>) {
        self.domains = domains;
    }

    /// Predict types for given clusters, mutating them in place.
    pub fn predict_types(&self, clusters: &mut [Cluster]) -> Result<()> {
        let model = self
            .model
            .as_ref()
            .ok_or_else(|| anyhow::anyhow!("TypeClassifier model not set"))?;

        // Compute domain compositions
        let comps: Vec<Vec<f64>> = clusters
            .iter()
            .map(|c| c.domain_composition(Some(&self.domains), true, false, true))
            .collect();

        // Predict probabilities
        let probas = model.predict_proba(&comps)?;

        // Assign types to clusters
        for (cluster, proba) in clusters.iter_mut().zip(probas.iter()) {
            // Threshold at 0.5
            let type_flags: Vec<bool> = proba.iter().map(|&p| p > 0.5).collect();
            let types = self.binarizer.inverse_transform(&[type_flags]);
            cluster.cluster_type = types.into_iter().next();
            cluster.type_probabilities = self
                .binarizer
                .classes
                .iter()
                .zip(proba.iter())
                .map(|(cls, &p)| (cls.clone(), p))
                .collect();
        }

        Ok(())
    }
}