def get_drawset(runset: RunSet, params: List[str] = None) -> pd.DataFrame:
"""
Returns the assembled sample as a pandas DataFrame consisting of
one column per parameter and one row per draw.
:param runset: record of completed run of NUTS sampler
:param params: list of model parameter names.
"""
pnames_base = [name.split('.')[0] for name in runset.column_names]
if params is not None:
for p in params:
if not (p in runset._column_names or p in pnames_base):
raise ValueError('unknown parameter: {}'.format(p))
runset.assemble_sample()
data = runset.sample.reshape((runset.draws * runset.chains),
len(runset.column_names),
order='A')
df = pd.DataFrame(data=data, columns=runset.column_names)
if params is None:
return df
mask = []
for p in params:
for name in runset.column_names:
if p == name or p == name.split('.')[0]:
mask.append(name)
return df[mask]
def test_sample_big(self):
# construct runset using existing sampler output
stan = os.path.join(datafiles_path, 'bernoulli.stan')
exe = os.path.join(datafiles_path, 'bernoulli')
model = Model(exe_file=exe, stan_file=stan)
output = os.path.join(datafiles_path, 'runset-big', 'output_icar_nyc')
args = SamplerArgs(model, chain_ids=[1, 2], output_file=output)
runset = RunSet(chains=2, args=args)
runset.validate_csv_files()
runset.assemble_sample()
sampler_state = [
'lp__', 'accept_stat__', 'stepsize__', 'treedepth__',
'n_leapfrog__', 'divergent__', 'energy__'
]
phis = ['phi.{}'.format(str(x + 1)) for x in range(2095)]
column_names = sampler_state + phis
self.assertEqual(runset.columns, len(column_names))
self.assertEqual(runset.column_names, tuple(column_names))
self.assertEqual(runset.metric_type, 'diag_e')
self.assertEqual(runset.stepsize.shape, (2, ))
self.assertEqual(runset.metric.shape, (2, 2095))
self.assertEqual((1000, 2, 2102), runset.sample.shape)
phis = get_drawset(runset, params=['phi'])
self.assertEqual((2000, 2095), phis.shape)
phi1 = get_drawset(runset, params=['phi.1'])
self.assertEqual((2000, 1), phi1.shape)
mo_phis = get_drawset(runset, params=['phi.1', 'phi.10', 'phi.100'])
self.assertEqual((2000, 3), mo_phis.shape)
phi2095 = get_drawset(runset, params=['phi.2095'])
self.assertEqual((2000, 1), phi2095.shape)
with self.assertRaises(Exception):
get_drawset(runset, params=['phi.2096'])
with self.assertRaises(Exception):
get_drawset(runset, params=['ph'])