Repository: stan-dev/cmdstanpy

Issue:
draws_pd() should include chain, iteration, and draw columns
#### Summary:
The method `draws_pd()` on a `cmdstanpy.CmdStanMCMC` object is analogous to CmdStanR's method `draws(..., format = "df")`.  CmdStanR also returns columns containing indices for chain, iteration, and draw, where chain indexes each chain, iteration indexes draws witihin each chain, and draw indexes across the chains.

It would be nice if `draws_pd()` gave such columns in the returned Pandas dataframe. 

#### Additional Information:

Assuming chains are stacked into a Pandas dataframe, It would probably be easiest to add two columns (chain and iteration) to each chain, stack the chains, and then add the column draw to the dataframe.  

#### Current Version:
```
cmdstanpy.show_versions()
INSTALLED VERSIONS
---------------------
python: 3.11.3 (main, Apr  7 2023, 19:29:16) [Clang 14.0.0 (clang-1400.0.29.202)]
python-bits: 64
OS: Darwin
OS-release: 21.6.0
machine: x86_64
processor: i386
byteorder: little
LC_ALL: None
LANG: en_US.UTF-8
LOCALE: ('en_US', 'UTF-8')
cmdstan_folder: /Users/ez/.cmdstan/cmdstan-2.32.2
cmdstan: (2, 32)
cmdstanpy: 1.1.0
pandas: 2.0.3
xarray: 2023.6.0
tqdm: 4.65.0
numpy: 1.23.5
```


Discussion:

Patch:
diff --git a/cmdstanpy/stanfit/gq.py b/cmdstanpy/stanfit/gq.py
index 5555ad7..8c6617d 100644
--- a/cmdstanpy/stanfit/gq.py
+++ b/cmdstanpy/stanfit/gq.py
@@ -344,25 +344,49 @@ class CmdStanGQ(Generic[Fit]):
 
         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 @@ class CmdStanGQ(Generic[Fit]):
             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(
@@ -661,7 +676,6 @@ class CmdStanGQ(Generic[Fit]):
         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
@@ -710,7 +724,6 @@ class CmdStanGQ(Generic[Fit]):
             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:
diff --git a/cmdstanpy/stanfit/mcmc.py b/cmdstanpy/stanfit/mcmc.py
index 8e0c9b5..2f53ac8 100644
--- a/cmdstanpy/stanfit/mcmc.py
+++ b/cmdstanpy/stanfit/mcmc.py
@@ -615,9 +615,31 @@ class CmdStanMCMC:
         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 patch:
diff --git a/test/test_generate_quantities.py b/test/test_generate_quantities.py
index d66fa90..c22a877 100644
--- a/test/test_generate_quantities.py
+++ b/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]
diff --git a/test/test_sample.py b/test/test_sample.py
index b8c0633..f600252 100644
--- a/test/test_sample.py
+++ b/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,14 +854,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,
     )
     # 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')
@@ -863,7 +869,26 @@ def test_instantiate_from_csvfiles() -> None:
     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,