def default_prior(model_desc_pre):
assert type(model_desc_pre) == ModelDescPre
family = model_desc_pre.response.family
assert type(family) == Family
assert family.link is not None
assert type(model_desc_pre.population) == PopulationPre
assert type(model_desc_pre.groups) == list
assert all(type(gm) == GroupPre for gm in model_desc_pre.groups)
b_children = [leaf(name) for name in model_desc_pre.population.coefs]
cor_children = [leaf(cols2str(group.columns)) for group in model_desc_pre.groups if group.corr]
sd_children = [Node(cols2str(gm.columns), None, False, [], [leaf(name) for name in gm.coefs]) for gm in
model_desc_pre.groups]
def mk_resp_prior_edit(param_name):
prior = get_response_prior(family.name, param_name)
if prior is not None:
return Prior(('resp', param_name), prior)
resp_children = [leaf(p.name, mk_resp_prior_edit(p.name), [chk_support(p.type)])
for p in model_desc_pre.response.nonlocparams]
return Node('root', None, False, [], [
Node('b', Prior(('b',), Cauchy(0., 1.)), False, [chk_support(Type['Real']())], b_children),
Node('sd', Prior(('sd',), HalfCauchy(3.)), False, [chk_support(Type['PosReal']())], sd_children),
Node('cor', Prior(('cor',), LKJ(1.)), False, [chk_lkj], cor_children),
Node('resp', None, False, [], resp_children)])
# Prior on coef of an interaction.
('y ~ x1:x2',
[Categorical('x1', list('ab')),
Categorical('x2', list('cd'))
], {}, Normal, [Prior(('b', 'x1[b]:x2[c]'),
Normal(0., 100.))], [('b_0', 'Cauchy', {}),
('b_1', 'Normal', {
'loc': 0.,
'scale': 100.
}), ('b_2', 'Cauchy', {}),
('sigma', 'HalfCauchy', {})]),
# Prior on group level `sd` choice.
('y ~ 1 + x2 + x3 | x1', [Categorical('x1', list('ab'))], {}, Normal,
[Prior(('sd', 'x1', 'intercept'),
HalfCauchy(4.))], [('sigma', 'HalfCauchy', {}),
('sd_0_0', 'HalfCauchy', {
'scale': 4.
}), ('sd_0_1', 'HalfCauchy', {}),
('z_0', 'Normal', {}), ('L_0', 'LKJ', {})]),
('y ~ 1 + x2 + x3 || x1', [Categorical('x1', list('ab'))], {}, Normal,
[Prior(('sd', 'x1', 'intercept'),
HalfNormal(4.))], [('sigma', 'HalfCauchy', {}),
('sd_0_0', 'HalfNormal', {
'scale': 4.
}), ('sd_0_1', 'HalfCauchy', {}),
('z_0', 'Normal', {})]),
('y ~ 1 + x || a:b',
[Categorical('a', ['a1', 'a2']),
Categorical('b', ['b1', 'b2'])], {}, Normal,
[Prior(('sd', 'a:b', 'intercept'),
from brmp.utils import join
# `is_param` indicates whether a node corresponds to a parameter in
# the model. (Nodes without this flag set exist only to add structure
# to the parameters.) This infomation is used when extracting
# information from the tree.
Node = namedtuple('Node', 'name prior_edit is_param checks children')
def leaf(name, prior_edit=None, checks=[]):
return Node(name, prior_edit, True, checks, [])
RESPONSE_PRIORS = {
'Normal': {
'sigma': HalfCauchy(3.)
},
}
def get_response_prior(family, parameter):
if family in RESPONSE_PRIORS:
return RESPONSE_PRIORS[family][parameter]
# This is similar to brms `set_prior`. (e.g. `set_prior('<prior>',
# coef='x1')` is similar to `Prior(['x1'], '<prior>)`.) By specifying
# paths (rather than class/group/coef) we're diverging from brms, but
# the hope is that a brms-like interface can be put in front of this.
Prior = namedtuple('Prior', 'path prior')