From 9cb1614d38dc29a35d9a40c8d6939b72928654b8 Mon Sep 17 00:00:00 2001
From: David Sinden <djps@users.noreply.github.com>
Date: Thu, 9 Nov 2023 17:18:47 +0100
Subject: [PATCH] refactor loops

---
 .../scale_source_terms_func.py                | 30 ++++++++++++-------
 1 file changed, 20 insertions(+), 10 deletions(-)

diff --git a/kwave/kWaveSimulation_helper/scale_source_terms_func.py b/kwave/kWaveSimulation_helper/scale_source_terms_func.py
index 428aa1c5..83a40e33 100644
--- a/kwave/kWaveSimulation_helper/scale_source_terms_func.py
+++ b/kwave/kWaveSimulation_helper/scale_source_terms_func.py
@@ -184,13 +184,13 @@ def scale_pressure_source_uniform_grid(source_p, c0, N, dx, dt, p_source_pos_ind
     else:
         # compute the scale parameter seperately for each source
         # position based on the sound speed at that position
-        #p_index = range(source_p[:, 0].size[0])
-        #mask = p_source_pos_index.flatten('F')[p_index]
-        #scale = (2.0 * dt) / (N * np.expand_dims(c0.ravel(order='F')[mask.ravel(order='F')], axis=-1) * dx)
-        #source_p[p_index, :] *= scale
+        p_index = range(source_p[:, 0].size)
+        mask = p_source_pos_index.flatten('F')[p_index]
+        scale = (2.0 * dt) / (N * np.expand_dims(c0.ravel(order='F')[mask.ravel(order='F')], axis=-1) * dx)
+        source_p[p_index, :] *= scale
 
-        for p_index in range(source_p[:, 0].size):
-            source_p[p_index, :] = source_p[p_index, :] * (2.0 * dt) / (N * c0[p_source_pos_index[p_index]] * dx)
+        # for p_index in range(source_p[:, 0].size):
+        #     source_p[p_index, :] = source_p[p_index, :] * (2.0 * dt) / (N * c0[p_source_pos_index[p_index]] * dx)
 
         return source_p
 
@@ -234,12 +234,19 @@ def scale_stress_source(source, c0, is_source_exists, is_p0_exists, source_val,
 
             else:
 
+                #for s_index in range(source_val.size[0]):
+                #    source_val[s_index, :] = source_val[s_index, :] * (2 * dt * c0[s_source_pos_index[s_index]] / (N * dx))
+
                 # compute the scale parameter seperately for each source
                 # position based on the sound speed at that position
+                # s_index = range(source_val.size[0])
+                # mask = s_source_pos_index.flatten('F')[s_index]
+                # scale = (2.0 * dt * np.expand_dims(c0.ravel(order='F')[mask.ravel(order='F')], axis=-1) ) / (N * dx)
+                # source_val[s_index, :] *= scale
 
                 s_index = range(source_val.size[0])
-                mask = s_source_pos_index.flatten('F')[s_index]
-                scale = (2.0 * dt * np.expand_dims(c0.ravel(order='F')[mask.ravel(order='F')], axis=-1) ) / (N * dx)
+                mask = s_source_pos_index[s_index]
+                scale = (2.0 * dt * c0[mask]) / (N * dx)
                 source_val[s_index, :] *= scale
                     
     return source_val
@@ -338,8 +345,11 @@ def scale_velocity_source(is_source, source_u_mode, source_val, c0, dt, u_source
         # scale = (2.0 * dt * np.expand_dims(c0.ravel(order='F')[mask.ravel(order='F')], axis=-1) ) / d_direction
         # source_val[u_index, :] *= scale
 
-        for u_index in range(source_val.size[0]):
-            source_val[u_index, :] = source_val[u_index, :] * (2 * c0(u_source_pos_index[u_index]) * dt / d_direction)
+        u_index = range(source_val.size[0])
+        source_val[u_index, :] = source_val[u_index, :] * (2 * c0[u_source_pos_index[u_index]] * dt / d_direction)
+
+        #for u_index in range(source_val.size[0]):
+        #    source_val[u_index, :] = source_val[u_index, :] * (2 * c0(u_source_pos_index[u_index]) * dt / d_direction)
                 
     return source_val