Repository: stan-dev/cmdstanpy
Pull request #418 (2021-08-03 15:23:15)

Expose write_stan_json utility
#### Submission Checklist

- [x] Run unit tests
- [x] Declare copyright holder and open-source license: see below

#### Summary
Addresses #410 by exposing `jsondump` as `write_stan_json` and **removing** `rdump`. 

As part of this I am allowing the more general `Mapping` type (which essentially means `.items()` is implemented) rather than strictly `Dict`, so you can now do things like `write_stan_json(path, dataframe)` or even `write_stan_json(path, xarray_dataset)`

I also added additional checks modeled on the cmdstanr function of the same name for clarity. These checks consist of:
* Verify the type of each value is numeric, a boolean, or a collection with a meaningful error if not.
* Verify that the values present are finite (meaning not NaN or inf/-inf)

Checking these two issues in particular mean that the call to `ujson.dump` should never itself fail

#### 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):
Simons Foundation

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/cmdstanpy/__init__.py b/cmdstanpy/__init__.py
index 537b9b0..28c4485 100644
--- a/cmdstanpy/__init__.py
+++ b/cmdstanpy/__init__.py
@@ -45,7 +45,7 @@ from .stanfit import (  # noqa
     from_csv,
 )
 from .utils import set_cmdstan_path  # noqa
