shear patch

Author: Sierd

aeolis/shear.py

@@ -576,31 +576,58 @@ def compute_shear(self, u0, nfilter=(1., 2.)):
 
         # Arrays in Fourier 
         k = np.sqrt(kx**2 + ky**2)
-        sigma = np.sqrt(1j * L * kx * z0new /l)
 
 
         time_start_perturbation = time.time()
-        
+
+
         # Shear stress perturbation
-        # Use safe division to avoid zero/invalid values at kx=0 or k=0
-        k_safe = np.where(k == 0, 1.0, k)
-        kx_safe = np.where(kx == 0, 1.0, kx)
-        
-        dtaux_t = hs * kx**2 / k_safe * 2 / ul**2 * \
-                  (-1. + (2. * np.log(l/z0new) + k**2/kx_safe**2) * sigma * \
-                    sc_kv(1., 2. * sigma) / sc_kv(0., 2. * sigma))
+        # Use masked computation to avoid division by zero and invalid special-function calls.
+        # Build boolean mask for valid Fourier modes where formula is defined.
+        valid = (k != 0) & (kx != 0)
+
+        # Pre-allocate zero arrays for Fourier-domain shear perturbations
+        dtaux_t = np.zeros_like(hs, dtype=complex)
+        dtauy_t = np.zeros_like(hs, dtype=complex)
+
+        if np.any(valid):
+            # Extract valid-mode arrays
+            k_v  = k[valid]
+            kx_v = kx[valid]
+            ky_v = ky[valid]
+            hs_v = hs[valid]
+
+            # z0new can be scalar or array; index accordingly
+            if np.size(z0new) == 1:
+                z0_v = z0new
+            else:
+                z0_v = z0new[valid]
 
-        dtauy_t = hs * kx * ky / k_safe * 2 / ul**2 * \
-                    2. * np.sqrt(2.) * sigma * sc_kv(1., 2. * np.sqrt(2.) * sigma)
+            # compute sigma on valid modes
+            sigma_v = np.sqrt(1j * L * kx_v * z0_v / l)
 
-        # Zero out invalid regions (kx=0 or k=0) where formulation is not valid
-        invalid_mask = (k == 0) | (kx == 0)
-        dtaux_t[invalid_mask] = 0.
-        dtauy_t[invalid_mask] = 0.
-        
+            # Evaluate Bessel K functions on valid arguments only
+            kv0 = sc_kv(0., 2. * sigma_v)
+            kv1 = sc_kv(1., 2. * sigma_v)
+
+            # main x-direction perturbation (vectorized on valid indices)
+            term_x = -1. + (2. * np.log(l / z0_v) + (k_v**2) / (kx_v**2)) * sigma_v * (kv1 / kv0)
+            dtaux_v = hs_v * (kx_v**2) / k_v * 2. / ul**2 * term_x
+
+            # y-direction perturbation (also vectorized)
+            kv1_y = sc_kv(1., 2. * np.sqrt(2.) * sigma_v)
+            dtauy_v = hs_v * (kx_v * ky_v) / k_v * 2. / ul**2 * 2. * np.sqrt(2.) * sigma_v * (kv1_y)
+
+            # store back into full arrays (other entries remain zero)
+            dtaux_t[valid] = dtaux_v
+            dtauy_t[valid] = dtauy_v
+
+        # invalid modes remain 0 (physically reasonable for k=0 or kx=0)
         gc['dtaux'] = np.real(np.fft.ifft2(dtaux_t))
         gc['dtauy'] = np.real(np.fft.ifft2(dtauy_t))
 
+
+        
 
     def separation_shear(self, hsep):
         '''Reduces the computed wind shear perturbation below the