1use 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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
25pub enum Rel {
26 IsA,
28 PartOf,
30}
31
32pub struct Ontology {
36 graph: DiGraph<Box<str>, Rel>,
37 idx: FxHashMap<Box<str>, NodeIndex>,
38 names: FxHashMap<Box<str>, Box<str>>,
39 defs: FxHashMap<Box<str>, Box<str>>,
41}
42
43struct ParsedTerm {
45 id: Box<str>,
46 name: Option<Box<str>>,
47 def: Option<Box<str>>,
48 is_a: Vec<Box<str>>,
50 part_of: Vec<Box<str>>,
52}
53
54impl Ontology {
55 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 let mut terms: Vec<ParsedTerm> = Vec::new();
63 for frame in doc.entities() {
64 let EntityFrame::Term(term) = frame else {
65 continue;
66 };
67 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 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 pub fn ids(&self) -> impl Iterator<Item = &str> + '_ {
141 self.idx.keys().map(|k| &**k)
142 }
143
144 #[must_use]
146 pub fn def(&self, id: &str) -> Option<&str> {
147 self.defs.get(id).map(|d| &**d)
148 }
149
150 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 #[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 #[must_use]
181 pub fn name(&self, id: &str) -> Option<&str> {
182 self.names.get(id).map(|n| &**n)
183 }
184
185 #[must_use]
187 pub fn ancestors_or_self(&self, id: &str) -> FxHashSet<Box<str>> {
188 self.ancestors_impl(id, false)
189 }
190
191 #[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 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 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 assert_eq!(onto.len(), 7);
256 assert!(!onto.contains("CL:9999999"));
257 assert_eq!(onto.name("CL:0000084"), Some("T cell"));
258
259 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 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 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}