From 3c1e57011af34eab411cba4f3dd66a1a710b978c Mon Sep 17 00:00:00 2001
From: Robert Izzard <r.izzard@surrey.ac.uk>
Date: Sat, 16 Oct 2021 14:22:11 +0100
Subject: [PATCH] add load_ensemble function
---
binarycpython/utils/functions.py | 58 +++++++++++++++++++-------------
1 file changed, 34 insertions(+), 24 deletions(-)
diff --git a/binarycpython/utils/functions.py b/binarycpython/utils/functions.py
index dbf90f113..60876ea0c 100644
--- a/binarycpython/utils/functions.py
+++ b/binarycpython/utils/functions.py
@@ -8,39 +8,36 @@ Tasks:
- TODO: change all prints to verbose_prints
"""
+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,
+)
+from colorama import Fore, Back, Style
+import copy
import datetime as dt
-import json
-import os
import gc
-import tempfile
-import copy
-import humanize
+import gzip
import inspect
+import h5py
+import humanize
+from io import StringIO
+import json
+import numpy as np
+import os
+import psutil
+import py_rinterpolate
import re
+import resource
import sys
import subprocess
+import tempfile
import time
import types
-import resource
-import psutil
-from colorama import Fore, Back, Style
-
-from io import StringIO
from typing import Union, Any
-from collections import (
- defaultdict,
- OrderedDict,
-)
-
-import h5py
-import numpy as np
-
-import astropy.units as u
-
-from binarycpython import _binary_c_bindings
-import binarycpython.utils.moe_di_stefano_2017_data as moe_di_stefano_2017_data
-
-import py_rinterpolate
########################################################
# Unsorted
@@ -2258,3 +2255,16 @@ class BinaryCEncoder(json.JSONEncoder):
# Let the base class default method raise the TypeError
return json.JSONEncoder.default(self, o)
+
+def load_ensemble(filename):
+ """
+ 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)
+ elif(filename.endswith('.gz')):
+ jfile = gzip.open(filename)
+ else:
+ jfile = open(filename)
+ return json.load(jfile)
--
GitLab