Skip to main content

data_beans/aux/
feature_types.rs

1//! The typed feature table that rides beside a mixed-type embedding:
2//! `{prefix}.feature_types.parquet`, string columns `feature` and `type`, one
3//! row per embedding row in the same order. `senna fne` and `gene-text` write
4//! it; a consumer that wants only one type of row (`senna bge` pinning gene
5//! rows, say) reads it. The type names are the shared vocabulary.
6
7use legume_numeric::matrix::parquet::{
8    read_parquet_string_columns_by_name, write_named_table, Column,
9};
10use std::path::Path;
11
12/// Nodes whose names are canonicalised as gene symbols.
13pub const GENE_TYPE: &str = "gene";
14/// Ontology terms and gene sets.
15pub const TERM_TYPE: &str = "term";
16/// Fixed genomic windows.
17pub const REGION_TYPE: &str = "region";
18/// Vocabulary words of a text relation.
19pub const WORD_TYPE: &str = "word";
20
21/// Whether rows of type `t` name a data feature — a gene, or a genomic window
22/// a peak can match — rather than a term, word or cell type.
23#[must_use]
24pub fn is_data_feature_type(t: &str) -> bool {
25    matches!(t, GENE_TYPE | REGION_TYPE)
26}
27
28pub fn feature_types_path(prefix: &str) -> String {
29    format!("{prefix}.feature_types.parquet")
30}
31
32/// Write the table for `names[i]` of type `types[i]`.
33pub fn write_feature_types(
34    prefix: &str,
35    names: &[Box<str>],
36    types: &[Box<str>],
37) -> anyhow::Result<()> {
38    anyhow::ensure!(
39        names.len() == types.len(),
40        "feature types: {} names for {} types",
41        names.len(),
42        types.len()
43    );
44    write_named_table(
45        &feature_types_path(prefix),
46        "feature",
47        names,
48        &[(Box::from("type"), Column::Str(types))],
49    )
50}
51
52/// One flag per row of a table (`names`, in order): true for a gene or a
53/// genomic window, by the table's types table `types`. `None` when `types`
54/// does not list `names` in order — written for another table under the
55/// same prefix. For [`crate::aux::frozen_features::load_frozen_feature_host_matching`].
56#[must_use]
57pub fn feature_rows(types: &[FeatureType], names: &[Box<str>]) -> Option<Vec<bool>> {
58    if !types.iter().map(|(n, _)| n).eq(names.iter()) {
59        return None;
60    }
61    Some(types.iter().map(|(_, t)| is_data_feature_type(t)).collect())
62}
63
64/// One row of the table: the feature's name and its type.
65pub type FeatureType = (Box<str>, Box<str>);
66
67/// The run's rows; `None` when the run wrote no table, which a caller reads as
68/// "every row is of the one type it expects".
69pub fn read_feature_types(prefix: &str) -> anyhow::Result<Option<Vec<FeatureType>>> {
70    let path = feature_types_path(prefix);
71    if !Path::new(&path).exists() {
72        return Ok(None);
73    }
74    let mut cols = read_parquet_string_columns_by_name(&path, &["feature", "type"])?;
75    let types = cols.pop().expect("two columns requested");
76    let names = cols.pop().expect("two columns requested");
77    Ok(Some(names.into_iter().zip(types).collect()))
78}
79
80#[cfg(test)]
81mod tests {
82    use super::*;
83
84    #[test]
85    fn feature_rows_are_marked_by_position_for_the_table_listed() {
86        let t = |n: &str, ty: &str| -> FeatureType { (n.into(), ty.into()) };
87        let types = [
88            t("CD4", "cell_type"),
89            t("CD4", "gene"),
90            t("chr1:0-5000", "region"),
91        ];
92        let names: Vec<Box<str>> = vec!["CD4".into(), "CD4".into(), "chr1:0-5000".into()];
93        assert_eq!(feature_rows(&types, &names), Some(vec![false, true, true]));
94        assert_eq!(feature_rows(&types, &names[..2]), None);
95        let other: Vec<Box<str>> = vec!["CD4".into(), "MYC".into(), "chr1:0-5000".into()];
96        assert_eq!(feature_rows(&types, &other), None);
97        assert!(is_data_feature_type(GENE_TYPE) && !is_data_feature_type(WORD_TYPE));
98    }
99
100    #[test]
101    fn round_trip_and_absence() {
102        let dir = tempfile::tempdir().unwrap();
103        let prefix = dir.path().join("run").to_string_lossy().into_owned();
104        assert!(read_feature_types(&prefix).unwrap().is_none());
105        let names: Vec<Box<str>> = vec!["TP53".into(), "GO:1".into()];
106        let types: Vec<Box<str>> = vec![GENE_TYPE.into(), TERM_TYPE.into()];
107        write_feature_types(&prefix, &names, &types).unwrap();
108        let rows = read_feature_types(&prefix).unwrap().unwrap();
109        assert_eq!(
110            rows,
111            vec![
112                ("TP53".into(), GENE_TYPE.into()),
113                ("GO:1".into(), TERM_TYPE.into())
114            ]
115        );
116        assert!(write_feature_types(&prefix, &names, &types[..1]).is_err());
117    }
118}