Coverage for kwave/kWaveSimulation_helper/scale_source_terms_func.py: 18%
130 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
1import math
2import numpy as np
3from warnings import warn
5from kwave import kWaveGrid
6from kwave.ksource import kSource
7from kwave.utils import matlab_mask
8from kwave.utils import dotdict
11def scale_source_terms_func(
12 c0, dt, kgrid: kWaveGrid, source,
13 p_source_pos_index, s_source_pos_index, u_source_pos_index,
14 transducer_input_signal,
15 flags: dotdict):
16 """
17 Subscript for the first-order k-Wave simulation functions to scale source terms to the correct units.
18 Args:
19 Returns:
21 """
22 dx, dy, dz = kgrid.dx, kgrid.dy, kgrid.dz,
24 # get the dimension size
25 N = kgrid.dim
27 if not check_conditions(flags.nonuniform_grid, flags.source_uy, flags.source_uz, flags.transducer_source):
28 return
30 # =========================================================================
31 # PRESSURE SOURCES
32 # =========================================================================
34 apply_pressure_source_correction(flags.source_p, flags.use_w_source_correction_p, source, dt)
36 scale_pressure_source(flags.source_p, source, kgrid, N, c0, dx, dt, p_source_pos_index, flags.nonuniform_grid)
38 # =========================================================================
39 # STRESS SOURCES
40 # =========================================================================
42 scale_stress_sources(source, c0, flags, dt, dx, N, s_source_pos_index)
44 # =========================================================================
45 # VELOCITY SOURCES
46 # =========================================================================
48 apply_velocity_source_corrections(flags.use_w_source_correction_u, flags.source_ux, flags.source_uy, flags.source_uz, source, dt)
50 scale_velocity_sources(flags, source, kgrid, c0, dt, dx, dy, dz, u_source_pos_index)
52 # =========================================================================
53 # TRANSDUCER SOURCE
54 # =========================================================================
55 transducer_input_signal = scale_transducer_source(flags.transducer_source, transducer_input_signal,
56 c0, dt, dx, u_source_pos_index)
57 return transducer_input_signal
60def check_conditions(is_nonuniform_grid, is_source_uy, is_source_uz, is_transducer_source):
61 """
62 check for non-uniform grid and give error for source terms that haven't yet been implemented
63 Returns:
65 """
66 if is_nonuniform_grid and (is_source_uy or is_source_uz or is_transducer_source):
67 warn('WARNING: source scaling not implemented for non-uniform grids with given source condition')
68 return False
69 return True
72def apply_pressure_source_correction(is_source_p, use_w_source_correction_p, source, dt):
73 """
74 apply k-space source correction expressed as a function of w
75 Args:
76 is_source_p:
77 use_w_source_correction_p:
78 source:
79 dt:
81 Returns:
83 """
84 if is_source_p and use_w_source_correction_p:
85 source.p = apply_source_correction(source.p, source.p_frequency_ref, dt)
88def scale_pressure_source(is_source_p, source, kgrid, N, c0, dx, dt, p_source_pos_index, is_nonuniform_grid):
89 """
90 scale the input pressure by 1/c0^2 (to convert to units of density), then
91 by 1/N (to split the input across the split density field). If the
92 pressure is injected as a mass source, also scale the pressure by
93 2*dt*c0/dx to account for the time step and convert to units of [kg/(m^3 s)]
94 Args:
95 is_source_p:
96 source:
97 kgrid:
98 N:
99 c0:
100 dx:
101 dt:
102 p_source_pos_index:
103 is_nonuniform_grid:
105 Returns:
107 """
108 if not is_source_p:
109 return
111 if source.p_mode == 'dirichlet':
112 source.p = scale_pressure_source_dirichlet(source.p, c0, N, p_source_pos_index)
113 else:
114 if is_nonuniform_grid:
115 source.p = scale_pressure_source_nonuniform_grid(source.p, kgrid, c0, N, dt, p_source_pos_index)
117 else:
118 source.p = scale_pressure_source_uniform_grid(source.p, c0, N, dx, dt, p_source_pos_index)
121def scale_pressure_source_dirichlet(source_p, c0, N, p_source_pos_index):
122 if c0.size == 1:
123 # compute the scale parameter based on the homogeneous
124 # sound speed
125 source_p = source_p / (N * (c0 ** 2))
127 else:
128 # compute the scale parameter seperately for each source
129 # position based on the sound speed at that position
130 for p_index in range(source_p.size[0]):
131 source_p[p_index, :] = source_p[p_index, :] / (N * (c0[p_source_pos_index[p_index]] ** 2))
132 return source_p
135def scale_pressure_source_nonuniform_grid(source_p, kgrid, c0, N, dt, p_source_pos_index):
136 x = kgrid.x
137 xn = kgrid.xn
138 yn = kgrid.yn
139 zn = kgrid.zn
140 x_size, y_size, z_size = kgrid.size
142 # create empty matrix
143 grid_point_sep = np.zeros(x.size)
145 # compute averaged grid point seperation map, the interior
146 # points are calculated using the average distance to all
147 # connected grid points (the edge values are not calculated
148 # assuming there are no source points in the PML)
149 if kgrid.dim == 1:
150 grid_point_sep[1:-1] = x_size * (xn[2:, 0] - xn[0:-2, 0]) / 2
151 elif kgrid.dim == 2:
152 grid_point_sep[1:-1, 1:-1] = x_size * (xn[2:, 1:-1] - xn[0:-2, 1:- 1]) / 4 + \
153 y_size * (yn[1:-1, 2:] - yn[1:-1, 0:-2]) / 4
154 elif kgrid.dim == 3:
155 grid_point_sep[1:-1, 1:-1, 1:-1] = x_size * (xn[2:, 1:-1, 1:-1] - xn[0:-2, 1:-1, 1:-1]) / 6 + \
156 y_size * (yn[1:-1, 2:, 1:-1] - yn[1:-1, 0:-2, 1:- 1]) / 6 + \
157 z_size * (zn[1:-1, 1:-1, 2:] - zn[1:-1, 1:-1, 0:-2]) / 6
159 # compute and apply scale parameter
160 for p_index in range(source_p.size[0]):
161 if c0.size == 1:
163 # compute the scale parameter based on the homogeneous sound speed
164 source_p[p_index, :] = source_p[p_index, :] * (2 * dt / (N * c0 * grid_point_sep[p_source_pos_index[p_index]]))
166 else:
168 # compute the scale parameter based on the sound speed at that position
169 source_p[p_index, :] = source_p[p_index, :] * (2 * dt / (N * c0[p_source_pos_index[p_index]] * grid_point_sep[p_source_pos_index[p_index]]))
170 return source_p
173def scale_pressure_source_uniform_grid(source_p, c0, N, dx, dt, p_source_pos_index):
174 if c0.size == 1:
175 # compute the scale parameter based on the homogeneous
176 # sound speed
177 source_p = source_p * (2 * dt / (N * c0 * dx))
179 else:
180 # compute the scale parameter seperately for each source
181 # position based on the sound speed at that position
182 for p_index in range(source_p[:, 0].size):
183 source_p[p_index, :] = source_p[p_index, :] * (2 * dt / (N * matlab_mask(c0, p_source_pos_index.flatten('F')[p_index]) * dx))
184 return source_p
187def scale_stress_sources(source, c0, flags, dt, dx, N, s_source_pos_index):
188 """
189 scale the stress source by 1/N to divide amoungst the split field
190 components, and if source.s_mode is not set to 'dirichlet', also scale by
191 2*dt*c0/dx to account for the time step and convert to units of
192 [kg/(m^3 s)] (note dx is used in all dimensions)
193 Args:
194 source:
195 c0:
196 flags:
197 dt:
198 dx:
199 N:
200 s_source_pos_index:
202 Returns:
204 """
205 source.sxx = scale_stress_source(source, c0, flags.source_sxx, flags.source_p0, source.sxx, dt, N, dx, s_source_pos_index)
206 source.syy = scale_stress_source(source, c0, flags.source_syy, flags.source_p0, source.syy, dt, N, dx, s_source_pos_index)
207 source.szz = scale_stress_source(source, c0, flags.source_szz, flags.source_p0, source.szz, dt, N, dx, s_source_pos_index)
208 source.sxy = scale_stress_source(source, c0, flags.source_sxy, True, source.sxy, dt, N, dx, s_source_pos_index)
209 source.sxz = scale_stress_source(source, c0, flags.source_sxz, True, source.sxz, dt, N, dx, s_source_pos_index)
210 source.syz = scale_stress_source(source, c0, flags.source_syz, True, source.syz, dt, N, dx, s_source_pos_index)
213def scale_stress_source(source, c0, is_source_exists, is_p0_exists, source_val, dt, N, dx, s_source_pos_index):
214 if is_source_exists:
215 if source.s_mode == 'dirichlet' or is_p0_exists:
216 source_val = source_val / N
217 else:
218 if c0.size == 1:
220 # compute the scale parameter based on the homogeneous sound
221 # speed
222 source_val = source_val * (2 * dt * c0 / (N * dx))
224 else:
226 # compute the scale parameter seperately for each source
227 # position based on the sound speed at that position
228 for s_index in range(source_val.size[0]):
229 source_val[s_index, :] = source_val[s_index, :] * (2 * dt * c0[s_source_pos_index[s_index]] / (N * dx))
230 return source_val
233def apply_velocity_source_corrections(
234 use_w_source_correction_u: bool, is_ux_exists: bool, is_uy_exists: bool, is_uz_exists: bool,
235 source: kSource, dt: float
236):
237 """
238 apply k-space source correction expressed as a function of w
239 Args:
240 use_w_source_correction_u:
241 is_ux_exists:
242 is_uy_exists:
243 is_uz_exists:
244 source:
245 dt:
247 Returns:
249 """
250 if not use_w_source_correction_u:
251 return
253 if is_ux_exists:
254 source.ux = apply_source_correction(source.ux, source.u_frequency_ref, dt)
256 if is_uy_exists:
257 source.uy = apply_source_correction(source.uy, source.u_frequency_ref, dt)
259 if is_uz_exists:
260 source.uz = apply_source_correction(source.uz, source.u_frequency_ref, dt)
263def apply_source_correction(source_val, frequency_ref, dt):
264 return source_val * math.cos(2 * math.pi * frequency_ref * dt/2)
267def scale_velocity_sources(flags, source, kgrid, c0, dt, dx, dy, dz, u_source_pos_index):
268 source.ux = scale_velocity_source_x(flags.source_ux, source.u_mode, source.ux, kgrid, c0, dt, dx, u_source_pos_index, flags.nonuniform_grid)
269 source.uy = scale_velocity_source(flags.source_uy, source.u_mode, source.uy, c0, dt, u_source_pos_index, dy)
270 source.uz = scale_velocity_source(flags.source_uz, source.u_mode, source.uz, c0, dt, u_source_pos_index, dz)
273def scale_velocity_source_x(is_source_ux, source_u_mode, source_val, kgrid, c0, dt, dx, u_source_pos_index, is_nonuniform_grid):
274 """
275 if source.u_mode is not set to 'dirichlet', scale the x-direction
276 velocity source terms by 2*dt*c0/dx to account for the time step and
277 convert to units of [m/s^2]
278 Returns:
280 """
281 if not is_source_ux or source_u_mode == 'dirichlet':
282 return
284 if is_nonuniform_grid:
285 source_val = scale_velocity_source_nonuniform(is_source_ux, source_u_mode, kgrid, source_val,
286 c0, dt, u_source_pos_index)
287 else:
288 source_val = scale_velocity_source(is_source_ux, source_u_mode, source_val, c0, dt, u_source_pos_index, dx)
289 return source_val
292def scale_velocity_source(is_source, source_u_mode, source_val, c0, dt, u_source_pos_index, d_direction):
293 """
294 if source.u_mode is not set to 'dirichlet', scale the d_direction
295 velocity source terms by 2*dt*c0/dz to account for the time step and
296 convert to units of [m/s^2]
297 Args:
298 is_source:
299 source_u_mode:
300 source_val:
301 c0:
302 dt:
303 u_source_pos_index:
304 d_direction:
306 Returns:
308 """
309 if not is_source or source_u_mode == 'dirichlet':
310 return source_val
312 if c0.size == 1:
313 # compute the scale parameter based on the homogeneous sound speed
314 source_val = source_val * (2 * c0 * dt / d_direction)
315 else:
316 # compute the scale parameter seperately for each source position
317 # based on the sound speed at that position
318 for u_index in range(source_val.size[0]):
319 source_val[u_index, :] = source_val[u_index, :] * (2 * c0(u_source_pos_index[u_index]) * dt / d_direction)
320 return source_val
323def scale_velocity_source_nonuniform(is_source, source_u_mode, kgrid, source_val, c0, dt, u_source_pos_index):
324 """
325 if source.u_mode is not set to 'dirichlet', scale the d_direction
326 velocity source terms by 2*dt*c0/dz to account for the time step and
327 convert to units of [m/s^2]
328 Args:
329 is_source:
330 source_u_mode:
331 kgrid:
332 source_val:
333 c0:
334 dt:
335 u_source_pos_index:
337 Returns:
339 """
340 if not is_source or source_u_mode == 'dirichlet':
341 return source_val
343 # create empty matrix
344 x = kgrid.x
345 xn = kgrid.xn
346 x_size = kgrid.size[0]
347 grid_point_sep = np.zeros_like(x)
349 # compute averaged grid point seperation map, the interior
350 # points are calculated using the average distance to all
351 # connected grid points (the edge values are not calculated
352 # assuming there are no source points in the PML)
353 grid_point_sep[1:- 1, :, :] = x_size * (xn[2:, :, :] - xn[1:- 2, :, :]) / 2
355 # compute and apply scale parameter
356 for u_index in range(source_val.size[0]):
357 if c0.size == 1:
358 # compute the scale parameter based on the homogeneous sound speed
359 source_val[u_index, :] = source_val[u_index, :] * (2 * c0 * dt / (grid_point_sep[u_source_pos_index[u_index]]))
360 else:
361 # compute the scale parameter based on the sound speed at that position
362 source_val[u_index, :] = source_val[u_index, :] * (2 * c0[u_source_pos_index[u_index]] * dt / (grid_point_sep[u_source_pos_index[u_index]]))
363 return source_val
366def scale_transducer_source(is_transducer_source, transducer_input_signal, c0, dt, dx, u_source_pos_index):
367 """
368 scale the transducer source term by 2*dt*c0/dx to account for the time
369 step and convert to units of [m/s^2]
370 Args:
371 is_transducer_source:
372 transducer_input_signal:
373 c0:
374 dt:
375 dx:
376 u_source_pos_index:
378 Returns:
380 """
381 if is_transducer_source:
382 if c0.size == 1:
383 transducer_input_signal = transducer_input_signal * (2 * c0 * dt / dx)
384 else:
385 # compute the scale parameter based on the average sound speed at the
386 # transducer positions (only one input signal is used to drive the transducer)
387 transducer_input_signal = transducer_input_signal * (2 * (np.mean(c0.flatten(order='F')[u_source_pos_index - 1])) * dt / dx)
388 return transducer_input_signal