Merge pull request #77 from pymc-devs/main

CAR random variables (pymc-devs#4596)

Author: sthagen

pymc3/distributions/multivariate.py

@@ -24,8 +24,10 @@
 import numpy as np
 import scipy
 
+from aesara.assert_op import Assert
 from aesara.graph.basic import Apply
 from aesara.graph.op import Op
+from aesara.sparse.basic import sp_sum
 from aesara.tensor import gammaln, sigmoid
 from aesara.tensor.nlinalg import det, eigh, matrix_inverse, trace
 from aesara.tensor.random.basic import MultinomialRV, dirichlet, multivariate_normal
@@ -397,7 +399,6 @@ def __new__(cls, name, *args, **kwargs):
 
     @classmethod
     def dist(cls, a, **kwargs):
-
         a = at.as_tensor_variable(a)
         # mean = a / at.sum(a)
         # mode = at.switch(at.all(a > 1), (a - 1) / at.sum(a - 1), np.nan)
@@ -491,7 +492,6 @@ class Multinomial(Discrete):
 
     @classmethod
     def dist(cls, n, p, *args, **kwargs):
-
         p = p / at.sum(p, axis=-1, keepdims=True)
         n = at.as_tensor_variable(n)
         p = at.as_tensor_variable(p)
@@ -1925,6 +1925,81 @@ def _distr_parameters_for_repr(self):
         return ["mu"]
 
 
+class CARRV(RandomVariable):
+    name = "car"
+    ndim_supp = 1
+    ndims_params = [1, 2, 0, 0]
+    dtype = "floatX"
+    _print_name = ("CAR", "\\operatorname{CAR}")
+
+    def make_node(self, rng, size, dtype, mu, W, alpha, tau):
+        mu = at.as_tensor_variable(floatX(mu))
+
+        W = aesara.sparse.as_sparse_or_tensor_variable(floatX(W))
+        if not W.ndim == 2:
+            raise ValueError("W must be a matrix (ndim=2).")
+
+        sparse = isinstance(W, aesara.sparse.SparseVariable)
+        msg = "W must be a symmetric adjacency matrix."
+        if sparse:
+            abs_diff = aesara.sparse.basic.mul(aesara.sparse.basic.sgn(W - W.T), W - W.T)
+            W = Assert(msg)(W, at.isclose(aesara.sparse.basic.sp_sum(abs_diff), 0))
+        else:
+            W = Assert(msg)(W, at.allclose(W, W.T))
+
+        tau = at.as_tensor_variable(floatX(tau))
+        alpha = at.as_tensor_variable(floatX(alpha))
+
+        return super().make_node(rng, size, dtype, mu, W, alpha, tau)
+
+    def _infer_shape(self, size, dist_params, param_shapes=None):
+        shape = tuple(size) + tuple(dist_params[0].shape)
+        return shape
+
+    @classmethod
+    def rng_fn(cls, rng: np.random.RandomState, mu, W, alpha, tau, size):
+        """
+        Implementation of algorithm from paper
+        Havard Rue, 2001. "Fast sampling of Gaussian Markov random fields,"
+        Journal of the Royal Statistical Society Series B, Royal Statistical Society,
+        vol. 63(2), pages 325-338. DOI: 10.1111/1467-9868.00288
+        """
+        if not scipy.sparse.issparse(W):
+            W = scipy.sparse.csr_matrix(W)
+        s = np.asarray(W.sum(axis=0))[0]
+        D = scipy.sparse.diags(s)
+        tau = scipy.sparse.csr_matrix(tau)
+        alpha = scipy.sparse.csr_matrix(alpha)
+
+        Q = tau.multiply(D - alpha.multiply(W))
+
+        perm_array = scipy.sparse.csgraph.reverse_cuthill_mckee(Q, symmetric_mode=True)
+        inv_perm = np.argsort(perm_array)
+
+        Q = Q[perm_array, :][:, perm_array]
+
+        Qb = Q.diagonal()
+        u = 1
+        while np.count_nonzero(Q.diagonal(u)) > 0:
+            Qb = np.vstack((np.pad(Q.diagonal(u), (u, 0), constant_values=(0, 0)), Qb))
+            u += 1
+
+        L = scipy.linalg.cholesky_banded(Qb, lower=False)
+
+        size = tuple(size or ())
+        if size:
+            mu = np.broadcast_to(mu, size + mu.shape)
+        z = rng.normal(size=mu.shape)
+        samples = np.empty(z.shape)
+        for idx in np.ndindex(mu.shape[:-1]):
+            samples[idx] = scipy.linalg.cho_solve_banded((L, False), z[idx]) + mu[idx][perm_array]
+        samples = samples[..., inv_perm]
+        return samples
+
+
+car = CARRV()
+
+
 class CAR(Continuous):
     r"""
     Likelihood for a conditional autoregression. This is a special case of the
@@ -1966,45 +2041,13 @@ class CAR(Continuous):
         "Generalized Hierarchical Multivariate CAR Models for Areal Data"
         Biometrics, Vol. 61, No. 4 (Dec., 2005), pp. 950-961
     """
+    rv_op = car
 
-    def __init__(self, mu, W, alpha, tau, sparse=False, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
-        D = W.sum(axis=0)
-        d, _ = W.shape
-
-        self.d = d
-        self.median = self.mode = self.mean = self.mu = at.as_tensor_variable(mu)
-        self.sparse = sparse
-
-        if not W.ndim == 2 or not np.allclose(W, W.T):
-            raise ValueError("W must be a symmetric adjacency matrix.")
-
-        if sparse:
-            W_sparse = scipy.sparse.csr_matrix(W)
-            self.W = aesara.sparse.as_sparse_variable(W_sparse)
-        else:
-            self.W = at.as_tensor_variable(W)
-
-        # eigenvalues of D^−1/2 * W * D^−1/2
-        Dinv_sqrt = np.diag(1 / np.sqrt(D))
-        DWD = np.matmul(np.matmul(Dinv_sqrt, W), Dinv_sqrt)
-        self.lam = scipy.linalg.eigvalsh(DWD)
-        self.D = at.as_tensor_variable(D)
-
-        tau = at.as_tensor_variable(tau)
-        if tau.ndim > 0:
-            self.tau = tau[:, None]
-        else:
-            self.tau = tau
-
-        alpha = at.as_tensor_variable(alpha)
-        if alpha.ndim > 0:
-            self.alpha = alpha[:, None]
-        else:
-            self.alpha = alpha
+    @classmethod
+    def dist(cls, mu, W, alpha, tau, *args, **kwargs):
+        return super().dist([mu, W, alpha, tau], **kwargs)
 
-    def logp(self, value):
+    def logp(value, mu, W, alpha, tau):
         """
         Calculate log-probability of a CAR-distributed vector
         at specified value. This log probability function differs from
@@ -2021,30 +2064,37 @@ def logp(self, value):
         TensorVariable
         """
 
+        sparse = isinstance(W, aesara.sparse.SparseVariable)
+
+        if sparse:
+            D = sp_sum(W, axis=0)
+            Dinv_sqrt = at.diag(1 / at.sqrt(D))
+            DWD = at.dot(aesara.sparse.dot(Dinv_sqrt, W), Dinv_sqrt)
+        else:
+            D = W.sum(axis=0)
+            Dinv_sqrt = at.diag(1 / at.sqrt(D))
+            DWD = at.dot(at.dot(Dinv_sqrt, W), Dinv_sqrt)
+        lam = at.slinalg.eigvalsh(DWD, at.eye(DWD.shape[0]))
+
+        d, _ = W.shape
+
         if value.ndim == 1:
             value = value[None, :]
 
-        logtau = self.d * at.log(self.tau).sum()
-        logdet = at.log(1 - self.alpha.T * self.lam[:, None]).sum()
-        delta = value - self.mu
+        logtau = d * at.log(tau).sum()
+        logdet = at.log(1 - alpha.T * lam[:, None]).sum()
+        delta = value - mu
 
-        if self.sparse:
-            Wdelta = aesara.sparse.dot(delta, self.W)
+        if sparse:
+            Wdelta = aesara.sparse.dot(delta, W)
         else:
-            Wdelta = at.dot(delta, self.W)
+            Wdelta = at.dot(delta, W)
 
-        tau_dot_delta = self.D[None, :] * delta - self.alpha * Wdelta
-        logquad = (self.tau * delta * tau_dot_delta).sum(axis=-1)
+        tau_dot_delta = D[None, :] * delta - alpha * Wdelta
+        logquad = (tau * delta * tau_dot_delta).sum(axis=-1)
         return bound(
             0.5 * (logtau + logdet - logquad),
-            self.alpha >= -1,
-            self.alpha <= 1,
-            self.tau > 0,
-            broadcast_conditions=False,
+            at.all(alpha <= 1),
+            at.all(alpha >= -1),
+            tau > 0,
         )
-
-    def random(self, point=None, size=None):
-        raise NotImplementedError("Sampling from a CAR distribution is not supported.")
-
-    def _distr_parameters_for_repr(self):
-        return ["mu", "W", "alpha", "tau"]

pymc3/tests/conftest.py

@@ -18,11 +18,12 @@
 
 import pymc3 as pm
 
-# @pytest.fixture(scope="function", autouse=True)
-# def aesara_config():
-#     config = aesara.config.change_flags(compute_test_value="raise")
-#     with config:
-#         yield
+
+@pytest.fixture(scope="function", autouse=True)
+def aesara_config():
+    config = aesara.config.change_flags(on_opt_error="raise")
+    with config:
+        yield
 
 
 @pytest.fixture(scope="function", autouse=True)

pymc3/tests/test_distributions.py

@@ -3025,41 +3025,87 @@ def test_ordered_probit_probs():
     assert isinstance(x, TensorVariable)
 
 
-@pytest.mark.xfail(reason="Distribution not refactored yet")
-@pytest.mark.parametrize("size", [(1,), (4,)], ids=str)
-def test_car_logp(size):
+@pytest.mark.parametrize(
+    "sparse, size",
+    [(False, ()), (False, (1,)), (False, (4,)), (False, (4, 4, 4)), (True, ()), (True, (4,))],
+    ids=str,
+)
+def test_car_logp(sparse, size):
     """
     Tests the log probability function for the CAR distribution by checking
     against Scipy's multivariate normal logpdf, up to an additive constant.
     The formula used by the CAR logp implementation omits several additive terms.
     """
     np.random.seed(1)
 
-    xs = np.random.randn(*size)
-
     # d x d adjacency matrix for a square (d=4) of rook-adjacent sites
     W = np.array(
         [[0.0, 1.0, 1.0, 0.0], [1.0, 0.0, 0.0, 1.0], [1.0, 0.0, 0.0, 1.0], [0.0, 1.0, 1.0, 0.0]]
     )
 
-    tau = 2.0
+    tau = 2
     alpha = 0.5
     mu = np.zeros(4)
 
+    xs = np.random.randn(*(size + mu.shape))
+
     # Compute CAR covariance matrix and resulting MVN logp
     D = W.sum(axis=0)
     prec = tau * (np.diag(D) - alpha * W)
     cov = np.linalg.inv(prec)
     scipy_logp = scipy.stats.multivariate_normal.logpdf(xs, mu, cov)
 
-    car_logp = logp(CAR.dist(mu, W, alpha, tau, size=size), xs).eval()
+    W = aesara.tensor.as_tensor_variable(W)
+    if sparse:
+        W = aesara.sparse.csr_from_dense(W)
+
+    car_dist = CAR.dist(mu, W, alpha, tau, size=size)
+    car_logp = logp(car_dist, xs).eval()
 
     # Check to make sure that the CAR and MVN log PDFs are equivalent
     # up to an additive constant which is independent of the CAR parameters
     delta_logp = scipy_logp - car_logp
 
     # Check to make sure all the delta values are identical.
-    assert np.allclose(delta_logp - delta_logp[0], 0.0)
+    tol = 1e-08
+    if aesara.config.floatX == "float32":
+        tol = 1e-5
+    assert np.allclose(delta_logp - delta_logp[0], 0.0, atol=tol)
+
+
+@pytest.mark.parametrize(
+    "sparse",
+    [False, True],
+    ids=str,
+)
+def test_car_matrix_check(sparse):
+    """
+    Tests the check of W matrix symmetry in CARRV.make_node.
+    """
+    np.random.seed(1)
+    tau = 2
+    alpha = 0.5
+    mu = np.zeros(4)
+    xs = np.random.randn(*mu.shape)
+
+    # non-symmetric matrix
+    W = np.array(
+        [[0.0, 1.0, 2.0, 0.0], [1.0, 0.0, 0.0, 1.0], [1.0, 0.0, 0.0, 1.0], [0.0, 1.0, 1.0, 0.0]]
+    )
+    W = aesara.tensor.as_tensor_variable(W)
+    if sparse:
+        W = aesara.sparse.csr_from_dense(W)
+
+    car_dist = CAR.dist(mu, W, alpha, tau)
+    with pytest.raises(AssertionError, match="W must be a symmetric adjacency matrix"):
+        logp(car_dist, xs).eval()
+
+    # W.ndim != 2
+    if not sparse:
+        W = np.array([0.0, 1.0, 2.0, 0.0])
+        W = aesara.tensor.as_tensor_variable(W)
+        with pytest.raises(ValueError, match="W must be a matrix"):
+            car_dist = CAR.dist(mu, W, alpha, tau)
 
 
 class TestBugfixes:

pymc3/tests/test_distributions_random.py

@@ -2380,3 +2380,43 @@ def test_with_cov_rv(self, sample_shape, dist_shape, mu_shape):
             prior = pm.sample_prior_predictive(samples=sample_shape)
 
         assert prior["mv"].shape == to_tuple(sample_shape) + dist_shape
+
+
+@pytest.mark.parametrize("sparse", [True, False])
+def test_car_rng_fn(sparse):
+    delta = 0.05  # limit for KS p-value
+    n_fails = 20  # Allows the KS fails a certain number of times
+    size = (100,)
+
+    W = np.array(
+        [[0.0, 1.0, 1.0, 0.0], [1.0, 0.0, 0.0, 1.0], [1.0, 0.0, 0.0, 1.0], [0.0, 1.0, 1.0, 0.0]]
+    )
+
+    tau = 2
+    alpha = 0.5
+    mu = np.array([1, 1, 1, 1])
+
+    D = W.sum(axis=0)
+    prec = tau * (np.diag(D) - alpha * W)
+    cov = np.linalg.inv(prec)
+    W = aesara.tensor.as_tensor_variable(W)
+    if sparse:
+        W = aesara.sparse.csr_from_dense(W)
+
+    with pm.Model(rng_seeder=1):
+        car = pm.CAR("car", mu, W, alpha, tau, size=size)
+        mn = pm.MvNormal("mn", mu, cov, size=size)
+        check = pm.sample_prior_predictive(n_fails)
+
+    p, f = delta, n_fails
+    while p <= delta and f > 0:
+        car_smp, mn_smp = check["car"][f - 1, :, :], check["mn"][f - 1, :, :]
+        p = min(
+            st.ks_2samp(
+                np.atleast_1d(car_smp[..., idx]).flatten(),
+                np.atleast_1d(mn_smp[..., idx]).flatten(),
+            )[1]
+            for idx in range(car_smp.shape[-1])
+        )
+        f -= 1
+    assert p > delta