def _posterior_to_trace(self, chain=0) -> NDArray:
"""Save results into a PyMC trace
This method should not be overwritten.
"""
lenght_pos = len(self.tempered_posterior)
varnames = [v.name for v in self.variables]
with self.model:
strace = NDArray(name=self.model.name)
strace.setup(lenght_pos, chain)
for i in range(lenght_pos):
value = []
size = 0
for varname in varnames:
shape, new_size = self.var_info[varname]
var_samples = self.tempered_posterior[i][size:size + new_size]
# Round discrete variable samples. The rounded values were the ones
# actually used in the logp evaluations (see logp_forw)
var = self.model[varname]
if var.dtype in discrete_types:
var_samples = np.round(var_samples).astype(var.dtype)
value.append(var_samples.reshape(shape))
size += new_size
strace.record(point={k: v for k, v in zip(varnames, value)})
return strace
def test_choose_chains(n_points, tune, expected_length, expected_n_traces):
with pm.Model() as model:
a = pm.Normal("a", mu=0, sigma=1)
trace_0 = NDArray(model)
trace_1 = NDArray(model)
trace_2 = NDArray(model)
trace_0.setup(n_points[0], 1)
trace_1.setup(n_points[1], 1)
trace_2.setup(n_points[2], 1)
for _ in range(n_points[0]):
trace_0.record({"a": 0})
for _ in range(n_points[1]):
trace_1.record({"a": 0})
for _ in range(n_points[2]):
trace_2.record({"a": 0})
traces, length = pm.sampling._choose_chains([trace_0, trace_1, trace_2], tune=tune)
assert length == expected_length
assert expected_n_traces == len(traces)