-from .utils import cmdstan_path, install_cmdstan, set_make_env
+from .utils import cmdstan_path, install_cmdstan, set_make_env, write_stan_json
 
 __all__ = [
     'set_cmdstan_path',
@@ -58,4 +58,5 @@ __all__ = [
     'CmdStanVB',
     'CmdStanModel',
     'from_csv',
+    'write_stan_json',
 ]
diff --git a/cmdstanpy/model.py b/cmdstanpy/model.py
index 957d78d..84c66da 100644
--- a/cmdstanpy/model.py
+++ b/cmdstanpy/model.py
@@ -10,7 +10,7 @@ from collections import OrderedDict
 from concurrent.futures import ThreadPoolExecutor
 from multiprocessing import cpu_count
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Mapping, Optional, Union
 
 from cmdstanpy.cmdstan_args import (
     CmdStanArgs,
@@ -360,7 +360,7 @@ class CmdStanModel:
 
     def optimize(
         self,
-        data: Union[Dict[str, Any], str, None] = None,
+        data: Union[Mapping[str, Any], str, None] = None,
         seed: Optional[int] = None,
         inits: Union[Dict[str, float], float, str, None] = None,
         output_dir: Optional[str] = None,
@@ -507,7 +507,7 @@ class CmdStanModel:
     # pylint: disable=too-many-arguments
     def sample(
         self,
-        data: Union[Dict[str, Any], str, None] = None,
+        data: Union[Mapping[str, Any], str, None] = None,
         chains: Optional[int] = None,
         parallel_chains: Optional[int] = None,
         threads_per_chain: Optional[int] = None,
@@ -868,7 +868,7 @@ class CmdStanModel:
 
     def generate_quantities(
         self,
-        data: Union[Dict[str, Any], str, None] = None,
+        data: Union[Mapping[str, Any], str, None] = None,
         mcmc_sample: Union[CmdStanMCMC, List[str], None] = None,
         seed: Optional[int] = None,
         gq_output_dir: Optional[str] = None,
@@ -1024,7 +1024,7 @@ class CmdStanModel:
 
     def variational(
         self,
-        data: Union[Dict[str, Any], str, None] = None,
+        data: Union[Mapping[str, Any], str, None] = None,
         seed: Optional[int] = None,
         inits: Optional[float] = None,
         output_dir: Optional[str] = None,
diff --git a/cmdstanpy/utils.py b/cmdstanpy/utils.py
index 989c1df..e607cc7 100644
--- a/cmdstanpy/utils.py
+++ b/cmdstanpy/utils.py
@@ -12,13 +12,14 @@ import subprocess
 import sys
 import tempfile
 from collections import OrderedDict
-from collections.abc import Collection, Sequence
+from collections.abc import Collection
 from typing import (
     Any,
     Callable,
     Dict,
     Iterator,
     List,
+    Mapping,
     MutableMapping,
     Optional,
     TextIO,
@@ -379,39 +380,61 @@ def cxx_toolchain_path(
     return compiler_path, tool_path
 
 
-def _rdump_array(key: str, val: np.ndarray) -> str:
-    """Flatten numpy ndarray, format as Rdump variable declaration."""
-    c = 'c(' + ', '.join(map(str, val.T.flat)) + ')'
-    if (val.size,) == val.shape:
-        return '{key} <- {c}'.format(key=key, c=c)
-    else:
-        dim = '.Dim = c{}'.format(val.shape)
-        struct = '{key} <- structure({c}, {dim})'.format(key=key, c=c, dim=dim)
-        return struct
+def write_stan_json(path: str, data: Mapping[str, Any]) -> None:
+    """
+    Dump a mapping of strings to data to a JSON file.
 
+    Values can be any numeric type, a boolean (converted to int),
+    or any collection compatible with ``numpy.asarray``, e.g a
+    pandas.Series.
 
-def jsondump(path: str, data: Dict[str, Any]) -> None:
-    """Dump a dict of data to a JSON file."""
-    data = data.copy()
+    Produces a file compatible with the
+    `Json Format for Cmdstan
+    <https://mc-stan.org/docs/2_27/cmdstan-guide/json.html>`__
+
+    :param path: File path for the created json. Will be overwritten if
+    already in existence.
+    :param data: A mapping from strings to values. This can be a dictionary
+    or something more exotic like an xarray.Dataset. This will be copied
+    before type conversion, not modified
+    """
+    data_out = {}
     for key, val in data.items():
+        if val is not None:
+            if isinstance(val, (str, bytes)) or (
+                type(val).__module__ != 'numpy'
+                and not isinstance(val, (Collection, bool, int, float))
+            ):
+                raise TypeError(
+                    f"Invalid type '{type(val)}' provided to "
+                    + f"write_stan_json for key '{key}'"
+                )
+            try:
+                if not np.all(np.isfinite(val)):
+                    raise ValueError(
+                        "Input to write_stan_json has nan or infinite "
+                        + f"values for key '{key}'"
+                    )
+            except TypeError:
+                # handles cases like val == ['hello']
+                # pylint: disable=raise-missing-from
+                raise ValueError(
+                    "Invalid type provided to "
+                    + f"write_stan_json for key '{key}' "
+                    + f"as part of collection {type(val)}"
+                )
+
         if type(val).__module__ == 'numpy':
-            data[key] = val.tolist()
+            data_out[key] = val.tolist()
         elif isinstance(val, Collection):
-            data[key] = np.asarray(val).tolist()
-
-    with open(path, 'w') as fd:
-        json.dump(data, fd)
-
+            data_out[key] = np.asarray(val).tolist()
+        elif isinstance(val, bool):
+            data_out[key] = int(val)
+        else:
+            data_out[key] = val
 
-def rdump(path: str, data: Dict[str, Any]) -> None:
-    """Dump a dict of data to a R dump format file."""
     with open(path, 'w') as fd:
-        for key, val in data.items():
-            if isinstance(val, (np.ndarray, Sequence)):
-                line = _rdump_array(key, np.asarray(val))
-            else:
-                line = '{} <- {}'.format(key, val)
-            print(line, file=fd)
+        json.dump(data_out, fd)
 
 
 def rload(fname: str) -> Optional[Dict[str, Union[int, float, np.ndarray]]]:
@@ -1113,27 +1136,20 @@ class MaybeDictToFilePath:
 
     def __init__(
         self,
-        *objs: Union[str, Dict[Any, Any], List[Any], int, float, None],
-        logger: Optional[logging.Logger] = None
+        *objs: Union[str, Mapping[str, Any], List[Any], int, float, None],
+        logger: Optional[logging.Logger] = None,
     ):
         self._unlink = [False] * len(objs)
         self._paths: List[Any] = [''] * len(objs)
         self._logger = logger or get_logger()
         i = 0
         for obj in objs:
-            if isinstance(obj, dict):
+            if isinstance(obj, Mapping):
                 data_file = create_named_text_file(
                     dir=_TMPDIR, prefix='', suffix='.json'
                 )
                 self._logger.debug('input tempfile: %s', data_file)
-                if any(
-                    not item
-                    for item in obj
-                    if isinstance(item, (Sequence, np.ndarray))
-                ):
-                    rdump(data_file, obj)
-                else:
-                    jsondump(data_file, obj)
+                write_stan_json(data_file, obj)
                 self._paths[i] = data_file
                 self._unlink[i] = True
             elif isinstance(obj, str):


Test patch:
diff --git a/test/test_utils.py b/test/test_utils.py
index 3336651..1a5a35c 100644
--- a/test/test_utils.py
+++ b/test/test_utils.py
@@ -24,17 +24,16 @@ from cmdstanpy.utils import (
     cmdstan_version_at,
     do_command,
     get_latest_cmdstan,
-    jsondump,
     parse_method_vars,
     parse_rdump_value,
     parse_stan_vars,
-    rdump,
     read_metric,
     rload,
     set_cmdstan_path,
     validate_cmdstan_path,
     validate_dir,
     windows_short_path,
+    write_stan_json,
 )
 
 HERE = os.path.dirname(os.path.abspath(__file__))
@@ -208,7 +207,7 @@ class DataFilesTest(unittest.TestCase):
             with MaybeDictToFilePath(123, dict_good) as (fg1, fg2):
                 pass
 
-    def test_jsondump(self):
+    def test_write_stan_json(self):
         def cmp(d1, d2):
             self.assertEqual(d1.keys(), d2.keys())
             for k in d1:
@@ -221,40 +220,58 @@ class DataFilesTest(unittest.TestCase):
 
         dict_list = {'a': [1.0, 2.0, 3.0]}
         file_list = os.path.join(_TMPDIR, 'list.json')
-        jsondump(file_list, dict_list)
+        write_stan_json(file_list, dict_list)
         with open(file_list) as fd:
             cmp(json.load(fd), dict_list)
 
         arr = np.repeat(3, 4)
         dict_vec = {'a': arr}
         file_vec = os.path.join(_TMPDIR, 'vec.json')
-        jsondump(file_vec, dict_vec)
+        write_stan_json(file_vec, dict_vec)
         with open(file_vec) as fd:
             cmp(json.load(fd), dict_vec)
 
+        dict_bool = {'a': False}
+        file_bool = os.path.join(_TMPDIR, 'bool.json')
+        write_stan_json(file_bool, dict_bool)
+        with open(file_bool) as fd:
+            cmp(json.load(fd), {'a': 0})
+
+        dict_none = {'a': None}
+        file_none = os.path.join(_TMPDIR, 'none.json')
+        write_stan_json(file_none, dict_none)
+        with open(file_none) as fd:
+            cmp(json.load(fd), dict_none)
+
         series = pd.Series(arr)
         dict_vec_pd = {'a': series}
         file_vec_pd = os.path.join(_TMPDIR, 'vec_pd.json')
-        jsondump(file_vec_pd, dict_vec_pd)
+        write_stan_json(file_vec_pd, dict_vec_pd)
         with open(file_vec_pd) as fd:
             cmp(json.load(fd), dict_vec_pd)
 
+        df_vec = pd.DataFrame(dict_list)
+        file_pd = os.path.join(_TMPDIR, 'pd.json')
+        write_stan_json(file_pd, df_vec)
+        with open(file_pd) as fd:
+            cmp(json.load(fd), dict_list)
+
         dict_zero_vec = {'a': []}
         file_zero_vec = os.path.join(_TMPDIR, 'empty_vec.json')
-        jsondump(file_zero_vec, dict_zero_vec)
+        write_stan_json(file_zero_vec, dict_zero_vec)
         with open(file_zero_vec) as fd:
             cmp(json.load(fd), dict_zero_vec)
 
         dict_zero_matrix = {'a': [[], [], []]}
         file_zero_matrix = os.path.join(_TMPDIR, 'empty_matrix.json')
-        jsondump(file_zero_matrix, dict_zero_matrix)
+        write_stan_json(file_zero_matrix, dict_zero_matrix)
         with open(file_zero_matrix) as fd:
             cmp(json.load(fd), dict_zero_matrix)
 
         arr = np.zeros(shape=(5, 0))
         dict_zero_matrix = {'a': arr}
         file_zero_matrix = os.path.join(_TMPDIR, 'empty_matrix.json')
-        jsondump(file_zero_matrix, dict_zero_matrix)
+        write_stan_json(file_zero_matrix, dict_zero_matrix)
         with open(file_zero_matrix) as fd:
             cmp(json.load(fd), dict_zero_matrix)
 
@@ -264,7 +281,7 @@ class DataFilesTest(unittest.TestCase):
 
         dict_3d_matrix = {'a': arr}
         file_3d_matrix = os.path.join(_TMPDIR, '3d_matrix.json')
-        jsondump(file_3d_matrix, dict_3d_matrix)
+        write_stan_json(file_3d_matrix, dict_3d_matrix)
         with open(file_3d_matrix) as fd:
             cmp(json.load(fd), dict_3d_matrix)
 
@@ -272,10 +289,29 @@ class DataFilesTest(unittest.TestCase):
         self.assertTrue(type(scalr).__module__ == 'numpy')
         dict_scalr = {'a': scalr}
         file_scalr = os.path.join(_TMPDIR, 'scalr.json')
-        jsondump(file_scalr, dict_scalr)
+        write_stan_json(file_scalr, dict_scalr)
         with open(file_scalr) as fd:
             cmp(json.load(fd), dict_scalr)
 
+    def test_write_stan_json_bad(self):
+        file_bad = os.path.join(_TMPDIR, 'bad.json')
+
+        dict_badtype = {'a': 'a string'}
+        with self.assertRaises(TypeError):
+            write_stan_json(file_bad, dict_badtype)
+
+        dict_badtype_nested = {'a': ['a string']}
+        with self.assertRaises(ValueError):
+            write_stan_json(file_bad, dict_badtype_nested)
+
+        dict_inf = {'a': [np.inf]}
+        with self.assertRaises(ValueError):
+            write_stan_json(file_bad, dict_inf)
+
+        dict_nan = {'a': np.nan}
+        with self.assertRaises(ValueError):
+            write_stan_json(file_bad, dict_nan)
+
 
 class ReadStanCsvTest(unittest.TestCase):
     def test_check_sampler_csv_1(self):
@@ -545,21 +581,6 @@ class RloadTest(unittest.TestCase):
         with self.assertRaises(ValueError):
             rload(dfile)
 
-    def test_roundtrip_metric(self):
-        dfile = os.path.join(DATAFILES_PATH, 'metric_diag.data.R')
-        data_dict_1 = rload(dfile)
-        self.assertEqual(data_dict_1['inv_metric'].shape, (3,))
-
-        dfile_tmp = os.path.join(DATAFILES_PATH, 'tmp.data.R')
-        rdump(dfile_tmp, data_dict_1)
-        data_dict_2 = rload(dfile_tmp)
-
-        self.assertTrue('inv_metric' in data_dict_2)
-        for i, x in enumerate(data_dict_2['inv_metric']):
-            self.assertEqual(x, data_dict_2['inv_metric'][i])
-
-        os.remove(dfile_tmp)
-
     def test_parse_rdump_value(self):
         struct1 = (
             'structure(c(1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16),.Dim=c(2,8))'