diff --git a/pystencils/datahandling/serial_datahandling.py b/pystencils/datahandling/serial_datahandling.py index 25a4d23ee40eb4fefb8ddcda3342eef4ba2f1e56..ce4629f6ab7f35c96cec437ab76b778dfec83e04 100644 --- a/pystencils/datahandling/serial_datahandling.py +++ b/pystencils/datahandling/serial_datahandling.py @@ -425,14 +425,15 @@ class SerialDataHandling(DataHandling): np.savez_compressed(file, **self.cpu_arrays) def load_all(self, file): + if '.npz' not in file: + file += '.npz' file_contents = np.load(file) for arr_name, arr_contents in self.cpu_arrays.items(): if arr_name not in file_contents: print(f"Skipping read data {arr_name} because there is no data with this name in data handling") continue if file_contents[arr_name].shape != arr_contents.shape: - print("Skipping read data {} because shapes don't match. " - "Read array shape {}, existing array shape {}".format(arr_name, file_contents[arr_name].shape, - arr_contents.shape)) + print(f"Skipping read data {arr_name} because shapes don't match. " + f"Read array shape {file_contents[arr_name].shape}, existing array shape {arr_contents.shape}") continue np.copyto(arr_contents, file_contents[arr_name]) diff --git a/pystencils_tests/test_data/datahandling_load_test.npz b/pystencils_tests/test_data/datahandling_load_test.npz new file mode 100644 index 0000000000000000000000000000000000000000..d363a8a0aba1bb78a06314a19b887eb4c4975334 Binary files /dev/null and b/pystencils_tests/test_data/datahandling_load_test.npz differ diff --git a/pystencils_tests/test_data/datahandling_save_test.npz b/pystencils_tests/test_data/datahandling_save_test.npz new file mode 100644 index 0000000000000000000000000000000000000000..d363a8a0aba1bb78a06314a19b887eb4c4975334 Binary files /dev/null and b/pystencils_tests/test_data/datahandling_save_test.npz differ diff --git a/pystencils_tests/test_datahandling.py b/pystencils_tests/test_datahandling.py index 6e53d1e8b8cfc89a9f5b71351a248396d69c1d1a..c18cfba98fb1542c0f7ff71584ed288c6c57e8e2 100644 --- a/pystencils_tests/test_datahandling.py +++ b/pystencils_tests/test_datahandling.py @@ -310,3 +310,44 @@ def test_log(): dh.log_on_root() assert dh.is_root assert dh.world_rank == 0 + + +def test_save_data(): + domain_shape = (2, 2) + + dh = create_data_handling(domain_size=domain_shape, default_ghost_layers=1) + dh.add_array("src", values_per_cell=9) + dh.fill("src", 1.0, ghost_layers=True) + dh.add_array("dst", values_per_cell=9) + dh.fill("dst", 1.0, ghost_layers=True) + + dh.save_all('test_data/datahandling_save_test') + + +def test_load_data(): + domain_shape = (2, 2) + + dh = create_data_handling(domain_size=domain_shape, default_ghost_layers=1) + dh.add_array("src", values_per_cell=9) + dh.fill("src", 0.0, ghost_layers=True) + dh.add_array("dst", values_per_cell=9) + dh.fill("dst", 0.0, ghost_layers=True) + + dh.load_all('test_data/datahandling_load_test') + assert np.all(dh.cpu_arrays['src']) == 1 + assert np.all(dh.cpu_arrays['dst']) == 1 + + domain_shape = (3, 3) + + dh = create_data_handling(domain_size=domain_shape, default_ghost_layers=1) + dh.add_array("src", values_per_cell=9) + dh.fill("src", 0.0, ghost_layers=True) + dh.add_array("dst", values_per_cell=9) + dh.fill("dst", 0.0, ghost_layers=True) + dh.add_array("dst2", values_per_cell=9) + dh.fill("dst2", 0.0, ghost_layers=True) + + dh.load_all('test_data/datahandling_load_test') + assert np.all(dh.cpu_arrays['src']) == 0 + assert np.all(dh.cpu_arrays['dst']) == 0 + assert np.all(dh.cpu_arrays['dst2']) == 0