Repository: stan-dev/cmdstanpy
Pull request #337 (2021-01-14 08:42:27)

Add extract function
#### Submission Checklist

- [x] Run unit tests
- [x] Declare copyright holder and open-source license: see below

#### Summary

Add option to extract ndim objects.

#### Copyright and Licensing

Please list the copyright holder for the work you are submitting (this will be you or your assignee, such as a university or company):

Ari Hartikainen

By submitting this pull request, the copyright holder is agreeing to license the submitted work under the following licenses:
- Code: BSD 3-clause (https://opensource.org/licenses/BSD-3-Clause)




Patch:
diff --git a/.pylintrc b/.pylintrc
index 5d42e56..7ae58f5 100644
--- a/.pylintrc
+++ b/.pylintrc
@@ -3,7 +3,7 @@
 # A comma-separated list of package or module names from where C extensions may
 # be loaded. Extensions are loading into the active Python interpreter and may
 # run arbitrary code.
-extension-pkg-whitelist=
+extension-pkg-whitelist=ujson
 
 # Add files or directories to the blacklist. They should be base names, not
 # paths.
diff --git a/cmdstanpy/__init__.py b/cmdstanpy/__init__.py
index 58ef9b8..2bdc473 100644
--- a/cmdstanpy/__init__.py
+++ b/cmdstanpy/__init__.py
@@ -40,10 +40,12 @@ from .model import CmdStanModel  # noqa
 from .stanfit import CmdStanGQ, CmdStanMCMC, CmdStanMLE, CmdStanVB  # noqa
 from .utils import set_cmdstan_path  # noqa
 from .utils import cmdstan_path, install_cmdstan, set_make_env
+from .utils import extract
 
 __all__ = [
     'set_cmdstan_path',
     'cmdstan_path',
+    'extract',
     'set_make_env',
     'install_cmdstan',
     'CmdStanMCMC',
diff --git a/cmdstanpy/utils.py b/cmdstanpy/utils.py
index aa58044..69a9c64 100644
--- a/cmdstanpy/utils.py
+++ b/cmdstanpy/utils.py
@@ -10,10 +10,10 @@ import shutil
 import subprocess
 import sys
 import tempfile
-from collections import OrderedDict
+from collections import OrderedDict, defaultdict
 from collections.abc import Sequence
 from numbers import Integral, Real
-from typing import Dict, List, TextIO, Tuple, Union
+from typing import Any, Dict, List, Optional, Union, TextIO, Tuple
 
 import numpy as np
 import pandas as pd
@@ -1061,3 +1061,74 @@ class TemporaryCopiedFile:
     def __exit__(self, exc_type, exc_val, exc_tb):
         if self._tmpdir:
             shutil.rmtree(self._tmpdir, ignore_errors=True)
+
+
+def extract(
+    fit: Any,
+    parameters: Optional[List] = None,
+    inc_diagnostics: Optional[bool] = False,
+    inc_warmup: Optional[bool] = False,
+) -> Dict[str, Any]:
+    """Extract ndim dictionary from CmdStanPy fit.
+
+    Parameters
+    ----------
+    fit: CmdStanMCMC
+    parameters: list
+    inc_diagnostics: bool
+        Don't include diagnostics with default paramaters.
+    inc_warmup: bool
+
+    Returns
+    -------
+    dict
+    """
+    if parameters is None:
+        parameters = fit.column_names
+        if not inc_diagnostics:
+            parameters = [
+                param for param in parameters if not param.endswith("__")
+            ]
+
+    if inc_warmup and not fit._save_warmup:
+        inc_warmup = False
+
+    data = fit.draws(inc_warmup=inc_warmup)
+    draws, chains, *_ = data.shape
+    column_groups = defaultdict(list)
+    column_locs = defaultdict(list)
+    # iterate flat column names
+    for i, col in enumerate(fit.column_names):
+        col_base, *col_tail = col.replace(']', '').replace('[', ',').split(',')
+        if col_tail:
+            # gather nD array locations
+            column_groups[col_base].append(tuple(map(int, col_tail)))
+        # gather raw data locations for each parameter
+        column_locs[col_base].append(i)
+    dims = {}
+    for colname, col_dims in column_groups.items():
+        # gather parameter dimensions (assumes dense arrays)
+        dims[colname] = tuple(np.array(col_dims).max(0))
+
+    sample = {}
+    valid_base_cols = []
+    # get list of parameters for extraction (basename) X[1,2] --> X
+    for col in parameters:
+        base_col = col.split('[')[0]
+        if base_col not in valid_base_cols:
+            valid_base_cols.append(base_col)
+
+    for key in valid_base_cols:
+        ndim = dims.get(key, None)
+        shape_location = column_groups.get(key, None)
+        if ndim is not None:
+            sample[key] = np.full((draws, chains, *ndim), np.nan)
+        if shape_location is None:
+            (i,) = column_locs[key]
+            sample[key] = data[..., i]
+        else:
+            for i, shape_loc in zip(column_locs[key], shape_location):
+                # location to insert extracted array
+                shape_loc = tuple([Ellipsis] + [j - 1 for j in shape_loc])
+                sample[key][shape_loc] = data[..., i]
+    return sample


Test patch:
diff --git a/test/test_utils.py b/test/test_utils.py
index c812894..9dd8440 100644
--- a/test/test_utils.py
+++ b/test/test_utils.py
@@ -20,6 +20,7 @@ from cmdstanpy.utils import (
     check_sampler_csv,
     cmdstan_path,
     cmdstan_version_at,
+    extract,
     get_latest_cmdstan,
     jsondump,
     parse_rdump_value,
@@ -620,6 +621,52 @@ class ParseVarDimsTest(unittest.TestCase):
         self.assertEqual(vars_dict['phi'], (2, 5))
         self.assertEqual(vars_dict['bar'], (2,))
 
+    def test_extract(self):
+        stan_file = os.path.join(_TMPDIR, "extract_test", "model.stan")
+        os.makedirs(os.path.dirname(stan_file), exist_ok=True)
+        model_code = """
+            parameters { matrix[2,3] y[4]; }
+            model {
+                for (n in 1:4) {
+                    for (m in 1:2) {
+                        for (l in 1:3) {
+                            y[n, m, l] ~ normal(n, 1);
+            }}}}
+        """
+        with open(stan_file, 'w') as fd:
+            print(model_code, file=fd)
+        model = CmdStanModel(stan_file=stan_file)
+        model.compile()
+        fit = model.sample(
+            chains=2,
+            parallel_chains=1,
+            seed=12345,
+            iter_sampling=490,
+            iter_warmup=490,
+            save_warmup=True,
+        )
+        sample = extract(fit)
+        assert 'y' in sample
+        assert 'lp__' not in sample
+        self.assertEqual(sample['y'].shape, (490, 2, 4, 2, 3))
+
+        sample = extract(fit, inc_warmup=True)
+        assert 'y' in sample
+        assert 'lp__' not in sample
+        self.assertEqual(sample['y'].shape, (490 + 490, 2, 4, 2, 3))
+
+        sample = extract(fit, inc_diagnostics=True)
+        assert 'y' in sample
+        assert 'lp__' in sample
+        self.assertEqual(sample['y'].shape, (490, 2, 4, 2, 3))
+        self.assertEqual(sample['lp__'].shape, (490, 2))
+
+        sample = extract(fit, inc_diagnostics=True, inc_warmup=True)
+        assert 'y' in sample
+        assert 'lp__' in sample
+        self.assertEqual(sample['y'].shape, (490 + 490, 2, 4, 2, 3))
+        self.assertEqual(sample['lp__'].shape, (490 + 490, 2))
+
 
 if __name__ == '__main__':
     unittest.main()