diff --git a/src/causalprog/_abc/labelled.py b/src/causalprog/_abc/labelled.py index 92250bc..b62770e 100644 --- a/src/causalprog/_abc/labelled.py +++ b/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) diff --git a/src/causalprog/algorithms/do.py b/src/causalprog/algorithms/do.py index 9b53b97..9c0d6f9 100644 --- a/src/causalprog/algorithms/do.py +++ b/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) diff --git a/src/causalprog/graph/__init__.py b/src/causalprog/graph/__init__.py index 0825556..02eb604 100644 --- a/src/causalprog/graph/__init__.py +++ b/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, +) diff --git a/src/causalprog/graph/node/__init__.py b/src/causalprog/graph/node/__init__.py index 1698224..a6acf66 100644 --- a/src/causalprog/graph/node/__init__.py +++ b/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 diff --git a/src/causalprog/graph/node/base.py b/src/causalprog/graph/node/base.py index ee6a89a..e7d4f54 100644 --- a/src/causalprog/graph/node/base.py +++ b/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 diff --git a/src/causalprog/graph/node/random_variables.py b/src/causalprog/graph/node/random_variables.py new file mode 100644 index 0000000..9f5dd38 --- /dev/null +++ b/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, + ) diff --git a/tests/test__abc/test_labelled.py b/tests/test__abc/test_labelled.py new file mode 100644 index 0000000..3f17704 --- /dev/null +++ b/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) diff --git a/tests/test_graph/test_component.py b/tests/test_graph/test_component.py index 9a417b8..3c6b356 100644 --- a/tests/test_graph/test_component.py +++ b/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,) diff --git a/tests/test_graph/test_model.py b/tests/test_graph/test_model.py index 7a18b0d..1f3e38d 100644 --- a/tests/test_graph/test_model.py +++ b/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): diff --git a/tests/test_node/test_random_variables.py b/tests/test_node/test_random_variables.py new file mode 100644 index 0000000..3905806 --- /dev/null +++ b/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) diff --git a/tests/test_node/test_vector_nodes.py b/tests/test_node/test_vector_nodes.py index 39230ef..44e848e 100644 --- a/tests/test_node/test_vector_nodes.py +++ b/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, )