Coverage for kwave/options/simulation_options.py: 39%
154 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-03-24 12:06 -0700
1from __future__ import annotations
3import os
4from dataclasses import dataclass, field
5from enum import Enum
6from tempfile import gettempdir
7from typing import TYPE_CHECKING, List, Optional
9import numpy as np
11if TYPE_CHECKING: 11 ↛ 13line 11 didn't jump to line 13 because the condition on line 11 was never true
12 # Found here: https://adamj.eu/tech/2021/05/13/python-type-hints-how-to-fix-circular-imports/
13 from kwave.kgrid import kWaveGrid
14from kwave.utils.data import get_date_string
15from kwave.utils.io import get_h5_literals
16from kwave.utils.pml import get_optimal_pml_size
19class SimulationType(Enum):
20 """
21 Enum for the simulation type
23 In the original matlab code, simulation type was determined
24 by looking at the calling function name and the user args.
26 Rules from the original matlab code:
27 AXISYMMETRIC => if the calling function name started with 'kspaceFirstOrderAS'
28 or if the userarg_axisymmetric is set to true
29 ELASTIC => if the calling function name started with 'pstdElastic' or 'kspaceElastic'
30 ELASTIC_WITH_KSPACE_CORRECTION => if the calling function name started with 'kspaceElastic'
31 """
33 FLUID = 1
34 AXISYMMETRIC = 2
35 ELASTIC = 3
36 ELASTIC_WITH_KSPACE_CORRECTION = 4
38 def is_elastic_simulation(self):
39 return self in [SimulationType.ELASTIC, SimulationType.ELASTIC_WITH_KSPACE_CORRECTION]
41 def is_axisymmetric(self):
42 return self == SimulationType.AXISYMMETRIC
45@dataclass
46class SimulationOptions(object):
47 """
48 Args:
49 axisymmetric: Flag that indicates whether axisymmetric simulation is used
50 cart_interp: Interpolation mode used to extract the pressure when a Cartesian sensor mask is given.
51 If set to 'nearest' and more than one Cartesian point maps to the same grid point,
52 duplicated data points are discarded and sensor_data will be returned
53 with less points than that specified by sensor.mask (default = 'linear').
54 pml_inside: put the PML inside the grid defined by the user
55 pml_alpha: Absorption within the perfectly matched layer in Nepers per grid point (default = 2).
56 save_to_disk: save the input data to a HDF5 file
57 save_to_disk_exit: exit the simulation after saving the HDF5 file
58 scale_source_terms: apply the source scaling term to time varying sources
59 smooth_c0: smooth the sound speed distribution
60 smooth_rho0: smooth the density distribution
61 smooth_p0: smooth the initial pressure distribution
62 use_kspace: use the k-space correction
63 use_sg: use a staggered grid
64 use_fd: Use finite difference gradients instead of spectral (in 1D)
65 pml_auto: automatically choose the PML size to give small prime factors
66 create_log: create a log using diary
67 use_finite_difference: use finite difference gradients instead of spectral (in 1D)
68 stream_to_disk: String containing a filename (including pathname if required).
69 If set, after the precomputation phase, the input variables used in the time loop are saved
70 the specified location in HDF5 format. The simulation then exits.
71 The saved variables can be used to run simulations using the C++ code.
72 data_recast: recast the sensor data back to double precision
73 cartesian_interp: interpolation mode for Cartesian sensor mask
74 hdf_compression_level: zip compression level for HDF5 input files
75 data_cast: data cast
76 pml_search_range: search range used when automatically determining PML size
77 radial_symmetry: radial symmetry used in axisymmetric code
78 multi_axial_PML_ratio: MPML settings
79 pml_x_alpha: PML Alpha for x-axis
80 pml_y_alpha: PML Alpha for y-axis
81 pml_z_alpha: PML Alpha for z-axis
82 pml_x_size: PML Size for x-axis
83 pml_y_size: PML Size for y-axis
84 pml_z_size: PML Size for z-axis
85 """
87 simulation_type: SimulationType = SimulationType.FLUID
88 cart_interp: str = "linear"
89 pml_inside: bool = True
90 pml_alpha: float = 2.0
91 save_to_disk: bool = False
92 save_to_disk_exit: bool = False
93 scale_source_terms: bool = True
94 smooth_c0: bool = False
95 smooth_rho0: bool = False
96 smooth_p0: bool = True
97 use_kspace: bool = True
98 use_sg: bool = True
99 use_fd: Optional[int] = None
100 pml_auto: bool = False
101 create_log: bool = False
102 use_finite_difference: bool = False
103 stream_to_disk: bool = False
104 data_recast: Optional[bool] = False
105 cartesian_interp: str = "linear"
106 hdf_compression_level: Optional[int] = None
107 data_cast: str = "off"
108 pml_search_range: List[int] = field(default_factory=lambda: [10, 40])
109 radial_symmetry: str = "WSWA-FFT"
110 multi_axial_PML_ratio: float = 0.1
111 data_path: Optional[str] = field(default_factory=lambda: gettempdir())
112 input_filename: Optional[str] = field(default_factory=lambda: f"{get_date_string()}_kwave_input.h5")
113 output_filename: Optional[str] = field(default_factory=lambda: f"{get_date_string()}_kwave_output.h5")
114 pml_x_alpha: Optional[float] = None
115 pml_y_alpha: Optional[float] = None
116 pml_z_alpha: Optional[float] = None
117 pml_size: Optional[List[int]] = None
118 pml_x_size: Optional[int] = None
119 pml_y_size: Optional[int] = None
120 pml_z_size: Optional[int] = None
122 def __post_init__(self):
123 assert self.cartesian_interp in [
124 "linear",
125 "nearest",
126 ], "Optional input ''cartesian_interp'' must be set to ''linear'' or ''nearest''."
128 assert isinstance(self.data_cast, str), "Optional input ''data_cast'' must be a string."
130 assert self.data_cast in ["off", "double", "single"], "Invalid input for ''data_cast''."
132 if self.data_cast == "double": 132 ↛ 133line 132 didn't jump to line 133 because the condition on line 132 was never true
133 self.data_cast = "off"
135 # load the HDF5 literals (for the default compression level)
136 h5_literals = get_h5_literals()
137 self.hdf_compression_level = h5_literals.HDF_COMPRESSION_LEVEL
138 # check value is an integer between 0 and 9
139 assert (
140 isinstance(self.hdf_compression_level, int) and 0 <= self.hdf_compression_level <= 9
141 ), "Optional input ''hdf_compression_level'' must be an integer between 0 and 9."
143 assert (
144 np.isscalar(self.multi_axial_PML_ratio) and self.multi_axial_PML_ratio >= 0
145 ), "Optional input ''multi_axial_PML_ratio'' must be a single positive value."
147 assert np.isscalar(self.stream_to_disk) or isinstance(
148 self.stream_to_disk, bool
149 ), "Optional input ''stream_to_disk'' must be a single scalar or Boolean value."
151 boolean_inputs = {
152 "use_sg": self.use_sg,
153 "data_recast": self.data_recast,
154 "save_to_disk_exit": self.save_to_disk_exit,
155 "use_kspace": self.use_kspace,
156 "save_to_disk": self.save_to_disk,
157 "pml_inside": self.pml_inside,
158 "create_log": self.create_log,
159 "scale_source_terms": self.scale_source_terms,
160 }
162 for key, val in boolean_inputs.items():
163 assert isinstance(val, bool), f"Optional input ''{key}'' must be Boolean."
165 assert self.radial_symmetry in [
166 "WSWA",
167 "WSWS",
168 "WSWA-FFT",
169 "WSWS-FFT",
170 ], "Optional input ''RadialSymmetry'' must be set to ''WSWA'', ''WSWS'', ''WSWA-FFT'', ''WSWS-FFT''."
172 # automatically assign the PML size to give small prime factors
173 if self.pml_auto and self.pml_inside: 173 ↛ 174line 173 didn't jump to line 174 because the condition on line 173 was never true
174 raise NotImplementedError("''pml_size'' set to ''auto'' is only supported with ''pml_inside'' set to false.")
176 if self.pml_size is not None: 176 ↛ 178line 176 didn't jump to line 178 because the condition on line 176 was never true
177 # TODO(walter): remove auto option in exchange for pml_auto=True
178 if isinstance(self.pml_size, int):
179 self.pml_size = np.array([self.pml_size])
180 if not isinstance(self.pml_size, (list, np.ndarray)):
181 raise ValueError("Optional input ''PMLSize'' must be a integer array of 1, 2 or 3 dimensions.")
183 # Check if each member variable is None, and set it to self.pml_alpha if it is
184 self.pml_x_alpha = self.pml_alpha if self.pml_x_alpha is None else self.pml_x_alpha
185 self.pml_y_alpha = self.pml_alpha if self.pml_y_alpha is None else self.pml_y_alpha
186 self.pml_z_alpha = self.pml_alpha if self.pml_z_alpha is None else self.pml_z_alpha
188 # add pathname to input and output filenames
189 self.input_filename = os.path.join(self.data_path, self.input_filename)
190 self.output_filename = os.path.join(self.data_path, self.output_filename)
192 assert self.use_fd is None or (
193 np.issubdtype(self.use_fd, np.number) and self.use_fd in [2, 4]
194 ), "Optional input ''UseFD'' can only be set to 2, 4."
196 @staticmethod
197 def option_factory(kgrid: "kWaveGrid", options: SimulationOptions):
198 """
199 Initialize the Simulation Options
201 Args:
202 kgrid: kWaveGrid instance
203 elastic_code: Flag that indicates whether elastic simulation is used
204 **kwargs: Dictionary that holds following optional simulation properties:
206 * cart_interp: Interpolation mode used to extract the pressure when a Cartesian sensor mask is given.
207 If set to 'nearest', duplicated data points are discarded and sensor_data
208 will be returned with fewer points than specified by sensor.mask (default = 'linear').
209 * create_log: Boolean controlling whether the command line output is saved using the diary function
210 with a date and time stamped filename (default = false).
211 * data_cast: String input of the data type that variables are cast to before computation.
212 For example, setting to 'single' will speed up the computation time
213 (due to the improved efficiency of fftn and ifftn for this data type) at the expense
214 of a loss in precision. This variable is also useful for utilising GPU parallelisation
215 through libraries such as the Parallel Computing Toolbox
216 by setting 'data_cast' to 'gpuArray-single' (default = 'off').
217 * data_recast: Boolean controlling whether the output data is cast back to double precision.
218 If set to false, sensor_data will be returned in
219 the data format set using the 'data_cast' option.
220 * hdf_compression_level: Compression level used for writing the input HDF5 file when using
221 'save_to_disk' or kspaceFirstOrder3DC. Can be set to an integer
222 between 0 (no compression, the default) and 9 (maximum compression).
223 The compression is lossless. Increasing the compression level will reduce
224 the file size if there are portions of the medium that are homogeneous,
225 but will also increase the time to create the HDF5 file.
226 * multi_axial_pml_ratio: MPML settings
227 * pml_alpha: Absorption within the perfectly matched layer in Nepers per grid point (default = 2).
228 * pml_inside: Boolean controlling whether the perfectly matched layer is inside or outside the grid.
229 If set to false, the input grids are enlarged by pml_size
230 before running the simulation (default = true).
231 * pml_range: Search range used when automatically determining PML size. Tuple of two elements
232 * pml_size: Size of the perfectly matched layer in grid points. By default, the PML is added evenly to
233 all sides of the grid, however, both pml_size and pml_alpha can be given as three element
234 arrays to specify the x, y, and z properties, respectively.
235 To remove the PML, set the appropriate pml_alpha to zero rather than forcing
236 the PML to be of zero size (default = 10).
237 * radial_symmetry: Radial symmetry used in axisymmetric code
238 * stream_to_disk: Boolean controlling whether sensor_data is periodically saved to disk to avoid storing
239 the complete matrix in memory. StreamToDisk may also be given as an integer which
240 specifies the number of time steps that are taken before the data
241 is saved to disk (default = 200).
242 * save_to_disk: String containing a filename (including pathname if required).
243 If set, after the precomputation phase, the input variables used in the time loop are
244 saved the specified location in HDF5 format. The simulation then exits.
245 The saved variables can be used to run simulations using the C++ code.
246 * save_to_disk_exit: Exit the simulation after saving the HDF5 file
247 * scale_source_terms: Apply the source scaling term to time varying sources
248 * use_fd: Use finite difference gradients instead of spectral (in 1D)
249 * use_k_space: use the k-space correction
250 * use_sg: Use a staggered grid
253 Returns:
254 SimulationOptions instance
255 """
257 STREAM_TO_DISK_STEPS_DEF = 200 # number of steps before streaming to disk
259 if options.pml_size is not None and not isinstance(options.pml_size, bool):
260 if len(options.pml_size) > kgrid.dim:
261 if kgrid.dim > 1:
262 raise ValueError(f"Optional input ''pml_size'' must be a 1 or {kgrid.dim} element numerical array.")
263 else:
264 raise ValueError("Optional input ''pml_size'' must be a single numerical value.")
266 if kgrid.dim == 1:
267 options.pml_x_size = options.pml_size if options.pml_size else 20
268 options.plot_scale = [-1.1, 1.1]
269 elif kgrid.dim == 2:
270 if options.pml_size is not None:
271 if len(options.pml_size) == kgrid.dim:
272 options.pml_x_size, options.pml_y_size = np.asarray(options.pml_size, dtype=int).ravel()
273 else:
274 options.pml_x_size, options.pml_y_size = (options.pml_size[0], options.pml_size[0])
275 else:
276 options.pml_x_size, options.pml_y_size = (20, 20)
277 options.plot_scale = [-1, 1]
278 elif kgrid.dim == 3:
279 if (options.pml_size is not None) and (len(options.pml_size) == kgrid.dim):
280 options.pml_x_size, options.pml_y_size, options.pml_z_size = np.asarray(options.pml_size).ravel()
281 else:
282 if options.pml_size is None:
283 options.pml_x_size = 10
284 options.pml_y_size = 10
285 options.pml_z_size = 10
286 else:
287 options.pml_x_size = options.pml_size[0]
288 options.pml_y_size = options.pml_x_size
289 options.pml_z_size = options.pml_x_size
290 options.plot_scale = [-1, 1]
292 # replace defaults with user defined values if provided and check inputs
293 if (val := options.pml_alpha) is not None and not isinstance(options.pml_alpha, str):
294 # check input is correct size
295 val = np.atleast_1d(val)
296 if val.size > kgrid.dim:
297 if kgrid.dim > 1:
298 raise ValueError(f"Optional input ''pml_alpha'' must be a 1 or {kgrid.dim} element numerical array.")
299 else:
300 raise ValueError("Optional input ''pml_alpha'' must be a single numerical value.")
302 # assign input based on number of dimensions
303 if kgrid.dim == 1:
304 options.pml_x_alpha = val
305 elif kgrid.dim == 2:
306 options.pml_x_alpha = val[0]
307 options.pml_y_alpha = val[-1]
308 elif kgrid.dim == 3:
309 options.pml_x_alpha = val[0]
310 options.pml_y_alpha = val[len(val) // 2]
311 options.pml_z_alpha = val[-1]
313 if options.save_to_disk_exit:
314 assert kgrid.dim != 1, "Optional input ''save_to_disk'' is not compatible with 1D simulations."
316 if options.stream_to_disk:
317 assert (
318 not options.simulation_type.is_elastic_simulation() and kgrid.dim == 3
319 ), "Optional input ''stream_to_disk'' is currently only compatible with 3D fluid simulations."
320 # if given as a Boolean, replace with the default number of time steps
321 if isinstance(options.stream_to_disk, bool) and options.stream_to_disk:
322 options.stream_to_disk = STREAM_TO_DISK_STEPS_DEF
324 if options.save_to_disk or options.save_to_disk_exit:
325 assert kgrid.dim != 1, "Optional input ''save_to_disk'' is not compatible with 1D simulations."
327 if options.use_fd:
328 # input only supported in 1D fluid code
329 assert kgrid.dim == 1 and not options.simulation_type.is_elastic_simulation(), "Optional input ''use_fd'' only supported in 1D."
330 # get optimal pml size
331 if options.simulation_type.is_axisymmetric() or options.pml_auto:
332 if options.simulation_type.is_axisymmetric():
333 pml_size_temp = get_optimal_pml_size(kgrid, options.pml_search_range, options.radial_symmetry[:4])
334 else:
335 pml_size_temp = get_optimal_pml_size(kgrid, options.pml_search_range)
337 # assign to individual variables
338 if kgrid.dim == 1:
339 options.pml_x_size = int(pml_size_temp[0])
340 elif kgrid.dim == 2:
341 options.pml_x_size = int(pml_size_temp[0])
342 options.pml_y_size = int(pml_size_temp[1])
343 elif kgrid.dim == 3:
344 options.pml_x_size = int(pml_size_temp[0])
345 options.pml_y_size = int(pml_size_temp[1])
346 options.pml_z_size = int(pml_size_temp[2])
348 # cleanup unused variables
349 del pml_size_temp
350 return options