Skip to content

Add moments for CAR distribution #5220

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Jan 26, 2022
Merged

Add moments for CAR distribution #5220

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions pymc/distributions/multivariate.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2115,6 +2115,13 @@ class CAR(Continuous):
def dist(cls, mu, W, alpha, tau, *args, **kwargs):
return super().dist([mu, W, alpha, tau], **kwargs)

def get_moment(rv, size, mu, W, alpha, tau):
moment = mu
if not rv_size_is_none(size):
moment_size = at.concatenate([size, moment.shape])
moment = at.full(moment_size, mu)
return moment

def logp(value, mu, W, alpha, tau):
"""
Calculate log-probability of a CAR-distributed vector
Expand Down
30 changes: 29 additions & 1 deletion pymc/tests/test_distributions_moments.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
import pymc as pm

from pymc.distributions import (
CAR,
AsymmetricLaplace,
Bernoulli,
Beta,
Expand Down Expand Up @@ -109,7 +110,6 @@ def test_all_distributions_have_moments():
# Distributions that have been refactored but don't yet have moments
not_implemented |= {
dist_module.discrete.DiscreteWeibull,
dist_module.multivariate.CAR,
dist_module.multivariate.DirichletMultinomial,
dist_module.multivariate.Wishart,
}
Expand Down Expand Up @@ -932,6 +932,34 @@ def test_mv_normal_moment(mu, cov, size, expected):
assert_moment_is_expected(model, expected, check_finite_logp=x.ndim < 3)


@pytest.mark.parametrize(
"mu, size, expected",
[
(
np.array([1, 0, 3.0, 4]),
None,
np.array([1, 0, 3.0, 4]),
),
(np.array([1, 0, 3.0, 4]), 6, np.full((6, 4), [1, 0, 3.0, 4])),
(np.array([1, 0, 3.0, 4]), (5, 3), np.full((5, 3, 4), [1, 0, 3.0, 4])),
(
np.array([[3.0, 5, 2, 1], [1, 4, 0.5, 9]]),
(4, 5),
np.full((4, 5, 2, 4), [[3.0, 5, 2, 1], [1, 4, 0.5, 9]]),
),
],
)
def test_car_moment(mu, size, expected):
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
Comment on lines +956 to +957
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Are tau and alpha constrained to be scalars?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

just noticed the class docstring allows them to be arrays. I just don't see it being tested anywhere in

pymc/pymc/tests/test_distributions.py

Lines 3162 to 3242 in 99ec0ff

@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)
# 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
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)
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.
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)

and

pymc/pymc/tests/test_distributions_random.py

Lines 2368 to 2405 in 99ec0ff

@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, return_inferencedata=False)
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

which is what this PR's tests are based on. I also don't remember any literature that treats tau is anything more than a scalar.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If the random/logp methods work with non-scalar inputs for those parameters, we need get_moment to do the same. If they are not really supposed to work with those being non-scalar we don't need to, but we might have to open a new issue similar to #5214

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

#5214 was fixed in #5241 - the changeset was rather simple. @zoj613 do you want to add the enforcement of scalar tau and alpha, or add them as vectors to the test cases?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

According to

pymc/pymc/distributions/multivariate.py

Line 1999 in b6f76e5

ndims_params = [1, 2, 0, 0]

it looks like tau and alpha are both scalars. The docstring seems wrong about allowing array input for those parameters.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ndims_params indicates the minimum/base dimensions, not whether they can't be of higher dimensionality. Idelly we wouldn't impose any limitations, but the rng_fn/logp methods might not broadcast/handle higer values than the minimum case properly

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@ckrapu is the tau parameter meant to handle array input? I have not seen that case in literature. I also do not see that case being tested in the implementation unittests. I'm asking because it appears you submitted the PR for the distribution. Would you be able to clarify if we need to account for non-scalar tau?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@zoj613 I think it's safe to go ahead and add that explicit limitation for the time being. We can always lift it later (and add tests). Do you want to include that in this PR?

Copy link
Contributor Author

@zoj613 zoj613 Jan 26, 2022 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@ricardoV94 I think it's best that array tau be handled as a separate PR. I believe there is no good motivation to account for array tau given its lack of application in literature. IIRC the paper linked in the doc-string only mentions tau as an array in the case an MCAR model, which I believe is different from what is implemented by CAR.

Copy link
Member

@ricardoV94 ricardoV94 Jan 26, 2022 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Lack of application in the literature shouldn't be a reason to restrict it, but instead whether we handling it properly (and have tests to give us confidence we are). I agree it can be done in a separate PR, would you be interested in doing that? I'll go ahead and merge this one after one last review

with Model() as model:
CAR("x", mu=mu, W=W, alpha=alpha, tau=tau, size=size)
assert_moment_is_expected(model, expected)


@pytest.mark.parametrize(
"mu, sigma, size, expected",
[
Expand Down