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_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