Make distributions and graph talk to each other

src/causalprog/algorithms/expectation.py

@@ -1,15 +1,17 @@
 """Algorithms for estimating the expectation and standard deviation."""
 
-import numpy as np
+import jax
 import numpy.typing as npt
 
 from causalprog.graph import Graph
 
 
 def sample(
     graph: Graph,
-    outcome_node_label: str | None = None,
-    samples: int = 1000,
+    outcome_node_label: str,
+    samples: int,
+    *,
+    rng_key: jax.Array,
 ) -> npt.NDArray[float]:
     """Sample data from a graph."""
     if outcome_node_label is None:
@@ -18,24 +20,30 @@ def sample(
     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)
+    keys = jax.random.split(rng_key, len(nodes))
+
+    for node, key in zip(nodes, keys, strict=False):
+        values[node.label] = node.sample(values, samples, rng_key=key)
     return values[outcome_node_label]
 
 
 def expectation(
     graph: Graph,
-    outcome_node_label: str | None = None,
-    samples: int = 1000,
+    outcome_node_label: str,
+    samples: int,
+    *,
+    rng_key: jax.Array,
 ) -> float:
     """Estimate the expectation of a graph."""
-    return sample(graph, outcome_node_label, samples).mean()
+    return sample(graph, outcome_node_label, samples, rng_key=rng_key).mean()
 
 
 def standard_deviation(
     graph: Graph,
-    outcome_node_label: str | None = None,
-    samples: int = 1000,
+    outcome_node_label: str,
+    samples: int,
+    *,
+    rng_key: jax.Array,
 ) -> float:
     """Estimate the standard deviation of a graph."""
-    return np.std(sample(graph, outcome_node_label, samples))
+    return sample(graph, outcome_node_label, samples, rng_key=rng_key).std()

src/causalprog/distribution/family.py

@@ -39,8 +39,8 @@ class DistributionFamily(Generic[CreatesDistribution], Labelled):
     @property
     def _member(self) -> Callable[..., Distribution]:
         """Constructor method for family members, given parameters."""
