Skip to main content

data_beans/aux/
ontology.rs

1//! Generic OBO ontology (`cl-basic.obo`, `go-basic.obo`, …) parsed with
2//! `fastobo` and stored as a `petgraph` directed graph (edges point
3//! **child → parent**, tagged with the relation kind).
4//!
5//! This is a thin, reusable wrapper: parse once, then query each term's
6//! ancestors via [`Ontology::ancestors_or_self`] (`is_a` only — the relation a
7//! collapsed-label tree needs) or [`Ontology::ancestors_or_self_with_part_of`]
8//! (`is_a` + `part_of`, the GO "true-path" closure); [`Ontology::edges`] hands
9//! the whole hierarchy to callers that embed or draw it rather than walk it.
10//!
11//! Prefix-agnostic: any term id is kept (`CL:`, `GO:`, …); edges to
12//! unknown/obsolete targets are dropped, as are obsolete terms. The
13//! `{is_inferred="true"}` qualifier blocks fastobo handles natively (a hand
14//! rolled parser missing them was the original prototype's bug).
15
16use anyhow::{Context, Result};
17use fastobo::ast::{EntityFrame, TermClause};
18use petgraph::graph::{DiGraph, NodeIndex};
19use petgraph::visit::EdgeRef;
20use petgraph::Direction::Outgoing;
21use rustc_hash::{FxHashMap, FxHashSet};
22
23/// Ontology relation kind carried on each `child → parent` edge.
24#[derive(Debug, Clone, Copy, PartialEq, Eq)]
25pub enum Rel {
26    /// `is_a` subsumption (subclass).
27    IsA,
28    /// `relationship: part_of` (mereological containment; GO true-path).
29    PartOf,
30}
31
32/// Parsed OBO DAG. Edges are `child → parent` (tagged [`Rel`]), so a node's
33/// ancestors are reachable along `Outgoing` edges and its children along
34/// `Incoming` edges.
35pub struct Ontology {
36    graph: DiGraph<Box<str>, Rel>,
37    idx: FxHashMap<Box<str>, NodeIndex>,
38    names: FxHashMap<Box<str>, Box<str>>,
39    /// `def:` text per term, unescaped, without the xref list.
40    defs: FxHashMap<Box<str>, Box<str>>,
41}
42
43/// One parsed, non-obsolete OBO term gathered in the first pass.
44struct ParsedTerm {
45    id: Box<str>,
46    name: Option<Box<str>>,
47    def: Option<Box<str>>,
48    /// `is_a` parent ids.
49    is_a: Vec<Box<str>>,
50    /// `relationship: part_of` parent ids.
51    part_of: Vec<Box<str>>,
52}
53
54impl Ontology {
55    /// Parse an OBO file into the `is_a` + `part_of` DAG. Obsolete terms and
56    /// edges to unknown/obsolete targets are skipped.
57    pub fn load_obo(path: &str) -> Result<Self> {
58        let doc = fastobo::from_file(path)
59            .with_context(|| format!("failed to parse OBO file: {path}"))?;
60
61        // Pass 1: collect non-obsolete terms.
62        let mut terms: Vec<ParsedTerm> = Vec::new();
63        for frame in doc.entities() {
64            let EntityFrame::Term(term) = frame else {
65                continue;
66            };
67            // fastobo's `Line<_>` Display can carry a trailing newline / inline
68            // comment, so trim every extracted token down to the bare value.
69            let id: Box<str> = term.id().to_string().trim().into();
70            let mut name: Option<Box<str>> = None;
71            let mut def: Option<Box<str>> = None;
72            let mut is_a: Vec<Box<str>> = Vec::new();
73            let mut part_of: Vec<Box<str>> = Vec::new();
74            let mut obsolete = false;
75            for line in term.clauses() {
76                match &**line {
77                    TermClause::Name(n) => name = Some(n.to_string().trim().into()),
78                    TermClause::Def(d) => {
79                        let text = d.text().as_str().trim();
80                        if !text.is_empty() {
81                            def = Some(text.into());
82                        }
83                    }
84                    TermClause::IsObsolete(b) => obsolete = obsolete || *b,
85                    TermClause::IsA(parent) => is_a.push(parent.to_string().trim().into()),
86                    TermClause::Relationship(rel, target) => {
87                        if rel.to_string().trim() == "part_of" {
88                            part_of.push(target.to_string().trim().into());
89                        }
90                    }
91                    _ => {}
92                }
93            }
94            if !obsolete {
95                terms.push(ParsedTerm {
96                    id,
97                    name,
98                    def,
99                    is_a,
100                    part_of,
101                });
102            }
103        }
104
105        // Pass 2: nodes first, then edges (only between known nodes).
106        let mut graph: DiGraph<Box<str>, Rel> = DiGraph::new();
107        let mut idx: FxHashMap<Box<str>, NodeIndex> = FxHashMap::default();
108        let mut names: FxHashMap<Box<str>, Box<str>> = FxHashMap::default();
109        let mut defs: FxHashMap<Box<str>, Box<str>> = FxHashMap::default();
110        for term in &terms {
111            let node = graph.add_node(term.id.clone());
112            idx.insert(term.id.clone(), node);
113            if let Some(n) = &term.name {
114                names.insert(term.id.clone(), n.clone());
115            }
116            if let Some(d) = &term.def {
117                defs.insert(term.id.clone(), d.clone());
118            }
119        }
120        for term in &terms {
121            let child = idx[&term.id];
122            for (parents, rel) in [(&term.is_a, Rel::IsA), (&term.part_of, Rel::PartOf)] {
123                for p in parents {
124                    if let Some(&parent) = idx.get(p) {
125                        graph.add_edge(child, parent, rel);
126                    }
127                }
128            }
129        }
130
131        Ok(Self {
132            graph,
133            idx,
134            names,
135            defs,
136        })
137    }
138
139    /// Every (non-obsolete) term id, in no particular order.
140    pub fn ids(&self) -> impl Iterator<Item = &str> + '_ {
141        self.idx.keys().map(|k| &**k)
142    }
143
144    /// The term's `def:` text (`None` if the term is unknown or undefined).
145    #[must_use]
146    pub fn def(&self, id: &str) -> Option<&str> {
147        self.defs.get(id).map(|d| &**d)
148    }
149
150    /// Every `child → parent` edge with its relation kind, in no particular
151    /// order — the hierarchy as a graph, for callers that embed or draw it
152    /// rather than walk it.
153    pub fn edges(&self) -> impl Iterator<Item = (&str, &str, Rel)> + '_ {
154        self.graph.edge_references().map(|e| {
155            (
156                &*self.graph[e.source()],
157                &*self.graph[e.target()],
158                *e.weight(),
159            )
160        })
161    }
162
163    /// Number of (non-obsolete) terms.
164    #[must_use]
165    pub fn len(&self) -> usize {
166        self.graph.node_count()
167    }
168
169    #[must_use]
170    pub fn is_empty(&self) -> bool {
171        self.graph.node_count() == 0
172    }
173
174    #[must_use]
175    pub fn contains(&self, id: &str) -> bool {
176        self.idx.contains_key(id)
177    }
178
179    /// Human-readable term name (`None` if the term is unknown or unnamed).
180    #[must_use]
181    pub fn name(&self, id: &str) -> Option<&str> {
182        self.names.get(id).map(|n| &**n)
183    }
184
185    /// All `is_a` ancestors of `id` plus `id` itself (empty if `id` is unknown).
186    #[must_use]
187    pub fn ancestors_or_self(&self, id: &str) -> FxHashSet<Box<str>> {
188        self.ancestors_impl(id, false)
189    }
190
191    /// All `is_a` + `part_of` ancestors of `id` plus `id` itself — the GO
192    /// "true-path" closure (empty if `id` is unknown).
193    #[must_use]
194    pub fn ancestors_or_self_with_part_of(&self, id: &str) -> FxHashSet<Box<str>> {
195        self.ancestors_impl(id, true)
196    }
197
198    /// Walk ancestors in `NodeIndex` space (Copy — no per-visit string clones),
199    /// following `is_a` edges and, when `with_part_of`, `part_of` edges too;
200    /// then materialize each id exactly once at the boundary.
201    fn ancestors_impl(&self, id: &str, with_part_of: bool) -> FxHashSet<Box<str>> {
202        let Some(&start) = self.idx.get(id) else {
203            return FxHashSet::default();
204        };
205        let mut seen: FxHashSet<NodeIndex> = FxHashSet::default();
206        seen.insert(start);
207        let mut stack = vec![start];
208        while let Some(n) = stack.pop() {
209            for edge in self.graph.edges_directed(n, Outgoing) {
210                let follow = matches!(edge.weight(), Rel::IsA)
211                    || (with_part_of && matches!(edge.weight(), Rel::PartOf));
212                if follow && seen.insert(edge.target()) {
213                    stack.push(edge.target());
214                }
215            }
216        }
217        seen.iter().map(|&n| self.graph[n].clone()).collect()
218    }
219}
220
221#[cfg(test)]
222mod tests {
223    use super::*;
224    use std::io::Write;
225
226    /// Minimal CL-shaped OBO: cell → lymphocyte → T → {CD4, CD8}; B sibling;
227    /// one obsolete term; one `{is_inferred="true"}` qualifier on an is_a line;
228    /// one `part_of` edge (CD8 T part_of an immune-system "compartment") to
229    /// exercise the relation-filtered traversal.
230    fn write_obo() -> tempfile::NamedTempFile {
231        let mut f = tempfile::NamedTempFile::new().unwrap();
232        writeln!(
233            f,
234            "format-version: 1.2\n\n\
235             [Term]\nid: CL:0000000\nname: cell\n\n\
236             [Term]\nid: CL:0000542\nname: lymphocyte\nis_a: CL:0000000 ! cell\n\n\
237             [Term]\nid: CL:0000084\nname: T cell\ndef: \"A lymphocyte with a \\\"TCR\\\", made in the thymus.\" [GOC:add]\nis_a: CL:0000542 {{is_inferred=\"true\"}} ! lymphocyte\n\n\
238             [Term]\nid: CL:0000624\nname: CD4 T\nis_a: CL:0000084 ! T cell\n\n\
239             [Term]\nid: CL:0000625\nname: CD8 T\nis_a: CL:0000084 ! T cell\nrelationship: part_of CL:1000000 ! compartment\n\n\
240             [Term]\nid: CL:1000000\nname: immune compartment\n\n\
241             [Term]\nid: CL:0000236\nname: B cell\nis_a: CL:0000542 ! lymphocyte\n\n\
242             [Term]\nid: CL:9999999\nname: dead\nis_obsolete: true\n"
243        )
244        .unwrap();
245        f.flush().unwrap();
246        f
247    }
248
249    #[test]
250    fn parses_and_resolves_ancestry() {
251        let f = write_obo();
252        let onto = Ontology::load_obo(f.path().to_str().unwrap()).unwrap();
253
254        // obsolete term dropped; 7 live terms (incl. the compartment).
255        assert_eq!(onto.len(), 7);
256        assert!(!onto.contains("CL:9999999"));
257        assert_eq!(onto.name("CL:0000084"), Some("T cell"));
258
259        // The {is_inferred} qualifier must NOT break the edge: CD4/CD8 under T,
260        // T under lymphocyte under cell.
261        let anc = onto.ancestors_or_self("CL:0000624");
262        for a in ["CL:0000624", "CL:0000084", "CL:0000542", "CL:0000000"] {
263            assert!(anc.contains(a), "missing ancestor {a}");
264        }
265        assert!(!anc.contains("CL:0000236"));
266    }
267
268    #[test]
269    fn definitions_are_kept_unescaped_and_the_hierarchy_is_exposed_as_edges() {
270        let f = write_obo();
271        let onto = Ontology::load_obo(f.path().to_str().unwrap()).unwrap();
272        assert_eq!(
273            onto.def("CL:0000084"),
274            Some("A lymphocyte with a \"TCR\", made in the thymus.")
275        );
276        assert_eq!(onto.def("CL:0000000"), None, "no def: line");
277        assert_eq!(onto.def("CL:9999999"), None, "obsolete");
278        let mut edges: Vec<(String, String, Rel)> = onto
279            .edges()
280            .map(|(c, p, r)| (c.to_string(), p.to_string(), r))
281            .collect();
282        edges.sort_by(|a, b| (&a.0, &a.1).cmp(&(&b.0, &b.1)));
283        assert_eq!(edges.len(), 6, "5 is_a + 1 part_of among live terms");
284        assert!(edges.contains(&("CL:0000625".into(), "CL:1000000".into(), Rel::PartOf)));
285        assert!(edges.contains(&("CL:0000084".into(), "CL:0000542".into(), Rel::IsA)));
286        assert!(edges
287            .iter()
288            .all(|(c, p, _)| onto.contains(c) && onto.contains(p)));
289    }
290
291    #[test]
292    fn part_of_only_followed_on_demand() {
293        let f = write_obo();
294        let onto = Ontology::load_obo(f.path().to_str().unwrap()).unwrap();
295
296        // is_a-only walk does NOT cross the part_of edge to the compartment.
297        let isa = onto.ancestors_or_self("CL:0000625");
298        assert!(isa.contains("CL:0000084"), "is_a ancestor missing");
299        assert!(
300            !isa.contains("CL:1000000"),
301            "part_of must not leak into is_a-only walk"
302        );
303
304        // is_a + part_of walk reaches the compartment.
305        let full = onto.ancestors_or_self_with_part_of("CL:0000625");
306        assert!(full.contains("CL:0000084"), "is_a ancestor missing");
307        assert!(full.contains("CL:1000000"), "part_of ancestor missing");
308    }
309}