Rename, add gq, fix tests

Author: WardBrian

cmdstanpy/stanfit/gq.py

@@ -344,17 +344,41 @@ 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]
@@ -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(

cmdstanpy/stanfit/mcmc.py

@@ -635,11 +635,11 @@ def draws_pd(
         )
         draws = np.concatenate([chains_col, iter_col, draw_col, draws], axis=2)
 
-        cols = ['chain', 'iter', 'draw'] + cols
+        cols = ['chain__', 'iter__', 'draw__'] + cols
 
         return pd.DataFrame(
             data=flatten_chains(draws),
-            columns=['chain', 'iter', 'draw'] + list(self.column_names),
+            columns=['chain__', 'iter__', 'draw__'] + list(self.column_names),
         )[cols]
 
     def draws_xr(

test/test_generate_quantities.py

@@ -77,14 +77,16 @@ 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]
     )
 
-    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_from_csv_files_bad() -> None:
@@ -153,7 +155,7 @@ 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]
@@ -267,14 +269,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 +525,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 +667,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,22 @@ 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,
     )
 
 
@@ -930,7 +936,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 +1298,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 +1377,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,