def extract_variables_from_df(dataframe):
"""
Extract random variables and their shapes from the pymc3-pandas data-frame
Parameters
----------
dataframe : :class:`pandas.DataFrame`
Returns
-------
flat_names : dict
with variable-names and respective flat-name indexes to data-frame
var_shapes : dict
with variable names and shapes
"""
all_df_indexes = [str(flatvar) for flatvar in dataframe.columns]
varnames = list(set([index.split('__')[0] for index in all_df_indexes]))
flat_names = {}
var_shapes = {}
for varname in varnames:
indexes = []
for index in all_df_indexes:
if index.split('__')[0] == varname:
indexes.append(index)
flat_names[varname] = indexes
var_shapes[varname] = ttab._create_shape(indexes)
return flat_names, var_shapes
def test_create_flat_names_2d():
shape = 2, 3
result = ttab.create_flat_names('x', shape)
expected = ['x__0_0', 'x__0_1', 'x__0_2',
'x__1_0', 'x__1_1', 'x__1_2']
assert result == expected
assert ttab._create_shape(result) == shape
def test_create_flat_names_3d():
shape = 2, 3, 4
assert ttab._create_shape(ttab.create_flat_names('x', shape)) == shape
def test_create_flat_names_2d():
shape = 2, 3
result = ttab.create_flat_names('x', shape)
expected = ['x__0_0', 'x__0_1', 'x__0_2', 'x__1_0', 'x__1_1', 'x__1_2']
assert result == expected
assert ttab._create_shape(result) == shape
def test_create_flat_names_0d():
shape = ()
result = ttab.create_flat_names('x', shape)
expected = ['x']
assert result == expected
assert ttab._create_shape(result) == shape
def test_create_flat_names_3d():
shape = 2, 3, 4
assert ttab._create_shape(ttab.create_flat_names('x', shape)) == shape
def test_create_flat_names_0d():
shape = ()
result = ttab.create_flat_names('x', shape)
expected = ['x']
assert result == expected
assert ttab._create_shape(result) == shape
def test_create_flat_names_2d():
shape = 2, 3
result = ttab.create_flat_names("x", shape)
expected = ["x__0_0", "x__0_1", "x__0_2", "x__1_0", "x__1_1", "x__1_2"]
assert result == expected
assert ttab._create_shape(result) == shape
def test_create_flat_names_1d():
shape = (2, )
result = ttab.create_flat_names("x", shape)
expected = ["x__0", "x__1"]
assert result == expected
assert ttab._create_shape(result) == shape
