subsume 0.8.2

Geometric region embeddings (boxes, cones, octagons, Gaussians, hyperbolic intervals, sheaf networks) for subsumption, entailment, and logical query answering
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81

//! Adapter for constructing [`Dataset`] from [`petgraph`] graphs.
//!
//! Converts a `petgraph::Graph` (or `DiGraph`) where nodes are entity names
//! and edges are relation names into subsume's [`Triple`] / [`Dataset`] format.

use crate::dataset::{Dataset, Triple};

/// Convert a directed petgraph into a [`Dataset`] with all triples in the training split.
///
/// Node weights become entity names (via `ToString`), edge weights become
/// relation names. All edges become training triples; validation and test
/// splits are empty (callers can split manually).
///
/// # Example
///
/// ```rust,ignore
/// use petgraph::graph::DiGraph;
/// use subsume::petgraph_adapter::from_graph;
///
/// let mut g = DiGraph::new();
/// let dog = g.add_node("dog");
/// let animal = g.add_node("animal");
/// g.add_edge(dog, animal, "is_a");
///
/// let dataset = from_graph(&g);
/// assert_eq!(dataset.train.len(), 1);
/// ```
pub fn from_graph<N, E>(graph: &petgraph::graph::DiGraph<N, E>) -> Dataset
where
    N: ToString,
    E: ToString,
{
    let triples: Vec<Triple> = graph
        .edge_indices()
        .filter_map(|e| {
            let (src, dst) = graph.edge_endpoints(e)?;
            let rel = graph.edge_weight(e)?;
            Some(Triple::new(
                graph[src].to_string(),
                rel.to_string(),
                graph[dst].to_string(),
            ))
        })
        .collect();

    Dataset::new(triples, Vec::new(), Vec::new())
}

/// Convert a directed petgraph into a [`Dataset`] with train/valid/test split.
///
/// Splits edges into train (80%), valid (10%), test (10%) by edge index order.
/// For reproducible random splits, shuffle edge indices before calling.
pub fn from_graph_with_split<N, E>(graph: &petgraph::graph::DiGraph<N, E>) -> Dataset
where
    N: ToString,
    E: ToString,
{
    let all_triples: Vec<Triple> = graph
        .edge_indices()
        .filter_map(|e| {
            let (src, dst) = graph.edge_endpoints(e)?;
            let rel = graph.edge_weight(e)?;
            Some(Triple::new(
                graph[src].to_string(),
                rel.to_string(),
                graph[dst].to_string(),
            ))
        })
        .collect();

    let n = all_triples.len();
    let train_end = (n * 8) / 10;
    let valid_end = train_end + n / 10;

    let mut iter = all_triples.into_iter();
    let train: Vec<Triple> = iter.by_ref().take(train_end).collect();
    let valid: Vec<Triple> = iter.by_ref().take(valid_end - train_end).collect();
    let test: Vec<Triple> = iter.collect();

    Dataset::new(train, valid, test)
}