Skeleton DistributionFamilys (#12)

* Ignore the .vscode folder before I inevitably commit it at least once

* Skeleton base class for distributions

* Base class for distribution families

* Distributions sample method is backend-agnostic

* Fix return type of sample method

* Immutable type for storing backend-compatibility info

* Remove distributon family for now

* Clean up imports and typevars

* Basic docstrings

* Give useful information in docstrings

* Write tests for SampleCompatibility class

* Write test for Distribution

* tests now require multiple backends

* Create a Translator class as you're going to need to do this a lot, Will

* Fix the tests that I broke

* Add some more detailed docstrings

* Test the translation method of translators

* Fix docstring linting issues

* Module name to singular to be consistent with #5

* Skeleton for distribution family

* Add test for distribution family builder method

* Expose getter method for Distribution backends

* Translator is optional

* Fix docstring

* Add note about how distribution families are created

Author: willGraham01

src/causalprog/distribution/base.py

@@ -64,6 +64,10 @@ def __init__(
         )
         self._backend_translator.validate_compatible(backend_distribution)
 
+    def get_dist(self) -> SupportsSampling:
+        """Access to the backend distribution."""
+        return self._dist
+
     def sample(self, rng_key: SupportsRNG, sample_shape: ArrayLike = ()) -> ArrayLike:
         """
         Draw samples from the distribution.

src/causalprog/distribution/family.py

@@ -0,0 +1,73 @@
+"""Parametrised groups of ``Distribution``s."""
+
+from collections.abc import Callable
+from typing import Generic, TypeVar
+
+from numpy.typing import ArrayLike
+
+from causalprog.distribution.base import Distribution, SupportsSampling
+from causalprog.utils.translator import Translator
+
+CreatesDistribution = TypeVar(
+    "CreatesDistribution", bound=Callable[..., SupportsSampling]
+)
+
+
+class DistributionFamily(Generic[CreatesDistribution]):
+    r"""
+    A family of ``Distributions``, that share the same parameters.
+
+    A ``DistributionFamily`` is essentially a ``Distribution`` that has not yet had its
+    parameter values explicitly specified. Explicit values for the parameters can be
+    passed to a ``DistributionFamily``'s ``construct`` method, which will then proceed
+    to construct a ``Distribution`` with those parameter values.
+
+    As an explicit example, the (possibly multivariate) normal distribution is
+    parametrised by two quantities - the (vector of) mean values $\mu$ and covariates
+    $\Sigma$. A ``DistributionFamily`` represents this general
+    $\mathcal{N}(\mu, \Sigma)$ parametrised form, however without explicit $\mu$ and
+    $\Sigma$ values we cannot perform operations like drawing samples. Specifying, for
+    example, $\mu = 0$ and $\Sigma = 1$ by invoking ``.construct(0., 1.)`` will return a
+    ``Distribution`` instance representing $\mathcal{N}(0., 1.)$, which can then have
+    samples drawn from it.
+    """
+
+    _family: CreatesDistribution
+    _family_translator: Translator | None
+
+    @property
+    def _member(self) -> Callable[..., Distribution]:
+        """Constructor method for family members, given parameters."""
+        return lambda *parameters: Distribution(
+            self._family(*parameters), backend_translator=self._family_translator
+        )
+
+    def __init__(
+        self,
+        backend_family: CreatesDistribution,
+        backend_translator: Translator | None = None,
+    ) -> None:
+        """
+        Create a new family of distributions.
+
+        Args:
+            backend_family (CreatesDistribution): Backend callable that assembles the
+                distribution, given explicit parameter values. Currently, this callable
+                can only accept the parameters as a sequence of positional arguments.
+            backend_translator (Translator): ``Translator`` instance that to be
+                passed to the ``Distribution`` constructor.
+
+        """
+        self._family = backend_family
+        self._family_translator = backend_translator
+
+    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,
+                passed as sequential arguments.
+
+        """
+        return self._member(*parameters)

tests/test_distributions/conftest.py

@@ -0,0 +1,27 @@
+import jax.numpy as jnp
+import jax.random as jrn
+import pytest
+from jax._src.basearray import Array
+
+
+@pytest.fixture
+def seed() -> int:
+    return 0
+
+
+@pytest.fixture
+def rng_key(seed: int):
+    return jrn.key(seed)
+
+
+@pytest.fixture
+def n_dim_std_normal(request) -> tuple[Array, Array]:
+    """
+    Mean and covariance matrix of the n-dimensional standard normal distribution.
+
+    ``request.param`` should be an integer corresponding to the number of dimensions.
+    """
+    n_dims = request.param
+    mean = jnp.array([0.0] * n_dims)
+    cov = jnp.diag(jnp.array([1.0] * n_dims))
+    return mean, cov

tests/test_distributions/test_different_backends.py

@@ -2,13 +2,12 @@
 
 import distrax
 import jax.numpy as jnp
-import jax.random as jrn
 from numpyro.distributions.continuous import MultivariateNormal
 
 from causalprog.distribution.base import Distribution, SampleTranslator
 
 
-def test_different_backends() -> None:
+def test_different_backends(rng_key) -> None:
     """
     Test that ``Distribution`` can use different (but equivalent) backends.
 
@@ -20,15 +19,14 @@ def test_different_backends() -> None:
     n_dims = 2
     mean = jnp.array([0.0] * n_dims)
     cov = jnp.diag(jnp.array([1.0] * n_dims))
-    rng = jrn.key(0)
     sample_size = (10, 5)
 
     distrax_normal = distrax.MultivariateNormalFullCovariance(mean, cov)
     distrax_dist = Distribution(distrax_normal, SampleTranslator(rng_key="seed"))
-    distrax_samples = distrax_dist.sample(rng, sample_size)
+    distrax_samples = distrax_dist.sample(rng_key, sample_size)
 
     npyo_normal = MultivariateNormal(mean, cov)
     npyo_dist = Distribution(npyo_normal, SampleTranslator(rng_key="key"))
-    npyo_samples = npyo_dist.sample(rng, sample_size)
+    npyo_samples = npyo_dist.sample(rng_key, sample_size)
 
     assert jnp.allclose(distrax_samples, npyo_samples)

tests/test_distributions/test_family.py

@@ -0,0 +1,26 @@
+import distrax
+import pytest
+
+from causalprog.distribution.base import SampleTranslator
+from causalprog.distribution.family import DistributionFamily
+
+
+@pytest.mark.parametrize(
+    ("n_dim_std_normal"),
+    [pytest.param(2, id="2D normal")],
+    indirect=["n_dim_std_normal"],
+)
+def test_builder_matches_backend(n_dim_std_normal) -> None:
+    """
+    Test that building from a family is equivalent
+    to building via the backend explicitly.
+
+    """
+    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)
+
+    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)