diff --git a/binarycpython/utils/dicts.py b/binarycpython/utils/dicts.py index dcdd189137ed0610a5dad0b28a5f23713f9c805e..c474b3dce0e8d94343bbd4a3841460b4724503d8 100644 --- a/binarycpython/utils/dicts.py +++ b/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])))) diff --git a/binarycpython/utils/ensemble.py b/binarycpython/utils/ensemble.py index 2a5d9db42ac7415fb64d686aff1f9d45c3afd716..d7f4fbcadb1f412c0f1a3bb8373865e69a8b83a7 100644 --- a/binarycpython/utils/ensemble.py +++ b/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