Allow use of stan_variables through .

Author: WardBrian

cmdstanpy/stanfit/mcmc.py

@@ -117,6 +117,14 @@ def __repr__(self) -> str:
         # TODO - hamiltonian, profiling files
         return repr
 
+    def __getattr__(self, attr: str) -> np.ndarray:
+        """Synonymous with ``fit.stan_variable(attr)"""
+        try:
+            return self.stan_variable(attr)
+        except ValueError as e:
+            # pylint: disable=raise-missing-from
+            raise AttributeError(*e.args)
+
     @property
     def chains(self) -> int:
         """Number of chains."""
@@ -647,6 +655,9 @@ def stan_variable(
         and the sample consists of 4 chains with 1000 post-warmup draws,
         this function will return a numpy.ndarray with shape (4000,3,3).
 
+        This functionaltiy is also available via a shortcut using ``.`` -
+        writing ``fit.a`` is a synonym for ``fit.stan_variable("a")``
+
         :param var: variable name
 
         :param inc_warmup: When ``True`` and the warmup draws are present in
@@ -769,6 +780,14 @@ def __repr__(self) -> str:
         )
         return repr
 
+    def __getattr__(self, attr: str) -> np.ndarray:
+        """Synonymous with ``fit.stan_variable(attr)"""
+        try:
+            return self.stan_variable(attr)
+        except ValueError as e:
+            # pylint: disable=raise-missing-from
+            raise AttributeError(*e.args)
+
     def _validate_csv_files(self) -> Dict[str, Any]:
         """
         Checks that Stan CSV output files for all chains are consistent
@@ -1160,6 +1179,9 @@ def stan_variable(
         and the sample consists of 4 chains with 1000 post-warmup draws,
         this function will return a numpy.ndarray with shape (4000,3,3).
 
+        This functionaltiy is also available via a shortcut using ``.`` -
+        writing ``fit.a`` is a synonym for ``fit.stan_variable("a")``
+
         :param var: variable name
 
         :param inc_warmup: When ``True`` and the warmup draws are present in

cmdstanpy/stanfit/mle.py

@@ -50,6 +50,14 @@ def __repr__(self) -> str:
             repr = '{} optimization failed to converge.'.format(repr)
         return repr
 
+    def __getattr__(self, attr: str) -> Union[np.ndarray, float]:
+        """Synonymous with ``fit.stan_variable(attr)"""
+        try:
+            return self.stan_variable(attr)
+        except ValueError as e:
+            # pylint: disable=raise-missing-from
+            raise AttributeError(*e.args)
+
     def _set_mle_attrs(self, sample_csv_0: str) -> None:
         meta = scan_optimize_csv(sample_csv_0, self._save_iterations)
         self._metadata = InferenceMetadata(meta)
@@ -165,6 +173,9 @@ def stan_variable(
         for the named Stan program variable where the dimensions of the
         numpy.ndarray match the shape of the Stan program variable.
 
+        This functionaltiy is also available via a shortcut using ``.`` -
+        writing ``fit.a`` is a synonym for ``fit.stan_variable("a")``
+
         :param var: variable name
 
         :param inc_iterations: When ``True`` and the intermediate estimates

cmdstanpy/stanfit/vb.py

@@ -41,6 +41,14 @@ def __repr__(self) -> str:
         # TODO - diagnostic, profiling files
         return repr
 
+    def __getattr__(self, attr: str) -> Union[np.ndarray, float]:
+        """Synonymous with ``fit.stan_variable(attr)"""
+        try:
+            return self.stan_variable(attr)
+        except ValueError as e:
+            # pylint: disable=raise-missing-from
+            raise AttributeError(*e.args)
+
     def _set_variational_attrs(self, sample_csv_0: str) -> None:
         meta = scan_variational_csv(sample_csv_0)
         self._metadata = InferenceMetadata(meta)
@@ -109,6 +117,9 @@ def stan_variable(self, var: str) -> Union[np.ndarray, float]:
         for the named Stan program variable where the dimensions of the
         numpy.ndarray match the shape of the Stan program variable.
 
+        This functionaltiy is also available via a shortcut using ``.`` -
+        writing ``fit.a`` is a synonym for ``fit.stan_variable("a")``
+
         :param var: variable name
 
         See Also

test/data/named_output.stan

@@ -0,0 +1,23 @@
+data {
+  int<lower=0> N;
+  int<lower=0,upper=1> y[N];
+}
+parameters {
+  real<lower=0,upper=1> theta;
+}
+model {
+  theta ~ beta(1,1);  // uniform prior on interval 0,1
+  y ~ bernoulli(theta);
+}
+
+generated quantities {
+   // these should be accessible via .
+   real a = 4.5;
+   array[3] real b = {1, 2.5, 4.5};
+
+   // these should not override built in properties/funs
+   real thin = 3.5;
+   int draws = 0;
+   int optimized_params_np = 0;
+   int variational_params_np = 0;
+}

test/test_generate_quantities.py

@@ -444,6 +444,26 @@ def test_complex_output(self):
             fit.draws_xr().z.isel(chain=0, draw=1).data[()], 3 + 4j
         )
 
+    def test_attrs(self):
+        stan_bern = os.path.join(DATAFILES_PATH, 'bernoulli.stan')
+        model_bern = CmdStanModel(stan_file=stan_bern)
+        jdata = os.path.join(DATAFILES_PATH, 'bernoulli.data.json')
+        fit_sampling = model_bern.sample(chains=1, iter_sampling=10, data=jdata)
+
+        stan = os.path.join(DATAFILES_PATH, 'named_output.stan')
+        model = CmdStanModel(stan_file=stan)
+        fit = model.generate_quantities(data=jdata, mcmc_sample=fit_sampling)
+
+        self.assertEqual(fit.a[0], 4.5)
+        self.assertEqual(fit.b.shape, (10, 3))
+        self.assertEqual(fit.theta.shape, (10,))
+
+        fit.draws()
+        self.assertEqual(fit.stan_variable('draws')[0], 0)
+
+        with self.assertRaisesRegex(AttributeError, 'Unknown variable name:'):
+            dummy = fit.c
+
 
 if __name__ == '__main__':
     unittest.main()

test/test_optimize.py

@@ -609,6 +609,24 @@ def test_complex_output(self):
         # make sure the name 'imag' isn't magic
         self.assertEqual(fit.stan_variable('imag').shape, (2,))
 
+    def test_attrs(self):
+        stan = os.path.join(DATAFILES_PATH, 'named_output.stan')
+        model = CmdStanModel(stan_file=stan)
+        jdata = os.path.join(DATAFILES_PATH, 'bernoulli.data.json')
+        fit = model.optimize(data=jdata)
+
+        self.assertEqual(fit.a, 4.5)
+        self.assertEqual(fit.b.shape, (3,))
+        self.assertIsInstance(fit.theta, float)
+
+        self.assertEqual(fit.stan_variable('thin'), 3.5)
+
+        self.assertIsInstance(fit.optimized_params_np, np.ndarray)
+        self.assertEqual(fit.stan_variable('optimized_params_np'), 0)
+
+        with self.assertRaisesRegex(AttributeError, 'Unknown variable name:'):
+            dummy = fit.c
+
 
 if __name__ == '__main__':
     unittest.main()

test/test_sample.py

@@ -1776,6 +1776,25 @@ def test_complex_output(self):
             fit.draws_xr().z.isel(chain=0, draw=1).data[()], 3 + 4j
         )
 
+    def test_attrs(self):
+        stan = os.path.join(DATAFILES_PATH, 'named_output.stan')
+        model = CmdStanModel(stan_file=stan)
+        jdata = os.path.join(DATAFILES_PATH, 'bernoulli.data.json')
+        fit = model.sample(chains=1, iter_sampling=10, data=jdata)
+
+        self.assertEqual(fit.a[0], 4.5)
+        self.assertEqual(fit.b.shape, (10, 3))
+        self.assertEqual(fit.theta.shape, (10,))
+
+        self.assertEqual(fit.thin, 1)
+        self.assertEqual(fit.stan_variable('thin')[0], 3.5)
+
+        fit.draws()
+        self.assertEqual(fit.stan_variable('draws')[0], 0)
+
+        with self.assertRaisesRegex(AttributeError, 'Unknown variable name:'):
+            dummy = fit.c
+
 
 if __name__ == '__main__':
     unittest.main()

test/test_variational.py

@@ -7,6 +7,7 @@
 import unittest
 from math import fabs
 
+import numpy as np
 from testfixtures import LogCapture
 
 from cmdstanpy.cmdstan_args import CmdStanArgs, VariationalArgs
@@ -264,6 +265,29 @@ def test_complex_output(self):
         # make sure the name 'imag' isn't magic
         self.assertEqual(fit.stan_variable('imag').shape, (2,))
 
+    def test_attrs(self):
+        stan = os.path.join(DATAFILES_PATH, 'named_output.stan')
+        model = CmdStanModel(stan_file=stan)
+        jdata = os.path.join(DATAFILES_PATH, 'bernoulli.data.json')
+        fit = model.variational(
+            data=jdata,
+            require_converged=False,
+            seed=12345,
+            algorithm='meanfield',
+        )
+
+        self.assertEqual(fit.a, 4.5)
+        self.assertEqual(fit.b.shape, (3,))
+        self.assertIsInstance(fit.theta, float)
+
+        self.assertEqual(fit.stan_variable('thin'), 3.5)
+
+        self.assertIsInstance(fit.variational_params_np, np.ndarray)
+        self.assertEqual(fit.stan_variable('variational_params_np'), 0)
+
+        with self.assertRaisesRegex(AttributeError, 'Unknown variable name:'):
+            dummy = fit.c
+
 
 if __name__ == '__main__':
     unittest.main()