Merge pull request #30 from s-kuberski/feature/json_io

covobs in json I/O

pyerrors/input/json.py

@@ -7,6 +7,7 @@ import datetime
 import platform
 import warnings
 from ..obs import Obs
+from ..covobs import Covobs
 from .. import version as pyerrorsversion
 
 
@@ -48,9 +49,25 @@ def create_json_string(ol, description='', indent=1):
         def __str__(self):
             return self.__repr__()
 
+    class Floatlist:
+        def __init__(self, li):
+            self.li = list(li)
+
+        def __repr__(self):
+            s = '['
+            for i in range(len(self.li)):
+                if i > 0:
+                    s += ', '
+                s += '%1.15e' % (self.li[i])
+            s += ']'
+            return s
+
+        def __str__(self):
+            return self.__repr__()
+
     def _gen_data_d_from_list(ol):
         dl = []
-        for name in ol[0].e_names:
+        for name in ol[0].mc_names:
             ed = {}
             ed['id'] = name
             ed['replica'] = []
@@ -69,6 +86,23 @@ def create_json_string(ol, description='', indent=1):
             dl.append(ed)
         return dl
 
+    def _gen_cdata_d_from_list(ol):
+        dl = []
+        for name in ol[0].cov_names:
+            ed = {}
+            ed['id'] = name
+            ed['layout'] = str(ol[0].covobs[name].cov.shape).lstrip('(').rstrip(')').rstrip(',')
+            ed['cov'] = Floatlist(np.ravel(ol[0].covobs[name].cov))
+            ncov = ol[0].covobs[name].cov.shape[0]
+            ed['grad'] = []
+            for i in range(ncov):
+                ed['grad'].append([])
+                for o in ol:
+                    ed['grad'][-1].append(o.covobs[name].grad[i][0])
+                ed['grad'][-1] = Floatlist(ed['grad'][-1])
+            dl.append(ed)
+        return dl
+
     def _assert_equal_properties(ol, otype=Obs):
         for o in ol:
             if not isinstance(o, otype):
@@ -92,7 +126,12 @@ def create_json_string(ol, description='', indent=1):
         if o.reweighted:
             d['reweighted'] = o.reweighted
         d['value'] = [o.value]
-        d['data'] = _gen_data_d_from_list([o])
+        data = _gen_data_d_from_list([o])
+        if len(data) > 0:
+            d['data'] = data
+        cdata = _gen_cdata_d_from_list([o])
+        if len(cdata) > 0:
+            d['cdata'] = cdata
         return d
 
     def write_List_to_dict(ol):
@@ -106,7 +145,12 @@ def create_json_string(ol, description='', indent=1):
         if ol[0].reweighted:
             d['reweighted'] = ol[0].reweighted
         d['value'] = [o.value for o in ol]
-        d['data'] = _gen_data_d_from_list(ol)
+        data = _gen_data_d_from_list(ol)
+        if len(data) > 0:
+            d['data'] = data
+        cdata = _gen_cdata_d_from_list(ol)
+        if len(cdata) > 0:
+            d['cdata'] = cdata
         return d
 
     def write_Array_to_dict(oa):
@@ -121,7 +165,12 @@ def create_json_string(ol, description='', indent=1):
         if ol[0].reweighted:
             d['reweighted'] = ol[0].reweighted
         d['value'] = [o.value for o in ol]
-        d['data'] = _gen_data_d_from_list(ol)
+        data = _gen_data_d_from_list(ol)
+        if len(data) > 0:
+            d['data'] = data
+        cdata = _gen_cdata_d_from_list(ol)
+        if len(cdata) > 0:
+            d['cdata'] = cdata
         return d
 
     if not isinstance(ol, list):
@@ -154,9 +203,9 @@ def create_json_string(ol, description='', indent=1):
         deltas = False
         split = s.split('\n')
         for i in range(len(split)):
-            if '"deltas":' in split[i]:
+            if '"deltas":' in split[i] or '"cov":' in split[i] or '"grad":' in split[i]:
                 deltas = True
-            elif deltas:
+            if deltas:
                 split[i] = split[i].replace('"[', '[').replace(']"', ']')
                 if split[i][-1] == ']':
                     deltas = False
@@ -222,16 +271,32 @@ def import_json_string(json_string, verbose=True, full_output=False):
 
     def _gen_obsd_from_datad(d):
         retd = {}
-        retd['names'] = []
+        if d:
-        retd['idl'] = []
+            retd['names'] = []
-        retd['deltas'] = []
+            retd['idl'] = []
-        retd['is_merged'] = {}
+            retd['deltas'] = []
+            retd['is_merged'] = {}
+            for ens in d:
+                for rep in ens['replica']:
+                    retd['names'].append(rep['name'])
+                    retd['idl'].append([di[0] for di in rep['deltas']])
+                    retd['deltas'].append(np.array([di[1:] for di in rep['deltas']]))
+                    retd['is_merged'][rep['name']] = rep.get('is_merged', False)
+        return retd
+
+    def _gen_covobsd_from_cdatad(d):
+        retd = {}
         for ens in d:
-            for rep in ens['replica']:
+            retl = []
-                retd['names'].append(rep['name'])
+            name = ens['id']
-                retd['idl'].append([di[0] for di in rep['deltas']])
+            layouts = ens.get('layout', '1').strip()
-                retd['deltas'].append(np.array([di[1:] for di in rep['deltas']]))
+            layout = [int(ls.strip()) for ls in layouts.split(',') if len(ls) > 0]
-                retd['is_merged'][rep['name']] = rep.get('is_merged', False)
+            cov = np.reshape(ens['cov'], layout)
+            grad = ens['grad']
+            nobs = len(grad[0])
+            for i in range(nobs):
+                retl.append({'name': name, 'cov': cov, 'grad': [g[i] for g in grad]})
+            retd[name] = retl
         return retd
 
     def get_Obs_from_dict(o):
