Add tests

Author: WardBrian

cmdstanpy/stanfit/gq.py

@@ -385,8 +385,8 @@ def draws_pd(
             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()
@@ -676,7 +676,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
@@ -725,7 +724,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:

test/test_generate_quantities.py

@@ -89,6 +89,45 @@ def test_from_csv_files(caplog: pytest.LogCaptureFixture) -> None:
     )
 
 
+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:
     # gq model
     stan = os.path.join(DATAFILES_PATH, 'bernoulli_ppc.stan')

test/test_sample.py

@@ -873,6 +873,25 @@ def test_instantiate_from_csvfiles() -> None:
     )
 
 
+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,
+    )
+
+
 def test_instantiate_from_csvfiles_fail(
     caplog: pytest.LogCaptureFixture,
 ) -> None: