Remove deprecated stan_variable behavior for optimization

Author: WardBrian

cmdstanpy/stanfit/mle.py

@@ -1,7 +1,7 @@
 """Container for the result of running optimization"""
 
 from collections import OrderedDict
-from typing import Optional, Union
+from typing import Optional
 
 import numpy as np
 import pandas as pd
@@ -95,7 +95,7 @@ def __repr__(self) -> str:
             repr = '{} optimization failed to converge.'.format(repr)
         return repr
 
-    def __getattr__(self, attr: str) -> Union[np.ndarray, float]:
+    def __getattr__(self, attr: str) -> np.ndarray:
         """Synonymous with ``fit.stan_variable(attr)"""
         if attr.startswith("_"):
             raise AttributeError(f"Unknown variable name {attr}")
@@ -206,7 +206,7 @@ def stan_variable(
         *,
         inc_iterations: bool = False,
         warn: bool = True,
-    ) -> Union[np.ndarray, float]:
+    ) -> np.ndarray:
         """
         Return a numpy.ndarray which contains the estimates for the
         for the named Stan program variable where the dimensions of the
@@ -254,14 +254,6 @@ def stan_variable(
             out: np.ndarray = self._metadata.stan_vars[var].extract_reshape(
                 data
             )
-            # TODO(2.0) remove
-            if out.shape == () or out.shape == (1,):
-                get_logger().warning(
-                    "The default behavior of CmdStanMLE.stan_variable() "
-                    "will change in a future release to always return a "
-                    "numpy.ndarray, even for scalar variables."
-                )
-                return out.item()  # type: ignore
             return out
         except KeyError:
             # pylint: disable=raise-missing-from
@@ -273,7 +265,7 @@ def stan_variable(
 
     def stan_variables(
         self, inc_iterations: bool = False
-    ) -> dict[str, Union[np.ndarray, float]]:
+    ) -> dict[str, np.ndarray]:
         """
         Return a dictionary mapping Stan program variables names
         to the corresponding numpy.ndarray containing the inferred values.

test/test_optimize.py

@@ -201,8 +201,7 @@ def test_variable_bern() -> None:
     assert 'theta' in bern_mle.metadata.stan_vars
     assert bern_mle.metadata.stan_vars['theta'].dimensions == ()
     theta = bern_mle.stan_variable(var='theta')
-    # TODO(2.0): remove before `or` clause
-    assert isinstance(theta, float) or theta.shape == ()
+    assert theta.shape == ()
     with pytest.raises(ValueError):
         bern_mle.stan_variable(var='eta')
     with pytest.raises(ValueError):
@@ -234,17 +233,15 @@ def test_variables_3d() -> None:
     var_beta = multidim_mle.stan_variable(var='beta')
     assert var_beta.shape == (2,)
     var_frac_60 = multidim_mle.stan_variable(var='frac_60')
-    # TODO(2.0): remove before `or` clause
-    assert isinstance(var_frac_60, float) or var_frac_60.shape == ()
+    assert var_frac_60.shape == ()
     vars = multidim_mle.stan_variables()
     assert len(vars) == len(multidim_mle.metadata.stan_vars)
     assert 'y_rep' in vars
     assert vars['y_rep'].shape == (5, 4, 3)
     assert 'beta' in vars
     assert vars['beta'].shape == (2,)
     assert 'frac_60' in vars
-    # TODO(2.0): remove before `or` clause
-    assert isinstance(vars['frac_60'], float) or vars['frac_60'].shape == ()
+    assert vars['frac_60'].shape == ()
 
     multidim_mle_iters = multidim_model.optimize(
         data=jdata,
@@ -565,8 +562,7 @@ def test_single_row_csv() -> None:
     model = CmdStanModel(stan_file=stan)
     mle = model.optimize()
     theta = mle.stan_variable('theta')
-    # TODO(2.0): remove before `or` clause
-    assert isinstance(theta, float) or theta.shape == ()
+    assert theta.shape == ()
     z_as_ndarray = mle.stan_variable(var="z")
     assert z_as_ndarray.shape == (4, 3)
     for i in range(4):
@@ -627,8 +623,7 @@ def test_attrs() -> None:
 
     assert fit.a == 4.5
     assert fit.b.shape == (3,)
-    # TODO(2.0) remove before `or` clause
-    assert isinstance(fit.theta, float) or fit.theta.shape == ()
+    assert fit.theta.shape == ()
 
     assert fit.stan_variable('thin') == 3.5
 
@@ -673,7 +668,7 @@ def test_serialization() -> None:
     )
 
 
-def test_optimize_create_inits():
+def test_optimize_create_inits() -> None:
     stan = os.path.join(DATAFILES_PATH, 'bernoulli.stan')
     bern_model = CmdStanModel(stan_file=stan)
     jdata = os.path.join(DATAFILES_PATH, 'bernoulli.data.json')
@@ -686,7 +681,7 @@ def test_optimize_create_inits():
     assert len(inits) == 1
 
 
-def test_optimize_init_sampling():
+def test_optimize_init_sampling() -> None:
     stan = os.path.join(DATAFILES_PATH, 'logistic.stan')
     logistic_model = CmdStanModel(stan_file=stan)
     logistic_data = os.path.join(DATAFILES_PATH, 'logistic.data.R')

test/test_variational.py

@@ -344,7 +344,7 @@ def test_serialization() -> None:
     )
 
 
-def test_variational_create_inits():
+def test_variational_create_inits() -> None:
     stan = os.path.join(DATAFILES_PATH, 'bernoulli.stan')
     bern_model = CmdStanModel(stan_file=stan)
     jdata = os.path.join(DATAFILES_PATH, 'bernoulli.data.json')
@@ -374,7 +374,7 @@ def test_variational_create_inits():
     )
 
 
-def test_variational_init_sampling():
+def test_variational_init_sampling() -> None:
     stan = os.path.join(DATAFILES_PATH, 'logistic.stan')
     logistic_model = CmdStanModel(stan_file=stan)
     logistic_data = os.path.join(DATAFILES_PATH, 'logistic.data.R')