in pyro/contrib/oed/glmm/glmm.py [0:0]
def bayesian_linear_model(design, w_means={}, w_sqrtlambdas={}, re_group_sizes={},
re_alphas={}, re_betas={}, obs_sd=None,
alpha_0=None, beta_0=None, response="normal",
response_label="y", k=None):
"""
A pyro model for Bayesian linear regression.
If :param:`response` is `"normal"` this corresponds to a linear regression
model
:math:`Y = Xw + \\epsilon`
with `\\epsilon`` i.i.d. zero-mean Gaussian. The observation standard deviation
(:param:`obs_sd`) may be known or unknown. If unknown, it is assumed to follow an
inverse Gamma distribution with parameters :param:`alpha_0` and :param:`beta_0`.
If the response type is `"bernoulli"` we instead have :math:`Y \\sim Bernoulli(p)`
with
:math:`logit(p) = Xw`
Given parameter groups in :param:`w_means` and :param:`w_sqrtlambda`, the fixed effects
regression coefficient is taken to be Gaussian with mean `w_mean` and standard deviation
given by
:math:`\\sigma / \\sqrt{\\lambda}`
corresponding to the normal inverse Gamma family.
The random effects coefficient is constructed as follows. For each random effect
group, standard deviations for that group are sampled from a normal inverse Gamma
distribution. For each group, a random effect coefficient is then sampled from a zero
mean Gaussian with those standard deviations.
:param torch.Tensor design: a tensor with last two dimensions `n` and `p`
corresponding to observations and features respectively.
:param OrderedDict w_means: map from variable names to tensors of fixed effect means.
:param OrderedDict w_sqrtlambdas: map from variable names to tensors of square root
:math:`\\lambda` values for fixed effects.
:param OrderedDict re_group_sizes: map from variable names to int representing the
group size
:param OrderedDict re_alphas: map from variable names to `torch.Tensor`, the tensor
consists of Gamma dist :math:`\\alpha` values
:param OrderedDict re_betas: map from variable names to `torch.Tensor`, the tensor
consists of Gamma dist :math:`\\beta` values
:param torch.Tensor obs_sd: the observation standard deviation (if assumed known).
This is still relevant in the case of Bernoulli observations when coefficeints
are sampled using `w_sqrtlambdas`.
:param torch.Tensor alpha_0: Gamma :math:`\\alpha` parameter for unknown observation
covariance.
:param torch.Tensor beta_0: Gamma :math:`\\beta` parameter for unknown observation
covariance.
:param str response: Emission distribution. May be `"normal"` or `"bernoulli"`.
:param str response_label: Variable label for response.
:param torch.Tensor k: Only used for a sigmoid response. The slope of the sigmoid
transformation.
"""
# design is size batch x n x p
# tau is size batch
batch_shape = design.shape[:-2]
with ExitStack() as stack:
for plate in iter_plates_to_shape(batch_shape):
stack.enter_context(plate)
if obs_sd is None:
# First, sample tau (observation precision)
tau_prior = dist.Gamma(alpha_0.unsqueeze(-1),
beta_0.unsqueeze(-1)).to_event(1)
tau = pyro.sample("tau", tau_prior)
obs_sd = 1./torch.sqrt(tau)
elif alpha_0 is not None or beta_0 is not None:
warnings.warn("Values of `alpha_0` and `beta_0` unused becased"
"`obs_sd` was specified already.")
obs_sd = obs_sd.expand(batch_shape + (1,))
# Build the regression coefficient
w = []
# Allow different names for different coefficient groups
# Process fixed effects
for name, w_sqrtlambda in w_sqrtlambdas.items():
w_mean = w_means[name]
# Place a normal prior on the regression coefficient
w_prior = dist.Normal(w_mean, obs_sd / w_sqrtlambda).to_event(1)
w.append(pyro.sample(name, w_prior))
# Process random effects
for name, group_size in re_group_sizes.items():
# Sample `G` once for this group
alpha, beta = re_alphas[name], re_betas[name]
G_prior = dist.Gamma(alpha, beta).to_event(1)
G = 1./torch.sqrt(pyro.sample("G_" + name, G_prior))
# Repeat `G` for each group
repeat_shape = tuple(1 for _ in batch_shape) + (group_size,)
u_prior = dist.Normal(torch.tensor(0.), G.repeat(repeat_shape)).to_event(1)
w.append(pyro.sample(name, u_prior))
# Regression coefficient `w` is batch x p
w = broadcast_cat(w)
# Run the regressor forward conditioned on inputs
prediction_mean = rmv(design, w)
if response == "normal":
# y is an n-vector: hence use .to_event(1)
return pyro.sample(response_label, dist.Normal(prediction_mean, obs_sd).to_event(1))
elif response == "bernoulli":
return pyro.sample(response_label, dist.Bernoulli(logits=prediction_mean).to_event(1))
elif response == "sigmoid":
base_dist = dist.Normal(prediction_mean, obs_sd).to_event(1)
# You can add loc via the linear model itself
k = k.expand(prediction_mean.shape)
transforms = [AffineTransform(loc=torch.tensor(0.), scale=k), SigmoidTransform()]
response_dist = dist.TransformedDistribution(base_dist, transforms)
return pyro.sample(response_label, response_dist)
else:
raise ValueError("Unknown response distribution: '{}'".format(response))