Remove deprecated stan_variable behavior for advi

Author: WardBrian

cmdstanpy/stanfit/vb.py

@@ -8,7 +8,6 @@
 
 from cmdstanpy.cmdstan_args import Method
 from cmdstanpy.utils import stancsv
-from cmdstanpy.utils.logging import get_logger
 
 from .metadata import InferenceMetadata
 from .runset import RunSet
@@ -100,7 +99,7 @@ def __repr__(self) -> str:
         # TODO - diagnostic, profiling files
         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}")
@@ -163,9 +162,7 @@ def metadata(self) -> InferenceMetadata:
         """
         return self._metadata
 
-    def stan_variable(
-        self, var: str, *, mean: Optional[bool] = None
-    ) -> Union[np.ndarray, float]:
+    def stan_variable(self, var: str, *, mean: bool = False) -> np.ndarray:
         """
         Return a numpy.ndarray which contains the estimates for the
         for the named Stan program variable where the dimensions of the
@@ -188,8 +185,7 @@ def stan_variable(
         :param var: variable name
 
         :param mean: if True, return the variational mean. Otherwise,
-            return the variational sample.  The default behavior will
-            change in a future release to return the variational sample.
+            return the variational sample. Defaults to False.
 
         See Also
         --------
@@ -200,16 +196,7 @@ def stan_variable(
         CmdStanGQ.stan_variable
         CmdStanLaplace.stan_variable
         """
-        # TODO(2.0): remove None case, make default `False`
-        if mean is None:
-            get_logger().warning(
-                "The default behavior of CmdStanVB.stan_variable() "
-                "will change in a future release to return the "
-                "variational sample, rather than the mean.\n"
-                "To maintain the current behavior, pass the argument "
-                "mean=True"
-            )
-            mean = True
+
         if mean:
             draws = self._variational_mean
         else:
@@ -219,16 +206,7 @@ def stan_variable(
             out: np.ndarray = self._metadata.stan_vars[var].extract_reshape(
                 draws
             )
-            # TODO(2.0): remove
-            if out.shape == () or out.shape == (1,):
-                if mean:
-                    get_logger().warning(
-                        "The default behavior of "
-                        "CmdStanVB.stan_variable(mean=True) 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
@@ -238,9 +216,7 @@ def stan_variable(
                 + ", ".join(self._metadata.stan_vars.keys())
             )
 
-    def stan_variables(
-        self, *, mean: Optional[bool] = None
-    ) -> dict[str, Union[np.ndarray, float]]:
+    def stan_variables(self, *, mean: bool = False) -> dict[str, np.ndarray]:
         """
         Return a dictionary mapping Stan program variables names
         to the corresponding numpy.ndarray containing the inferred values.

test/test_variational.py

@@ -282,9 +282,8 @@ def test_complex_output() -> None:
     )
 
     assert fit.stan_variable('zs', mean=False).shape == (1000, 2, 3)
-    # TODO(2.0): change
-    np.testing.assert_equal(fit.z, 3 + 4j)
-    # np.testing.assert_equal(fit.z, np.repeat(3 + 4j, 1000))
+    np.testing.assert_equal(fit.stan_variable("z", mean=True), 3 + 4j)
+    np.testing.assert_equal(fit.z, np.repeat(3 + 4j, 1000))
 
     np.testing.assert_allclose(
         fit.stan_variable('zs', mean=False)[0],
@@ -309,14 +308,9 @@ def test_attrs() -> None:
         algorithm='meanfield',
     )
 
-    # TODO(2.0): swap tests
-    np.testing.assert_equal(fit.a, 4.5)
-    assert fit.b.shape == (3,)
-    assert isinstance(fit.theta, float)
-
-    # np.testing.assert_equal(fit.a, np.repeat(4.5, 1000))
-    # assert fit.b.shape == (1000, 3)
-    # assert fit.theta.shape == (1000,)
+    np.testing.assert_equal(fit.a, np.repeat(4.5, 1000))
+    assert fit.b.shape == (1000, 3)
+    assert fit.theta.shape == (1000,)
 
     assert fit.stan_variable('thin', mean=True) == 3.5