-        return lambda *parameters: Distribution(
-            self._family(*parameters),
+        return lambda **parameters: Distribution(
+            self._family(**parameters),
             backend_translator=self._family_translator,
         )
 
@@ -67,13 +67,13 @@ def __init__(
         self._family = backend_family
         self._family_translator = backend_translator
 
-    def construct(self, *parameters: ArrayLike) -> Distribution:
+    def construct(self, **parameters: ArrayLike) -> Distribution:
         """
         Create a distribution from an explicit set of parameters.
 
         Args:
-            *parameters (ArrayLike): Parameters that define a member of this family,
+            **parameters (ArrayLike): Parameters that define a member of this family,
                 passed as sequential arguments.
 
         """
-        return self._member(*parameters)
+        return self._member(**parameters)

src/causalprog/distribution/normal.py

@@ -58,6 +58,8 @@ def __init__(self, mean: ArrayCompatible, cov: ArrayCompatible) -> None:
             cov (ArrayCompatible): Matrix of covariates, $\Sigma$.
 
         """
+        mean = jnp.atleast_1d(mean)
+        cov = jnp.atleast_2d(cov)
         super().__init__(_Normal(mean, cov), label=f"({mean.ndim}-dim) Normal")
 
 
@@ -76,7 +78,7 @@ def __init__(self) -> None:
         """Create a family of normal distributions."""
         super().__init__(Normal, family_name="Normal")
 
-    def construct(self, mean: ArrayCompatible, cov: ArrayCompatible) -> Normal:
+    def construct(self, mean: ArrayCompatible, cov: ArrayCompatible) -> Normal:  # type: ignore # noqa: PGH003
         r"""
         Construct a normal distribution with the given mean and covariates.
 
@@ -85,4 +87,4 @@ def construct(self, mean: ArrayCompatible, cov: ArrayCompatible) -> Normal:
             cov (ArrayCompatible): Matrix of covariates, $\Sigma$.
 
         """
-        return super().construct(mean, cov)
+        return super().construct(mean=mean, cov=cov)

src/causalprog/graph/node.py

@@ -3,48 +3,17 @@
 from __future__ import annotations
 
 import typing
-from abc import ABC, abstractmethod
+from abc import abstractmethod
 
+import jax
 import numpy as np
 
 if typing.TYPE_CHECKING:
     import numpy.typing as npt
 
-from causalprog._abc.labelled import Labelled
-
-
-class Distribution(ABC):
-    """Placeholder class."""
-
-    @abstractmethod
-    def sample(
-        self, sampled_dependencies: dict[str, npt.NDArray[float]], samples: int
-    ) -> npt.NDArray[float]:
-        """Sample."""
-
-
-class NormalDistribution(Distribution):
-    """Normal distribution."""
+    from causalprog.distribution.family import DistributionFamily
 
-    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
+from causalprog._abc.labelled import Labelled
 
 
 class Node(Labelled):
@@ -57,7 +26,10 @@ def __init__(self, label: str, *, is_outcome: bool = False) -> None:
 
     @abstractmethod
     def sample(
-        self, sampled_dependencies: dict[str, npt.NDArray[float]], samples: int
+        self,
+        sampled_dependencies: dict[str, npt.NDArray[float]],
+        samples: int,
+        rng_key: jax.Array,
     ) -> float:
         """Sample a value from the node."""
 
@@ -72,20 +44,40 @@ class DistributionNode(Node):
 
     def __init__(
         self,
-        distribution: Distribution,
+        distribution: DistributionFamily,
         label: str,
         *,
+        parameters: dict[str, str] | None = None,
+        constant_parameters: dict[str, float] | None = None,
         is_outcome: bool = False,
     ) -> None:
         """Initialise."""
         self._dist = distribution
+        self._constant_parameters = constant_parameters if constant_parameters else {}
+        self._parameters = parameters if parameters else {}
         super().__init__(label, is_outcome=is_outcome)
 
     def sample(
-        self, sampled_dependencies: dict[str, npt.NDArray[float]], samples: int
-    ) -> float:
+        self,
+        sampled_dependencies: dict[str, npt.NDArray[float]],
+        samples: int,
+        rng_key: jax.Array,
+    ) -> npt.NDArray[float]:
         """Sample a value from the node."""
-        return self._dist.sample(sampled_dependencies, samples)
+        if not self._parameters:
+            concrete_dist = self._dist.construct(**self._constant_parameters)
+            return concrete_dist.sample(rng_key, samples)
+        output = np.zeros(samples)
+        new_key = jax.random.split(rng_key, samples)
+        for sample in range(samples):
+            parameters = {
+                i: sampled_dependencies[j][sample] for i, j in self._parameters.items()
+            }
+            concrete_dist = self._dist.construct(
+                **parameters, **self._constant_parameters
+            )
+            output[sample] = concrete_dist.sample(new_key[sample], 1)[0][0]
+        return output
 
     def __repr__(self) -> str:
         return f'DistributionNode("{self.label}")'

tests/test_distributions/conftest.py → tests/conftest.py

@@ -15,7 +15,7 @@ def rng_key(seed: int):
 
 
 @pytest.fixture
-def n_dim_std_normal(request) -> tuple[Array, Array]:
+def n_dim_std_normal(request) -> dict[str, Array]:
     """
     Mean and covariance matrix of the n-dimensional standard normal distribution.
 
@@ -24,4 +24,4 @@ def n_dim_std_normal(request) -> tuple[Array, Array]:
     n_dims = request.param
     mean = jnp.array([0.0] * n_dims)
     cov = jnp.diag(jnp.array([1.0] * n_dims))
-    return mean, cov
+    return {"mean": mean, "cov": cov}

tests/test_distributions/test_family.py

@@ -19,8 +19,10 @@ def test_builder_matches_backend(n_dim_std_normal) -> None:
     mnv = distrax.MultivariateNormalFullCovariance
 
     mnv_family = DistributionFamily(mnv, SampleTranslator(rng_key="seed"))
-    via_family = mnv_family.construct(*n_dim_std_normal)
-    via_backend = mnv(*n_dim_std_normal)
+    via_family = mnv_family.construct(
+        loc=n_dim_std_normal["mean"], covariance_matrix=n_dim_std_normal["cov"]
+    )
+    via_backend = mnv(n_dim_std_normal["mean"], n_dim_std_normal["cov"])
 
     assert via_backend.kl_divergence(via_family.get_dist()) == pytest.approx(0.0)
     assert via_family.get_dist().kl_divergence(via_backend) == pytest.approx(0.0)