data_beans/aux/
frozen_features.rs

1//! Load a pre-trained per-gene embedding (and optional per-gene bias)
2//! from parquet and strictly intersect its row axis against a caller's
3//! target feature axis.
4//!
5//! Used by `senna gbe / topic / cell-embedded-topic` to freeze the
6//! gene-side parameter table (`E_feat` in gbe, ρ in the ETM topic models)
7//! so cells train on a shared, pre-fit gene-relation space.
8//!
9//! Source formats supported via [`FrozenLoadArgs`]:
10//! - **gbe**: `{prefix}.dictionary.parquet` + `{prefix}.feature_bias.parquet`
11//!   (gene × H plus gene × 1 bias).
12//! - **topic / cell-embedded-topic**: `{prefix}.feature_embedding.parquet`
13//!   alone — bias defaults to zeros, which is what the topic models use
14//!   internally (no per-gene additive bias on ρ).
15//!
16//! Name resolution goes through [`FeatureNameKind`] so `TGFB1` and
17//! `ENSG00000105329_TGFB1` resolve to the same row.
18
19use crate::aux::feature_names::FeatureNameKind;
20use legume_numeric::matrix::traits::IoOps;
21use nalgebra::DMatrix;
22use rustc_hash::{FxHashMap, FxHashSet};
23
24/// Loaded + aligned frozen feature side ready to hand off to a candle
25/// engine (`graph-embedding-util` or one of the topic-model encoders).
26///
27/// Rows are reordered to follow the *target* axis. `keep_target_indices`
28/// records which positions in the caller's target feature axis survived
29/// the intersection — the caller MUST restrict its data (triplets,
30/// encoder D, decoder β) to these indices, otherwise the row order
31/// disagrees with the embedding rows.
32pub struct FrozenFeatureHost {
33    /// `[|keep|, H]`, rows in the same order as `keep_target_indices`.
34    pub e_feat: DMatrix<f32>,
35    /// `[|keep|]`. Zeros when no `bias_path` was given.
36    pub b_feat: Vec<f32>,
37    /// Indices into the *target* feature axis that matched a source row.
38    /// Length equals `e_feat.nrows()`.
39    pub keep_target_indices: Vec<usize>,
40    /// The *source* (dictionary) row each kept target row came from, parallel to
41    /// `keep_target_indices` — what a caller needs to look up any other table
42    /// keyed on the dictionary's rows (a module membership, say).
43    pub keep_src_indices: Vec<usize>,
44    /// The whole source table `[n_src, H]` and its row names, as read — so a
45    /// caller that needs the unmatched rows too (to place a gene by its
46    /// neighbours' rows) does not decode the file a second time.
47    pub src_e_feat: DMatrix<f32>,
48    pub src_names: Vec<Box<str>>,
49    /// Rows in the dictionary file, i.e. how many features the MODEL has.
50    ///
51    /// The only field that survives the intersection unfiltered, and the reason
52    /// it exists: `e_feat` and `keep_target_indices` are both already restricted
53    /// to the matched features, so a coverage fraction built from them is
54    /// identically 1 and tells a caller nothing. This is the denominator.
55    pub n_src: usize,
56    pub h: usize,
57}
58
59/// A rename of source row names, see [`FrozenLoadArgs::source_name_map`].
60pub type SourceNameMap<'a> = &'a dyn Fn(&str) -> Box<str>;
61
62pub struct FrozenLoadArgs<'a> {
63    /// Path to the `[D_src, H]` parquet (gbe `dictionary.parquet` or
64    /// topic `feature_embedding.parquet`). Row column 0 is the gene name.
65    pub dictionary_path: &'a str,
66    /// Optional path to a `[D_src, 1]` per-gene bias parquet (gbe
67    /// `feature_bias.parquet`). `None` → bias filled with zeros, which
68    /// matches the topic models' implicit "no per-gene bias on ρ".
69    pub bias_path: Option<&'a str>,
70    /// Caller's feature axis (e.g. `unified.feature_names` for gbe;
71    /// the topic models' `gene_names`). Output rows follow this order
72    /// after dropping unmatched entries.
73    pub target_feature_names: &'a [Box<str>],
74    /// Per-name canonicalization rule applied to both source and target
75    /// names before intersection. [`FeatureNameKind::Exact`] for strict
76    /// matching; [`FeatureNameKind::Gene { delim: '_' }`] is the typical
77    /// choice for scRNA gene IDs.
78    pub name_kind: FeatureNameKind,
79    /// Applied to every SOURCE row name before canonicalization, and kept as
80    /// the host's `src_names`: how a caller whose axis carries a row grammar
81    /// (`{gene}/count/spliced`) reads a plain gene table, lifting each bare
82    /// name into the grammar first. `None` = the names as read.
83    pub source_name_map: Option<SourceNameMap<'a>>,
84}
85
86pub fn load_frozen_feature_host(args: FrozenLoadArgs) -> anyhow::Result<FrozenFeatureHost> {
87    let dict = <DMatrix<f32> as IoOps>::from_parquet(args.dictionary_path)?;
88    let n_src = dict.rows.len();
89    let h = dict.mat.ncols();
90    anyhow::ensure!(
91        h > 0 && dict.mat.nrows() == n_src,
92        "{}: malformed dictionary (rows={}, mat dims={}x{})",
93        args.dictionary_path,
94        n_src,
95        dict.mat.nrows(),
96        h
97    );
98
99    let src_bias: Vec<f32> = match args.bias_path {
100        None => vec![0.0; n_src],
101        Some(p) => {
102            let bias = <DMatrix<f32> as IoOps>::from_parquet(p)?;
103            anyhow::ensure!(
104                bias.rows == dict.rows,
105                "{} row names disagree with {} (both files must come from the same training run)",
106                p,
107                args.dictionary_path
108            );
109            anyhow::ensure!(
110                bias.mat.ncols() == 1,
111                "{}: expected 1 data column (bias), got {}",
112                p,
113                bias.mat.ncols()
114            );
115            (0..n_src).map(|i| bias.mat[(i, 0)]).collect()
116        }
117    };
118
119    let src_names: Vec<Box<str>> = match args.source_name_map {
120        Some(f) => dict.rows.iter().map(|n| f(n)).collect(),
121        None => dict.rows,
122    };
123    let mut src_by_canon: FxHashMap<Box<str>, usize> = FxHashMap::default();
124    let mut src_dupes = 0usize;
125    for (i, name) in src_names.iter().enumerate() {
126        let canon = args.name_kind.canonicalize(name);
127        // First occurrence wins, as documented; `insert` would keep the last.
128        if let std::collections::hash_map::Entry::Vacant(e) = src_by_canon.entry(canon) {
129            e.insert(i);
130        } else {
131            src_dupes += 1;
132        }
133    }
134    if src_dupes > 0 {
135        log::warn!(
136            "{}: {} source rows had duplicate canonical names — kept first occurrence",
137            args.dictionary_path,
138            src_dupes
139        );
140    }
141
142    let mut keep_target_indices = Vec::new();
143    let mut keep_src_indices = Vec::new();
144    for (target_i, name) in args.target_feature_names.iter().enumerate() {
145        let canon = args.name_kind.canonicalize(name);
146        if let Some(&src_i) = src_by_canon.get(&canon) {
147            keep_target_indices.push(target_i);
148            keep_src_indices.push(src_i);
149        }
150    }
151    anyhow::ensure!(
152        !keep_target_indices.is_empty(),
153        "No feature names matched between {} (n={}) and target axis (n={}) under {:?} \
154         — check the gene-name kind (Exact / Gene / Locus / Mixed) and source axis",
155        args.dictionary_path,
156        n_src,
157        args.target_feature_names.len(),
158        args.name_kind
159    );
160
161    let unique_src_used: FxHashSet<usize> = keep_src_indices.iter().copied().collect();
162    // A dictionary is a plain gene table. Source rows carrying the channelized
163    // row grammar ({gene}/{modality}/... ) that matched nothing usually mean
164    // the caller fed a channelized or co-embedding artifact; a PARTIAL match
165    // would otherwise proceed silently on the plain-name subset.
166    let channelized_unmatched = src_names
167        .iter()
168        .enumerate()
169        .filter(|(i, r)| {
170            !unique_src_used.contains(i) && crate::aux::feature_rows::parse_feature_row(r).is_some()
171        })
172        .count();
173    if channelized_unmatched > 0 {
174        log::warn!(
175            "{}: {} unmatched source rows carry the channelized row grammar —              is this a raw gene dictionary, or a channelized/co-embedding output?",
176            args.dictionary_path,
177            channelized_unmatched
178        );
179    }
180    log::info!(
181        "Frozen feature side from {}: {}/{} target features matched (H={}, {} of {} source rows reused, kind={:?})",
182        args.dictionary_path,
183        keep_target_indices.len(),
184        args.target_feature_names.len(),
185        h,
186        unique_src_used.len(),
187        n_src,
188        args.name_kind
189    );
190
191    let k = keep_target_indices.len();
192    let mut e_feat = DMatrix::<f32>::zeros(k, h);
193    let mut b_feat = Vec::with_capacity(k);
194    for (out_i, &src_i) in keep_src_indices.iter().enumerate() {
195        for j in 0..h {
196            e_feat[(out_i, j)] = dict.mat[(src_i, j)];
197        }
198        b_feat.push(src_bias[src_i]);
199    }
200
201    Ok(FrozenFeatureHost {
202        e_feat,
203        b_feat,
204        keep_target_indices,
205        keep_src_indices,
206        src_e_feat: dict.mat,
207        src_names,
208        n_src,
209        h,
210    })
211}
212
213#[cfg(test)]
214mod tests {
215    use super::*;
216    use legume_numeric::matrix::traits::IoOps;
217
218    fn write_test_parquet(
219        path: &str,
220        rows: &[&str],
221        row_axis: &str,
222        cols: &[&str],
223        data: &DMatrix<f32>,
224    ) {
225        let row_names: Vec<Box<str>> = rows.iter().map(|s| (*s).into()).collect();
226        let col_names: Vec<Box<str>> = cols.iter().map(|s| (*s).into()).collect();
227        data.to_parquet_with_names(path, (Some(&row_names), Some(row_axis)), Some(&col_names))
228            .unwrap();
229    }
230
231    #[test]
232    fn strict_intersection_drops_unmatched_and_preserves_target_order() {
233        let dir = tempfile::tempdir().unwrap();
234        let dict_path = dir.path().join("d.parquet").to_str().unwrap().to_string();
235
236        // Source: 4 genes × H=3. Source row "ENSG_DROP" has no target match.
237        let src = DMatrix::<f32>::from_row_slice(
238            4,
239            3,
240            &[
241                1.0, 2.0, 3.0, // TGFB1
242                4.0, 5.0, 6.0, // MYC
243                7.0, 8.0, 9.0, // ENSG_DROP (unmatched)
244                10.0, 11.0, 12.0, // TP53
245            ],
246        );
247        write_test_parquet(
248            &dict_path,
249            &["TGFB1", "MYC", "ENSG_DROP", "TP53"],
250            "gene",
251            &["h0", "h1", "h2"],
252            &src,
253        );
254
255        // Target: 5 genes; "FOO" and "BAR" don't appear in source.
256        let target: Vec<Box<str>> = ["FOO", "TP53", "TGFB1", "BAR", "MYC"]
257            .iter()
258            .map(|s| (*s).into())
259            .collect();
260
261        let host = load_frozen_feature_host(FrozenLoadArgs {
262            dictionary_path: &dict_path,
263            bias_path: None,
264            target_feature_names: &target,
265            name_kind: FeatureNameKind::Exact,
266            source_name_map: None,
267        })
268        .unwrap();
269
270        // Kept target indices = positions of TP53, TGFB1, MYC in target order.
271        assert_eq!(host.keep_target_indices, vec![1, 2, 4]);
272        assert_eq!(host.h, 3);
273        assert_eq!(host.e_feat.nrows(), 3);
274        assert_eq!(host.b_feat, vec![0.0, 0.0, 0.0]);
275
276        // Row 0 of e_feat should be source row for TP53 (= source row 3).
277        assert_eq!(host.e_feat[(0, 0)], 10.0);
278        assert_eq!(host.e_feat[(0, 2)], 12.0);
279        // Row 1: TGFB1 → source row 0.
280        assert_eq!(host.e_feat[(1, 0)], 1.0);
281        // Row 2: MYC → source row 1.
282        assert_eq!(host.e_feat[(2, 1)], 5.0);
283    }
284
285    /// A source of bare gene names read onto an axis that carries the row
286    /// grammar: the map lifts each source name into the grammar before the
287    /// canonical match, and the host reports the lifted names.
288    #[test]
289    fn a_source_name_map_is_applied_before_matching_and_kept_in_src_names() {
290        let dir = tempfile::tempdir().unwrap();
291        let dict_path = dir.path().join("d.parquet").to_str().unwrap().to_string();
292        let src = DMatrix::<f32>::from_row_slice(2, 2, &[1.0, 2.0, 3.0, 4.0]);
293        write_test_parquet(
294            &dict_path,
295            &["TGFB1", "MYC/count/unspliced"],
296            "gene",
297            &["h0", "h1"],
298            &src,
299        );
300        let target: Vec<Box<str>> = [
301            "ENSG_TGFB1/count/spliced",
302            "ENSG_MYC/count/spliced",
303            "ENSG_MYC/count/unspliced",
304        ]
305        .iter()
306        .map(|s| (*s).into())
307        .collect();
308        let lift = |n: &str| -> Box<str> {
309            if n.contains('/') {
310                n.into()
311            } else {
312                format!("{n}/count/spliced").into()
313            }
314        };
315        let host = load_frozen_feature_host(FrozenLoadArgs {
316            dictionary_path: &dict_path,
317            bias_path: None,
318            target_feature_names: &target,
319            name_kind: FeatureNameKind::Gene { delim: '_' },
320            source_name_map: Some(&lift),
321        })
322        .unwrap();
323        assert_eq!(host.keep_target_indices, vec![0, 2]);
324        assert_eq!(host.keep_src_indices, vec![0, 1]);
325        assert_eq!(
326            host.src_names,
327            vec![
328                Box::<str>::from("TGFB1/count/spliced"),
329                Box::<str>::from("MYC/count/unspliced")
330            ]
331        );
332        assert_eq!(host.e_feat[(0, 0)], 1.0);
333        assert_eq!(host.e_feat[(1, 1)], 4.0);
334    }
335
336    #[test]
337    fn gene_canon_matches_across_delim_variants() {
338        let dir = tempfile::tempdir().unwrap();
339        let dict_path = dir.path().join("d.parquet").to_str().unwrap().to_string();
340
341        // Source uses ENSG-prefixed; target uses bare gene symbols.
342        let src = DMatrix::<f32>::from_row_slice(2, 2, &[1.0, 2.0, 3.0, 4.0]);
343        write_test_parquet(
344            &dict_path,
345            &["ENSG00000105329_TGFB1", "ENSG00000141510_TP53"],
346            "gene",
347            &["h0", "h1"],
348            &src,
349        );
350        let target: Vec<Box<str>> = ["TP53", "TGFB1"].iter().map(|s| (*s).into()).collect();
351
352        let host = load_frozen_feature_host(FrozenLoadArgs {
353            dictionary_path: &dict_path,
354            bias_path: None,
355            target_feature_names: &target,
356            name_kind: FeatureNameKind::Gene { delim: '_' },
357            source_name_map: None,
358        })
359        .unwrap();
360
361        assert_eq!(host.keep_target_indices, vec![0, 1]);
362        // Row 0 (target TP53) ← source row 1.
363        assert_eq!(host.e_feat[(0, 0)], 3.0);
364        // Row 1 (target TGFB1) ← source row 0.
365        assert_eq!(host.e_feat[(1, 0)], 1.0);
366    }
367
368    #[test]
369    fn empty_intersection_errors() {
370        let dir = tempfile::tempdir().unwrap();
371        let dict_path = dir.path().join("d.parquet").to_str().unwrap().to_string();
372        let src = DMatrix::<f32>::from_row_slice(2, 2, &[1.0, 2.0, 3.0, 4.0]);
373        write_test_parquet(&dict_path, &["A", "B"], "gene", &["h0", "h1"], &src);
374        let target: Vec<Box<str>> = ["C", "D"].iter().map(|s| (*s).into()).collect();
375        let result = load_frozen_feature_host(FrozenLoadArgs {
376            dictionary_path: &dict_path,
377            bias_path: None,
378            target_feature_names: &target,
379            name_kind: FeatureNameKind::Exact,
380            source_name_map: None,
381        });
382        let err = match result {
383            Ok(_) => panic!("expected empty-intersection error"),
384            Err(e) => e,
385        };
386        assert!(err.to_string().contains("No feature names matched"));
387    }
388
389    #[test]
390    fn bias_loaded_when_provided() {
391        let dir = tempfile::tempdir().unwrap();
392        let dict_path = dir.path().join("d.parquet").to_str().unwrap().to_string();
393        let bias_path = dir.path().join("b.parquet").to_str().unwrap().to_string();
394        let src = DMatrix::<f32>::from_row_slice(2, 2, &[1.0, 2.0, 3.0, 4.0]);
395        write_test_parquet(&dict_path, &["A", "B"], "gene", &["h0", "h1"], &src);
396        let bias = DMatrix::<f32>::from_row_slice(2, 1, &[0.5, -0.3]);
397        write_test_parquet(&bias_path, &["A", "B"], "gene", &["bias"], &bias);
398
399        let target: Vec<Box<str>> = ["B", "A"].iter().map(|s| (*s).into()).collect();
400        let host = load_frozen_feature_host(FrozenLoadArgs {
401            dictionary_path: &dict_path,
402            bias_path: Some(&bias_path),
403            target_feature_names: &target,
404            name_kind: FeatureNameKind::Exact,
405            source_name_map: None,
406        })
407        .unwrap();
408        // Row 0 of output = target "B" = source row 1.
409        assert_eq!(host.b_feat, vec![-0.3, 0.5]);
410    }
411}
data_beans/aux/frozen_features.rs

data_beans/aux/
frozen_features.rs
