add functions to dicts to handle 2D dicts:

keys_of_2d_dict
fill_2d_dict
and fillgrid

and ordered() to return an ordered dict

in the ensemble.py module, ensemble_dist_to_2D_dict() converts an ensemble-style nested dataset

...->{xkeys}->{xvalue}->{ykeys}->{yvalue} = z

    to a dictionary like

    {xvalue}->{yvalue} = z

which is handy for other purposes.

binarycpython/utils/dicts.py

@@ -2,13 +2,15 @@
 Module containing functions that binary_c-python uses to modify
 dictionaries.
 """
+import astropy.units as u
 import copy
 import collections
-import astropy.units as u
+import math
 import numpy as np
 from collections import (
     OrderedDict,
 )
+import sys
 
 # we need to convert keys to floats:
 # this is ~ a factor 10 faster than David's
@@ -826,3 +828,135 @@ def normalize_dict(result_dict, verbosity=0):
     for key in result_dict.keys():
         result_dict[key] = result_dict[key] / sum_result
     return result_dict
+
+def fill_2d_dict(data,fill):
+    """
+    Fill a 2D dict with whatever fill is set to when keys are missing.
+    """
+    if fill is not None:
+        _ys = {}
+        for x in data.keys():
+            for y in data[x].keys():
+                _ys[y] = 1
+        xs = data.keys()
+        ys = _ys.keys()
+        for x in xs:
+            for y in ys:
+                if not y in data[x]:
+                    data[x][y] = fill
+
+def format_2d_dict_keys(data,format):
+    """
+    Reformat float keys in a 2d dict according to the given format
+    """
+    xs = list(data.keys())
+    for x in xs:
+        h = data[x]
+        if isinstance(x,float):
+            _y = data[x]
+            del data[x]
+            x = format.format(x)
+            h = data[x] = _y
+        ys = list(h.keys())
+        for y in ys:
+            if isinstance(y,float):
+                z = h[y]
+                del h[y]
+                y = format.format(y)
+                data[x][y] = z
+    return data
+
+def keys_of_2d_dict(data,numpy=False):
+    """
+    Given a 2D dict data[x][y] = z, return the list of keys as tuples, (x,y),
+    and equivalent z list.
+
+    Args:
+        data: the input 2D dict
+        numpy: if True convert to numpy arrays
+    """
+    xys = []
+    zs = []
+    for x in data.keys():
+        for y in data[x].keys():
+            xys.append((x,y))
+            zs.append(data[x][y])
+    if numpy:
+        xys = np.array(xys)
+        zs = np.array(zs)
+    return xys,zs
+
+def _mindiff(list,tol=1e-10):
+    """
+    Given a list of unique values, within tol(=1e-7), find the
+    minimum difference between any two elements.
+
+    Return None on error.
+    """
+    if len(list)<2:
+        return None
+
+    for index,x in enumerate(sorted(list)):
+        if index == 0:
+            mindiff = sys.float_info.max
+        else:
+            diff = list[index]-list[index-1]
+            if index > 1 and math.fabs(1.0-diff/mindiff) > tol:
+                mindiff = min(mindiff,math.fabs(mindiff - diff))
+            else:
+                mindiff = min(mindiff,diff)
+
+    if mindiff == sys.float_info.max:
+        return None
+    return mindiff
+
+def _find_nearest_index(array,value):
+    """
+    Find index of nearest item in array to value.
+    """
+    array = np.asarray(array)
+    idx = (np.abs(array-value)).argmin()
+    return idx
+
+
+def fillgrid(data,tol=1e-10,format=None):
+    """
+    Given a data[x][y] dict, fill it to the appropriate bin widths.
+    """
+
+    data = ordered(data)
+    xdata = list(data.keys())
+    ydata = list(data.values())
+    dx = _mindiff(xdata)
+
+    if not dx:
+        print("Warning: mindiff failed in fill() on {} data items. Returning original data.".format(len(xdata)))
+        return data
+
+    maxx = xdata[-1]
+    x = xdata[0]
+    newdata = {}
+    index = 0
+    while x<=maxx:
+        matchi = _find_nearest_index(xdata,x)
+        matchx = xdata[matchi]
+
+        if (x==0.0 and matchx==0.0) or abs(1.0-matchx/x)<=tol:
+            newdata[x] = ydata[index]
+            index += 1
+        else:
+            # missing data
+            newdata[x] = 0.0
+        x += dx
+        if format:
+            x = float(format.format(x))
+
+    return ordered(newdata)
+
+def ordered(dict):
+    """
+    Return an ordered version of dict
+    """
+    return collections.OrderedDict(
+        sorted(dict.items(),
+               key = lambda x:(float(x[0]))))

binarycpython/utils/ensemble.py

@@ -3,32 +3,31 @@ Module containing functions to interact with binary_c's
 population ensemble using the binarycpython package
 
 """
-
-import sys
+from binarycpython.utils.dicts import (
+    AutoVivificationDict,
+    custom_sort_dict,
+    fill_2d_dict,
+    keys_to_floats,
+    merge_dicts,
+    recursive_change_key_to_float,
+    recursive_change_key_to_string,
+)
+from binarycpython.utils.functions import verbose_print
 import bz2
-import time
-import json
-from typing import Any
+from collections.abc import Iterable, Mapping
 import gc
 import gzip
-import inspect
-
-
 from halo import Halo
+import inspect
+import json
 import msgpack
+# import orjson # not required any more?
 import py_rinterpolate
-
 import simplejson
+import sys
+import time
+from typing import Any
 
-# import orjson
-
-from binarycpython.utils.dicts import (
-    keys_to_floats,
-    recursive_change_key_to_float,
-    custom_sort_dict,
-    recursive_change_key_to_string,
-)
-from binarycpython.utils.functions import verbose_print
 
 def ensemble_setting(ensemble, parameter_name):
     """
@@ -370,3 +369,61 @@ def format_ensemble_results(ensemble_dictionary):
 
     # Put back in the dictionary
     return reformatted_ensemble_results
+
+def ensemble_dist_to_2D_dict(ensemble,xkeys,ykeys,fill=None,format=None):
+    """
+    Function to convert an ensemble dictionary like
+
+    ...->{xkeys}->{xvalue}->{ykeys}->{yvalue} = z
+
+    to a dictionary like
+
+    {xvalue}->{yvalue} = z
+
+    Note that xkeys and ykeys can be lists, in which case
+    we loop over several keys
+
+    Args:
+        ensemble : the base of the ensemble data
+        xkeys : scalar or list of x keys
+        ykeys : scalar or list of y keys
+        fill : if not None, fill out the data with whatever fill is set to
+
+    """
+    data = AutoVivificationDict()
+    if not isinstance(xkeys,list):
+        xkeys = [xkeys]
+    if not isinstance(ykeys,list):
+        ykeys = [ykeys]
+    h = ensemble
+    for xk in xkeys:
+        h = h[xk]
+    for x in h.keys():
+        h2 = h[x]
+        for yk in ykeys:
+            h2 = h2[yk]
+        for y in h2.keys():
+            data[x][y] += h2[y]
+    if fill is not None:
+        fill_2d_dict(data,fill)
+
+    return data
+
+def ensemble_flatten(ensemble):
+    """
+    Given a piece of an ensemble which is like this
+
+    {key1}->{more data1}-> ...
+    {key2}->{more data2}-> ...
+     ...
+
+    flatten the key list [key1, key2, ...] by merging all data that lies deeper.
+
+    Returns:
+
+    {more data1 + more data2 + ...}
+    """
+    data = {}
+    for k in ensemble:
+        data = merge_dicts(data,ensemble[k])
+    return data