Merge pull request #598 from stan-dev/update/column_order

Keep column order for dataframe

Author: mitzimorris

cmdstanpy/stanfit/gq.py

@@ -268,6 +268,7 @@ def draws_pd(
                 vars_list = [vars]
             else:
                 vars_list = vars
+            vars_list = list(dict.fromkeys(vars_list))
         if (
             inc_warmup
             and not self.mcmc_sample.metadata.cmdstan_config['save_warmup']
@@ -282,7 +283,7 @@ def draws_pd(
         gq_cols = []
         mcmc_vars = []
         if vars is not None:
-            for var in set(vars_list):
+            for var in vars_list:
                 if var in self.metadata.stan_vars_cols:
                     for idx in self.metadata.stan_vars_cols[var]:
                         gq_cols.append(self.column_names[idx])
@@ -295,6 +296,7 @@ def draws_pd(
                     raise ValueError('Unknown variable: {}'.format(var))
         else:
             gq_cols = list(self.column_names)
+            vars_list = gq_cols
 
         if inc_sample and mcmc_vars:
             if gq_cols:
@@ -311,7 +313,7 @@ def draws_pd(
                         )[gq_cols],
                     ],
                     axis='columns',
-                )
+                )[vars_list]
             else:
                 return self.mcmc_sample.draws_pd(
                     vars=mcmc_vars, inc_warmup=inc_warmup

cmdstanpy/stanfit/mcmc.py

@@ -589,7 +589,7 @@ def draws_pd(
             self._assemble_draws()
         cols = []
         if vars is not None:
-            for var in set(vars_list):
+            for var in dict.fromkeys(vars_list):
                 if (
                     var not in self.metadata.method_vars_cols
                     and var not in self.metadata.stan_vars_cols

test/test_generate_quantities.py

@@ -84,6 +84,11 @@ def test_from_csv_files(self):
             + bern_gqs.draws_pd().shape[1],
         )
 
+        self.assertEqual(
+            list(bern_gqs.draws_pd(vars=['y_rep']).columns),
+            column_names,
+        )
+
     def test_from_csv_files_bad(self):
         # gq model
         stan = os.path.join(DATAFILES_PATH, 'bernoulli_ppc.stan')

test/test_sample.py

@@ -741,6 +741,15 @@ def test_validate_good_run(self):
         self.assertEqual(fit.draws_pd(vars=['theta', 'lp__']).shape, (400, 2))
         self.assertEqual(fit.draws_pd(vars='theta').shape, (400, 1))
 
+        self.assertEqual(
+            list(fit.draws_pd(vars=['theta', 'lp__']).columns),
+            ['theta', 'lp__']
+        )
+        self.assertEqual(
+            list(fit.draws_pd(vars=['lp__', 'theta']).columns),
+            ['lp__', 'theta']
+        )
+
         summary = fit.summary()
         self.assertIn('5%', list(summary.columns))
         self.assertIn('50%', list(summary.columns))