Fix ustarn calculation: initialization and FFT shear formula bugs (#265)

* Initial plan

* Fix ustars0 and ustarn0 initialization bug in wind.py

Fixed bug where ustars0 and ustarn0 were incorrectly set to ustar magnitude
instead of their respective directional components ustars and ustarn.

Co-authored-by: Sierd <[email protected]>

* input files for debugging

* Fix missing division in dtauy FFT shear calculation

The dtauy_t formula in the FFT shear method was missing the division by
sc_kv(0., 2.*sqrt(2.)*sigma), causing incorrect y-direction shear stress
perturbations. This resulted in non-zero ustarn values even when the bed
had no y-direction variability.

The formula now matches the structure of dtaux_t calculation.

Co-authored-by: Sierd <[email protected]>

* Fix frequency array alignment in FFT shear calculation

The kx and ky frequency arrays were misaligned with the FFT output. The code
was creating frequency arrays with fftfreq(n+1)[1:] which skips the DC
component, but the FFT includes all frequencies. This caused incorrect
mapping between frequencies and FFT coefficients, leading to wrong shear
perturbations.

Fixed by using fftfreq(n) which properly matches the FFT output dimensions
and includes the DC component at the correct index.

Co-authored-by: Sierd <[email protected]>

* Handle division by zero at DC component in FFT shear

The previous fix included the DC component (kx=0, ky=0) in the frequency
arrays, which caused division by zero errors in the shear calculations.

Fixed by:
- Using safe division with np.where to replace zeros with 1.0 temporarily
- Explicitly setting DC component of perturbations to 0 after calculation
- Applying same fix to filter_highfrequencies function

The DC component represents the mean value and doesn't contribute to
perturbations, so setting it to zero is physically correct.

Co-authored-by: Sierd <[email protected]>

* Use masked computation to avoid invalid values in FFT shear

The previous fix still allowed invalid values when kx=0 (but ky≠0) because
sigma depends on kx. When kx=0, sigma=0 which causes issues in Bessel
function ratios.

New approach:
- Create a mask for valid computations: (k > 0) & (abs(kx) > 0)
- Initialize perturbation arrays with zeros
- Only compute perturbations for valid frequency combinations
- All invalid combinations (kx=0 or k=0) automatically get zero perturbation

This is physically correct: perturbations at kx=0 (pure y-direction
frequencies) should be zero for the x-direction flow formulation.

Co-authored-by: Sierd <[email protected]>

* Simplify FFT shear edge case handling

Simplified the masked computation approach to a cleaner implementation:
- Use np.where for safe division (replace zeros with 1.0 temporarily)
- Compute formulas normally with safe arrays
- Apply invalid_mask at the end to zero out problematic regions

This achieves the same result with much simpler, more readable code.

Co-authored-by: Sierd <[email protected]>

* Update shear.py

reverted the suggested division in dtauy_t calculation

* deleted wind tests

* shear patch

---------

Co-authored-by: copilot-swe-agent[bot] <[email protected]>
Co-authored-by: Sierd <[email protected]>
Co-authored-by: Sierd <[email protected]>

Author: Copilot

aeolis/shear.py

@@ -544,8 +544,8 @@ def compute_shear(self, u0, nfilter=(1., 2.)):
             return
 
         ny, nx = gc['z'].shape
-        kx, ky = np.meshgrid(2. * np.pi * np.fft.fftfreq(nx+1, gc['dx'])[1:],
-                             2. * np.pi * np.fft.fftfreq(ny+1, gc['dy'])[1:])
+        kx, ky = np.meshgrid(2. * np.pi * np.fft.fftfreq(nx, gc['dx']),
+                             2. * np.pi * np.fft.fftfreq(ny, gc['dy']))
 
         hs = np.fft.fft2(gc['z'])
         hs = self.filter_highfrequenies(kx, ky, hs, nfilter)
@@ -576,25 +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
-        
-        dtaux_t = hs * kx**2 / k * 2 / ul**2 * \
-                  (-1. + (2. * np.log(l/z0new) + k**2/kx**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 * 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)
 
-        
+            # 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 
@@ -668,8 +701,11 @@ def filter_highfrequenies(self, kx, ky, hs, nfilter=(1, 2)):
         if nfilter is not None:
             n1 = np.min(nfilter)
             n2 = np.max(nfilter)
-            px = 2 * np.pi / self.cgrid['dx'] / np.abs(kx)
-            py = 2 * np.pi / self.cgrid['dy'] / np.abs(ky)
+            # Avoid division by zero at DC component (kx=0, ky=0)
+            kx_safe = np.where(kx == 0, 1.0, kx)
+            ky_safe = np.where(ky == 0, 1.0, ky)
+            px = 2 * np.pi / self.cgrid['dx'] / np.abs(kx_safe)
+            py = 2 * np.pi / self.cgrid['dy'] / np.abs(ky_safe)
             s1 =  n1 / np.log(1. / .01 - 1.)
             s2 = -n2 / np.log(1. / .99 - 1.)
             f1 = 1. / (1. + np.exp(-(px + n1 - n2) / s1))
@@ -882,4 +918,4 @@ def interpolate(self, x, y, z, xi, yi, z0):
 
 
 
-        
+        

aeolis/wind.py

@@ -148,8 +148,8 @@ def interpolate(s, p, t):
     s = velocity_stress(s,p)
 
     s['ustar0'] = s['ustar'].copy()
-    s['ustars0'] = s['ustar'].copy()
-    s['ustarn0'] = s['ustar'].copy()
+    s['ustars0'] = s['ustars'].copy()
+    s['ustarn0'] = s['ustarn'].copy()
 
     s['tau0'] = s['tau'].copy()
     s['taus0'] = s['taus'].copy()