Coverage for kwave/kWaveSimulation_helper/expand_grid_matrices.py: 14%
121 statements
« prev ^ index » next coverage.py v6.5.0, created at 2022-10-24 11:55 -0700
« prev ^ index » next coverage.py v6.5.0, created at 2022-10-24 11:55 -0700
1from kwave import kWaveGrid, kWaveMedium, SimulationOptions, NotATransducer
2from kwave.data import Array
3from kwave.utils import matlab_find, expand_matrix, get_smallest_possible_type
4from kwave.utils import dotdict
5import numpy as np
8def expand_grid_matrices(
9 kgrid: kWaveGrid, medium: kWaveMedium, source, sensor, opt: SimulationOptions,
10 values: dotdict, flags: dotdict):
11 # update command line status
12 print(' expanding computational grid...')
14 #####################
15 # Grab values
16 #####################
18 # retaining the values for kgrid time array
19 pml_size = [opt.pml_x_size, opt.pml_y_size, opt.pml_z_size]
20 pml_size = Array(pml_size[:kgrid.dim])
22 p_source_pos_index = values.p_source_pos_index
23 u_source_pos_index = values.u_source_pos_index
24 s_source_pos_index = values.s_source_pos_index
26 is_source_sensor_same = (isinstance(sensor, NotATransducer) and sensor == source)
28 #####################
29 # Expand Structures
30 #####################
32 # expand the computational grid, replacing the original grid
33 kgrid = expand_kgrid(kgrid, flags.axisymmetric, pml_size.numpy())
35 expand_size = calculate_expand_size(kgrid, flags.axisymmetric, pml_size)
37 # update the data type in case adding the PML requires additional index precision
38 total_grid_points = kgrid.total_grid_points
39 index_data_type = get_smallest_possible_type(total_grid_points, 'uint', default='double')
41 expand_sensor(sensor, expand_size, flags.use_sensor, flags.blank_sensor)
43 # TODO why it is not self.record ? "self"
44 record = expand_cuboid_corner_list(flags.cuboid_corners, kgrid, pml_size)
46 expand_medium(medium, expand_size)
48 p_source_pos_index, u_source_pos_index, s_source_pos_index = expand_source(
49 source, is_source_sensor_same, flags, expand_size, index_data_type,
50 p_source_pos_index, u_source_pos_index, s_source_pos_index
51 )
53 expand_directivity_angle(kgrid, sensor, expand_size, flags.use_sensor, flags.compute_directivity)
55 print_grid_size(kgrid)
57 return kgrid, index_data_type, p_source_pos_index, u_source_pos_index, s_source_pos_index
60def expand_kgrid(kgrid, is_axisymmetric, pml_size):
61 Nt_temp, dt_temp = kgrid.Nt, kgrid.dt
63 if kgrid.dim == 1:
64 new_size = kgrid.N + 2 * pml_size
65 elif kgrid.dim == 2:
66 if is_axisymmetric:
67 new_size = [kgrid.Nx + 2 * pml_size[0], kgrid.Ny + pml_size[1]]
68 else:
69 new_size = kgrid.N + 2 * pml_size
70 elif kgrid.dim == 3:
71 new_size = kgrid.N + 2 * pml_size
72 else:
73 raise NotImplementedError
75 kgrid = kWaveGrid(new_size, kgrid.spacing)
76 # re-assign original time array
77 kgrid.setTime(Nt_temp, dt_temp)
79 return kgrid
82def calculate_expand_size(kgrid, is_axisymmetric, pml_size):
83 # set the PML size for use with expandMatrix, don't expand the inner radial
84 # dimension if using the axisymmetric code
85 if kgrid.dim == 1:
86 expand_size = pml_size[0]
87 elif kgrid.dim == 2:
88 if is_axisymmetric:
89 expand_size = [pml_size[0], pml_size[0], 0, pml_size[1]]
90 else:
91 expand_size = pml_size
92 elif kgrid.dim == 3:
93 expand_size = pml_size
94 else:
95 raise NotImplementedError
96 return np.array(expand_size)
99def expand_medium(medium: kWaveMedium, expand_size):
100 # enlarge the sound speed grids by exting the edge values into the expanded grid
101 medium.sound_speed = np.atleast_1d(medium.sound_speed)
102 if medium.sound_speed.size > 1:
103 medium.sound_speed = expand_matrix(medium.sound_speed, expand_size)
105 # enlarge the grid of density by exting the edge values into the expanded grid
106 medium.density = np.atleast_1d(medium.density)
107 if medium.density.size > 1:
108 medium.density = expand_matrix(medium.density, expand_size)
110 # for key in ['alpha_coeff', 'alpha_coeff_compression', 'alpha_coeff_shear', 'BonA']:
111 for key in ['alpha_coeff', 'BonA']:
112 # enlarge the grid of medium[key] if given
113 attr = getattr(medium, key)
114 if attr is not None and np.atleast_1d(attr).size > 1:
115 attr = expand_matrix(np.atleast_1d(attr), expand_size)
116 setattr(medium, key, attr)
118 # enlarge the absorption filter mask if given
119 if medium.alpha_filter is not None:
120 medium.alpha_filter = expand_matrix(medium.alpha_filter, expand_size, 0)
123def expand_source(
124 source, is_source_sensor_same, flags, expand_size, index_data_type,
125 p_source_pos_index, u_source_pos_index, s_source_pos_index):
127 p_source_pos_index = expand_pressure_sources(source, expand_size, flags.source_p0, flags.source_p,
128 index_data_type, p_source_pos_index)
130 u_source_pos_index = expand_velocity_sources(
131 source, expand_size, is_source_sensor_same, index_data_type, u_source_pos_index,
132 flags.source_ux, flags.source_uy, flags.source_uz, flags.transducer_source
133 )
135 s_source_pos_index = expand_stress_sources(source, expand_size, flags, index_data_type, s_source_pos_index)
137 return p_source_pos_index, u_source_pos_index, s_source_pos_index
140def expand_pressure_sources(source, expand_size, is_source_p0, is_source_p, index_data_type, p_source_pos_index):
141 # enlarge the initial pressure if given
142 if is_source_p0:
143 source.p0 = expand_matrix(source.p0, expand_size, 0)
145 # enlarge the pressure source mask if given
146 if is_source_p:
148 # enlarge the pressure source mask
149 source.p_mask = expand_matrix(source.p_mask, expand_size, 0)
151 # create an indexing variable corresponding to the source elements
152 # and convert the data type deping on the number of indices
153 p_source_pos_index = matlab_find(source.p_mask).astype(index_data_type)
154 return p_source_pos_index
157def expand_velocity_sources(
158 source, expand_size, is_source_sensor_same, index_data_type, u_source_pos_index,
159 is_source_ux, is_source_uy, is_source_uz, is_transducer_source
160):
161 """
162 enlarge the velocity source mask if given
163 Args:
164 source:
165 expand_size:
166 is_source_sensor_same:
167 index_data_type:
168 u_source_pos_index:
169 is_source_ux:
170 is_source_uy:
171 is_source_uz:
172 is_transducer_source:
174 Returns:
176 """
177 if is_source_ux or is_source_uy or is_source_uz or is_transducer_source:
179 # update the source indexing variable
180 if isinstance(source, NotATransducer):
182 # check if the sensor is also the same transducer, if so, don't expand the grid again
183 if not is_source_sensor_same:
185 # expand the transducer mask
186 source.expand_grid(expand_size)
188 # get the new active elements mask
189 active_elements_mask = source.active_elements_mask
191 # update the indexing variable corresponding to the active elements
192 u_source_pos_index = matlab_find(active_elements_mask)
193 else:
195 # enlarge the velocity source mask
196 source.u_mask = expand_matrix(source.u_mask, expand_size, 0)
198 # create an indexing variable corresponding to the source elements
199 u_source_pos_index = matlab_find(source.u_mask)
201 # convert the data type deping on the number of indices
202 u_source_pos_index = u_source_pos_index.astype(index_data_type)
203 return u_source_pos_index
206def expand_stress_sources(source, expand_size, flags, index_data_type, s_source_pos_index):
207 # enlarge the stress source mask if given
208 if flags.source_sxx or flags.source_syy or flags.source_szz or flags.source_sxy or flags.source_sxz or flags.source_syz:
209 # enlarge the velocity source mask
210 source.s_mask = expand_matrix(source.s_mask, expand_size, 0)
212 # create an indexing variable corresponding to the source elements
213 s_source_pos_index = matlab_find(source.s_mask != 0)
215 # convert the data type deping on the number of indices
216 s_source_pos_index = s_source_pos_index.astype(index_data_type)
217 return s_source_pos_index
220def expand_directivity_angle(kgrid, sensor, expand_size, is_use_sensor, is_compute_directivity):
221 """
222 enlarge the directivity angle if given (2D only)
223 Args:
224 kgrid:
225 sensor:
226 expand_size:
227 is_use_sensor:
228 is_compute_directivity:
230 Returns:
232 """
233 if is_use_sensor and kgrid.dim == 2 and is_compute_directivity:
234 # enlarge the directivity angle
235 sensor.directivity.angle = expand_matrix(sensor.directivity.angle, expand_size, 0)
236 # re-assign the wavenumber vectors
237 sensor.directivity.wavenumbers = np.hstack((kgrid.ky.T, kgrid.kx.T))
240def print_grid_size(kgrid):
241 """
242 update command line status
243 Args:
244 kgrid:
246 Returns:
248 """
249 k_Nx, k_Ny, k_Nz = kgrid.Nx, kgrid.Ny, kgrid.Nz
250 if kgrid.dim == 1:
251 print(' computational grid size:', int(k_Nx), 'grid points')
252 elif kgrid.dim == 2:
253 print(' computational grid size:', int(k_Nx), 'by', int(k_Ny), 'grid points')
254 elif kgrid.dim == 3:
255 print(' computational grid size:', int(k_Nx), 'by', int(k_Ny), 'by', int(k_Nz), 'grid points')
258def expand_cuboid_corner_list(is_cuboid_list, kgrid, pml_size: Array):
259 """
260 add the PML size to cuboid corner indices if using a cuboid sensor mask
261 Args:
262 is_cuboid_list:
263 kgrid:
265 Returns:
267 """
268 if not is_cuboid_list:
269 return
271 record = dotdict()
272 if kgrid.dim == 1:
273 record.cuboid_corners_list = record.cuboid_corners_list + pml_size.x
274 elif kgrid.dim == 2:
275 record.cuboid_corners_list[[0, 2], :] = record.cuboid_corners_list[[0, 2], :] + pml_size.x
276 record.cuboid_corners_list[[1, 3], :] = record.cuboid_corners_list[[1, 3], :] + pml_size.y
277 elif kgrid.dim == 3:
278 record.cuboid_corners_list[[0, 3], :] = record.cuboid_corners_list[[0, 3], :] + pml_size.x
279 record.cuboid_corners_list[[1, 4], :] = record.cuboid_corners_list[[1, 4], :] + pml_size.y
280 record.cuboid_corners_list[[2, 5], :] = record.cuboid_corners_list[[2, 5], :] + pml_size.z
281 return record
284def expand_sensor(sensor, expand_size, is_use_sensor, is_blank_sensor):
285 """
286 enlarge the sensor mask (for Cartesian sensor masks and cuboid corners,
287 this has already been converted to a binary mask for display in inputChecking)
288 Args:
289 sensor:
290 expand_size:
291 is_use_sensor:
292 is_blank_sensor:
294 Returns:
296 """
297 if is_use_sensor and not is_blank_sensor:
298 sensor.expand_grid(expand_size)