more unit tests; fixed gq bug; different test model

Author: mitzimorris

cmdstanpy/stanfit.py

@@ -1983,16 +1983,26 @@ def stan_variable(
             return self.mcmc_sample.stan_variable(var, inc_warmup=inc_warmup)
         else:  # is gq variable
             self._assemble_generated_quantities()
-            col_idxs = self._metadata.stan_vars_cols[var]
+            draw1 = 0
             if (
                 not inc_warmup
                 and self.mcmc_sample.metadata.cmdstan_config['save_warmup']
             ):
-                draw1 = self.mcmc_sample.num_draws_warmup * self.chains
-                return flatten_chains(self._draws)[  # type: ignore
-                    draw1:, col_idxs
-                ]
-            return flatten_chains(self._draws)[:, col_idxs]  # type: ignore
+                draw1 = self.mcmc_sample.num_draws_warmup
+            num_draws = self.mcmc_sample.num_draws_sampling
+            if (
+                inc_warmup
+                and self.mcmc_sample.metadata.cmdstan_config['save_warmup']
+            ):
+                num_draws += self.mcmc_sample.num_draws_warmup
+            dims = [num_draws * self.chains]
+            col_idxs = self._metadata.stan_vars_cols[var]
+            if len(col_idxs) > 0:
+                dims.extend(self._metadata.stan_vars_dims[var])
+            # pylint: disable=redundant-keyword-arg
+            return self._draws[draw1:, :, col_idxs].reshape(  # type: ignore
+                dims, order='F'
+            )
 
     def stan_variables(self, inc_warmup: bool = False) -> Dict[str, np.ndarray]:
         """

test/data/matrix_var.stan

@@ -0,0 +1,19 @@
+transformed data {
+  int y[10] = {0,1,0,0,0,0,0,0,0,1};
+}
+parameters {
+  real<lower=0,upper=1> theta;
+}
+model {
+  theta ~ beta(1,1);  // uniform prior on interval 0,1
+  y ~ bernoulli(theta);
+}
+generated quantities {
+  # x is a 4 x 3 matrix where i,j entry == rownum
+  matrix[4, 3] z;
+  for (row_num in 1:4) {
+    for (col_num in 1:3) {
+      z[row_num, col_num] = row_num;
+    }
+  }
+}

test/test_generate_quantities.py

@@ -426,6 +426,28 @@ def test_no_xarray(self):
             with self.assertRaises(RuntimeError):
                 bern_gqs.draws_xr()
 
+    def test_bug_455(self):
+        stan = os.path.join(DATAFILES_PATH, 'bernoulli.stan')
+        bern_model = CmdStanModel(stan_file=stan)
+        jdata = os.path.join(DATAFILES_PATH, 'bernoulli.data.json')
+        bern_fit = bern_model.sample(
+            data=jdata,
+            chains=1,
+            seed=12345,
+            iter_sampling=1,
+        )
+        stan = os.path.join(DATAFILES_PATH, 'matrix_var.stan')
+        model = CmdStanModel(stan_file=stan)
+        gqs = model.generate_quantities(mcmc_sample=bern_fit)
+        z_as_ndarray = gqs.stan_variable(var="z")
+        self.assertEqual(z_as_ndarray.shape, (1, 4, 3))  # flattens chains
+        z_as_xr = gqs.draws_xr(vars="z")
+        self.assertEqual(z_as_xr.z.data.shape, (1, 1, 4, 3))  # keeps chains
+        for i in range(4):
+            for j in range(3):
+                self.assertEqual(int(z_as_ndarray[0, i, j]), i + 1)
+                self.assertEqual(int(z_as_xr.z.data[0, 0, i, j]), i + 1)
+
 
 if __name__ == '__main__':
     unittest.main()

test/test_optimize.py

@@ -287,36 +287,6 @@ def test_variables_3d(self):
         self.assertTrue('frac_60' in vars_iters)
         self.assertEqual(vars_iters['frac_60'].shape, (8,))
 
-    def test_variables_shape(self):
-        stan = os.path.join(DATAFILES_PATH, 'shape.stan')
-        model = CmdStanModel(stan_file=stan)
-        no_data = {}
-        mle = model.optimize(
-            seed=1239812093,
-            algorithm='LBFGS',
-            init_alpha=0.001,
-            iter=100,
-            tol_obj=1e-12,
-            tol_rel_obj=1e4,
-            tol_grad=1e-8,
-            tol_rel_grad=1e7,
-            tol_param=1e-8,
-            history_size=5,
-        )
-        for var, shape in {
-            "x": float,
-            "a": float,
-            "b": (2,),
-            "c": (2,3),
-            "d": (2,),
-            "e":(2,3),
-            "f": (4,2,3),
-            "g":(4,5,2,3),
-        }.items():
-            if isinstance(shape, tuple):
-                assert mle.stan_variable(var).shape == shape
-            else:
-                assert isinstance(mle.stan_variable(var), shape)
 
 class OptimizeTest(unittest.TestCase):
     def test_optimize_good(self):
@@ -611,6 +581,17 @@ def test_optimize_bad(self):
                 data=no_data, seed=1239812093, inits=None, algorithm='BFGS'
             )
 
+    def test_bug_455(self):
+        stan = os.path.join(DATAFILES_PATH, 'matrix_var.stan')
+        model = CmdStanModel(stan_file=stan)
+        mle = model.optimize()
+        self.assertTrue(isinstance(mle.stan_variable('theta'), float))
+        z_as_ndarray = mle.stan_variable(var="z")
+        self.assertEqual(z_as_ndarray.shape, (4, 3))
+        for i in range(4):
+            for j in range(3):
+                self.assertEqual(int(z_as_ndarray[i, j]), i + 1)
+
 
 if __name__ == '__main__':
     unittest.main()

test/test_sample.py

@@ -1668,6 +1668,19 @@ def test_no_xarray(self):
             with self.assertRaises(RuntimeError):
                 bern_fit.draws_xr()
 
+    def test_bug_455(self):
+        stan = os.path.join(DATAFILES_PATH, 'matrix_var.stan')
+        bug_455_model = CmdStanModel(stan_file=stan)
+        bug_455_fit = bug_455_model.sample(iter_sampling=1, chains=1)
+        z_as_ndarray = bug_455_fit.stan_variable(var="z")
+        self.assertEqual(z_as_ndarray.shape, (1, 4, 3))  # flattens chains
+        z_as_xr = bug_455_fit.draws_xr(vars="z")
+        self.assertEqual(z_as_xr.z.data.shape, (1, 1, 4, 3))  # keeps chains
+        for i in range(4):
+            for j in range(3):
+                self.assertEqual(int(z_as_ndarray[0, i, j]), i + 1)
+                self.assertEqual(int(z_as_xr.z.data[0, 0, i, j]), i + 1)
+
 
 if __name__ == '__main__':
     unittest.main()

test/test_variational.py

@@ -253,6 +253,17 @@ def test_variational_eta_fail(self):
             )
         )
 
+    def test_bug_455(self):
+        stan = os.path.join(DATAFILES_PATH, 'matrix_var.stan')
+        model = CmdStanModel(stan_file=stan)
+        vb_fit = model.variational()
+        self.assertTrue(isinstance(vb_fit.stan_variable('theta'), float))
+        z_as_ndarray = vb_fit.stan_variable(var="z")
+        self.assertEqual(z_as_ndarray.shape, (4, 3))
+        for i in range(4):
+            for j in range(3):
+                self.assertEqual(int(z_as_ndarray[i, j]), i + 1)
+
 
 if __name__ == '__main__':
     unittest.main()