Handle inf/nan in json

Author: WardBrian

cmdstanpy/utils.py

@@ -409,6 +409,21 @@ def cxx_toolchain_path(
     return compiler_path, tool_path
 
 
+def rewrite_inf_nan(
+    data: Union[float, int, List[Any]]
+) -> Union[str, int, float, List[Any]]:
+    if isinstance(data, float):
+        if math.isnan(data):
+            return 'NaN'
+        if math.isinf(data):
+            return ('+' if data > 0 else '-') + 'inf'
+        return data
+    elif isinstance(data, list):
+        return [rewrite_inf_nan(item) for item in data]
+    else:
+        return data
+
+
 def write_stan_json(path: str, data: Mapping[str, Any]) -> None:
     """
     Dump a mapping of strings to data to a JSON file.
@@ -427,9 +442,13 @@ def write_stan_json(path: str, data: Mapping[str, Any]) -> None:
     :param data: A mapping from strings to values. This can be a dictionary
         or something more exotic like an :class:`xarray.Dataset`. This will be
         copied before type conversion, not modified
+
+    :param handle_nan_inf: If enabled, perform the (Slow!) checks necessary to
+        output NaN and inf as required for Stan
     """
     data_out = {}
     for key, val in data.items():
+        handle_nan_inf = False
         if val is not None:
             if isinstance(val, (str, bytes)) or (
                 type(val).__module__ != 'numpy'
@@ -440,18 +459,14 @@ def write_stan_json(path: str, data: Mapping[str, Any]) -> None:
                     + 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}'"
-                    )
+                handle_nan_inf = not np.all(np.isfinite(val))
             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)}"
+                    f"write_stan_json for key '{key}' "
+                    f"as part of collection {type(val)}"
                 )
 
         if type(val).__module__ == 'numpy':
@@ -463,6 +478,9 @@ def write_stan_json(path: str, data: Mapping[str, Any]) -> None:
         else:
             data_out[key] = val
 
+        if handle_nan_inf:
+            data_out[key] = rewrite_inf_nan(data_out[key])
+
     with open(path, 'w') as fd:
         json.dump(data_out, fd)
 

test/test_utils.py

@@ -338,6 +338,17 @@ def cmp(d1, d2):
         with open(file_scalr) as fd:
             cmp(json.load(fd), dict_scalr)
 
+        # custom Stan serialization
+
+        dict_inf_nan = {
+            'a': np.array([[-np.inf, np.inf, np.NaN, float('NaN')]])
+        }
+        dict_inf_nan_exp = {'a': np.array([["-inf", "+inf", "NaN", "NaN"]])}
+        file_fin = os.path.join(_TMPDIR, 'inf.json')
+        write_stan_json(file_fin, dict_inf_nan)
+        with open(file_fin) as fd:
+            cmp(json.load(fd), dict_inf_nan_exp)
+
     def test_write_stan_json_bad(self):
         file_bad = os.path.join(_TMPDIR, 'bad.json')
 
@@ -349,14 +360,6 @@ def test_write_stan_json_bad(self):
         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):