Merge pull request #543 from stan-dev/stanvariables-dot

Support `.` syntax for Stan variable access

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."""
@@ -619,7 +627,7 @@ def draws_xr(
 
     def stan_variable(
         self,
-        var: Optional[str] = None,
+        var: str,
         inc_warmup: bool = False,
     ) -> np.ndarray:
         """
@@ -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
@@ -660,10 +671,12 @@ def stan_variable(
         CmdStanVB.stan_variable
         CmdStanGQ.stan_variable
         """
-        if var is None:
-            raise ValueError('No variable name specified.')
         if var not in self._metadata.stan_vars_dims:
-            raise ValueError('Unknown variable name: {}'.format(var))
+            raise ValueError(
+                f'Unknown variable name: {var}\n'
+                'Available variables are '
+                + ", ".join(self._metadata.stan_vars_dims)
+            )
         if self._draws.shape == (0,):
             self._assemble_draws()
         draw1 = 0
@@ -767,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
@@ -1130,7 +1151,7 @@ def draws_xr(
 
     def stan_variable(
         self,
-        var: Optional[str] = None,
+        var: str,
         inc_warmup: bool = False,
     ) -> np.ndarray:
         """
@@ -1158,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
@@ -1171,12 +1195,14 @@ def stan_variable(
         CmdStanMLE.stan_variable
         CmdStanVB.stan_variable
         """
-        if var is None:
-            raise ValueError('No variable name specified.')
         model_var_names = self.mcmc_sample.metadata.stan_vars_cols.keys()
         gq_var_names = self.metadata.stan_vars_cols.keys()
         if not (var in model_var_names or var in gq_var_names):
-            raise ValueError('Unknown variable name: {}'.format(var))
+            raise ValueError(
+                f'Unknown variable name: {var}\n'
+                'Available variables are '
+                + ", ".join(model_var_names | gq_var_names)
+            )
         if var not in gq_var_names:
             return self.mcmc_sample.stan_variable(var, inc_warmup=inc_warmup)
         else:  # is gq variable

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)
@@ -155,7 +163,7 @@ def optimized_params_dict(self) -> Dict[str, float]:
 
     def stan_variable(
         self,
-        var: Optional[str] = None,
+        var: str,
         *,
         inc_iterations: bool = False,
         warn: bool = True,
@@ -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
@@ -179,10 +190,12 @@ def stan_variable(
         CmdStanVB.stan_variable
         CmdStanGQ.stan_variable
         """
-        if var is None:
-            raise ValueError('no variable name specified.')
         if var not in self._metadata.stan_vars_dims:
-            raise ValueError('unknown variable name: {}'.format(var))
+            raise ValueError(
+                f'Unknown variable name: {var}\n'
+                'Available variables are '
+                + ", ".join(self._metadata.stan_vars_dims)
+            )
         if warn and inc_iterations and not self._save_iterations:
             get_logger().warning(
                 'Intermediate iterations not saved to CSV output file. '

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)
@@ -103,14 +111,15 @@ def metadata(self) -> InferenceMetadata:
         """
         return self._metadata
 
-    def stan_variable(
-        self, var: Optional[str] = None
-    ) -> Union[np.ndarray, float]:
+    def stan_variable(self, var: str) -> Union[np.ndarray, float]:
         """
         Return a numpy.ndarray which contains the estimates for the
         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
@@ -123,7 +132,11 @@ def stan_variable(
         if var is None:
             raise ValueError('No variable name specified.')
         if var not in self._metadata.stan_vars_dims:
-            raise ValueError('Unknown variable name: {}'.format(var))
+            raise ValueError(
+                f'Unknown variable name: {var}\n'
+                'Available variables are '
+                + ", ".join(self._metadata.stan_vars_dims)
+            )
         col_idxs = list(self._metadata.stan_vars_cols[var])
         shape: Tuple[int, ...] = ()
         if len(col_idxs) > 1:

cmdstanpy/utils.py

@@ -730,7 +730,7 @@ def scan_variational_csv(path: str) -> Dict[str, Any]:
             lineno += 1
         xs = line.split(',')
         variational_mean = [float(x) for x in xs]
-        dict['variational_mean'] = variational_mean
+        dict['variational_mean'] = np.array(variational_mean)
         dict['variational_sample'] = pd.read_csv(
             path,
             comment='#',

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()