def test_proceed_deterministic_failure_on_unnamed_deterministic():
@pm.model
def model():
x = yield pm.Normal("x", 0, 1)
yield pm.Deterministic(None, x)
with pytest.raises(EvaluationError):
pm.evaluate_model(model())
def test_raise_if_return_distribution():
def invalid_model():
yield pm.distributions.Normal("n1", 0, 1)
return pm.distributions.Normal("n2", 0, 1)
with pytest.raises(pm.flow.executor.EvaluationError) as e:
pm.evaluate_model(invalid_model())
assert e.match("should not contain")
def test_deterministic_with_distribution_name_fails():
@pm.model
def model():
x = yield pm.Normal("x", 0, 1)
det = yield pm.Deterministic("x", x)
return det
with pytest.raises(pm.flow.executor.EvaluationError):
pm.evaluate_model(model())
def test_evaluate_model_failure_on_state_and_values():
values = {"model/x": 1}
st = pm.flow.executor.SamplingState(observed_values=values)
@pm.model
def model():
yield pm.Normal("x", 0, 1)
with pytest.raises(ValueError):
pm.evaluate_model(model(), state=st, values=values)
def test_unable_to_create_duplicate_variable():
def invdalid_model():
yield pm.distributions.HalfNormal("n", 1, transform=pm.distributions.transforms.Log())
yield pm.distributions.Normal("n", 0, 1)
with pytest.raises(pm.flow.executor.EvaluationError) as e:
pm.evaluate_model(invdalid_model())
assert e.match("duplicate")
with pytest.raises(pm.flow.executor.EvaluationError) as e:
pm.evaluate_model_transformed(invdalid_model())
assert e.match("duplicate")
def test_incompatible_observed_shape():
@pm.model
def model(observed):
a = yield pm.Normal("a", 0, [1, 2], observed=observed)
observed_value = np.arange(3, dtype="float32")
with pytest.raises(EvaluationError):
pm.evaluate_model(model(None), observed={"model/a": observed_value})
with pytest.raises(EvaluationError):
pm.evaluate_model(model(observed_value))
def test_unnamed_return_2():
@pm.model(name=None)
def a_model():
return (yield pm.HalfNormal("n", 1, transform=pm.distributions.transforms.Log()))
_, state = pm.evaluate_model(a_model(name="b_model"))
assert "b_model" in state.deterministics_values
with pytest.raises(pm.flow.executor.EvaluationError) as e:
pm.evaluate_model(a_model())
assert e.match("unnamed")
with pytest.raises(pm.flow.executor.EvaluationError) as e:
pm.evaluate_model_transformed(a_model())
assert e.match("unnamed")
def test_uncatched_exception_works():
@pm.model
def a_model():
try:
yield 1
except:
pass
yield pm.distributions.HalfNormal("n", 1, transform=pm.distributions.transforms.Log())
with pytest.raises(pm.flow.executor.StopExecution) as e:
pm.evaluate_model(a_model())
assert e.match("something_bad")
with pytest.raises(pm.flow.executor.StopExecution) as e:
pm.evaluate_model_transformed(a_model())
assert e.match("something_bad")
def test_unreduced_log_prob(fixture_batch_shapes):
observed_value = np.ones(10, dtype="float32")
@pm.model
def model():
a = yield pm.Normal("a", 0, 1)
b = yield pm.HalfNormal("b", 1)
c = yield pm.Normal("c",
loc=a,
scale=b,
event_stack=len(observed_value))
values = {
"model/a": np.zeros(fixture_batch_shapes, dtype="float32"),
"model/b": np.ones(fixture_batch_shapes, dtype="float32"),
}
observed = {
"model/c":
np.broadcast_to(observed_value,
fixture_batch_shapes + observed_value.shape)
}
state = pm.evaluate_model(model(), values=values, observed=observed)[1]
unreduced_log_prob = state.collect_unreduced_log_prob()
assert unreduced_log_prob.numpy().shape == fixture_batch_shapes
np.testing.assert_allclose(tf.reduce_sum(unreduced_log_prob),
state.collect_log_prob())
def test_as_sampling_state_works_observed_is_constrained(
complex_model_with_observed):
_, state = pm.evaluate_model(complex_model_with_observed())
sampling_state = state.as_sampling_state()
assert not sampling_state.transformed_values
assert set(sampling_state.observed_values) == {"complex_model/a/n"}
assert set(sampling_state.untransformed_values) == {"complex_model/n"}
def test_complex_model_no_keep_return(complex_model):
_, state = pm.evaluate_model(complex_model())
assert set(state.untransformed_values) == {"complex_model/n", "complex_model/a/n"}
assert set(state.deterministics_values) == {"complex_model/a"}
assert not state.transformed_values # we call untransformed executor
assert not state.observed_values
def test_differently_shaped_logp():
def model():
yield dist.Normal("n1", np.zeros(10), np.ones(10))
yield dist.Normal("n2", np.zeros(3), np.ones(3))
_, state = pm.evaluate_model(model())
state.collect_log_prob() # this should work
def test_observed_cant_mix_with_untransformed_and_raises_an_error_case_untransformed_executor(
transformed_model_with_observed, ):
with pytest.raises(pm.flow.executor.EvaluationError) as e:
_, state = pm.evaluate_model(transformed_model_with_observed(),
values=dict(n=0.0))
assert e.match("{'n': None}")
assert e.match("'n' from untransformed values")
def test_complex_model_keep_return():
@pm.model
def nested_model(cond):
norm = yield dist.HalfNormal("n",
cond**2,
transform=dist.transforms.Log())
return norm
@pm.model()
def complex_model():
norm = yield dist.Normal("n", 0, 1)
result = yield nested_model(norm, name="a")
return result
_, state = pm.evaluate_model(complex_model())
assert set(state.untransformed_values) == {
"complex_model/n", "complex_model/a/n"
}
assert set(state.deterministics) == {
"complex_model",
"complex_model/a",
}
assert not state.transformed_values # we call untransformed executor
assert not state.observed_values
def test_observed_are_passed_correctly(complex_model_with_observed):
_, state = pm.evaluate_model(complex_model_with_observed())
assert set(state.untransformed_values) == {"complex_model/n", "complex_model/a"}
assert not state.transformed_values # we call untransformed executor
assert set(state.observed_values) == {"complex_model/a/n"}
assert np.allclose(state.all_values["complex_model/a/n"], np.ones(10))
def test_executor_logp_tensorflow(transformed_model):
norm = tfd.HalfNormal(1)
_, state = pm.evaluate_model(transformed_model(), values=dict(n=math.pi))
np.testing.assert_allclose(state.collect_log_prob(),
norm.log_prob(math.pi),
equal_nan=False)
def test_executor_logp_tensorflow(transformed_model):
tfp = pytest.importorskip("tensorflow_probability")
tfd = tfp.distributions
norm = tfd.HalfNormal(1)
_, state = pm.evaluate_model(transformed_model(), values=dict(n=math.pi))
np.testing.assert_allclose(state.collect_log_prob(), norm.log_prob(math.pi), equal_nan=False)
def test_executor_on_conditionally_independent(fixture_batch_shapes):
@pm.model
def model():
a = yield pm.Normal("a", 0, 1, conditionally_independent=True)
b = yield pm.Normal("b", a, 1)
_, state = pm.evaluate_model(model(), sample_shape=fixture_batch_shapes)
assert state.untransformed_values["model/a"].shape == fixture_batch_shapes
assert state.untransformed_values["model/b"].shape == fixture_batch_shapes
def test_sampling_state_clone(deterministics_in_nested_models):
model = deterministics_in_nested_models[0]
observed = {"model/nested_model/x": 0.0}
_, state = pm.evaluate_model(model(), observed=observed)
clone = state.clone()
assert set(state.all_values) == set(clone.all_values)
assert all((state.all_values[k] == v for k, v in clone.all_values.items()))
assert set(state.deterministics) == set(clone.deterministics)
assert all((state.deterministics[k] == v
for k, v in clone.deterministics.items()))
assert state.posterior_predictives == clone.posterior_predictives
def test_observed_are_set_to_none_for_posterior_predictive_correctly(
complex_model_with_observed):
_, state = pm.evaluate_model(complex_model_with_observed(),
observed={"complex_model/a/n": None})
assert set(state.untransformed_values) == {
"complex_model/n", "complex_model/a/n"
}
assert set(state.deterministics) == {"complex_model/a"}
assert not state.transformed_values # we call untransformed executor
assert not state.observed_values
assert not np.allclose(state.all_values["complex_model/a/n"], np.ones(10))
def test_deterministics(model_with_deterministics):
model, expected_deterministics, expected_ops, expected_ops_inputs = model_with_deterministics
_, state = pm.evaluate_model(model())
assert len(state.deterministics) == len(expected_deterministics)
assert set(expected_deterministics) <= set(state.deterministics)
for expected_deterministic, op, op_inputs in zip(expected_deterministics,
expected_ops,
expected_ops_inputs):
inputs = [v for k, v in state.all_values.items() if k in op_inputs]
out = op(*inputs)
np.testing.assert_allclose(
state.deterministics[expected_deterministic], out)
def test_log_prob_elemwise(fixture_model_with_plates):
model, expected_rv_shapes = fixture_model_with_plates
_, state = pm.evaluate_model(model())
log_prob_elemwise = dict(
zip(state.distributions, state.collect_log_prob_elemwise())
) # This will discard potentials in log_prob_elemwise
log_prob = state.collect_log_prob()
assert len(log_prob_elemwise) == len(expected_rv_shapes)
assert all(rv in log_prob_elemwise for rv in expected_rv_shapes)
assert all(log_prob_elemwise[rv].shape == shape
for rv, shape in expected_rv_shapes.items())
assert log_prob.numpy() == sum(
map(tf.reduce_sum, log_prob_elemwise.values())).numpy()
def test_simple_model(simple_model):
_, state = pm.evaluate_model(simple_model())
assert "n" in state.untransformed_values
assert not state.observed_values
assert not state.transformed_values
def test_single_distribution():
_, state = pm.evaluate_model(pm.distributions.Normal("n", 0, 1))
assert "n" in state.all_values
def test_class_model(class_model):
"""Test that model can be defined as method in an object definition"""
_, state = pm.evaluate_model(class_model.class_model_method())
assert "class_model_method/n" in state.untransformed_values
assert not state.observed_values
assert not state.transformed_values
def test_unnamed_distribution_to_prior():
f = lambda: (yield pm.distributions.Normal.dist(0, 1).prior("n"))
_, state = pm.evaluate_model(f())
assert "n" in state.untransformed_values
def test_unnamed_distribution():
f = lambda: (yield pm.distributions.Normal.dist(0, 1))
with pytest.raises(pm.flow.executor.EvaluationError) as e:
pm.evaluate_model(f())
assert e.match("anonymous Distribution")
def test_data_from_model(model, modelParams, params_dict, to_return=None):
"""
Generates a test dataset from our model! Needs some pregenerated
data file which is quite strange but only uses it to cast the right dimensions!
Parameters
----------
model: pymc4 model
The model to sample from
modelParams: :py:class:`covid19_npis.ModelParams`
Instance of modelParams, mainly used for number of age groups and
number of countries.
params_dict: dictionary
Parameters for the test run.
to_return: list str
Return these variables
Returns
-------
: new_cases_inferred, R_t, interventions
"""
# ------------------------------------------------------------------------------ #
# Set params for the test dataset
# ------------------------------------------------------------------------------ #
len_gen_interv_kernel = 12
num_interventions = 2
if to_return is None:
to_return = []
model_name = model(modelParams).name
dict_with_model_name = {
f"{model_name}/{key}": tf.cast(value, "float32")[tf.newaxis,
tf.newaxis]
for key, value in params_dict.items()
}
trace = az.from_dict(posterior=dict_with_model_name, )
var_names = []
# variables = list(
# set(
# # ["R_0", "new_cases_inferred", "R_t", "g", "d_i_c_p", "new_I_t", "h_0_t",]
# to_return
# )
# )
variables = to_return
for var in variables:
var_names.append(f"{model_name}/{var}")
# Sample
trace = pm.sample_posterior_predictive(
model(modelParams),
trace,
var_names=var_names,
use_auto_batching=False,
)
# Convert to pandas
_, sample_state = pm.evaluate_model(model(modelParams))
def convert_to_pandas(key):
df = data.convert_trace_to_dataframe(
trace,
sample_state=sample_state,
key=key,
data_type="posterior_predictive",
)
df.index = df.index.droplevel(["chain", "draw"])
if "time" in df.index.names:
df = df.stack().unstack(level="time").T
df.columns = df.columns.droplevel(
-1) # stack adds an additional dimension
df.index = df.index.rename("date")
return df
"""
new_cases_inferred = convert_to_pandas("new_cases_inferred")
R_t = convert_to_pandas("R_t")
d = data.convert_trace_to_dataframe(
trace,
sample_state=sample_state,
key="d_i_c_p",
data_type="posterior_predictive",
)
# Take intervention array from model params
for c, country in enumerate(modelParams.countries):
if c == 0:
interv = country.data_interventions
else:
interv = interv.join(country.data_interventions)
extra = []
for var in to_return:
extra.append(convert_to_pandas(var))
"""
dfs = [convert_to_pandas(var) for var in variables]
return dfs
def test_transformed_model_untransformed_executor(transformed_model):
_, state = pm.evaluate_model(transformed_model())
assert set(state.untransformed_values) == {"n"}
assert not state.transformed_values # we call untransformed executor
assert not state.observed_values
def test_as_sampling_state_does_not_works_if_untransformed_exec(complex_model):
_, state = pm.evaluate_model(complex_model())
with pytest.raises(TypeError) as e:
state.as_sampling_state()
e.match("'complex_model/a/__log_n' is not found")
