Merge branch 'wgraham/signature-converting' into wgraham/distributions-are-backend-agnostic

Author: willGraham01

pyproject.toml

@@ -15,7 +15,11 @@ classifiers = [
     "Programming Language :: Python :: 3.13",
     "Typing :: Typed",
 ]
-dependencies = ["jax", "networkx"]
+dependencies = [
+    "jax",
+    "networkx",
+    "numpy",
+]
 description = "A Python package for causal modelling and inference with stochastic causal programming"
 dynamic = ["version"]
 keywords = []
@@ -35,6 +39,7 @@ optional-dependencies = {dev = [
     "mkdocstrings-python",
 ], test = [
     "distrax",
+    "numpy",
     "numpyro",
     "pytest",
     "pytest-cov",
@@ -75,6 +80,7 @@ lint.per-file-ignores = {"__init__.py" = [
     "ANN",
     "D",
     "INP001", # File is part of an implicit namespace package.
+    "PLR0913", # Too many arguments in function definition
     "S101", # Use of `assert` detected
 ]}
 lint.select = ["ALL"]

src/causalprog/__init__.py

@@ -1,4 +1,4 @@
 """causalprog package."""
 
-from . import graph
+from . import algorithms, distribution, graph, utils
 from ._version import __version__

src/causalprog/algorithms/__init__.py

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

src/causalprog/algorithms/expectation.py

@@ -0,0 +1,41 @@
+"""Algorithms for estimating the expectation and standard deviation."""
+
+import numpy as np
+import numpy.typing as npt
+
+from causalprog.graph import Graph
+
+
+def sample(
+    graph: Graph,
+    outcome_node_label: str | None = None,
+    samples: int = 1000,
+) -> npt.NDArray[float]:
+    """Sample data from a graph."""
+    if outcome_node_label is None:
+        outcome_node_label = graph.outcome.label
+
+    nodes = graph.roots_down_to_outcome(outcome_node_label)
+
+    values: dict[str, npt.NDArray[float]] = {}
+    for node in nodes:
+        values[node.label] = node.sample(values, 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/graph/__init__.py

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

src/causalprog/graph/graph.py

@@ -10,24 +10,98 @@
 class Graph(Labelled):
     """A directed acyclic graph that represents a causality tree."""
 
-    def __init__(self, graph: nx.Graph, label: str) -> None:
-        """Initialise a graph from a NetworkX graph."""
+    _nodes_by_label: dict[str, Node]
+
+    def __init__(self, label: str) -> None:
+        """Create end empty graph."""
         super().__init__(label=label)
+        self._graph = nx.DiGraph()
+        self._nodes_by_label = {}
+
+    def get_node(self, label: str) -> Node:
+        """Get a node from its label."""
+        node = self._nodes_by_label.get(label, None)
+        if not node:
+            msg = f'Node not found with label "{label}"'
+            raise KeyError(msg)
+        return node
+
+    def add_node(self, node: Node) -> None:
+        """Add a node to the graph."""
+        if node.label in self._nodes_by_label:
+            msg = f"Duplicate node label: {node.label}"
+            raise ValueError(msg)
+        self._nodes_by_label[node.label] = node
+        self._graph.add_node(node)
+
+    def add_edge(self, first_node: Node | str, second_node: Node | str) -> None:
+        """
+        Add an edge to the graph.
 
-        for node in graph.nodes:
-            if not isinstance(node, Node):
-                msg = f"Invalid node: {node}"
-                raise TypeError(msg)
+        Adding an edge between nodes not currently in the graph,
+        will cause said nodes to be added to the graph along with
+        the edge.
+        """
+        if isinstance(first_node, str):
+            first_node = self.get_node(first_node)
+        if isinstance(second_node, str):
+            second_node = self.get_node(second_node)
+        if first_node.label not in self._nodes_by_label:
+            self.add_node(first_node)
+        if second_node.label not in self._nodes_by_label:
+            self.add_node(second_node)
+        for node_to_check in (first_node, second_node):
+            if node_to_check != self._nodes_by_label[node_to_check.label]:
+                msg = "Invalid node: {node_to_check}"
+                raise ValueError(msg)
+        self._graph.add_edge(first_node, second_node)
 
-        self._graph = graph.copy()
-        self._nodes = list(graph.nodes())
-        self._depth_first_nodes = list(nx.algorithms.dfs_postorder_nodes(graph))
+    @property
+    def predecessors(self) -> dict[Node, Node]:
+        """Get predecessors of every node."""
+        return nx.algorithms.dfs_predecessors(self._graph)
 
-        outcomes = [node for node in self._nodes if node.is_outcome]
+    @property
+    def successors(self) -> dict[Node, list[Node]]:
+        """Get successors of every node."""
+        return nx.algorithms.dfs_successors(self._graph)
+
+    @property
+    def outcome(self) -> Node:
+        """The outcome node of the graph."""
+        outcomes = [node for node in self.nodes if node.is_outcome]
         if len(outcomes) == 0:
             msg = "Cannot create graph with no outcome nodes"
             raise ValueError(msg)
         if len(outcomes) > 1:
             msg = "Cannot yet create graph with multiple outcome nodes"
             raise ValueError(msg)
-        self._outcome = outcomes[0]
+        return outcomes[0]
+
+    @property
+    def nodes(self) -> list[Node]:
+        """The nodes of the graph."""
+        return list(self._graph.nodes())
+
+    @property
+    def ordered_nodes(self) -> list[Node]:
+        """Nodes ordered so that each node appears after its dependencies."""
+        if not nx.is_directed_acyclic_graph(self._graph):
+            msg = "Graph is not acyclic."
+            raise RuntimeError(msg)
+        return list(nx.topological_sort(self._graph))
+
+    def roots_down_to_outcome(
+        self,
+        outcome_node_label: str,
+    ) -> list[Node]:
+        """
+        Get ordered list of nodes that outcome depends on.
+
+        Nodes are ordered so that each node appears after its dependencies.
+        """
+        outcome = self.get_node(outcome_node_label)
+        ancestors = nx.ancestors(self._graph, outcome)
+        return [
+            node for node in self.ordered_nodes if node == outcome or node in ancestors
+        ]

src/causalprog/graph/node.py

@@ -2,97 +2,90 @@
 
 from __future__ import annotations
 
-from abc import abstractmethod
-from typing import Protocol, runtime_checkable
+import typing
+from abc import ABC, abstractmethod
 
-from causalprog._abc.labelled import Labelled
+import numpy as np
 
+if typing.TYPE_CHECKING:
+    import numpy.typing as npt
 
-class DistributionFamily:
-    """Placeholder class."""
+from causalprog._abc.labelled import Labelled
 
 
-class Distribution:
+class Distribution(ABC):
     """Placeholder class."""
 
-
-@runtime_checkable
-class Node(Protocol):
-    """An abstract node in a graph."""
-
-    @property
-    @abstractmethod
-    def label(self) -> str:
-        """The label of the node."""
-
-    @property
     @abstractmethod
-    def is_root(self) -> bool:
-        """Identify if the node is a root."""
-
-    @property
-    @abstractmethod
-    def is_outcome(self) -> bool:
-        """Identify if the node is an outcome."""
-
-
-class RootDistributionNode(Labelled):
-    """A root node containing a distribution family."""
+    def sample(
+        self, sampled_dependencies: dict[str, npt.NDArray[float]], samples: int
+    ) -> npt.NDArray[float]:
+        """Sample."""
+
+
+class NormalDistribution(Distribution):
+    """Normal distribution."""
+
+    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."""
+        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(Labelled):
+    """An abstract node in a graph."""
 
-    def __init__(
-        self,
-        family: DistributionFamily,
-        label: str,
-        *,
-        is_outcome: bool = False,
-    ) -> None:
-        """Initialise the node."""
+    def __init__(self, label: str, *, is_outcome: bool = False) -> None:
+        """Initialise."""
         super().__init__(label=label)
+        self._is_outcome = is_outcome
 
-        self._dfamily = family
-        self._outcome = is_outcome
-
-    def __repr__(self) -> str:
-        """Representation."""
-        return f'RootDistributionNode("{self._label}")'
-
-    @property
-    def is_root(self) -> bool:
-        """Identify if the node is a root."""
-        return True
+    @abstractmethod
+    def sample(
+        self, sampled_dependencies: dict[str, npt.NDArray[float]], samples: int
+    ) -> float:
+        """Sample a value from the node."""
 
     @property
     def is_outcome(self) -> bool:
         """Identify if the node is an outcome."""
-        return self._outcome
+        return self._is_outcome
 
 
-class DistributionNode(Labelled):
-    """A node containing a distribution family that depends on its parents."""
+class DistributionNode(Node):
+    """A node containing a distribution."""
 
     def __init__(
         self,
-        family: DistributionFamily,
+        distribution: Distribution,
         label: str,
         *,
         is_outcome: bool = False,
     ) -> None:
-        """Initialise the node."""
-        super().__init__(label=label)
+        """Initialise."""
+        self._dist = distribution
+        super().__init__(label, is_outcome=is_outcome)
 
-        self._dfamily = family
-        self._outcome = is_outcome
+    def sample(
+        self, sampled_dependencies: dict[str, npt.NDArray[float]], samples: int
+    ) -> float:
+        """Sample a value from the node."""
+        return self._dist.sample(sampled_dependencies, samples)
 
     def __repr__(self) -> str:
-        """Representation."""
-        return f'DistributionNode("{self._label}")'
-
-    @property
-    def is_root(self) -> bool:
-        """Identify if the node is a root."""
-        return False
-
-    @property
-    def is_outcome(self) -> bool:
-        """Identify if the node is an outcome."""
-        return self._outcome
+        return f'DistributionNode("{self.label}")'

tests/test_backend/test_signature_can_be_cast.py

@@ -160,7 +160,7 @@
         ),
     ],
 )
-def test_signature_can_be_cast(  # noqa: PLR0913
+def test_signature_can_be_cast(
     signature_to_convert: Signature,
     new_signature: Signature,
     old_to_new_names: ParamNameMap,