@@ -240,11 +305,21 @@ def import_json_string(json_string, verbose=True, full_output=False):
             raise Exception("layout is %s has to be 1 for type Obs." % (layouts), RuntimeWarning)
 
         values = o['value']
-        od = _gen_obsd_from_datad(o['data'])
+        od = _gen_obsd_from_datad(o.get('data', {}))
+        cd = _gen_covobsd_from_cdatad(o.get('cdata', {}))
+
+        if od:
+            ret = Obs([[ddi[0] + values[0] for ddi in di] for di in od['deltas']], od['names'], idl=od['idl'])
+            ret.is_merged = od['is_merged']
+        else:
+            ret = Obs([], [])
+            ret._value = values[0]
+        for name in cd:
+            co = cd[name][0]
+            ret._covobs[name] = Covobs(None, co['cov'], co['name'], grad=co['grad'])
+            ret.names.append(co['name'])
 
-        ret = Obs([[ddi[0] + values[0] for ddi in di] for di in od['deltas']], od['names'], idl=od['idl'])
         ret.reweighted = o.get('reweighted', False)
-        ret.is_merged = od['is_merged']
         ret.tag = o.get('tag', [None])[0]
         return ret
 
@@ -252,14 +327,25 @@ def import_json_string(json_string, verbose=True, full_output=False):
         layouts = o.get('layout', '1').strip()
         layout = int(layouts)
         values = o['value']
-        od = _gen_obsd_from_datad(o['data'])
+        od = _gen_obsd_from_datad(o.get('data', {}))
+        cd = _gen_covobsd_from_cdatad(o.get('cdata', {}))
 
         ret = []
         taglist = o.get('tag', layout * [None])
         for i in range(layout):
-            ret.append(Obs([list(di[:, i] + values[i]) for di in od['deltas']], od['names'], idl=od['idl']))
+            if od:
+                ret.append(Obs([list(di[:, i] + values[i]) for di in od['deltas']], od['names'], idl=od['idl']))
+                ret[-1].is_merged = od['is_merged']
+            else:
+                ret.append(Obs([], []))
+                ret[-1]._value = values[i]
+                print('Created Obs with means= ', values[i])
+            for name in cd:
+                co = cd[name][i]
+                ret[-1]._covobs[name] = Covobs(None, co['cov'], co['name'], grad=co['grad'])
+                ret[-1].names.append(co['name'])
+
             ret[-1].reweighted = o.get('reweighted', False)
-            ret[-1].is_merged = od['is_merged']
             ret[-1].tag = taglist[i]
         return ret
 
@@ -268,14 +354,23 @@ def import_json_string(json_string, verbose=True, full_output=False):
         layout = [int(ls.strip()) for ls in layouts.split(',') if len(ls) > 0]
         N = np.prod(layout)
         values = o['value']
-        od = _gen_obsd_from_datad(o['data'])
+        od = _gen_obsd_from_datad(o.get('data', {}))
+        cd = _gen_covobsd_from_cdatad(o.get('cdata', {}))
 
         ret = []
         taglist = o.get('tag', N * [None])
         for i in range(N):
-            ret.append(Obs([di[:, i] + values[i] for di in od['deltas']], od['names'], idl=od['idl']))
+            if od:
+                ret.append(Obs([di[:, i] + values[i] for di in od['deltas']], od['names'], idl=od['idl']))
+                ret[-1].is_merged = od['is_merged']
+            else:
+                ret.append(Obs([], []))
+                ret[-1]._value = values[i]
+            for name in cd:
+                co = cd[name][i]
+                ret[-1]._covobs[name] = Covobs(None, co['cov'], co['name'], grad=co['grad'])
+                ret[-1].names.append(co['name'])
             ret[-1].reweighted = o.get('reweighted', False)
-            ret[-1].is_merged = od['is_merged']
             ret[-1].tag = taglist[i]
         return np.reshape(ret, layout)
 

tests/io_test.py

@@ -13,10 +13,16 @@ def test_jsonio():
     otag = 'This has been merged!'
     o4.tag = otag
     do = o - .2 * o4
+    co1 = pe.cov_Obs(1., .123, 'cov1')
+    co3 = pe.cov_Obs(4., .1 ** 2, 'cov3')
+    do *= co1 / co3
     do.tag = {'A': 2}
 
     o5 = pe.pseudo_Obs(0.8, .1, 'two|r2')
-    o5.tag = 2*otag
+    co2 = pe.cov_Obs([1, 2], [[.12, .004], [.004, .02]], 'cov2')
+    o5 /= co2[0]
+    o3 /= co2[1]
+    o5.tag = 2 * otag
     testl = [o3, o5]
 
     arr = np.array([o3, o5])
@@ -33,7 +39,7 @@ def test_jsonio():
 
     tt.tag = 'Test Obs: Ä'
 
-    ol = [o4, do, testl, mat, arr, np.array([o]), np.array([tt, tt]), [tt, tt]]
+    ol = [o4, do, testl, mat, arr, np.array([o]), np.array([tt, tt]), [tt, tt], co1, co2, np.array(co2), co1 / co2[0]]
     fname = 'test_rw'
 
     jsonio.dump_to_json(ol, fname, indent=1, description='[I am a tricky description]')
@@ -45,7 +51,7 @@ def test_jsonio():
     for o, r in zip(ol, rl):
         assert np.all(o == r)
 
-    for i in range(len(rl)):
+    for i in range(len(ol)):
         if isinstance(ol[i], pe.Obs):
             o = ol[i] - rl[i]
             assert(o.is_zero())