diff --git a/badges/test_coverage.svg b/badges/test_coverage.svg
index 179c6a1788ecba3678a21ca99b50643f3ec81e57..6d68f472b0f880260c3f22f1fc057e768f16f432 100644
--- a/badges/test_coverage.svg
+++ b/badges/test_coverage.svg
@@ -15,7 +15,7 @@
<g fill="#fff" text-anchor="middle" font-family="DejaVu Sans,Verdana,Geneva,sans-serif" font-size="11">
<text x="31.5" y="15" fill="#010101" fill-opacity=".3">coverage</text>
<text x="31.5" y="14">coverage</text>
- <text x="80" y="15" fill="#010101" fill-opacity=".3">63%</text>
- <text x="80" y="14">63%</text>
+ <text x="80" y="15" fill="#010101" fill-opacity=".3">64%</text>
+ <text x="80" y="14">64%</text>
</g>
</svg>
diff --git a/binarycpython/tests/main.py b/binarycpython/tests/main.py
index e59ecd3b863d371ec163d4d230574b899fe293ce..da398439bc840f2e30c0f56d76022b2632a54545 100755
--- a/binarycpython/tests/main.py
+++ b/binarycpython/tests/main.py
@@ -43,6 +43,12 @@ from binarycpython.tests.test_dicts import (
from binarycpython.tests.test_ensemble import (
test_binaryc_json_serializer,
test_handle_ensemble_string_to_json,
+ test_BinarycEncoder,
+ test_BinarycDecoder,
+ test_extract_ensemble_json_from_string,
+ test_load_ensemble,
+ test_ensemble_file_type,
+ test_open_ensemble
)
from binarycpython.tests.test_functions import (
test_verbose_print,
diff --git a/binarycpython/tests/test_c_bindings.py b/binarycpython/tests/test_c_bindings.py
index 69200067075168f77e26f6e6654a9472bd4068ee..65ddd3252142ea891717704ff1ad8d416b91b247 100644
--- a/binarycpython/tests/test_c_bindings.py
+++ b/binarycpython/tests/test_c_bindings.py
@@ -19,7 +19,7 @@ from binarycpython.utils.functions import (
Capturing,
)
from binarycpython.utils.ensemble import (
- binarycDecoder,
+ BinarycDecoder,
handle_ensemble_string_to_json,
extract_ensemble_json_from_string,
)
diff --git a/binarycpython/tests/test_ensemble.py b/binarycpython/tests/test_ensemble.py
index ec19b99cc70c812bfab069549f57424057d1486e..32128c72ff3fb364763c7ace470c2e3d54ee009b 100644
--- a/binarycpython/tests/test_ensemble.py
+++ b/binarycpython/tests/test_ensemble.py
@@ -3,19 +3,32 @@ Unit tests for the ensemble module
"""
import os
+import sys
+import json
import unittest
+import io
+from io import StringIO
+
from binarycpython.utils.functions import (
temp_dir,
Capturing,
)
from binarycpython.utils.ensemble import (
binaryc_json_serializer,
- handle_ensemble_string_to_json
+ handle_ensemble_string_to_json,
+ BinarycEncoder,
+ BinarycDecoder,
+ extract_ensemble_json_from_string,
+ load_ensemble,
+ ensemble_file_type,
+ open_ensemble
+
)
+from binarycpython.utils.grid import Population
TMP_DIR = temp_dir("tests", "test_ensemble")
-
+TEST_VERBOSITY = 1
class test_binaryc_json_serializer(unittest.TestCase):
"""
@@ -71,5 +84,381 @@ class test_handle_ensemble_string_to_json(unittest.TestCase):
self.assertTrue(output_dict["ding"] == 10)
self.assertTrue(output_dict["list_example"] == [1, 2, 3])
+
+class test_BinarycEncoder(unittest.TestCase):
+ """
+ Unittests for class BinarycEncoder
+ """
+
+ def test_1(self):
+ with Capturing() as _:
+ self._test_1()
+
+ def _test_1(self):
+ """
+ Test that the object is converted to strings
+ """
+
+ input_1 = {'a': BinarycEncoder}
+ output_1 = json.dumps(input_1, cls=BinarycEncoder)
+ self.assertTrue(isinstance(output_1, str))
+
+ dict_output_1 = json.loads(output_1)
+ self.assertTrue(isinstance(dict_output_1['a'], str))
+
+
+class test_BinarycDecoder(unittest.TestCase):
+ """
+ Unittests for class BinarycDecoder
+ """
+
+ def test_1(self):
+ with Capturing() as _:
+ self._test_1()
+
+ def _test_1(self):
+ """
+ Test that the object is converted to floats
+ """
+
+
+ input_1 = '{"a": "10.0"}'
+ output_1 = json.loads(input_1)
+ output_2 = json.loads(input_1, cls=BinarycDecoder)
+
+ self.assertTrue(isinstance(output_1["a"], str))
+ self.assertTrue(isinstance(output_2["a"], float))
+
+
+class test_extract_ensemble_json_from_string(unittest.TestCase):
+ """
+ Unittests for class extract_ensemble_json_from_string
+ """
+
+ def test_1(self):
+ with Capturing() as _:
+ self._test_1()
+
+ def _test_1(self):
+ """
+ Simple test without errors
+ """
+
+ input_1 = "ENSEMBLE_JSON {\"a\": 10}"
+ output_1 = extract_ensemble_json_from_string(input_1)
+
+ self.assertTrue(isinstance(output_1, dict))
+ self.assertEqual(output_1, {'a': 10})
+
+ def test_2(self):
+ with Capturing() as _:
+ self._test_1()
+
+ def _test_2(self):
+ """
+ Simple test with 2 lines
+ """
+
+ input_1 = "ENSEMBLE_JSON {\"a\": 10}\nENSEMBLE_JSON {\"b\": 20}"
+
+ capturedOutput = StringIO() # Create StringIO object
+ sys.stdout = capturedOutput # and redirect stdout.
+ _ = extract_ensemble_json_from_string(input_1)
+ sys.stdout = sys.__stdout__ # Reset redirect.
+
+ self.assertTrue(capturedOutput.getvalue().startswith("Warning:"))
+
+ def test_3(self):
+ with Capturing() as _:
+ self._test_1()
+
+ def _test_3(self):
+ """
+ Simple test with empty input
+ """
+
+ input_1 = ""
+ output_1 = extract_ensemble_json_from_string(input_1)
+
+ self.assertTrue(isinstance(output_1, dict))
+ self.assertEqual(output_1, {})
+
+ def test_3(self):
+ with Capturing() as _:
+ self._test_1()
+
+ def test_3(self):
+ """
+ Simple test with missing starting string
+ """
+
+ input_1 = " {\"a\": 10}"
+
+ #
+ capturedOutput = StringIO() # Create StringIO object
+ sys.stdout = capturedOutput # and redirect stdout.
+ _ = extract_ensemble_json_from_string(input_1)
+ sys.stdout = sys.__stdout__ # Reset redirect.
+
+ self.assertTrue(capturedOutput.getvalue().startswith("Error:"))
+
+
+class test_load_ensemble(unittest.TestCase):
+ """
+ Unittests for class extract_ensemble_json_from_string
+ """
+
+ def __init__(self, *args, **kwargs):
+ """
+ init function
+ """
+ super(test_load_ensemble, self).__init__(*args, **kwargs)
+
+ # First
+ test_pop_1 = Population()
+ test_pop_1.set(
+ num_cores=2,
+ verbosity=TEST_VERBOSITY,
+ M_2=1,
+ orbital_period=100000,
+ ensemble=1,
+ ensemble_defer=1,
+ ensemble_filters_off=1,
+ ensemble_filter_STELLAR_TYPE_COUNTS=1,
+ ensemble_dt=1000,
+ )
+ test_pop_1.set(
+ data_dir=TMP_DIR,
+ combine_ensemble_with_thread_joining=True,
+ )
+
+ resolution = {"M_1": 10}
+
+ test_pop_1.add_grid_variable(
+ name="lnm1",
+ longname="Primary mass",
+ valuerange=[1, 100],
+ samplerfunc="self.const_linear(math.log(1), math.log(100), {})".format(
+ resolution["M_1"]
+ ),
+ precode="M_1=math.exp(lnm1)",
+ probdist="self.three_part_powerlaw(M_1, 0.1, 0.5, 1.0, 100, -1.3, -2.3, -2.3)*M_1",
+ dphasevol="dlnm1",
+ parameter_name="M_1",
+ condition="", # Impose a condition on this grid variable. Mostly for a check for yourself
+ )
+
+ _ = test_pop_1.evolve()
+ ensemble_output_1 = test_pop_1.grid_ensemble_results
+
+ self.normal_ensemble_output_name = os.path.join(TMP_DIR, 'test_load_ensemble_ensemble_output.json')
+ self.bzip2_ensemble_output_name = os.path.join(TMP_DIR, 'test_load_ensemble_ensemble_output.json.bz2')
+ self.no_extension_ensemble_output_name = os.path.join(TMP_DIR, 'test_load_ensemble_ensemble_output')
+
+ # Write ensemble to json with normal write
+ test_pop_1.write_ensemble(self.normal_ensemble_output_name)
+
+ # Write ensemble to bzip
+ test_pop_1.write_ensemble(self.bzip2_ensemble_output_name)
+
+ # Write ensemble without extension
+ with open(self.no_extension_ensemble_output_name, 'w') as f:
+ f.write(json.dumps(ensemble_output_1))
+
+ def test_1(self):
+ with Capturing() as _:
+ self._test_1()
+
+ def _test_1(self):
+ """
+ Simple test to load ensemble with normal filetype
+ """
+
+ # load data
+ loaded_data_1 = load_ensemble(self.normal_ensemble_output_name)
+
+ self.assertTrue(isinstance(loaded_data_1, dict))
+
+ def test_2(self):
+ with Capturing() as _:
+ self._test_2()
+
+ def _test_2(self):
+ """
+ Simple test to load ensemble with msgpack type
+ """
+
+ # load data
+ loaded_data_1 = load_ensemble(self.bzip2_ensemble_output_name)
+
+ self.assertTrue(isinstance(loaded_data_1, dict))
+
+
+ def test_3(self):
+ with Capturing() as _:
+ self._test_3()
+
+ def _test_3(self):
+ """
+ Simple test to load ensemble with timing output
+ """
+
+
+ #
+ capturedOutput = StringIO() # Create StringIO object
+ sys.stdout = capturedOutput # and redirect stdout.
+ _ = load_ensemble(self.normal_ensemble_output_name, timing=True)
+ sys.stdout = sys.__stdout__ # Reset redirect.
+
+ self.assertTrue("Took" in capturedOutput.getvalue())
+
+ def test_4(self):
+ with Capturing() as _:
+ self._test_4()
+
+ def _test_4(self):
+ """
+ Simple test to load ensemble conveting to floats and timing
+ """
+
+ #
+ capturedOutput = StringIO() # Create StringIO object
+ sys.stdout = capturedOutput # and redirect stdout.
+ loaded_data_1 = load_ensemble(self.normal_ensemble_output_name, timing=True, convert_float_keys=True)
+ sys.stdout = sys.__stdout__ # Reset redirect.
+
+ self.assertTrue("Took" in capturedOutput.getvalue())
+
+
+class test_ensemble_file_type(unittest.TestCase):
+ """
+ Unittests for class ensemble_file_type
+ """
+
+ def test_1(self):
+ with Capturing() as _:
+ self._test_1()
+
+ def _test_1(self):
+ """
+ filetype tests
+ """
+
+ # Json test
+ input_1 = "/tmp/test.json"
+ output_1 = ensemble_file_type(input_1)
+
+ self.assertEqual(output_1, "JSON")
+
+ # Msgpack
+ input_2 = "/tmp/test.msgpack"
+ output_2 = ensemble_file_type(input_2)
+
+ self.assertEqual(output_2, "msgpack")
+
+ # None
+ input_3 = "/tmp/test"
+ output_3 = ensemble_file_type(input_3)
+
+ self.assertIsNone(output_3)
+
+
+class test_open_ensemble(unittest.TestCase):
+ """
+ Unittests for class open_ensemble:
+ """
+
+ def __init__(self, *args, **kwargs):
+ """
+ init function
+ """
+ super(test_open_ensemble, self).__init__(*args, **kwargs)
+
+ # First
+ test_pop_1 = Population()
+ test_pop_1.set(
+ num_cores=2,
+ verbosity=TEST_VERBOSITY,
+ M_2=1,
+ orbital_period=100000,
+ ensemble=1,
+ ensemble_defer=1,
+ ensemble_filters_off=1,
+ ensemble_filter_STELLAR_TYPE_COUNTS=1,
+ ensemble_dt=1000,
+ )
+ test_pop_1.set(
+ data_dir=TMP_DIR,
+ combine_ensemble_with_thread_joining=True,
+ )
+
+ resolution = {"M_1": 10}
+
+ test_pop_1.add_grid_variable(
+ name="lnm1",
+ longname="Primary mass",
+ valuerange=[1, 100],
+ samplerfunc="self.const_linear(math.log(1), math.log(100), {})".format(
+ resolution["M_1"]
+ ),
+ precode="M_1=math.exp(lnm1)",
+ probdist="self.three_part_powerlaw(M_1, 0.1, 0.5, 1.0, 100, -1.3, -2.3, -2.3)*M_1",
+ dphasevol="dlnm1",
+ parameter_name="M_1",
+ condition="", # Impose a condition on this grid variable. Mostly for a check for yourself
+ )
+
+ _ = test_pop_1.evolve()
+ ensemble_output_1 = test_pop_1.grid_ensemble_results
+
+ self.normal_ensemble_output_name = os.path.join(TMP_DIR, 'test_open_ensemble_ensemble_output.json')
+ self.bzip2_ensemble_output_name = os.path.join(TMP_DIR, 'test_open_ensemble_ensemble_output.json.bz2')
+ self.gzip_ensemble_output_name = os.path.join(TMP_DIR, 'test_open_ensemble_ensemble_output.json.gz')
+ # self.msgpack_ensemble_output_name = os.path.join(TMP_DIR, 'test_open_ensemble_ensemble_output.msgpack')
+
+ # Write ensemble to json with normal write
+ test_pop_1.write_ensemble(self.normal_ensemble_output_name)
+
+ # Write ensemble to bzip
+ test_pop_1.write_ensemble(self.bzip2_ensemble_output_name)
+
+ # gzip
+ test_pop_1.write_ensemble(self.gzip_ensemble_output_name)
+
+ # # Msgpack
+ # test_pop_1.write_ensemble(self.msgpack_ensemble_output_name)
+
+ def test_1(self):
+ with Capturing() as _:
+ self._test_1()
+
+ def _test_1(self):
+ """
+ filetype tests
+ """
+
+ self.msgpack_ensemble_output_name = os.path.join(TMP_DIR, 'test_open_ensemble_ensemble_output.msgpack.gz')
+
+ #
+ handle_1 = open_ensemble(self.normal_ensemble_output_name)
+ self.assertTrue(isinstance(handle_1, io._io.TextIOWrapper))
+ handle_1.close()
+
+ #
+ handle_2 = open_ensemble(self.bzip2_ensemble_output_name)
+ self.assertTrue(isinstance(handle_2, io._io.TextIOWrapper))
+ handle_2.close()
+
+ #
+ handle_3 = open_ensemble(self.gzip_ensemble_output_name)
+ self.assertTrue(isinstance(handle_3, io._io.TextIOWrapper))
+ handle_3.close()
+
+ # # TODO: implement this again
+ # handle_4 = open_ensemble(self.msgpack_ensemble_output_name)
+ # self.assertTrue(isinstance(handle_4, io._io.TextIOWrapper))
+ # handle_4.close()
+
+
if __name__ == "__main__":
unittest.main()
diff --git a/binarycpython/tests/test_grid.py b/binarycpython/tests/test_grid.py
index 1e146ec34fad370a69ef9ed6135e604756940e5a..37e303d112595bb035667a8c05b4153a80c9a8b8 100644
--- a/binarycpython/tests/test_grid.py
+++ b/binarycpython/tests/test_grid.py
@@ -6,7 +6,6 @@ TODO: exit
TODO: _set_nprocesses
TODO: _pre_run_setup
TODO: clean
-TODO: evolve
TODO: _evolve_population
TODO: _system_queue_filler
TODO: _evolve_population_grid
@@ -660,9 +659,6 @@ class test_resultdict(unittest.TestCase):
test_case_dict, dict(example_pop.grid_results["example"]["mass"])
)
-
-
-
class test_grid_evolve(unittest.TestCase):
"""
Unittests for function Population.evolve()
@@ -955,7 +951,6 @@ Printf("TEST_CUSTOM_LOGGING_1 %30.12e %g %g %g %g\\n",
)
test_pop.set(
data_dir=TMP_DIR,
- ensemble_output_name="ensemble_output.json",
combine_ensemble_with_thread_joining=False,
)
@@ -1027,7 +1022,6 @@ Printf("TEST_CUSTOM_LOGGING_1 %30.12e %g %g %g %g\\n",
test_pop.set(
data_dir=TMP_DIR,
combine_ensemble_with_thread_joining=True,
- ensemble_output_name="ensemble_output.json",
)
resolution = {"M_1": 10}
@@ -1084,7 +1078,6 @@ Printf("TEST_CUSTOM_LOGGING_1 %30.12e %g %g %g %g\\n",
test_pop_1.set(
data_dir=TMP_DIR,
combine_ensemble_with_thread_joining=True,
- ensemble_output_name="ensemble_output.json",
)
resolution = {"M_1": 10}
@@ -1121,7 +1114,6 @@ Printf("TEST_CUSTOM_LOGGING_1 %30.12e %g %g %g %g\\n",
)
test_pop_2.set(
data_dir=TMP_DIR,
- ensemble_output_name="ensemble_output.json",
combine_ensemble_with_thread_joining=False,
)
diff --git a/binarycpython/utils/ensemble.py b/binarycpython/utils/ensemble.py
index 7573fc6b58ba1d1c52d3a78df62d4ad4b204e38c..1808ed7c31595b021a4c43fb5d1b377af3d26674 100644
--- a/binarycpython/utils/ensemble.py
+++ b/binarycpython/utils/ensemble.py
@@ -197,11 +197,11 @@ def load_ensemble(
# data = recursive_change_key_to_float(data) # 61s
data = keys_to_floats(data) # 6.94s
- if timing:
- print(
- "\n\nTook {} s to convert floats\n\n".format(time.time() - tstart),
- flush=True,
- )
+ if timing:
+ print(
+ "\n\nTook {} s to convert floats\n\n".format(time.time() - tstart),
+ flush=True,
+ )
# return data
return data
@@ -264,10 +264,10 @@ def handle_ensemble_string_to_json(raw_output):
raw_output: raw output of the ensemble dump by binary_c
Returns:
- json.loads(raw_output, cls=binarycDecoder)
+ json.loads(raw_output, cls=BinarycDecoder)
"""
- return json.loads(raw_output, cls=binarycDecoder)
+ return json.loads(raw_output, cls=BinarycDecoder)
def binaryc_json_serializer(obj: Any) -> Any:
@@ -289,7 +289,7 @@ def binaryc_json_serializer(obj: Any) -> Any:
return obj
-class binarycDecoder(json.JSONDecoder):
+class BinarycDecoder(json.JSONDecoder):
"""
Custom decoder to transform the numbers that are strings to actual floats
"""
@@ -327,7 +327,7 @@ class binarycDecoder(json.JSONDecoder):
return o
-class BinaryCEncoder(json.JSONEncoder):
+class BinarycEncoder(json.JSONEncoder):
"""
Encoding class function to attempt to convert things to strings.
"""
@@ -336,6 +336,7 @@ class BinaryCEncoder(json.JSONEncoder):
"""
Converting function. Well, could be more precise. look at the JSON module
"""
+
try:
str_repr = str(o)
except TypeError:
diff --git a/binarycpython/utils/population_extensions/dataIO.py b/binarycpython/utils/population_extensions/dataIO.py
index bd1da06bc59924dfe5d4ad2907d267eaffd6ce47..0723da759559885ee7c17fe68390e28f0eae2090 100644
--- a/binarycpython/utils/population_extensions/dataIO.py
+++ b/binarycpython/utils/population_extensions/dataIO.py
@@ -20,6 +20,7 @@ import compress_pickle
from binarycpython.utils.ensemble import (
ensemble_file_type,
+ ensemble_compression
)
from binarycpython.utils.dicts import (
merge_dicts,
@@ -358,8 +359,8 @@ class dataIO:
)
self.exit(code=1)
else:
- f = self.open(output_file, "wt", encoding=encoding)
if file_type == "JSON":
+ f = self.open(output_file, "wt", encoding=encoding)
# JSON output
f.write(
json.dumps(
@@ -370,8 +371,11 @@ class dataIO:
)
)
elif file_type == "msgpack":
+ f = self.open(output_file, "w") # TODO: i think something is going wrong here. not sure but doing msgpack and .gz e.g gives an error about str input rather than bytes. i think this is because the self.open does not take into account that the msgpack stream requires different properties.
+
# msgpack output
msgpack.dump(data, f)
+
f.close()
print(