diff --git a/kwave/kWaveSimulation_helper/save_to_disk_func.py b/kwave/kWaveSimulation_helper/save_to_disk_func.py index 2fbd25e1..f4771034 100644 --- a/kwave/kWaveSimulation_helper/save_to_disk_func.py +++ b/kwave/kWaveSimulation_helper/save_to_disk_func.py @@ -16,7 +16,8 @@ def save_to_disk_func( kgrid: kWaveGrid, medium: kWaveMedium, source, - opt: SimulationOptions, values: dotdict, flags: dotdict): + opt: SimulationOptions, auto_chunk: bool, + values: dotdict, flags: dotdict): # update command line status logging.log(logging.INFO, ' precomputation completed in ', scale_time(TicToc.toc())) TicToc.tic() @@ -56,7 +57,8 @@ def save_to_disk_func( # ========================================================================= remove_z_dimension(float_variables, kgrid.dim) - save_file(opt.input_filename, integer_variables, float_variables, opt.hdf_compression_level) + save_file(opt.input_filename, integer_variables, float_variables, opt.hdf_compression_level, + auto_chunk=auto_chunk) # update command line status logging.log(logging.INFO, ' completed in ', scale_time(TicToc.toc())) @@ -445,12 +447,12 @@ def enforce_filename_standards(filepath): return filepath, filename_ext -def save_file(filepath, integer_variables, float_variables, hdf_compression_level): +def save_file(filepath, integer_variables, float_variables, hdf_compression_level, auto_chunk): filepath, filename_ext = enforce_filename_standards(filepath) # save file if filename_ext == '.h5': - save_h5_file(filepath, integer_variables, float_variables, hdf_compression_level) + save_h5_file(filepath, integer_variables, float_variables, hdf_compression_level, auto_chunk) elif filename_ext == '.mat': save_mat_file(filepath, integer_variables, float_variables) @@ -459,7 +461,7 @@ def save_file(filepath, integer_variables, float_variables, hdf_compression_leve raise NotImplementedError('unknown file extension for ''save_to_disk'' filename') -def save_h5_file(filepath, integer_variables, float_variables, hdf_compression_level): +def save_h5_file(filepath, integer_variables, float_variables, hdf_compression_level, auto_chunk): # ---------------- # SAVE HDF5 FILE # ---------------- @@ -474,7 +476,7 @@ def save_h5_file(filepath, integer_variables, float_variables, hdf_compression_l for key, value in float_variables.items(): # cast matrix to single precision value = np.array(value, dtype=np.float32) - write_matrix(filepath, value, key, hdf_compression_level) + write_matrix(filepath, value, key, hdf_compression_level, auto_chunk) del value # change all the index variables to be in 64-bit unsigned integers @@ -482,7 +484,7 @@ def save_h5_file(filepath, integer_variables, float_variables, hdf_compression_l for key, value in integer_variables.items(): # cast matrix to 64-bit unsigned integer value = np.array(value, dtype=np.uint64) - write_matrix(filepath, value, key, hdf_compression_level) + write_matrix(filepath, value, key, hdf_compression_level, auto_chunk) del value # set additional file attributes diff --git a/kwave/kspaceFirstOrder2D.py b/kwave/kspaceFirstOrder2D.py index 5cba1506..3b151bc5 100644 --- a/kwave/kspaceFirstOrder2D.py +++ b/kwave/kspaceFirstOrder2D.py @@ -391,7 +391,7 @@ def kspaceFirstOrder2D( retract_size = [[options.pml_x_size, options.pml_y_size, options.pml_z_size]] # run subscript to save files to disk - save_to_disk_func(k_sim.kgrid, k_sim.medium, k_sim.source, k_sim.options, + save_to_disk_func(k_sim.kgrid, k_sim.medium, k_sim.source, k_sim.options, execution_options.auto_chunking, dotdict({ 'ddx_k_shift_pos': k_sim.ddx_k_shift_pos, 'ddx_k_shift_neg': k_sim.ddx_k_shift_neg, diff --git a/kwave/kspaceFirstOrder3D.py b/kwave/kspaceFirstOrder3D.py index b4a51a14..59144603 100644 --- a/kwave/kspaceFirstOrder3D.py +++ b/kwave/kspaceFirstOrder3D.py @@ -413,7 +413,7 @@ def kspaceFirstOrder3D( retract_size = [[options.pml_x_size, options.pml_y_size, options.pml_z_size]] # run subscript to save files to disk - save_to_disk_func(k_sim.kgrid, k_sim.medium, k_sim.source, k_sim.options, + save_to_disk_func(k_sim.kgrid, k_sim.medium, k_sim.source, k_sim.options, execution_options.auto_chunking, dotdict({ 'ddx_k_shift_pos': k_sim.ddx_k_shift_pos, 'ddx_k_shift_neg': k_sim.ddx_k_shift_neg, diff --git a/kwave/kspaceFirstOrderAS.py b/kwave/kspaceFirstOrderAS.py index b3f51d7b..3258b427 100644 --- a/kwave/kspaceFirstOrderAS.py +++ b/kwave/kspaceFirstOrderAS.py @@ -314,7 +314,7 @@ def kspaceFirstOrderAS( retract_size = [[options.pml_x_size, options.pml_y_size, options.pml_z_size]] # run subscript to save files to disk - save_to_disk_func(k_sim.kgrid, k_sim.medium, k_sim.source, k_sim.options, + save_to_disk_func(k_sim.kgrid, k_sim.medium, k_sim.source, k_sim.options, execution_options.auto_chunking, dotdict({ 'ddx_k_shift_pos': k_sim.ddx_k_shift_pos, 'ddx_k_shift_neg': k_sim.ddx_k_shift_neg, diff --git a/kwave/options/simulation_execution_options.py b/kwave/options/simulation_execution_options.py index 4b3bc833..4736665a 100644 --- a/kwave/options/simulation_execution_options.py +++ b/kwave/options/simulation_execution_options.py @@ -34,6 +34,9 @@ class SimulationExecutionOptions: system_call: Optional[str] = None verbose_level: int = 0 + # determine whether chunking is handled automatically (the default), or manually + auto_chunking: Optional[bool] = True + # show simulation log show_sim_log: bool = True diff --git a/kwave/utils/io.py b/kwave/utils/io.py index 6a3a7a84..da9f8237 100644 --- a/kwave/utils/io.py +++ b/kwave/utils/io.py @@ -52,10 +52,12 @@ def get_h5_literals(): return literals -def write_matrix(filename, matrix: np.ndarray, matrix_name, compression_level=None): +def write_matrix(filename, matrix: np.ndarray, matrix_name: str, compression_level:int =None, auto_chunk: bool =True): # get literals h5_literals = get_h5_literals() + assert isinstance(auto_chunk, bool), "auto_chunk must be a boolean." + if compression_level is None: compression_level = h5_literals.HDF_COMPRESSION_LEVEL @@ -181,8 +183,9 @@ def write_matrix(filename, matrix: np.ndarray, matrix_name, compression_level=No # allocate a holder for the new matrix within the file opts = { 'dtype': data_type_matlab, - 'chunks': tuple(chunk_size) + 'chunks': auto_chunk if auto_chunk is True else tuple(chunk_size) } + if compression_level != 0: # use compression opts['compression'] = compression_level