diff --git a/binarycpython/utils/functions.py b/binarycpython/utils/functions.py index f70db74f9d11cf4c88ee34487aa13aa9b002759b..8d7e271611dad01d4e8aab75e273017763618702 100644 --- a/binarycpython/utils/functions.py +++ b/binarycpython/utils/functions.py @@ -12,10 +12,7 @@ import astropy.units as u import binarycpython.utils.moe_di_stefano_2017_data as moe_di_stefano_2017_data from binarycpython import _binary_c_bindings import bz2 -from collections import ( - defaultdict, - OrderedDict, -) +import collections from colorama import Fore, Back, Style import copy import datetime as dt @@ -1149,7 +1146,7 @@ def example_parse_output(output: str, selected_header: str) -> dict: keys = value_dicts[0].keys() # Construct final dict. - final_values_dict = defaultdict(list) + final_values_dict = collections.defaultdict(list) for value_dict in value_dicts: for key in keys: final_values_dict[key].append(value_dict[key]) @@ -1724,7 +1721,7 @@ def inspect_dict( type(input_dict[key]) (except if the value is a dict) """ - structure_dict = OrderedDict() # TODO: check if this still works + structure_dict = collections.OrderedDict() # TODO: check if this still works # for key, value in input_dict.items(): @@ -1749,7 +1746,7 @@ def count_keys_recursive(input_dict): local_count = 0 for key in input_dict.keys(): local_count += 1 - if isinstance(input_dict[key], (dict, OrderedDict)): + if isinstance(input_dict[key], (dict, collections.OrderedDict)): local_count += count_keys_recursive(input_dict[key]) return local_count @@ -1779,7 +1776,7 @@ def merge_dicts(dict_1: dict, dict_2: dict) -> dict: """ # Set up new dict - new_dict = OrderedDict() # TODO: check if this still necessary + new_dict = collectsions.OrderedDict() # TODO: check if this still necessary # keys_1 = dict_1.keys() @@ -1824,9 +1821,9 @@ def merge_dicts(dict_1: dict, dict_2: dict) -> dict: # Exceptions: versions of dicts can be merged elif isinstance( - dict_1[key], (dict, OrderedDict, type(AutoVivificationDict)) + dict_1[key], (dict, collections.OrderedDict, type(AutoVivificationDict)) ) and isinstance( - dict_2[key], (dict, OrderedDict, type(AutoVivificationDict)) + dict_2[key], (dict, collections.OrderedDict, type(AutoVivificationDict)) ): new_dict[key] = merge_dicts(dict_1[key], dict_2[key]) @@ -1996,7 +1993,7 @@ def multiply_values_dict(input_dict, factor): for key in input_dict: if not key == "general_info": - if isinstance(input_dict[key], (dict, OrderedDict)): + if isinstance(input_dict[key], (dict, collections.OrderedDict)): input_dict[key] = multiply_values_dict(input_dict[key], factor) else: if isinstance(input_dict[key], (int, float)): @@ -2018,8 +2015,8 @@ def custom_sort_dict(input_dict): """ # If the new input is a dictionary, then try to sort it - if isinstance(input_dict, (dict, OrderedDict)): - new_dict = OrderedDict() + if isinstance(input_dict, (dict, collections.OrderedDict)): + new_dict = collections.OrderedDict() keys = input_dict.keys() @@ -2064,10 +2061,10 @@ def recursive_change_key_to_float(input_dict): Does not work with lists as values """ - new_dict = OrderedDict() # TODO: check if this still works + new_dict = collections.OrderedDict() # TODO: check if this still works for key in input_dict: - if isinstance(input_dict[key], (dict, OrderedDict)): + if isinstance(input_dict[key], (dict, collections.OrderedDict)): try: num_key = float(key) new_dict[num_key] = recursive_change_key_to_float( @@ -2092,10 +2089,10 @@ def recursive_change_key_to_string(input_dict): Function to recursively change the key back to a string but this time in a format that we decide """ - new_dict = OrderedDict() # TODO: check if this still works + new_dict = collections.OrderedDict() # TODO: check if this still works for key in input_dict: - if isinstance(input_dict[key], (dict, OrderedDict)): + if isinstance(input_dict[key], (dict, collections.OrderedDict)): if isinstance(key, (int, float)): string_key = "{:g}".format(key) new_dict[string_key] = recursive_change_key_to_string( @@ -2256,18 +2253,39 @@ class BinaryCEncoder(json.JSONEncoder): # Let the base class default method raise the TypeError return json.JSONEncoder.default(self, o) -def load_ensemble(filename): +def load_ensemble(filename,convert_float_keys=True): """ Function to load an ensemeble file, even if it is compressed, and return its contents to as a Python dictionary. """ if(filename.endswith('.bz2')): - jfile = bz2.open(filename) + jfile = bz2.open(filename,'rt') elif(filename.endswith('.gz')): - jfile = gzip.open(filename) + jfile = gzip.open(filename,'rt') else: - jfile = open(filename) - return json.load(jfile) + jfile = open(filename,'rt') + data = json.load(jfile) + if convert_float_keys == False: + return data + else: + # we need to convert keys to floats + def _to_float(json_data): + new_data = {} + for k,v in json_data.items(): + if isinstance(v, list): + v = [ _to_float(item) if isinstance(item, dict) else item for item in v ] + elif isinstance(v, collections.abc.Mapping): + # dict, ordereddict, etc. + v = _to_float(v) + try: + f = float(k) + new_data[f] = json_data[k] + except: + new_data[k] = v + return new_data + + data = _to_float(data) + return data def ensemble_setting(ensemble,parameter_name): """