Test fixes

Author: WardBrian

cmdstanpy/stanfit/gq.py

@@ -664,11 +664,14 @@ def _draws_start(self, inc_warmup: bool) -> tuple[int, int]:
                     num_draws = opt_iters
                 else:
                     draw1 = opt_iters - 1
-        else:  # CmdStanVB:
+        elif isinstance(p_fit, CmdStanVB):
             draw1 = 1  # skip mean
             num_draws = p_fit.variational_sample.shape[0]
             if inc_warmup:
                 num_draws += 1
+        else:
+            num_draws = p_fit.draws().shape[0]
+            draw1 = 0
 
         return draw1, num_draws
 

test/test_generate_quantities.py

@@ -8,7 +8,6 @@
 import pickle
 import shutil
 from test import check_present, without_import
-from unittest.mock import MagicMock, patch
 
 import numpy as np
 import pandas as pd
@@ -578,7 +577,7 @@ def test_from_optimization() -> None:
 
     # stan_variable
     theta = bern_gqs.stan_variable(var='theta')
-    assert theta.shape == (1,)
+    assert theta.shape == ()
     y_rep = bern_gqs.stan_variable(var='y_rep')
     assert y_rep.shape == (1, 10)
 
@@ -772,34 +771,3 @@ def test_vb_xarray() -> None:
     bern_gqs = model.generate_quantities(data=jdata, previous_fit=bern_fit)
     with pytest.raises(RuntimeError, match="via Sampling"):
         _ = bern_gqs.draws_xr()
-
-
-@patch(
-    'cmdstanpy.utils.cmdstan.cmdstan_version',
-    MagicMock(return_value=(2, 27)),
-)
-def test_from_non_hmc_old() -> None:
-    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_v = bern_model.variational(
-        data=jdata,
-        show_console=True,
-        require_converged=False,
-        seed=12345,
-    )
-
-    # gq_model
-    stan = os.path.join(DATAFILES_PATH, 'bernoulli_ppc.stan')
-    model = CmdStanModel(stan_file=stan)
-
-    with pytest.raises(RuntimeError, match="2.31"):
-        model.generate_quantities(data=jdata, previous_fit=bern_fit_v)
-
-    bern_fit_opt = bern_model.optimize(
-        data=jdata,
-        seed=12345,
-    )
-
-    with pytest.raises(RuntimeError, match="2.31"):
-        model.generate_quantities(data=jdata, previous_fit=bern_fit_opt)