Merge pull request #510 from qres/develop

assemble draws and generated quantities only once

Author: WardBrian

cmdstanpy/stanfit/mcmc.py

@@ -231,7 +231,7 @@ def draws(
         CmdStanMCMC.draws_xr
         CmdStanGQ.draws
         """
-        if self._draws.size == 0:
+        if self._draws.shape == (0,):
             self._assemble_draws()
 
         if inc_warmup and not self._save_warmup:
@@ -309,9 +309,6 @@ def _assemble_draws(self) -> None:
         Allocates and populates the step size, metric, and sample arrays
         by parsing the validated stan_csv files.
         """
-        if self._draws.shape != (0,):
-            return
-
         num_draws = self.num_draws_sampling
         sampling_iter_start = 0
         if self._save_warmup:
@@ -527,7 +524,8 @@ def draws_pd(
                 ' must run sampler with "save_warmup=True".'
             )
 
-        self._assemble_draws()
+        if self._draws.shape == (0,):
+            self._assemble_draws()
         cols = []
         if vars is not None:
             for var in set(vars_list):
@@ -583,7 +581,8 @@ def draws_xr(
         else:
             vars_list = vars
 
-        self._assemble_draws()
+        if self._draws.shape == (0,):
+            self._assemble_draws()
 
         num_draws = self.num_draws_sampling
         meta = self._metadata.cmdstan_config
@@ -663,7 +662,8 @@ def stan_variable(
             raise ValueError('No variable name specified.')
         if var not in self._metadata.stan_vars_dims:
             raise ValueError('Unknown variable name: {}'.format(var))
-        self._assemble_draws()
+        if self._draws.shape == (0,):
+            self._assemble_draws()
         draw1 = 0
         if not inc_warmup and self._save_warmup:
             draw1 = self.num_draws_warmup
@@ -705,7 +705,8 @@ def method_variables(self) -> Dict[str, np.ndarray]:
         containing per-draw diagnostic values.
         """
         result = {}
-        self._assemble_draws()
+        if self._draws.shape == (0,):
+            self._assemble_draws()
         for idxs in self.metadata.method_vars_cols.values():
             for idx in idxs:
                 result[self.column_names[idx]] = self._draws[:, :, idx]
@@ -868,7 +869,7 @@ def draws(
         CmdStanGQ.draws_xr
         CmdStanMCMC.draws
         """
-        if self._draws.size == 0:
+        if self._draws.shape == (0,):
             self._assemble_generated_quantities()
         if (
             inc_warmup
@@ -955,7 +956,8 @@ def draws_pd(
                 'Draws from warmup iterations not available,'
                 ' must run sampler with "save_warmup=True".'
             )
-        self._assemble_generated_quantities()
+        if self._draws.shape == (0,):
+            self._assemble_generated_quantities()
 
         gq_cols = []
         mcmc_vars = []
@@ -1076,7 +1078,8 @@ def draws_xr(
         for var in dup_vars:
             vars_list.remove(var)
 
-        self._assemble_generated_quantities()
+        if self._draws.shape == (0,):
+            self._assemble_generated_quantities()
 
         num_draws = self.mcmc_sample.num_draws_sampling
         sample_config = self.mcmc_sample.metadata.cmdstan_config
@@ -1173,7 +1176,8 @@ def stan_variable(
         if var not in gq_var_names:
             return self.mcmc_sample.stan_variable(var, inc_warmup=inc_warmup)
         else:  # is gq variable
-            self._assemble_generated_quantities()
+            if self._draws.shape == (0,):
+                self._assemble_generated_quantities()
             draw1 = 0
             if (
                 not inc_warmup
@@ -1222,7 +1226,7 @@ def stan_variables(self, inc_warmup: bool = False) -> Dict[str, np.ndarray]:
         return result
 
     def _assemble_generated_quantities(self) -> None:
-        # use numpy genfromtext
+        # use numpy loadtxt
         warmup = self.mcmc_sample.metadata.cmdstan_config['save_warmup']
         num_draws = self.mcmc_sample.draws(inc_warmup=warmup).shape[0]
         gq_sample = np.empty(