Merge pull request #677 from stan-dev/feature/676-pandas-columns

Add "chain__", "iter__", and "draw__" columns to draws_pd

Author: WardBrian

cmdstanpy/stanfit/gq.py

@@ -344,25 +344,49 @@ def draws_pd(
 
         previous_draws_pd = self._previous_draws_pd(mcmc_vars, inc_warmup)
 
+        draws = self.draws(inc_warmup=inc_warmup)
+        # add long-form columns for chain, iteration, draw
+        n_draws, n_chains, _ = draws.shape
+        chains_col = (
+            np.repeat(np.arange(1, n_chains + 1), n_draws)
+            .reshape(1, n_chains, n_draws)
+            .T
+        )
+        iter_col = (
+            np.tile(np.arange(1, n_draws + 1), n_chains)
+            .reshape(1, n_chains, n_draws)
+            .T
+        )
+        draw_col = (
+            np.arange(1, (n_draws * n_chains) + 1)
+            .reshape(1, n_chains, n_draws)
+            .T
+        )
+        draws = np.concatenate([chains_col, iter_col, draw_col, draws], axis=2)
+
+        vars_list = ['chain__', 'iter__', 'draw__'] + vars_list
+        if gq_cols:
+            gq_cols = ['chain__', 'iter__', 'draw__'] + gq_cols
+
+        draws_pd = pd.DataFrame(
+            data=flatten_chains(draws),
+            columns=['chain__', 'iter__', 'draw__'] + list(self.column_names),
+        )
+
         if inc_sample and mcmc_vars:
             if gq_cols:
                 return pd.concat(
                     [
                         previous_draws_pd,
-                        pd.DataFrame(
-                            data=flatten_chains(
-                                self.draws(inc_warmup=inc_warmup)
-                            ),
-                            columns=self.column_names,
-                        )[gq_cols],
+                        draws_pd[gq_cols],
                     ],
                     axis='columns',
                 )[vars_list]
             else:
                 return previous_draws_pd
         elif inc_sample and vars is None:
-            cols_1 = self.previous_fit.column_names
-            cols_2 = self.column_names
+            cols_1 = list(previous_draws_pd.columns)
+            cols_2 = list(draws_pd.columns)
             dups = [
                 item
                 for item, count in Counter(cols_1 + cols_2).items()
@@ -371,23 +395,14 @@ def draws_pd(
             return pd.concat(
                 [
                     previous_draws_pd.drop(columns=dups).reset_index(drop=True),
-                    pd.DataFrame(
-                        data=flatten_chains(self.draws(inc_warmup=inc_warmup)),
-                        columns=self.column_names,
-                    ),
+                    draws_pd,
                 ],
                 axis=1,
             )
         elif gq_cols:
-            return pd.DataFrame(
-                data=flatten_chains(self.draws(inc_warmup=inc_warmup)),
-                columns=self.column_names,
-            )[gq_cols]
-
-        return pd.DataFrame(
-            data=flatten_chains(self.draws(inc_warmup=inc_warmup)),
-            columns=self.column_names,
-        )
+            return draws_pd[gq_cols]
+
+        return draws_pd
 
     @overload
     def draws_xr(
@@ -657,7 +672,6 @@ def _draws_start(self, inc_warmup: bool) -> Tuple[int, int]:
         elif isinstance(p_fit, CmdStanMLE):
             num_draws = 1
             if p_fit._save_iterations:
-
                 opt_iters = len(p_fit.optimized_iterations_np)  # type: ignore
                 if inc_warmup:
                     num_draws = opt_iters
@@ -706,7 +720,6 @@ def _previous_draws_pd(
             return p_fit.draws_pd(vars or None, inc_warmup=inc_warmup)
 
         elif isinstance(p_fit, CmdStanMLE):
-
             if inc_warmup and p_fit._save_iterations:
                 return p_fit.optimized_iterations_pd[sel]  # type: ignore
             else:

cmdstanpy/stanfit/mcmc.py

@@ -615,9 +615,31 @@ def draws_pd(
         else:
             cols = list(self.column_names)
 
+        draws = self.draws(inc_warmup=inc_warmup)
+        # add long-form columns for chain, iteration, draw
+        n_draws, n_chains, _ = draws.shape
+        chains_col = (
+            np.repeat(np.arange(1, n_chains + 1), n_draws)
+            .reshape(1, n_chains, n_draws)
+            .T
+        )
+        iter_col = (
+            np.tile(np.arange(1, n_draws + 1), n_chains)
+            .reshape(1, n_chains, n_draws)
+            .T
+        )
+        draw_col = (
+            np.arange(1, (n_draws * n_chains) + 1)
+            .reshape(1, n_chains, n_draws)
+            .T
+        )
+        draws = np.concatenate([chains_col, iter_col, draw_col, draws], axis=2)
+
+        cols = ['chain__', 'iter__', 'draw__'] + cols
+
         return pd.DataFrame(
-            data=flatten_chains(self.draws(inc_warmup=inc_warmup)),
-            columns=self.column_names,
+            data=flatten_chains(draws),
+            columns=['chain__', 'iter__', 'draw__'] + list(self.column_names),
         )[cols]
 
     def draws_xr(

test/test_generate_quantities.py

@@ -77,14 +77,56 @@ def test_from_csv_files(caplog: pytest.LogCaptureFixture) -> None:
     )
 
     # draws_pd()
-    assert bern_gqs.draws_pd().shape == (400, 10)
+    assert bern_gqs.draws_pd().shape == (400, 13)
     assert (
         bern_gqs.draws_pd(inc_sample=True).shape[1]
         == bern_gqs.previous_fit.draws_pd().shape[1]
         + bern_gqs.draws_pd().shape[1]
+        - 3  # chain, iter, draw duplicates
     )
 
-    assert list(bern_gqs.draws_pd(vars=['y_rep']).columns) == column_names
+    assert list(bern_gqs.draws_pd(vars=['y_rep']).columns) == (
+        ["chain__", "iter__", "draw__"] + column_names
+    )
+
+
+def test_pd_xr_agreement():
+    # fitted_params sample - list of filenames
+    goodfiles_path = os.path.join(DATAFILES_PATH, 'runset-good', 'bern')
+    csv_files = []
+    for i in range(4):
+        csv_files.append('{}-{}.csv'.format(goodfiles_path, i + 1))
+
+    # gq_model
+    stan = os.path.join(DATAFILES_PATH, 'bernoulli_ppc.stan')
+    model = CmdStanModel(stan_file=stan)
+    jdata = os.path.join(DATAFILES_PATH, 'bernoulli.data.json')
+
+    bern_gqs = model.generate_quantities(data=jdata, previous_fit=csv_files)
+
+    draws_pd = bern_gqs.draws_pd(inc_sample=True)
+    draws_xr = bern_gqs.draws_xr(inc_sample=True)
+
+    # check that the indexing is the same between the two
+    np.testing.assert_equal(
+        draws_pd[draws_pd['chain__'] == 2]['y_rep[1]'],
+        draws_xr.y_rep.sel(chain=2).isel(y_rep_dim_0=0).values,
+    )
+    # "draw" is 0-indexed in xarray, equiv. "iter__" is 1-indexed in pandas
+    np.testing.assert_equal(
+        draws_pd[draws_pd['iter__'] == 100]['y_rep[1]'],
+        draws_xr.y_rep.sel(draw=99).isel(y_rep_dim_0=0).values,
+    )
+
+    # check for included sample as well
+    np.testing.assert_equal(
+        draws_pd[draws_pd['chain__'] == 2]['theta'],
+        draws_xr.theta.sel(chain=2).values,
+    )
+    np.testing.assert_equal(
+        draws_pd[draws_pd['iter__'] == 100]['theta'],
+        draws_xr.theta.sel(draw=99).values,
+    )
 
 
 def test_from_csv_files_bad() -> None:
@@ -153,16 +195,17 @@ def test_from_previous_fit_draws() -> None:
 
     bern_gqs = model.generate_quantities(data=jdata, previous_fit=bern_fit)
 
-    assert bern_gqs.draws_pd().shape == (400, 10)
+    assert bern_gqs.draws_pd().shape == (400, 13)
     assert (
         bern_gqs.draws_pd(inc_sample=True).shape[1]
         == bern_gqs.previous_fit.draws_pd().shape[1]
         + bern_gqs.draws_pd().shape[1]
+        - 3  # duplicates of chain, iter, and draw
     )
     row1_sample_pd = bern_fit.draws_pd().iloc[0]
     row1_gqs_pd = bern_gqs.draws_pd().iloc[0]
     np.testing.assert_array_equal(
-        pd.concat((row1_sample_pd, row1_gqs_pd), axis=0).values,
+        pd.concat((row1_sample_pd, row1_gqs_pd), axis=0).values[3:],
         bern_gqs.draws_pd(inc_sample=True).iloc[0].values,
     )
     # draws_xr
@@ -267,14 +310,14 @@ def test_save_warmup(caplog: pytest.LogCaptureFixture) -> None:
         10,
     )
 
-    assert bern_gqs.draws_pd().shape == (400, 10)
-    assert bern_gqs.draws_pd(inc_warmup=False).shape == (400, 10)
-    assert bern_gqs.draws_pd(inc_warmup=True).shape == (800, 10)
+    assert bern_gqs.draws_pd().shape == (400, 13)
+    assert bern_gqs.draws_pd(inc_warmup=False).shape == (400, 13)
+    assert bern_gqs.draws_pd(inc_warmup=True).shape == (800, 13)
     assert bern_gqs.draws_pd(vars=['y_rep'], inc_warmup=False).shape == (
         400,
-        10,
+        13,
     )
-    assert bern_gqs.draws_pd(vars='y_rep', inc_warmup=False).shape == (400, 10)
+    assert bern_gqs.draws_pd(vars='y_rep', inc_warmup=False).shape == (400, 13)
 
     theta = bern_gqs.stan_variable(var='theta')
     assert theta.shape == (400,)
@@ -523,7 +566,7 @@ def test_from_optimization() -> None:
     assert bern_gqs.draws(inc_sample=True).shape == (1, 1, 12)
 
     # draws_pd()
-    assert bern_gqs.draws_pd().shape == (1, 10)
+    assert bern_gqs.draws_pd().shape == (1, 13)
     assert (
         bern_gqs.draws_pd(inc_sample=True).shape[1]
         == bern_gqs.previous_fit.optimized_params_pd.shape[1]
@@ -665,7 +708,7 @@ def test_from_vb():
     assert bern_gqs.draws(inc_sample=True).shape == (1000, 1, 14)
 
     # draws_pd()
-    assert bern_gqs.draws_pd().shape == (1000, 10)
+    assert bern_gqs.draws_pd().shape == (1000, 13)
     assert (
         bern_gqs.draws_pd(inc_sample=True).shape[1]
         == bern_gqs.previous_fit.variational_sample_pd.shape[1]

test/test_sample.py

@@ -745,18 +745,24 @@ def test_validate_good_run() -> None:
     draws_pd = fit.draws_pd()
     assert draws_pd.shape == (
         fit.runset.chains * fit.num_draws_sampling,
-        len(fit.column_names),
+        len(fit.column_names) + 3,
     )
-    assert fit.draws_pd(vars=['theta']).shape == (400, 1)
-    assert fit.draws_pd(vars=['lp__', 'theta']).shape == (400, 2)
-    assert fit.draws_pd(vars=['theta', 'lp__']).shape == (400, 2)
-    assert fit.draws_pd(vars='theta').shape == (400, 1)
+    assert fit.draws_pd(vars=['theta']).shape == (400, 4)
+    assert fit.draws_pd(vars=['lp__', 'theta']).shape == (400, 5)
+    assert fit.draws_pd(vars=['theta', 'lp__']).shape == (400, 5)
+    assert fit.draws_pd(vars='theta').shape == (400, 4)
 
     assert list(fit.draws_pd(vars=['theta', 'lp__']).columns) == [
+        'chain__',
+        'iter__',
+        'draw__',
         'theta',
         'lp__',
     ]
     assert list(fit.draws_pd(vars=['lp__', 'theta']).columns) == [
+        'chain__',
+        'iter__',
+        'draw__',
         'lp__',
         'theta',
     ]
@@ -817,7 +823,7 @@ def test_validate_big_run() -> None:
     assert fit.step_size.shape == (2,)
     assert fit.metric.shape == (2, 2095)
     assert fit.draws().shape == (1000, 2, 2102)
-    assert fit.draws_pd(vars=['phi']).shape == (2000, 2095)
+    assert fit.draws_pd(vars=['phi']).shape == (2000, 2098)
     with raises_nested(ValueError, r'Unknown variable: gamma'):
         fit.draws_pd(vars=['gamma'])
 
@@ -828,14 +834,14 @@ def test_instantiate_from_csvfiles() -> None:
     draws_pd = bern_fit.draws_pd()
     assert draws_pd.shape == (
         bern_fit.runset.chains * bern_fit.num_draws_sampling,
-        len(bern_fit.column_names),
+        len(bern_fit.column_names) + 3,
     )
     csvfiles_path = os.path.join(DATAFILES_PATH, 'runset-big')
     big_fit = from_csv(path=csvfiles_path)
     draws_pd = big_fit.draws_pd()
     assert draws_pd.shape == (
         big_fit.runset.chains * big_fit.num_draws_sampling,
-        len(big_fit.column_names),
+        len(big_fit.column_names) + 3,
     )
     # list
     csvfiles_path = os.path.join(DATAFILES_PATH, 'runset-good')
@@ -848,22 +854,41 @@ def test_instantiate_from_csvfiles() -> None:
     draws_pd = bern_fit.draws_pd()
     assert draws_pd.shape == (
         bern_fit.runset.chains * bern_fit.num_draws_sampling,
-        len(bern_fit.column_names),
+        len(bern_fit.column_names) + 3,
     )
     # single csvfile
     bern_fit = from_csv(path=csvfiles[0])
     draws_pd = bern_fit.draws_pd()
     assert draws_pd.shape == (
         bern_fit.num_draws_sampling,
-        len(bern_fit.column_names),
+        len(bern_fit.column_names) + 3,
     )
     # glob
     csvfiles_path = os.path.join(csvfiles_path, '*.csv')
     big_fit = from_csv(path=csvfiles_path)
     draws_pd = big_fit.draws_pd()
     assert draws_pd.shape == (
         big_fit.runset.chains * big_fit.num_draws_sampling,
-        len(big_fit.column_names),
+        len(big_fit.column_names) + 3,
+    )
+
+
+def test_pd_xr_agreement():
+    csvfiles_path = os.path.join(DATAFILES_PATH, 'runset-good', '*.csv')
+    bern_fit = from_csv(path=csvfiles_path)
+
+    draws_pd = bern_fit.draws_pd()
+    draws_xr = bern_fit.draws_xr()
+
+    # check that the indexing is the same between the two
+    np.testing.assert_equal(
+        draws_pd[draws_pd['chain__'] == 2]['theta'],
+        draws_xr.theta.sel(chain=2).values,
+    )
+    # "draw" is 0-indexed in xarray, equiv. "iter__" is 1-indexed in pandas
+    np.testing.assert_equal(
+        draws_pd[draws_pd['iter__'] == 100]['theta'],
+        draws_xr.theta.sel(draw=99).values,
     )
 
 
@@ -930,7 +955,7 @@ def test_instantiate_from_csvfiles_fail(
 def test_from_csv_fixed_param() -> None:
     csv_path = os.path.join(DATAFILES_PATH, 'fixed_param_sample.csv')
     fixed_param_sample = from_csv(path=csv_path)
-    assert fixed_param_sample.draws_pd().shape == (100, 85)
+    assert fixed_param_sample.draws_pd().shape == (100, 88)
 
 
 def test_custom_metric() -> None:
@@ -1292,14 +1317,14 @@ def test_save_warmup() -> None:
         len(BERNOULLI_COLS),
     )
 
-    assert bern_fit.draws_pd().shape == (200, len(BERNOULLI_COLS))
+    assert bern_fit.draws_pd().shape == (200, len(BERNOULLI_COLS) + 3)
     assert bern_fit.draws_pd(inc_warmup=False).shape == (
         200,
-        len(BERNOULLI_COLS),
+        len(BERNOULLI_COLS) + 3,
     )
     assert bern_fit.draws_pd(inc_warmup=True).shape == (
         600,
-        len(BERNOULLI_COLS),
+        len(BERNOULLI_COLS) + 3,
     )
 
 
@@ -1371,7 +1396,7 @@ def test_dont_save_warmup(caplog: pytest.LogCaptureFixture) -> None:
     with caplog.at_level(logging.WARNING):
         assert bern_fit.draws_pd(inc_warmup=True).shape == (
             200,
-            len(BERNOULLI_COLS),
+            len(BERNOULLI_COLS) + 3,
         )
     check_present(
         caplog,