prodigal_rs/api/
meta_predictor.rs

1//! Reusable metagenomic gene predictor.
2//!
3//! Caches the 50 metagenomic models and uses a rayon thread pool
4//! with a large stack for prediction.
5
6use std::alloc::{alloc_zeroed, handle_alloc_error, Layout};
7use std::os::raw::c_int;
8use std::sync::Arc;
9
10use rayon::prelude::*;
11use rayon::ThreadPool;
12
13use super::convert::gene_to_predicted;
14use super::encode::SequenceBuffer;
15use super::types::{PredictedGene, ProdigalConfig, ProdigalError};
16use crate::types::{Gene, Training, MAX_GENES, MAX_SEQ, NUM_META};
17
18use super::predict::validate_config;
19
20/// Recommended stack size for Rayon pools used by `MetaPredictor`.
21///
22/// The low-level translated prediction routines use deep call stacks, so the
23/// default Rust thread stack can be too small for worker threads.
24pub const META_PREDICTOR_STACK_SIZE: usize = 32 * 1024 * 1024; // 32 MB
25
26use crate::dprog::{dprog, eliminate_bad_genes};
27use crate::gene::{add_genes, record_gene_data, tweak_final_starts};
28use crate::node::{add_nodes, record_overlapping_starts, reset_node_scores, score_nodes};
29
30/// Reusable metagenomic gene predictor.
31///
32/// Pre-loads 50 metagenomic models and evaluates qualifying models with
33/// Prodigal-compatible state preservation across adjacent model bins.
34pub struct MetaPredictor {
35    pool: Arc<ThreadPool>,
36    models: Arc<Vec<Box<Training>>>,
37    config: ProdigalConfig,
38}
39
40impl MetaPredictor {
41    /// Create a new predictor with default config.
42    pub fn new() -> Result<Self, ProdigalError> {
43        Self::with_config(ProdigalConfig::default())
44    }
45
46    /// Create a new predictor with custom config.
47    pub fn with_config(config: ProdigalConfig) -> Result<Self, ProdigalError> {
48        validate_config(&config)?;
49
50        let pool = rayon::ThreadPoolBuilder::new()
51            .stack_size(META_PREDICTOR_STACK_SIZE)
52            .build()
53            .map_err(|e| {
54                ProdigalError::Io(std::io::Error::new(
55                    std::io::ErrorKind::Other,
56                    e.to_string(),
57                ))
58            })?;
59
60        Self::with_config_and_thread_pool(config, Arc::new(pool))
61    }
62
63    /// Create a new predictor with default config and a shared Rayon thread pool.
64    ///
65    /// The supplied pool is reused for parallel metagenomic model evaluation.
66    /// For large inputs, build the pool with at least `META_PREDICTOR_STACK_SIZE`
67    /// stack bytes per worker thread.
68    pub fn with_thread_pool(pool: Arc<ThreadPool>) -> Result<Self, ProdigalError> {
69        Self::with_config_and_thread_pool(ProdigalConfig::default(), pool)
70    }
71
72    /// Create a new predictor with custom config and a shared Rayon thread pool.
73    ///
74    /// The supplied pool is reused for parallel metagenomic model evaluation.
75    /// For large inputs, build the pool with at least `META_PREDICTOR_STACK_SIZE`
76    /// stack bytes per worker thread.
77    pub fn with_config_and_thread_pool(
78        config: ProdigalConfig,
79        pool: Arc<ThreadPool>,
80    ) -> Result<Self, ProdigalError> {
81        validate_config(&config)?;
82        let models = Arc::new(load_meta_models());
83
84        Ok(MetaPredictor {
85            pool,
86            models,
87            config,
88        })
89    }
90
91    /// Predict genes in the given sequence.
92    pub fn predict(&self, seq: &[u8]) -> Result<Vec<PredictedGene>, ProdigalError> {
93        validate_sequence(seq)?;
94
95        self.pool
96            .install(|| predict_parallel(seq, &self.models, &self.config))
97    }
98
99    /// Predict genes in a batch of sequences, preserving input order.
100    pub fn predict_batch<S>(&self, seqs: &[S]) -> Result<Vec<Vec<PredictedGene>>, ProdigalError>
101    where
102        S: AsRef<[u8]> + Sync,
103    {
104        for seq in seqs {
105            validate_sequence(seq.as_ref())?;
106        }
107
108        self.pool.install(|| {
109            seqs.par_iter()
110                .map(|seq| predict_parallel(seq.as_ref(), &self.models, &self.config))
111                .collect()
112        })
113    }
114}
115
116/// Reject empty or oversized input sequences before encoding.
117fn validate_sequence(seq: &[u8]) -> Result<(), ProdigalError> {
118    if seq.is_empty() {
119        return Err(ProdigalError::EmptySequence);
120    }
121    if seq.len() > MAX_SEQ {
122        return Err(ProdigalError::SequenceTooLong {
123            length: seq.len(),
124            max: MAX_SEQ,
125        });
126    }
127    Ok(())
128}
129
130/// Allocate and initialize all `NUM_META` pre-trained metagenomic models on the heap.
131fn load_meta_models() -> Vec<Box<Training>> {
132    let mut models: Vec<Box<Training>> = Vec::with_capacity(NUM_META);
133    for i in 0..NUM_META {
134        unsafe {
135            let layout = Layout::new::<Training>();
136            let ptr = alloc_zeroed(layout) as *mut Training;
137            if ptr.is_null() {
138                handle_alloc_error(layout);
139            }
140            crate::training_data::load_metagenome(i, ptr);
141            models.push(Box::from_raw(ptr));
142        }
143    }
144    models
145}
146
147/// Run the metagenomic gene-prediction pipeline on a single sequence.
148///
149/// Encodes the input, rebuilds node arrays when the translation table changes,
150/// scores all GC-compatible models, keeps the highest-scoring solution, and
151/// converts its genes to `PredictedGene` values.
152fn predict_parallel(
153    seq: &[u8],
154    models: &[Box<Training>],
155    config: &ProdigalConfig,
156) -> Result<Vec<PredictedGene>, ProdigalError> {
157    let closed = if config.closed_ends { 1 } else { 0 };
158
159    let mut buf = SequenceBuffer::with_base_capacity(seq.len());
160    let (slen, gc) = unsafe { buf.encode(seq, config.mask_n_runs) };
161    if slen == 0 {
162        return Err(ProdigalError::EmptySequence);
163    }
164    buf.ensure_node_capacity(slen);
165
166    // GC window for model selection
167    let mut low = 0.88495 * gc - 0.0102337;
168    if low > 0.65 {
169        low = 0.65;
170    }
171    let mut high = 0.86596 * gc + 0.1131991;
172    if high < 0.35 {
173        high = 0.35;
174    }
175
176    let mut best_score = f64::NEG_INFINITY;
177    let mut best_nodes = Vec::new();
178    let mut best_genes: Vec<Gene> = Vec::new();
179    let mut best_tinf: Option<Training> = None;
180    let mut nn: c_int = 0;
181
182    unsafe {
183        for i in 0..NUM_META {
184            let need_rebuild = i == 0 || models[i].trans_table != models[i - 1].trans_table;
185            let mut tinf = (*models[i]).clone();
186
187            if need_rebuild {
188                buf.clear_nodes(nn);
189                nn = add_nodes(
190                    buf.seq.as_mut_ptr(),
191                    buf.rseq.as_mut_ptr(),
192                    slen,
193                    buf.nodes.as_mut_ptr(),
194                    closed,
195                    buf.masks.as_mut_ptr(),
196                    buf.nmask,
197                    &mut tinf,
198                );
199                buf.nodes[..nn as usize]
200                    .sort_unstable_by(|a, b| a.ndx.cmp(&b.ndx).then(b.strand.cmp(&a.strand)));
201            }
202
203            if tinf.gc < low || tinf.gc > high {
204                continue;
205            }
206
207            reset_node_scores(buf.nodes.as_mut_ptr(), nn);
208            score_nodes(
209                buf.seq.as_mut_ptr(),
210                buf.rseq.as_mut_ptr(),
211                slen,
212                buf.nodes.as_mut_ptr(),
213                nn,
214                &mut tinf,
215                closed,
216                1,
217            );
218            record_overlapping_starts(buf.nodes.as_mut_ptr(), nn, &mut tinf, 1);
219            let ipath = dprog(buf.nodes.as_mut_ptr(), nn, &mut tinf, 1);
220            if ipath < 0 || ipath >= nn {
221                continue;
222            }
223
224            let score = buf.nodes[ipath as usize].score;
225            if score > best_score {
226                best_score = score;
227                eliminate_bad_genes(buf.nodes.as_mut_ptr(), ipath, &mut tinf);
228
229                let mut genes: Vec<Gene> = vec![std::mem::zeroed(); MAX_GENES];
230                let ng = add_genes(genes.as_mut_ptr(), buf.nodes.as_mut_ptr(), ipath);
231                tweak_final_starts(
232                    genes.as_mut_ptr(),
233                    ng,
234                    buf.nodes.as_mut_ptr(),
235                    nn,
236                    &mut tinf,
237                );
238                genes.truncate(ng as usize);
239
240                best_nodes = buf.nodes[..nn as usize].to_vec();
241                best_genes = genes;
242                best_tinf = Some(tinf);
243            }
244        }
245    }
246
247    let Some(mut tinf) = best_tinf else {
248        return Ok(Vec::new());
249    };
250
251    unsafe {
252        record_gene_data(
253            best_genes.as_mut_ptr(),
254            best_genes.len() as c_int,
255            best_nodes.as_mut_ptr(),
256            &mut tinf,
257            1,
258        );
259
260        let mut result = Vec::with_capacity(best_genes.len());
261        for gene in &best_genes {
262            result.push(gene_to_predicted(
263                gene,
264                best_nodes.as_ptr(),
265                &tinf,
266                slen as usize,
267            ));
268        }
269        Ok(result)
270    }
271}
prodigal_rs/api/meta_predictor.rs

prodigal_rs/api/
meta_predictor.rs
