Properly handle complex outputs (#537)

* parse complex names, dims

* stan program for unit tests

* save var type info in metadata

* complex conversion OK, needs unit tests

* logic fix

* Start testing

* Add MLE, VB and tests

* Cleanup

* Fix pytest errors, test parsing

* Fix test

* Fix typo

Co-authored-by: Mitzi Morris <mitzi@panix.com>

Author: WardBrian

cmdstanpy/stanfit/mcmc.py

@@ -32,6 +32,7 @@
 from cmdstanpy.cmdstan_args import Method, SamplerArgs
 from cmdstanpy.utils import (
     EXTENSION,
+    BaseType,
     check_sampler_csv,
     cmdstan_path,
     cmdstan_version_before,
@@ -610,6 +611,7 @@ def draws_xr(
                 self._metadata.stan_vars_cols[var],
                 0,
                 self.draws(inc_warmup=inc_warmup),
+                self._metadata.stan_vars_types[var],
             )
         return xr.Dataset(data, coords=coordinates, attrs=attrs).transpose(
             'chain', 'draw', ...
@@ -674,8 +676,10 @@ def stan_variable(
         col_idxs = self._metadata.stan_vars_cols[var]
         if len(col_idxs) > 0:
             dims.extend(self._metadata.stan_vars_dims[var])
-        # pylint: disable=redundant-keyword-arg
-        return self._draws[draw1:, :, col_idxs].reshape(dims, order='F')
+        draws = self._draws[draw1:, :, col_idxs].reshape(dims, order='F')
+        if self._metadata.stan_vars_types[var] == BaseType.COMPLEX:
+            draws = draws[..., 0] + 1j * draws[..., 1]
+        return draws
 
     def stan_variables(self) -> Dict[str, np.ndarray]:
         """
@@ -1106,6 +1110,7 @@ def draws_xr(
                 self._metadata.stan_vars_cols[var],
                 0,
                 self.draws(inc_warmup=inc_warmup),
+                self._metadata.stan_vars_types[var],
             )
         if inc_sample:
             for var in mcmc_vars_list:
@@ -1116,6 +1121,7 @@ def draws_xr(
                     self.mcmc_sample.metadata.stan_vars_cols[var],
                     0,
                     self.mcmc_sample.draws(inc_warmup=inc_warmup),
+                    self.mcmc_sample._metadata.stan_vars_types[var],
                 )
 
         return xr.Dataset(data, coords=coordinates, attrs=attrs).transpose(
@@ -1193,7 +1199,10 @@ def stan_variable(
             if len(col_idxs) > 0:
                 dims.extend(self._metadata.stan_vars_dims[var])
             # pylint: disable=redundant-keyword-arg
-            return self._draws[draw1:, :, col_idxs].reshape(dims, order='F')
+            draws = self._draws[draw1:, :, col_idxs].reshape(dims, order='F')
+            if self._metadata.stan_vars_types[var] == BaseType.COMPLEX:
+                draws = draws[..., 0] + 1j * draws[..., 1]
+            return draws
 
     def stan_variables(self, inc_warmup: bool = False) -> Dict[str, np.ndarray]:
         """
@@ -1262,6 +1271,7 @@ def build_xarray_data(
     col_idxs: Tuple[int, ...],
     start_row: int,
     drawset: np.ndarray,
+    var_type: BaseType,
 ) -> None:
     """
     Adds Stan variable name, labels, and values to a dictionary
@@ -1270,12 +1280,19 @@ def build_xarray_data(
     var_dims: Tuple[str, ...] = ('draw', 'chain')
     if dims:
         var_dims += tuple(f"{var_name}_dim_{i}" for i in range(len(dims)))
+
+        draws = drawset[start_row:, :, col_idxs].reshape(
+            *drawset.shape[:2], *dims, order="F"
+        )
+        if var_type == BaseType.COMPLEX:
+            draws = draws[..., 0] + 1j * draws[..., 1]
+            var_dims = var_dims[:-1]
+
         data[var_name] = (
             var_dims,
-            drawset[start_row:, :, col_idxs].reshape(
-                *drawset.shape[:2], *dims, order="F"
-            ),
+            draws,
         )
+
     else:
         data[var_name] = (
             var_dims,

cmdstanpy/stanfit/metadata.py

@@ -3,7 +3,7 @@
 import copy
 from typing import Any, Dict, Tuple
 
-from cmdstanpy.utils import parse_method_vars, parse_stan_vars
+from cmdstanpy.utils import parse_method_vars, parse_stan_vars, BaseType
 
 
 class InferenceMetadata:
@@ -17,11 +17,12 @@ def __init__(self, config: Dict[str, Any]) -> None:
         """Initialize object from CSV headers"""
         self._cmdstan_config = config
         self._method_vars_cols = parse_method_vars(names=config['column_names'])
-        stan_vars_dims, stan_vars_cols = parse_stan_vars(
+        stan_vars_dims, stan_vars_cols, stan_vars_types = parse_stan_vars(
             names=config['column_names']
         )
         self._stan_vars_dims = stan_vars_dims
         self._stan_vars_cols = stan_vars_cols
+        self._stan_vars_types = stan_vars_types
 
     def __repr__(self) -> str:
         return 'Metadata:\n{}\n'.format(self._cmdstan_config)
@@ -66,3 +67,11 @@ def stan_vars_dims(self) -> Dict[str, Tuple[int, ...]]:
         Uses deepcopy for immutability.
         """
         return copy.deepcopy(self._stan_vars_dims)
+
+    @property
+    def stan_vars_types(self) -> Dict[str, BaseType]:
+        """
+        Returns map from Stan program variable names to variable base type.
+        Uses deepcopy for immutability.
+        """
+        return copy.deepcopy(self._stan_vars_types)

cmdstanpy/stanfit/mle.py

@@ -7,7 +7,7 @@
 import pandas as pd
 
 from cmdstanpy.cmdstan_args import Method, OptimizeArgs
-from cmdstanpy.utils import get_logger, scan_optimize_csv
+from cmdstanpy.utils import BaseType, get_logger, scan_optimize_csv
 
 from .metadata import InferenceMetadata
 from .runset import RunSet
@@ -199,22 +199,24 @@ def stan_variable(
         else:
             num_rows = 1
 
-        result: Union[np.ndarray, float]
         if len(col_idxs) > 1:  # container var
             dims = (num_rows,) + self._metadata.stan_vars_dims[var]
             # pylint: disable=redundant-keyword-arg
             if num_rows > 1:
                 result = self._all_iters[:, col_idxs].reshape(dims, order='F')
             else:
                 result = self._mle[col_idxs].reshape(dims[1:], order="F")
+
+            if self._metadata.stan_vars_types[var] == BaseType.COMPLEX:
+                result = result[..., 0] + 1j * result[..., 1]
+            return result
+
         else:  # scalar var
             col_idx = col_idxs[0]
             if num_rows > 1:
-                result = self._all_iters[:, col_idx]
+                return self._all_iters[:, col_idx]
             else:
-                result = float(self._mle[col_idx])
-
-        return result
+                return float(self._mle[col_idx])
 
     def stan_variables(
         self, inc_iterations: bool = False

cmdstanpy/stanfit/vb.py

@@ -7,7 +7,7 @@
 import pandas as pd
 
 from cmdstanpy.cmdstan_args import Method
-from cmdstanpy.utils import scan_variational_csv
+from cmdstanpy.utils import BaseType, scan_variational_csv
 
 from .metadata import InferenceMetadata
 from .runset import RunSet
@@ -126,15 +126,17 @@ def stan_variable(
             raise ValueError('Unknown variable name: {}'.format(var))
         col_idxs = list(self._metadata.stan_vars_cols[var])
         shape: Tuple[int, ...] = ()
-        result: Union[np.ndarray, float]
         if len(col_idxs) > 1:
             shape = self._metadata.stan_vars_dims[var]
-            result = np.asarray(self._variational_mean)[col_idxs].reshape(
-                shape, order="F"
-            )
+            result: np.ndarray = np.asarray(self._variational_mean)[
+                col_idxs
+            ].reshape(shape, order="F")
+
+            if self._metadata.stan_vars_types[var] == BaseType.COMPLEX:
+                result = result[..., 0] + 1j * result[..., 1]
+            return result
         else:
-            result = float(self._variational_mean[col_idxs[0]])
-        return result
+            return float(self._variational_mean[col_idxs[0]])
 
     def stan_variables(self) -> Dict[str, Union[np.ndarray, float]]:
         """

cmdstanpy/utils.py

@@ -14,6 +14,7 @@
 import tempfile
 from collections import OrderedDict
 from collections.abc import Collection
+from enum import Enum, auto
 from typing import (
     Any,
     Callable,
@@ -46,6 +47,16 @@
 EXTENSION = '.exe' if platform.system() == 'Windows' else ''
 
 
+class BaseType(Enum):
+    """Stan langauge base type"""
+
+    COMPLEX = auto()
+    PRIM = auto()  # future: int / real
+
+    def __repr__(self) -> str:
+        return '<%s.%s>' % (self.__class__.__name__, self.name)
+
+
 @functools.lru_cache(maxsize=None)
 def get_logger() -> logging.Logger:
     """cmdstanpy logger"""
@@ -794,10 +805,14 @@ def munge_varnames(names: List[str]) -> List[str]:
     """
     if names is None:
         raise ValueError('missing argument "names"')
-    return [
-        re.sub(r',([\d,]+)$', r'[\1]', column.replace('.', ','))
-        for column in names
-    ]
+    result = []
+    for name in names:
+        if '.' not in name:
+            result.append(name)
+        else:
+            head, *rest = name.split('.')
+            result.append(''.join([head, '[', ','.join(rest), ']']))
+    return result
 
 
 def parse_method_vars(names: Tuple[str, ...]) -> Dict[str, Tuple[int, ...]]:
@@ -816,38 +831,52 @@ def parse_method_vars(names: Tuple[str, ...]) -> Dict[str, Tuple[int, ...]]:
 
 def parse_stan_vars(
     names: Tuple[str, ...]
-) -> Tuple[Dict[str, Tuple[int, ...]], Dict[str, Tuple[int, ...]]]:
+) -> Tuple[
+    Dict[str, Tuple[int, ...]], Dict[str, Tuple[int, ...]], Dict[str, BaseType]
+]:
     """
     Parses out Stan variable names (i.e., names not ending in `__`)
     from list of CSV file column names.
-    Returns a pair of dicts which map variable names to dimensions and
-    variable names to columns, respectively, using zero-based column indexing.
+    Returns three dicts which map variable names to base type, dimensions and
+    CSV file columns, respectively, using zero-based column indexing.
     Note: assumes: (a) munged varnames and (b) container vars are non-ragged
-    and dense; no checks size, indices.
+    and dense; no checks on size, indices.
     """
     if names is None:
         raise ValueError('missing argument "names"')
     dims_map: Dict[str, Tuple[int, ...]] = {}
     cols_map: Dict[str, Tuple[int, ...]] = {}
+    types_map: Dict[str, BaseType] = {}
     idxs = []
     dims: Union[List[str], List[int]]
     for (idx, name) in enumerate(names):
+        if name.endswith('real]') or name.endswith('imag]'):
+            basetype = BaseType.COMPLEX
+        else:
+            basetype = BaseType.PRIM
         idxs.append(idx)
         var, *dims = name.split('[')
         if var.endswith('__'):
             idxs = []
         elif len(dims) == 0:
             dims_map[var] = ()
             cols_map[var] = tuple(idxs)
+            types_map[var] = basetype
             idxs = []
         else:
             if idx < len(names) - 1 and names[idx + 1].split('[')[0] == var:
                 continue
-            dims = [int(x) for x in dims[0][:-1].split(',')]
+            coords = dims[0][:-1].split(',')
+            if coords[-1] == 'imag':
+                dims = [int(x) for x in coords[:-1]]
+                dims.append(2)
+            else:
+                dims = [int(x) for x in coords]
             dims_map[var] = tuple(dims)
             cols_map[var] = tuple(idxs)
+            types_map[var] = basetype
             idxs = []
-    return (dims_map, cols_map)
+    return (dims_map, cols_map, types_map)
 
 
 def scan_hmc_params(

test/data/complex_var.stan

@@ -0,0 +1,21 @@
+transformed data {
+  array[10] int<lower=0,upper=1> y = {0,1,0,0,0,0,0,0,0,1};
+}
+parameters {
+  real<lower=0,upper=1> theta;
+}
+model {
+  theta ~ beta(1,1);
+  y ~ bernoulli(theta);
+}
+
+// model segment is just so that VB works
+
+generated quantities {
+  int a  = 1;
+  array[2,3,2] int ys = {{{3,0},{0,4}, {5,0}}, {{0,1}, {0,2}, {0,3}}};
+  array[2,3] complex zs = {{3,4i,5},{1i,2i,3i}};
+  complex z = 3 + 4i;
+
+  array[2] int imag = {3,4};
+}

test/test_generate_quantities.py

@@ -423,6 +423,27 @@ def test_show_console(self):
         self.assertTrue('Chain [3] method = generate' in console)
         self.assertTrue('Chain [4] method = generate' in console)
 
+    def test_complex_output(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, 'complex_var.stan')
+        model = CmdStanModel(stan_file=stan)
+        fit = model.generate_quantities(mcmc_sample=fit_sampling)
+
+        self.assertEqual(fit.stan_variable('zs').shape, (10, 2, 3))
+        self.assertEqual(fit.stan_variable('z')[0], 3 + 4j)
+        # make sure the name 'imag' isn't magic
+        self.assertEqual(fit.stan_variable('imag').shape, (10, 2))
+
+        self.assertNotIn("zs_dim_2", fit.draws_xr())
+        # getting a raw scalar out of xarray is heavy
+        self.assertEqual(
+            fit.draws_xr().z.isel(chain=0, draw=1).data[()], 3 + 4j
+        )
+
 
 if __name__ == '__main__':
     unittest.main()

test/test_optimize.py

@@ -599,6 +599,16 @@ def test_exe_only(self):
             mle.optimized_params_np[1], mle.optimized_params_dict['theta']
         )
 
+    def test_complex_output(self):
+        stan = os.path.join(DATAFILES_PATH, 'complex_var.stan')
+        model = CmdStanModel(stan_file=stan)
+        fit = model.optimize()
+
+        self.assertEqual(fit.stan_variable('zs').shape, (2, 3))
+        self.assertEqual(fit.stan_variable('z'), 3 + 4j)
+        # make sure the name 'imag' isn't magic
+        self.assertEqual(fit.stan_variable('imag').shape, (2,))
+
 
 if __name__ == '__main__':
     unittest.main()

test/test_sample.py

@@ -1762,6 +1762,22 @@ def test_single_row_csv(self):
                 self.assertEqual(int(z_as_ndarray[0, i, j]), i + 1)
                 self.assertEqual(int(z_as_xr.z.data[0, 0, i, j]), i + 1)
 
+    def test_complex_output(self):
+        stan = os.path.join(DATAFILES_PATH, 'complex_var.stan')
+        model = CmdStanModel(stan_file=stan)
+        fit = model.sample(chains=1, iter_sampling=10)
+
+        self.assertEqual(fit.stan_variable('zs').shape, (10, 2, 3))
+        self.assertEqual(fit.stan_variable('z')[0], 3 + 4j)
+        # make sure the name 'imag' isn't magic
+        self.assertEqual(fit.stan_variable('imag').shape, (10, 2))
+
+        self.assertNotIn("zs_dim_2", fit.draws_xr())
+        # getting a raw scalar out of xarray is heavy
+        self.assertEqual(
+            fit.draws_xr().z.isel(chain=0, draw=1).data[()], 3 + 4j
+        )
+
 
 if __name__ == '__main__':
     unittest.main()