make two distribution example

src/causalprog/algorithms/__init__.py

@@ -1,3 +1,3 @@
 """Algorithms."""
 
-from .expectation import expectation
+from .expectation import expectation, standard_deviation

src/causalprog/algorithms/expectation.py

@@ -1,21 +1,19 @@
-"""Algorithm for estimating the expectation of a process represented by a graph."""
+"""Algorithms for estimating the expectation and standard deviation."""
 
-import typing
-
-if typing.TYPE_CHECKING:
-    import numpy.typing as npt
+import numpy as np
+import numpy.typing as npt
 
 from causalprog.graph import Graph
 
 from .iteration import roots_down_to_outcome
 
 
-def expectation(
+def sample(
     graph: Graph,
     outcome_node_label: str | None = None,
     samples: int = 1000,
-) -> float:
-    """Estimate the expectation of a graph."""
+) -> npt.NDArray[float]:
+    """Sample data from a graph."""
     if outcome_node_label is None:
         outcome_node_label = graph.outcome.label
 
@@ -24,4 +22,22 @@ def expectation(
     values: dict[str, npt.NDArray[float]] = {}
     for node in nodes:
         values[node.label] = node.sample(values, samples)
-    return values[outcome_node_label].sum() / samples
+    return values[outcome_node_label]
+
+
+def expectation(
+    graph: Graph,
+    outcome_node_label: str | None = None,
+    samples: int = 1000,
+) -> float:
+    """Estimate the expectation of a graph."""
+    return sample(graph, outcome_node_label, samples).mean()
+
+
+def standard_deviation(
+    graph: Graph,
+    outcome_node_label: str | None = None,
+    samples: int = 1000,
+) -> float:
+    """Estimate the standard deviation of a graph."""
+    return np.std(sample(graph, outcome_node_label, samples))

src/causalprog/algorithms/iteration.py

@@ -19,10 +19,10 @@ def roots_down_to_outcome(
     while n < len(nodes_need_sampling):
         new_n = len(nodes_need_sampling)
         for node in nodes_need_sampling[n:]:
-            if node in pre:
-                for parent in pre[node]:
-                    if parent not in nodes_need_sampling:
-                        nodes_need_sampling.append(parent)
+            if node in pre and pre[node] not in nodes_need_sampling:
+                nodes_need_sampling.append(pre[node])
         n = new_n
 
-    return [node for node in graph.depth_first_nodes if node in nodes_need_sampling]
+    return [
+        node for node in graph.depth_first_nodes[::-1] if node in nodes_need_sampling
+    ]

src/causalprog/graph/__init__.py

@@ -1,4 +1,4 @@
 """Creation and storage of graphs."""
 
 from .graph import Graph
-from .node import DistributionNode, Node, RootDistributionNode
+from .node import DistributionNode, Node

src/causalprog/graph/graph.py

@@ -13,7 +13,7 @@ class Graph:
     def __init__(self, label: str) -> None:
         """Create end empty graph."""
         self._label = label
-        self._graph = nx.Graph()
+        self._graph = nx.DiGraph()
         self._nodes_by_label = {}
         self._node_index = 0
 

src/causalprog/graph/node.py

@@ -2,21 +2,47 @@
 
 from __future__ import annotations
 
+import typing
 from abc import ABC, abstractmethod
 
 import numpy as np
 
+if typing.TYPE_CHECKING:
+    import numpy.typing as npt
 
-class DistributionFamily:
+
+class Distribution(ABC):
     """Placeholder class."""
 
+    @abstractmethod
+    def sample(
+        self, sampled_dependencies: dict[str, npt.NDArray[float]], samples: int
+    ) -> npt.NDArray[float]:
+        """Sample."""
+
 
-class Distribution:
-    """Placeholder class."""
+class NormalDistribution(Distribution):
+    """Normal distribution."""
 
-    def sample(self, samples: int) -> float:
+    def __init__(self, mean: str | float = 0.0, std_dev: str | float = 1.0) -> None:
+        """Initialise."""
+        self.mean = mean
+        self.std_dev = std_dev
+
+    def sample(
+        self, sampled_dependencies: dict[str, npt.NDArray[float]], samples: int
+    ) -> npt.NDArray[float]:
         """Sample a normal distribution with mean 1."""
-        return np.random.normal(1.0, 1.0, samples)  # noqa: NPY002
+        values = np.random.normal(0.0, 1.0, samples)  # noqa: NPY002
+        if isinstance(self.std_dev, str):
+            values *= sampled_dependencies[self.std_dev]
+        else:
+            values *= self.std_dev
+        if isinstance(self.mean, str):
+            values += sampled_dependencies[self.mean]
+        else:
+            values += self.mean
+        return values
 
 
 class Node(ABC):
@@ -36,22 +62,19 @@ def label(self) -> str:
         return self._label
 
     @abstractmethod
-    def sample(self, sampled_dependencies: dict[str, float], samples: int) -> float:
+    def sample(
+        self, sampled_dependencies: dict[str, npt.NDArray[float]], samples: int
+    ) -> float:
         """Sample a value from the node."""
 
-    @property
-    @abstractmethod
-    def is_root(self) -> bool:
-        """Identify if the node is a root."""
-
     @property
     def is_outcome(self) -> bool:
         """Identify if the node is an outcome."""
         return self._is_outcome
 
 
-class RootDistributionNode(Node):
-    """A root node containing a distribution family."""
+class DistributionNode(Node):
+    """A node containing a distribution."""
 
     def __init__(
         self,
@@ -64,41 +87,11 @@ def __init__(
         self._dist = distribution
         super().__init__(label, is_outcome=is_outcome)
 
-    def sample(self, _sampled_dependencies: dict[str, float], samples: int) -> float:
-        """Sample a value from the node."""
-        return self._dist.sample(samples)
-
-    def __repr__(self) -> str:
-        return f'RootDistributionNode("{self.label}")'
-
-    @property
-    def is_root(self) -> bool:
-        """Identify if the node is a root."""
-        return True
-
-
-class DistributionNode(Node):
-    """A node containing a distribution family that depends on its parents."""
-
-    def __init__(
-        self,
-        family: DistributionFamily,
-        label: str | None = None,
-        *,
-        is_outcome: bool = False,
-    ) -> None:
-        """Initialise."""
-        self._dfamily = family
-        super().__init__(label, is_outcome=is_outcome)
-
-    def sample(self, sampled_dependencies: dict[str, float], samples: int) -> float:
+    def sample(
+        self, sampled_dependencies: dict[str, npt.NDArray[float]], samples: int
+    ) -> float:
         """Sample a value from the node."""
-        raise NotImplementedError
+        return self._dist.sample(sampled_dependencies, samples)
 
     def __repr__(self) -> str:
         return f'DistributionNode("{self.label}")'
-
-    @property
-    def is_root(self) -> bool:
-        """Identify if the node is a root."""
-        return False

tests/test_graph.py

@@ -9,11 +9,11 @@
 
 
 def test_label():
-    family = causalprog.graph.node.DistributionFamily()
-    node = causalprog.graph.RootDistributionNode(family)
-    node2 = causalprog.graph.RootDistributionNode(family, "node1")
-    node3 = causalprog.graph.RootDistributionNode(family, "Y")
-    node4 = causalprog.graph.DistributionNode(family)
+    d = causalprog.graph.node.NormalDistribution()
+    node = causalprog.graph.DistributionNode(d)
+    node2 = causalprog.graph.DistributionNode(d, "node1")
+    node3 = causalprog.graph.DistributionNode(d, "Y")
+    node4 = causalprog.graph.DistributionNode(d)
     node_copy = node
 
     assert node._label == node_copy._label  # noqa: SLF001
@@ -41,22 +41,22 @@ def test_label():
 
 
 def test_duplicate_label():
-    family = causalprog.graph.node.DistributionFamily()
+    d = causalprog.graph.node.NormalDistribution()
 
     graph = causalprog.graph.Graph("G0")
-    graph.add_node(causalprog.graph.RootDistributionNode(family, "X"))
+    graph.add_node(causalprog.graph.DistributionNode(d, "X"))
     with pytest.raises(ValueError, match=re.escape("Duplicate node label: X")):
-        graph.add_node(causalprog.graph.RootDistributionNode(family, "X"))
+        graph.add_node(causalprog.graph.DistributionNode(d, "X"))
 
 
 def test_simple_graph():
-    family = causalprog.graph.node.DistributionFamily()
-    n_x = causalprog.graph.RootDistributionNode(family, "N_X")
-    n_m = causalprog.graph.RootDistributionNode(family, "N_M")
-    u_y = causalprog.graph.RootDistributionNode(family, "U_Y")
-    x = causalprog.graph.DistributionNode(family, "X")
-    m = causalprog.graph.DistributionNode(family, "M")
-    y = causalprog.graph.DistributionNode(family, "Y", is_outcome=True)
+    d = causalprog.graph.node.NormalDistribution()
+    n_x = causalprog.graph.DistributionNode(d, "N_X")
+    n_m = causalprog.graph.DistributionNode(d, "N_M")
+    u_y = causalprog.graph.DistributionNode(d, "U_Y")
+    x = causalprog.graph.DistributionNode(d, "X")
+    m = causalprog.graph.DistributionNode(d, "M")
+    y = causalprog.graph.DistributionNode(d, "Y", is_outcome=True)
 
     graph = causalprog.graph.Graph("G0")
     graph.add_edge(n_x, x)
@@ -69,15 +69,15 @@ def test_simple_graph():
 
 
 def test_simple_graph_build_using_labels():
-    family = causalprog.graph.node.DistributionFamily()
+    d = causalprog.graph.node.NormalDistribution()
 
     graph = causalprog.graph.Graph("G0")
-    graph.add_node(causalprog.graph.RootDistributionNode(family, "N_X"))
-    graph.add_node(causalprog.graph.RootDistributionNode(family, "N_M"))
-    graph.add_node(causalprog.graph.RootDistributionNode(family, "U_Y"))
-    graph.add_node(causalprog.graph.DistributionNode(family, "X"))
-    graph.add_node(causalprog.graph.DistributionNode(family, "M"))
-    graph.add_node(causalprog.graph.DistributionNode(family, "Y", is_outcome=True))
+    graph.add_node(causalprog.graph.DistributionNode(d, "N_X"))
+    graph.add_node(causalprog.graph.DistributionNode(d, "N_M"))
+    graph.add_node(causalprog.graph.DistributionNode(d, "U_Y"))
+    graph.add_node(causalprog.graph.DistributionNode(d, "X"))
+    graph.add_node(causalprog.graph.DistributionNode(d, "M"))
+    graph.add_node(causalprog.graph.DistributionNode(d, "Y", is_outcome=True))
 
     graph.add_edge("N_X", "X")
     graph.add_edge("N_M", "M")
@@ -88,22 +88,57 @@ def test_simple_graph_build_using_labels():
     assert graph.label == "G0"
 
 
+@pytest.mark.parametrize("mean", [1.0, 2.0])
+@pytest.mark.parametrize("stdev", [0.8, 1.0])
 @pytest.mark.parametrize(
     ("samples", "rtol"),
     [
         (10, 1),
         (1000, 1e-1),
         (100000, 1e-2),
-        (100000, 1e-3),
+        (10000000, 1e-3),
     ],
 )
-def test_single_normal_node(samples, rtol):
-    normal = causalprog.graph.node.Distribution()
-    node = causalprog.graph.RootDistributionNode(normal, "X", is_outcome=True)
+def test_single_normal_node(samples, rtol, mean, stdev):
+    normal = causalprog.graph.node.NormalDistribution(mean, stdev)
+    node = causalprog.graph.DistributionNode(normal, "X", is_outcome=True)
 
     graph = causalprog.graph.Graph("G0")
     graph.add_node(node)
 
     assert np.isclose(
-        causalprog.algorithms.expectation(graph, samples=samples), 1.0, rtol=rtol
+        causalprog.algorithms.expectation(graph, samples=samples), mean, rtol=rtol
+    )
+    assert np.isclose(
+        causalprog.algorithms.standard_deviation(graph, samples=samples),
+        stdev,
+        rtol=rtol,
+    )
+
+
+@pytest.mark.parametrize("mean", [1.0, 2.0])
+@pytest.mark.parametrize("stdev", [0.8, 1.0])
+@pytest.mark.parametrize("stdev2", [0.8, 1.0])
+@pytest.mark.parametrize(
+    ("samples", "rtol"),
+    [
+        (100, 1),
+        (10000, 1e-1),
+        (1000000, 1e-2),
+    ],
+)
+def test_two_node_graph(samples, rtol, mean, stdev, stdev2):
+    normal = causalprog.graph.node.NormalDistribution(mean, stdev)
+    normal2 = causalprog.graph.node.NormalDistribution("UX", stdev2)
+
+    graph = causalprog.graph.Graph("G0")
+    graph.add_node(causalprog.graph.DistributionNode(normal, "UX"))
+    graph.add_node(causalprog.graph.DistributionNode(normal, "Y"))
+    graph.add_node(causalprog.graph.DistributionNode(normal, "Z"))
+    graph.add_node(causalprog.graph.DistributionNode(normal2, "X", is_outcome=True))
+    graph.add_edge("UX", "X")
+    graph.add_edge("Y", "Z")
+
+    assert np.isclose(
+        causalprog.algorithms.expectation(graph, samples=samples), mean, rtol=rtol
     )