rename sampling_fn argument to sample_fn in GraphicalSimulator.add_node method

Author: daniel-habermann

bayesflow/experimental/graphical_simulator/example_simulators.py

@@ -1,4 +1,5 @@
 import numpy as np
+
 from .graphical_simulator import GraphicalSimulator
 
 
@@ -60,21 +61,21 @@ def meta_fn():
     simulator = GraphicalSimulator(meta_fn=meta_fn)
     simulator.add_node(
         "schools",
-        sampling_fn=sample_school,
+        sample_fn=sample_school,
     )
     simulator.add_node(
         "exams",
-        sampling_fn=sample_exam,
+        sample_fn=sample_exam,
         reps="num_exams",
     )
     simulator.add_node(
         "questions",
-        sampling_fn=sample_question,
+        sample_fn=sample_question,
         reps="num_questions",
     )
     simulator.add_node(
         "students",
-        sampling_fn=sample_student,
+        sample_fn=sample_student,
         reps="num_students",
     )
 
@@ -109,8 +110,8 @@ def meta():
 
     simulator = GraphicalSimulator(meta_fn=meta)
 
-    simulator.add_node("prior", sampling_fn=prior)
-    simulator.add_node("likelihood", sampling_fn=likelihood)
+    simulator.add_node("prior", sample_fn=prior)
+    simulator.add_node("likelihood", sample_fn=likelihood)
 
     simulator.add_edge("prior", "likelihood")
 
@@ -140,15 +141,15 @@ def sample_y(local_mean, shared_std):
         return {"y": float(y)}
 
     simulator = GraphicalSimulator()
-    simulator.add_node("hypers", sampling_fn=sample_hypers, reps=1)
+    simulator.add_node("hypers", sample_fn=sample_hypers, reps=5)
 
     simulator.add_node(
         "locals",
         sampling_fn=sample_locals,
         reps=6,
     )
-    simulator.add_node("shared", sampling_fn=sample_shared, reps=1)
-    simulator.add_node("y", sampling_fn=sample_y, reps=10)
+    simulator.add_node("shared", sample_fn=sample_shared, reps=1)
+    simulator.add_node("y", sample_fn=sample_y, reps=10)
 
     simulator.add_edge("hypers", "locals")
     simulator.add_edge("locals", "y")
@@ -184,19 +185,19 @@ def sample_y(level_3_mean, shared_std):
         return {"y": y}
 
     simulator = GraphicalSimulator()
-    simulator.add_node("level1", sampling_fn=sample_level_1)
+    simulator.add_node("level1", sample_fn=sample_level_1)
     simulator.add_node(
         "level2",
-        sampling_fn=sample_level_2,
+        sample_fn=sample_level_2,
         reps=10,
     )
     simulator.add_node(
         "level3",
-        sampling_fn=sample_level_3,
+        sample_fn=sample_level_3,
         reps=20,
     )
-    simulator.add_node("shared", sampling_fn=sample_shared)
-    simulator.add_node("y", sampling_fn=sample_y, reps=10)
+    simulator.add_node("shared", sample_fn=sample_shared)
+    simulator.add_node("y", sample_fn=sample_y, reps=10)
 
     simulator.add_edge("level1", "level2")
     simulator.add_edge("level2", "level3")

bayesflow/experimental/graphical_simulator/graphical_simulator.py

@@ -29,8 +29,8 @@ def __init__(self, meta_fn: Optional[Callable[[], dict[str, Any]]] = None, *args
         self.graph = nx.DiGraph()
         self.meta_fn = meta_fn
 
-    def add_node(self, node: str, sampling_fn: Callable[..., dict[str, Any]], reps: int | str = 1):
-        self.graph.add_node(node, sampling_fn=sampling_fn, reps=reps)
+    def add_node(self, node: str, sample_fn: Callable[..., dict[str, Any]], reps: int | str = 1):
+        self.graph.add_node(node, sample_fn=sample_fn, reps=reps)
 
     def add_edge(self, from_node: str, to_node: str):
         self.graph.add_edge(from_node, to_node)
@@ -62,7 +62,7 @@ def sample(self, batch_shape: Shape, **kwargs) -> dict[str, np.ndarray]:
                 node_samples = []
 
                 parent_nodes = list(self.graph.predecessors(node))
-                sampling_fn = self.graph.nodes[node]["sampling_fn"]
+                sampling_fn = self.graph.nodes[node]["sample_fn"]
                 reps_field = self.graph.nodes[node]["reps"]
                 reps = reps_field if isinstance(reps_field, int) else meta_dict[reps_field]