Add random variable nodes

src/causalprog/_abc/labelled.py

@@ -20,4 +20,7 @@ def label(self) -> str:
         return self._label
 
     def __init__(self, *, label: str) -> None:
+        if not str.isidentifier(label):
+            msg = f"Label is not valid Python variable name: {label}"
+            raise ValueError(msg)
         self._label = str(label)

src/causalprog/algorithms/do.py

@@ -68,7 +68,7 @@ def do(graph: Graph, node: str, value: float, label: str | None = None) -> Graph
 
     """
     if label is None:
-        label = f"{graph.label}|do({node}={value})"
+        label = f"{graph.label}_do_{node}__" + f"{value}".replace(".", "_")
 
     nodes = {n.label: deepcopy(n) for n in graph.nodes if n.label != node}
 
@@ -92,7 +92,7 @@ def do(graph: Graph, node: str, value: float, label: str | None = None) -> Graph
 
     nodes[node] = ConstantNode(label=node, value=value)
 
-    g = Graph(label=f"{label}|do[{node}={value}]")
+    g = Graph(label=f"{label}_do_{node}__" + f"{value}".replace(".", "_"))
     for n in nodes.values():
         g.add_node(n)
 

src/causalprog/graph/__init__.py

@@ -1,4 +1,12 @@
 """Creation and storage of graphs."""
 
 from .graph import Graph
-from .node import ComponentNode, ConstantNode, DataNode, DistributionNode, Node
+from .node import (
+    ComponentNode,
+    ConstantNode,
+    ContinuousRandomVariableNode,
+    DataNode,
+    DiscreteRandomVariableNode,
+    DistributionNode,
+    Node,
+)

src/causalprog/graph/node/__init__.py

@@ -5,3 +5,4 @@
 from .constant import ConstantNode
 from .data import DataNode
 from .distribution import DistributionNode
+from .random_variables import ContinuousRandomVariableNode, DiscreteRandomVariableNode

src/causalprog/graph/node/base.py

@@ -15,18 +15,18 @@
 def _to_string(indices: int | slice | tuple[int | slice, ...]) -> str:
     """Convert getitem indices to a string."""
     if isinstance(indices, tuple):
-        return ", ".join(_to_string(i) for i in indices)
+        return "_".join(_to_string(i) for i in indices)
     if isinstance(indices, int):
         return f"{indices}"
     if isinstance(indices, slice):
         s = ""
         if indices.start is not None:
             s += f"{indices.start}"
-        s += ":"
+        s += "__"
         if indices.stop is not None:
             s += f"{indices.stop}"
         if indices.step is not None:
-            s += f":{indices.step}"
+            s += f"__{indices.step}"
         return s
     e = f"Invalid indices: {indices}"
     raise TypeError(e)
@@ -91,7 +91,7 @@ def __getitem__(self, indices: int | slice | tuple[int | slice, ...]) -> Node:
             self.label,
             indices,
             shape=shape,
-            label=f"{self.label}[{_to_string(indices)}]",
+            label=f"{self.label}_{_to_string(indices)}",
         )
 
     @abstractmethod

src/causalprog/graph/node/random_variables.py

@@ -0,0 +1,150 @@
+"""Graph nodes representing random variables."""
+
+import inspect
+import typing
+from abc import abstractmethod
+
+import jax
+import numpy as np
+import numpy.typing as npt
+from typing_extensions import override
+
+from .base import Node
+
+
+class RandomVariableNode(Node):
+    """A node containing a random variable (RV)."""
+
+    def __init__(
+        self,
+        *,
+        shape: tuple[int, ...] = (),
+        label: str,
+        compute: typing.Callable | None = None,
+    ) -> None:
+        """
+        Initialise.
+
+        Args:
+            shape: The shape of the output of the RV
+            label: A unique label to identify the node
+            compute: A function to compute node's value from given values of parents
+
+        """
+        super().__init__(label=label, shape=shape)
+        if compute is None:
+            self._parents = []
+        else:
+            self._parents = list(inspect.signature(compute).parameters.keys())
+        self._compute = compute
+
+    @override
+    def sample(
+        self,
+        parameter_values: dict[str, float],
+        sampled_dependencies: dict[str, npt.ArrayLike],
+        samples: int,
+        *,
+        rng_key: jax.Array,
+    ) -> npt.ArrayLike:
+        raise NotImplementedError
+
+    @override
+    def evaluate(
+        self,
+        **given_values: float | npt.NDArray[float],
+    ) -> float | npt.NDArray[float]:
+        if self.label in given_values:
+            value = given_values[self.label]
+            self.assert_is_valid_value(value)
+            return value
+
+        if self._compute is None:
+            msg = f"Missing input for node: {self.label}."
+            raise ValueError(msg)
+        return self._compute(**{p: given_values[p] for p in self._parents})
+
+    @override
+    @property
+    def parents(self) -> list[str]:
+        return self._parents
+
+    @abstractmethod
+    def is_valid_value(self, value: float | npt.NDArray[float]) -> bool:
+        """Check if a value is valid for this node."""
+
+    def assert_is_valid_value(self, value: float | npt.NDArray[float]) -> None:
+        """Check if a value is valid for this node."""
+        if not self.is_valid_value(value):
+            msg = (
+                f"Invalid value for {self.__class__.__name__}: "
+                f"{self.label} cannot be {value}"
+            )
+            raise ValueError(msg)
+        if self.shape != (value.shape if hasattr(value, "shape") else ()):
+            msg = f"Invalid value for node: {self.label}"
+            raise ValueError(msg)
+
+
+class ContinuousRandomVariableNode(RandomVariableNode):
+    """A node containing a continuous random variable (RV)."""
+
+    @override
+    def __repr__(self) -> str:
+        return f'ContinuousRandomVariableNode(label="{self.label}")'
+
+    @override
+    def is_valid_value(self, value: float | npt.NDArray[float]) -> bool:
+        return True
+
+    @override
+    def copy(self) -> Node:
+        return ContinuousRandomVariableNode(
+            shape=self.shape, label=self.label, compute=self._compute
+        )
+
+
+class DiscreteRandomVariableNode(RandomVariableNode):
+    """A node containing a discrete random variable (RV)."""
+
+    def __init__(
+        self,
+        *,
+        values: list[float] | list[npt.NDArray[float]],
+        shape: tuple[int, ...] = (),
+        label: str,
+        compute: typing.Callable | None = None,
+    ) -> None:
+        """
+        Initialise.
+
+        Args:
+            shape: The shape of the output of the RV
+            label: A unique label to identify the node
+            compute: A function to compute node's value from given values of parents
+
+        """
+        super().__init__(label=label, shape=shape, compute=compute)
+        self._values = values
+
+    @property
+    def possible_values(self) -> list[float] | list[npt.NDArray[float]]:
+        """The values that this RV can take."""
+        return self._values
+
+    @override
+    def __repr__(self) -> str:
+        return f'DiscreteRandomVariableNode(label="{self.label}")'
+
+    @override
+    def is_valid_value(self, value: float | npt.NDArray[float]) -> bool:
+        return any(np.allclose(v, value) for v in self._values)
+
+    @override
+    def copy(self) -> Node:
+        return DiscreteRandomVariableNode(
+            values=self._values,
+            shape=self.shape,
+            label=self.label,
+            compute=self._compute,
+        )

tests/test__abc/test_labelled.py

@@ -0,0 +1,36 @@
+import pytest
+
+from causalprog._abc.labelled import Labelled
+
+
+@pytest.mark.parametrize(
+    "label",
+    [
+        "1",
+        " ",
+        "a b",
+        "0a",
+        "a.b",
+        "a-b",
+        "a+b",
+        "a*b",
+        "a/b",
+    ],
+)
+def test_invalid_label(label, raises_context):
+    with raises_context(ValueError("Label is not valid Python variable name")):
+        Labelled(label=label)
+
+
+@pytest.mark.parametrize(
+    "label",
+    [
+        "a",
+        "A",
+        "a0",
+        "a_b",
+        "_a",
+    ],
+)
+def test_valid_label(label):
+    Labelled(label=label)

tests/test_graph/test_component.py

@@ -26,6 +26,6 @@ def test_component_node(param_values):
     graph.add_node(graph.get_node("X")[1, :, 2])
 
     assert graph.get_node("X").shape == (3, 1, 4)
-    assert graph.get_node("X[0]").shape == (1, 4)
-    assert graph.get_node("X[1, 0, 2]").shape == ()
-    assert graph.get_node("X[1, :, 2]").shape == (1,)
+    assert graph.get_node("X_0").shape == (1, 4)
+    assert graph.get_node("X_1_0_2").shape == ()
+    assert graph.get_node("X_1____2").shape == (1,)

tests/test_graph/test_model.py

@@ -84,7 +84,7 @@ def test_model_extension(
         parameters={"loc": "mean"},
         constant_parameters={"scale": 1.0},
     )
-    one_normal_graph = Graph(label="One normal")
+    one_normal_graph = Graph(label="One_normal")
     one_normal_graph.add_edge(mean, x)
 
     def extended_model(*, cov2, **parameter_values):

tests/test_node/test_random_variables.py

@@ -0,0 +1,35 @@
+import numpy as np
+
+from causalprog.graph import (
+    ContinuousRandomVariableNode,
+    DiscreteRandomVariableNode,
+    Graph,
+)
+
+
+def test_random_variable_node():
+    node = ContinuousRandomVariableNode(label="X")
+    assert np.isclose(node.evaluate(X=2.0), 2.0)
+
+
+def test_missing_input(raises_context):
+    node = ContinuousRandomVariableNode(label="X")
+
+    with raises_context(ValueError("Missing input for node")):
+        node.evaluate()
+
+
+def test_invalid_discrete_node_value(raises_context):
+    node = DiscreteRandomVariableNode(label="Y", values=[-0.5, 0.0, 0.5])
+    with raises_context(ValueError("Invalid value for DiscreteRandomVariableNode")):
+        node.evaluate(Y=-1.0)
+
+
+def test_evaluate_down_graph():
+    graph = Graph(label="G")
+
+    graph.add_node(ContinuousRandomVariableNode(label="x"))
+    graph.add_node(DiscreteRandomVariableNode(label="y", values=[-0.5, 0.0, 0.5]))
+    graph.add_node(ContinuousRandomVariableNode(label="z", compute=lambda x, y: x + y))
+
+    assert np.isclose(graph.get_node("z").evaluate(x=3.0, y=-0.5), 2.5)

tests/test_node/test_vector_nodes.py

@@ -45,6 +45,6 @@ def test_get_component_node(rng_key):
     assert s0.shape == (100, 5)
     assert node[:, 1].sample({}, {"X": s}, 100, rng_key=rng_key).shape == (100, 4)
     assert node[0, 1].sample({}, {"X": s}, 100, rng_key=rng_key).shape == (100,)
-    assert node[0][1].sample({}, {"X": s, "X[0]": s0}, 100, rng_key=rng_key).shape == (
+    assert node[0][1].sample({}, {"X": s, "X_0": s0}, 100, rng_key=rng_key).shape == (
         100,
     )