Serialize complex numbers to json

Author: WardBrian

cmdstanpy/utils.py

@@ -3,6 +3,7 @@
 """
 import contextlib
 import functools
+import json
 import logging
 import math
 import os
@@ -30,7 +31,7 @@
 
 import numpy as np
 import pandas as pd
-import ujson as json
+import ujson
 from tqdm.auto import tqdm
 
 from cmdstanpy import (
@@ -439,6 +440,13 @@ def rewrite_inf_nan(
         return data
 
 
+def serialize_complex(c: Any) -> List[float]:
+    if isinstance(c, complex):
+        return [c.real, c.imag]
+    else:
+        raise TypeError(f"Unserializable type: {type(c)}")
+
+
 def write_stan_json(path: str, data: Mapping[str, Any]) -> None:
     """
     Dump a mapping of strings to data to a JSON file.
@@ -494,7 +502,11 @@ def write_stan_json(path: str, data: Mapping[str, Any]) -> None:
             data_out[key] = rewrite_inf_nan(data_out[key])
 
     with open(path, 'w') as fd:
-        json.dump(data_out, fd)
+        try:
+            ujson.dump(data_out, fd)
+        except TypeError as e:
+            get_logger().debug(e)
+            json.dump(data_out, fd, default=serialize_complex)
 
 
 def rload(fname: str) -> Optional[Dict[str, Union[int, float, np.ndarray]]]:
@@ -948,7 +960,7 @@ def read_metric(path: str) -> List[int]:
     """
     if path.endswith('.json'):
         with open(path, 'r') as fd:
-            metric_dict = json.load(fd)
+            metric_dict = ujson.load(fd)
         if 'inv_metric' in metric_dict:
             dims_np: np.ndarray = np.asarray(metric_dict['inv_metric'])
             return list(dims_np.shape)

test/test_utils.py

@@ -38,14 +38,14 @@
     parse_method_vars,
     parse_rdump_value,
     parse_stan_vars,
+    pushd,
     read_metric,
     rload,
     set_cmdstan_path,
     validate_cmdstan_path,
     validate_dir,
     windows_short_path,
     write_stan_json,
-    pushd,
 )
 
 HERE = os.path.dirname(os.path.abspath(__file__))
@@ -347,6 +347,13 @@ def cmp(d1, d2):
         with open(file_fin) as fd:
             cmp(json.load(fd), dict_inf_nan_exp)
 
+        dict_complex = {'a': np.array([np.complex64(3), 3 + 4j])}
+        dict_complex_exp = {'a': [[3, 0], [3, 4]]}
+        file_complex = os.path.join(_TMPDIR, 'complex.json')
+        write_stan_json(file_complex, dict_complex)
+        with open(file_complex) as fd:
+            cmp(json.load(fd), dict_complex_exp)
+
     def test_write_stan_json_bad(self):
         file_bad = os.path.join(_TMPDIR, 'bad.json')
 
@@ -809,7 +816,6 @@ def test_exit(self):
 
 
 class PushdTest(unittest.TestCase):
-
     def test_restore_cwd(self):
         "Ensure do_command in a different cwd restores cwd after error."
         cwd = os.getcwd()