fix shaping issues

Author: ahartikainen

cmdstanpy/stanfit.py

@@ -1074,7 +1074,6 @@ def draws_xr(
                 0,
                 self.draws(inc_warmup=inc_warmup),
             )
-
         return xr.Dataset(data, coords=coordinates, attrs=attrs).transpose(
             'chain', 'draw', ...
         )
@@ -1373,7 +1372,7 @@ def stan_variable(
         inc_iterations: bool = False,
         warn: bool = True,
         name: Optional[str] = None,
-    ) -> np.ndarray:
+    ) -> Union[np.ndarray, float]:
         """
         Return a numpy.ndarray which contains the estimates for the
         for the named Stan program variable where the dimensions of the
@@ -1416,33 +1415,29 @@ def stan_variable(
                 'Invalid estimate, optimization failed to converge.'
             )
 
-        col_idxs = self._metadata.stan_vars_cols[var]
+        col_idxs = list(self._metadata.stan_vars_cols[var])
         if inc_iterations and self._save_iterations:
             num_rows = self._all_iters.shape[0]
         else:
             num_rows = 1
-
-        if len(col_idxs) > 0:  # container var
+        if len(col_idxs) > 1:  # container var
             dims = (num_rows,) + self._metadata.stan_vars_dims[var]
             # pylint: disable=redundant-keyword-arg
             if num_rows > 1:
                 result = self._all_iters[:, col_idxs].reshape(  # type: ignore
                     dims, order='F'
                 )
             else:
-                mle = np.expand_dims(self._mle, axis=0)  # hack for col indexing
-                result = (
-                    mle[0, col_idxs]
-                    .reshape(dims, order='F')  # type: ignore
-                    .squeeze(axis=0)
-                )
+                result = self._mle[col_idxs].reshape(dims[1:], order="F")
         else:  # scalar var
+            col_idx = col_idxs[0]
             if num_rows > 1:
-                result = self._all_iters[:, col_idxs]
+                result = self._all_iters[:, col_idx]
             else:
-                result = np.atleast_1d(mle[0, col_idxs])
-
-        assert isinstance(result, np.ndarray)  # make the typechecker happy
+                result = float(self._mle[col_idx])
+        assert isinstance(
+            result, (np.ndarray, float)
+        )  # make the typechecker happy
         return result
 
     def stan_variables(
@@ -2143,7 +2138,7 @@ def metadata(self) -> InferenceMetadata:
 
     def stan_variable(
         self, var: Optional[str] = None, *, name: Optional[str] = None
-    ) -> np.ndarray:
+    ) -> Union[np.ndarray, float]:
         """
         Return a numpy.ndarray which contains the estimates for the
         for the named Stan program variable where the dimensions of the
@@ -2172,12 +2167,16 @@ def stan_variable(
         if var not in self._metadata.stan_vars_dims:
             raise ValueError('Unknown variable name: {}'.format(var))
         col_idxs = list(self._metadata.stan_vars_cols[var])
-        vals = list(self._variational_mean)
-        xs = [vals[x] for x in col_idxs]
         shape: Tuple[int, ...] = ()
-        if len(col_idxs) > 0:
+        if len(col_idxs) > 1:
             shape = self._metadata.stan_vars_dims[var]
-        return np.array(xs).reshape(shape)
+            result = np.asarray(self._variational_mean)[col_idxs].reshape(
+                shape, order="F"
+            )
+        else:
+            result = float(self._variational_mean[col_idxs[0]])
+        assert isinstance(result, (np.ndarray, float))
+        return result
 
     def stan_variables(self) -> Dict[str, np.ndarray]:
         """
@@ -2424,7 +2423,12 @@ def build_xarray_data(
     var_dims: Tuple[str, ...] = ('draw', 'chain')
     if dims:
         var_dims += tuple(f"{var_name}_dim_{i}" for i in range(len(dims)))
-        data[var_name] = (var_dims, drawset[start_row:, :, col_idxs])
+        data[var_name] = (
+            var_dims,
+            drawset[start_row:, :, col_idxs].reshape(
+                *drawset.shape[:2], *dims, order="F"
+            ),
+        )
     else:
         data[var_name] = (
             var_dims,

docsrc/conf.py

@@ -441,5 +441,7 @@ def emit(self, record):
 # }
 
 #  Makes the copying behavior on code examples cleaner by removing things like In [10]: from the text to be copied
-copybutton_prompt_text = r">>> |\.\.\. |\$ |In \[\d*\]: | {2,5}\.\.\.: | {5,8}: "
+copybutton_prompt_text = (
+    r">>> |\.\.\. |\$ |In \[\d*\]: | {2,5}\.\.\.: | {5,8}: "
+)
 copybutton_prompt_is_regexp = True

test/test_optimize.py

@@ -208,13 +208,15 @@ def test_variable_bern(self):
         self.assertTrue('theta' in bern_mle.metadata.stan_vars_dims)
         self.assertEqual(bern_mle.metadata.stan_vars_dims['theta'], ())
         theta = bern_mle.stan_variable(var='theta')
-        self.assertEqual(theta.shape, ())
+        self.assertTrue(isinstance(theta, float))
         with self.assertRaises(ValueError):
             bern_mle.stan_variable(var='eta')
         with self.assertRaises(ValueError):
             bern_mle.stan_variable(var='lp__')
         with LogCapture() as log:
-            self.assertEqual(bern_mle.stan_variable(name='theta').shape, ())
+            self.assertTrue(
+                isinstance(bern_mle.stan_variable(name='theta'), float)
+            )
         log.check_present(
             (
                 'cmdstanpy',
@@ -250,15 +252,15 @@ def test_variables_3d(self):
         var_beta = multidim_mle.stan_variable(var='beta')
         self.assertEqual(var_beta.shape, (2,))  # 1-element tuple
         var_frac_60 = multidim_mle.stan_variable(var='frac_60')
-        self.assertEqual(var_frac_60.shape, ())
+        self.assertTrue(isinstance(var_frac_60, float))
         vars = multidim_mle.stan_variables()
         self.assertEqual(len(vars), len(multidim_mle.metadata.stan_vars_dims))
         self.assertTrue('y_rep' in vars)
         self.assertEqual(vars['y_rep'].shape, (5, 4, 3))
         self.assertTrue('beta' in vars)
         self.assertEqual(vars['beta'].shape, (2,))
         self.assertTrue('frac_60' in vars)
-        self.assertEqual(vars['frac_60'].shape, ())
+        self.assertTrue(isinstance(vars['frac_60'], float))
 
         multidim_mle_iters = multidim_model.optimize(
             data=jdata,

test/test_variational.py

@@ -126,7 +126,7 @@ def test_variables_3d(self):
         var_beta = multidim_variational.stan_variable(var='beta')
         self.assertEqual(var_beta.shape, (2,))  # 1-element tuple
         var_frac_60 = multidim_variational.stan_variable(var='frac_60')
-        self.assertEqual(var_frac_60.shape, ())
+        self.assertTrue(isinstance(var_frac_60, float))
         vars = multidim_variational.stan_variables()
         self.assertEqual(
             len(vars), len(multidim_variational.metadata.stan_vars_dims)
@@ -136,7 +136,7 @@ def test_variables_3d(self):
         self.assertTrue('beta' in vars)
         self.assertEqual(vars['beta'].shape, (2,))
         self.assertTrue('frac_60' in vars)
-        self.assertEqual(vars['frac_60'].shape, ())
+        self.assertTrue(isinstance(vars['frac_60'], float))
         with self.assertRaises(ValueError):
             multidim_variational.stan_variable(var='beta', name='yrep')
         with LogCapture() as log: