Merge pull request #535 from stan-dev/serialize-complex

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
@@ -31,7 +32,7 @@
 
 import numpy as np
 import pandas as pd
-import ujson as json
+import ujson
 from tqdm.auto import tqdm
 
 from cmdstanpy import (
@@ -450,6 +451,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.
@@ -505,7 +513,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]]]:
@@ -977,7 +989,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

@@ -348,6 +348,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')