put the eval context into a struct and expose methods for getting arb…

…itrary eval nodes and edges
Pometry · May 27, 2024 · e0a3b6a · e0a3b6a
1 parent b4daff8
commit e0a3b6a
Show file tree

Hide file tree

Showing 15 changed files with 217 additions and 183 deletions.
diff --git a/examples/rust/src/bin/hulongbay/main.rs b/examples/rust/src/bin/hulongbay/main.rs
@@ -2,7 +2,13 @@
 #![allow(dead_code)]
 use itertools::Itertools;
 use raphtory::{
-    algorithms::{components::weakly_connected_components, motifs::triangle_count::triangle_count},
+    algorithms::{
+        components::weakly_connected_components,
+        motifs::{
+            global_temporal_three_node_motifs::global_temporal_three_node_motif,
+            triangle_count::triangle_count,
+        },
+    },
     graph_loader::source::csv_loader::CsvLoader,
     prelude::*,
 };
@@ -207,8 +213,16 @@ fn try_main_bm() -> Result<(), Box<dyn Error>> {
     Ok(())
 }
 
+fn try_motif() -> Result<(), Box<dyn Error>> {
+    let args: Vec<String> = env::args().collect();
+    let data_dir = Path::new(args.get(1).ok_or(MissingArgumentError)?);
+    let graph = loader(data_dir)?;
+    global_temporal_three_node_motif(&graph, 3600, None);
+    Ok(())
+}
+
 fn main() {
-    if let Err(e) = try_main_bm() {
+    if let Err(e) = try_motif() {
         eprintln!("Failed: {}", e);
         std::process::exit(1)
     }

diff --git a/raphtory/src/algorithms/motifs/global_temporal_three_node_motifs.rs b/raphtory/src/algorithms/motifs/global_temporal_three_node_motifs.rs
@@ -23,15 +23,11 @@ use std::collections::HashMap;
 
 ///////////////////////////////////////////////////////
 
-pub fn star_motif_count<G>(
-    graph: &G,
-    evv: &EvalNodeView<G, ()>,
-    deltas: Vec<i64>,
-) -> Vec<[usize; 32]>
+pub fn star_motif_count<G>(evv: &EvalNodeView<G, ()>, deltas: Vec<i64>) -> Vec<[usize; 32]>
 where
     G: StaticGraphViewOps,
 {
-    let two_n_c = twonode_motif_count(graph, evv, deltas.clone());
+    let two_n_c = twonode_motif_count(evv, deltas.clone());
     let neigh_map: HashMap<u64, usize> = evv
         .neighbours()
         .into_iter()
@@ -75,20 +71,16 @@ where
 
 ///////////////////////////////////////////////////////
 
-pub fn twonode_motif_count<G>(
-    graph: &G,
-    evv: &EvalNodeView<G, ()>,
-    deltas: Vec<i64>,
-) -> Vec<[usize; 8]>
+pub fn twonode_motif_count<G>(evv: &EvalNodeView<G, ()>, deltas: Vec<i64>) -> Vec<[usize; 8]>
 where
     G: StaticGraphViewOps,
 {
     let mut results = deltas.iter().map(|_| [0; 8]).collect::<Vec<[usize; 8]>>();
 
     for nb in evv.neighbours().into_iter() {
         let nb_id = nb.id();
-        let out = graph.edge(evv.id(), nb_id);
-        let inc = graph.edge(nb_id, evv.id());
+        let out = evv.graph().edge(evv.id(), nb_id);
+        let inc = evv.graph().edge(nb_id, evv.id());
         let events: Vec<TwoNodeEvent> = out
             .iter()
             .flat_map(|e| e.explode())
@@ -265,8 +257,7 @@ where
     let out1 = triangle_motifs(g, deltas.clone(), threads);
 
     let step1 = ATask::new(move |evv: &mut EvalNodeView<G, _>| {
-        let g = evv.graph();
-        let star_nodes = star_motif_count(g, evv, deltas.clone());
+        let star_nodes = star_motif_count(evv, deltas.clone());
         for (i, star) in star_nodes.iter().enumerate() {
             evv.global_update(&star_mc[i], *star);
         }

diff --git a/raphtory/src/algorithms/motifs/local_temporal_three_node_motifs.rs b/raphtory/src/algorithms/motifs/local_temporal_three_node_motifs.rs
@@ -126,7 +126,6 @@ where
 ///////////////////////////////////////////////////////
 
 pub fn twonode_motif_count<'a, 'b, G, GH>(
-    graph: &'a G,
     evv: &'a EvalNodeView<'b, '_, G, MotifCounter, GH>,
     deltas: Vec<i64>,
 ) -> Vec<[usize; 8]>
@@ -144,8 +143,8 @@ where
 
     for nb in evv.neighbours().into_iter() {
         let nb_id = nb.id();
-        let out = graph.edge(evv.id(), nb_id);
-        let inc = graph.edge(nb_id, evv.id());
+        let out = evv.graph().edge(evv.id(), nb_id);
+        let inc = evv.graph().edge(nb_id, evv.id());
         let events: Vec<TwoNodeEvent> = out
             .iter()
             .flat_map(|e| e.explode())
@@ -337,8 +336,7 @@ where
     let out1 = triangle_motifs(g, deltas.clone(), motifs_counter, threads);
 
     let step1 = ATask::new(move |evv: &mut EvalNodeView<G, MotifCounter>| {
-        let g = evv.graph();
-        let two_nodes = twonode_motif_count(g, evv, deltas.clone());
+        let two_nodes = twonode_motif_count(evv, deltas.clone());
         let star_nodes = star_motif_count(evv, deltas.clone());
 
         *evv.get_mut() = MotifCounter::new(

diff --git a/raphtory/src/db/api/storage/nodes/node_entry.rs b/raphtory/src/db/api/storage/nodes/node_entry.rs
@@ -97,6 +97,10 @@ impl<'a, 'b: 'a> NodeStorageOps<'a> for &'a NodeStorageEntry<'b> {
         for_all!(self, node => node.vid())
     }
 
+    fn id(self) -> u64 {
+        for_all!(self, node => node.id())
+    }
+
     fn name(self) -> Option<&'a str> {
         for_all!(self, node => node.name())
     }

diff --git a/raphtory/src/db/api/storage/nodes/node_owned_entry.rs b/raphtory/src/db/api/storage/nodes/node_owned_entry.rs
@@ -94,6 +94,10 @@ impl<'a> NodeStorageOps<'a> for &'a NodeOwnedEntry {
         for_all!(self, node => node.vid())
     }
 
+    fn id(self) -> u64 {
+        for_all!(self, node => node.id())
+    }
+
     fn name(self) -> Option<&'a str> {
         for_all!(self, node => node.name())
     }

diff --git a/raphtory/src/db/api/storage/nodes/node_ref.rs b/raphtory/src/db/api/storage/nodes/node_ref.rs
@@ -91,6 +91,10 @@ impl<'a> NodeStorageOps<'a> for NodeStorageRef<'a> {
         for_all!(self, node => node.vid())
     }
 
+    fn id(self) -> u64 {
+        for_all!(self, node => node.id())
+    }
+
     fn name(self) -> Option<&'a str> {
         for_all!(self, node => node.name())
     }

diff --git a/raphtory/src/db/api/storage/nodes/node_storage_ops.rs b/raphtory/src/db/api/storage/nodes/node_storage_ops.rs
@@ -25,6 +25,8 @@ pub trait NodeStorageOps<'a>: Sized {
 
     fn vid(self) -> VID;
 
+    fn id(self) -> u64;
+
     fn name(self) -> Option<&'a str>;
 
     fn find_edge(self, dst: VID, layer_ids: &LayerIds) -> Option<EdgeRef>;
@@ -59,6 +61,10 @@ impl<'a> NodeStorageOps<'a> for &'a NodeStore {
         self.vid
     }
 
+    fn id(self) -> u64 {
+        self.global_id
+    }
+
     fn name(self) -> Option<&'a str> {
         self.name.as_str()
     }

diff --git a/raphtory/src/db/api/view/node.rs b/raphtory/src/db/api/view/node.rs
@@ -6,7 +6,7 @@ use crate::{
     },
     db::api::{
         properties::{internal::PropertiesOps, Properties},
-        storage::storage_ops::GraphStorage,
+        storage::{nodes::node_storage_ops::NodeStorageOps, storage_ops::GraphStorage},
         view::{
             internal::{CoreGraphOps, OneHopFilter, TimeSemantics},
             reset_filter::ResetFilter,
@@ -175,7 +175,7 @@ impl<'graph, V: BaseNodeViewOps<'graph> + 'graph> NodeViewOps<'graph> for V {
 
     #[inline]
     fn id(&self) -> Self::ValueType<u64> {
-        self.map(|_cg, g, v| g.node_id(v))
+        self.map(|cg, _g, v| cg.node(v).id())
     }
     #[inline]
     fn name(&self) -> Self::ValueType<String> {

diff --git a/raphtory/src/db/task/edge/eval_edge.rs b/raphtory/src/db/task/edge/eval_edge.rs
@@ -18,7 +18,7 @@ use crate::{
 
 use crate::db::task::edge::eval_edges::EvalEdges;
 
-use crate::db::api::storage::storage_ops::GraphStorage;
+use crate::db::{api::storage::storage_ops::GraphStorage, task::eval_graph::EvalGraph};
 use std::{cell::RefCell, rc::Rc};
 
 pub struct EvalEdgeView<'graph, 'a, G, GH, CS: Clone, S> {
@@ -102,15 +102,19 @@ impl<
         let node_state = self.node_state.clone();
         let local_state_prev = self.local_state_prev;
         let storage = self.storage;
-        EvalNodeView {
+        let base_graph = self.edge.base_graph;
+        let eval_graph = EvalGraph {
             ss,
-            node: node.node,
-            graph: node.base_graph,
-            base_graph: node.base_graph,
+            base_graph,
             storage,
-            local_state: None,
             local_state_prev,
             node_state,
+        };
+        EvalNodeView {
+            node: node.node,
+            graph: node.base_graph,
+            eval_graph,
+            local_state: None,
         }
     }
 

diff --git a/raphtory/src/db/task/edge/eval_edges.rs b/raphtory/src/db/task/edge/eval_edges.rs
@@ -12,6 +12,7 @@ use crate::{
         graph::edges::Edges,
         task::{
             edge::eval_edge::EvalEdgeView,
+            eval_graph::EvalGraph,
             node::{eval_node::EvalPathFromNode, eval_node_state::EVState},
             task_state::PrevLocalState,
         },
@@ -168,14 +169,18 @@ impl<
         let local_state_prev = self.local_state_prev;
         let path = self.edges.map_nodes(op);
         let base_graph = self.edges.base_graph;
+        let storage = self.storage;
+        let eval_graph = EvalGraph {
+            ss,
+            base_graph,
+            storage,
+            local_state_prev,
+            node_state,
+        };
         EvalPathFromNode {
             graph: base_graph,
-            base_graph: base_graph,
+            base_graph: eval_graph,
             op: path.op,
-            ss,
-            node_state,
-            local_state_prev,
-            storage: self.storage,
         }
     }
 

diff --git a/raphtory/src/db/task/eval_graph.rs b/raphtory/src/db/task/eval_graph.rs
@@ -0,0 +1,61 @@
+use crate::{
+    core::{
+        entities::nodes::node_ref::AsNodeRef,
+        state::compute_state::{ComputeState, ComputeStateVec},
+    },
+    db::{
+        api::storage::storage_ops::GraphStorage,
+        task::{
+            edge::eval_edge::EvalEdgeView,
+            node::{eval_node::EvalNodeView, eval_node_state::EVState},
+            task_state::PrevLocalState,
+        },
+    },
+    prelude::GraphViewOps,
+};
+use std::{cell::RefCell, rc::Rc};
+
+#[derive(Debug)]
+pub struct EvalGraph<'graph, 'a, G, S, CS: Clone = ComputeStateVec> {
+    pub(crate) ss: usize,
+    pub(crate) base_graph: &'graph G,
+    pub(crate) storage: &'graph GraphStorage,
+    pub(crate) local_state_prev: &'graph PrevLocalState<'a, S>,
+    pub(crate) node_state: Rc<RefCell<EVState<'a, CS>>>,
+}
+
+impl<'graph, 'a, G, S, CS: Clone> Clone for EvalGraph<'graph, 'a, G, S, CS> {
+    fn clone(&self) -> Self {
+        Self {
+            ss: self.ss,
+            base_graph: self.base_graph,
+            storage: self.storage,
+            local_state_prev: self.local_state_prev,
+            node_state: self.node_state.clone(),
+        }
+    }
+}
+
+impl<'graph, 'a: 'graph, G: GraphViewOps<'graph>, S: 'static, CS: ComputeState + Clone>
+    EvalGraph<'graph, 'a, G, S, CS>
+{
+    pub fn node(&self, n: impl AsNodeRef) -> Option<EvalNodeView<'graph, 'a, G, S, &'graph G, CS>> {
+        let node = (&self.base_graph).node(n)?;
+        Some(EvalNodeView::new_local(node.node, self.clone(), None))
+    }
+
+    pub fn edge<N: AsNodeRef>(
+        &self,
+        src: N,
+        dst: N,
+    ) -> Option<EvalEdgeView<'graph, 'a, G, &'graph G, CS, S>> {
+        let edge = (&self.base_graph).edge(src, dst)?;
+        Some(EvalEdgeView::new(
+            self.ss,
+            edge,
+            self.storage,
+            self.node_state.clone(),
+            self.local_state_prev,
+        ))
+    }
+}
diff --git a/raphtory/src/db/task/mod.rs b/raphtory/src/db/task/mod.rs
@@ -4,6 +4,7 @@ use std::sync::Arc;
 
 pub mod context;
 pub mod edge;
+mod eval_graph;
 pub mod node;
 pub mod task;
 pub mod task_runner;