Refactor sweep functions and remove unused code

aeolis/avalanching.py

@@ -33,6 +33,8 @@
 # package modules
 from aeolis.utils import *
 
+from numba import njit
+
 # initialize logger
 logger = logging.getLogger(__name__)
 
@@ -60,10 +62,10 @@ def angele_of_repose(s,p):
     # comment Lisa: dependence on moisture content is not yet implemented 
     # Can we do something with theta dependent on vegetation cover (larger rhoveg = larger theta?)    
 
-    theta_stat = p['theta_stat']
+    # theta_stat = p['theta_stat']
     theta_dyn  = p['theta_dyn']
 
-    s['theta_stat'] = theta_stat
+    # s['theta_stat'] = theta_stat
     s['theta_dyn'] = theta_dyn
 
     return s
@@ -92,176 +94,150 @@ def avalanche(s, p):
     '''
 
     if p['process_avalanche']:
+        nx = p['nx'] + 1
+        ny = p['ny'] + 1
 
-        nx = p['nx']+1
-        ny = p['ny']+1
-
-        #parameters
-
-        tan_stat = np.tan(np.deg2rad(s['theta_stat']))
+        # parameters - only dynamic angle used in loop for now. 
+        # Static angle can be used for more complex criteria in later
+        # tan_stat = np.tan(np.deg2rad(s['theta_stat']))
         tan_dyn = np.tan(np.deg2rad(s['theta_dyn']))
-
 
-
-        E = 0.2
-
-        grad_h_down = np.zeros((ny,nx,4))
-        flux_down = np.zeros((ny,nx,4))
-        slope_diff = np.zeros((ny,nx))
-        grad_h = np.zeros((ny,nx))
+        E = 0.1
 
         max_iter_ava = p['max_iter_ava']
-
-        max_grad_h, grad_h, grad_h_down = calc_gradients(s['zb'], nx, ny, s['ds'], s['dn'], s['zne'])
-
-        s['gradh'] = grad_h.copy()
-
-        initiate_avalanche = (max_grad_h > tan_stat) 
-
-        if initiate_avalanche:
-
-            for i in range(0,max_iter_ava):
-
-                grad_h_down *= 0.
-                flux_down *= 0.
-                slope_diff *= 0.
-                grad_h *= 0.
-
-                max_grad_h, grad_h, grad_h_down = calc_gradients(s['zb'], nx, ny, s['ds'], s['dn'], s[ 'zne'])
-
-                if max_grad_h < tan_dyn:
-                    break
-
-                # Calculation of flux
-
-                grad_h_nonerod = (s['zb'] - s['zne']) / s['ds'] # HAS TO BE ADJUSTED!    
-
-                ix = np.logical_and(grad_h > tan_dyn, grad_h_nonerod > 0)
-                slope_diff[ix] = np.tanh(grad_h[ix]) - np.tanh(0.9*tan_dyn)    
-
-                ix = grad_h_nonerod < grad_h - tan_dyn 
-                slope_diff[ix] = np.tanh(grad_h_nonerod[ix])				                    
-
-                ix = grad_h != 0
-
-                if ny == 1:
-                    #1D interpretation
-                    flux_down[:,:,0][ix] = slope_diff[ix] * grad_h_down[:,:,0][ix] / grad_h[ix]
-                    flux_down[:,:,2][ix] = slope_diff[ix] * grad_h_down[:,:,2][ix] / grad_h[ix]
-
-                    # Calculation of change in bed level
-
-                    q_in = np.zeros((ny,nx))
-
-                    q_out = 0.5*np.abs(flux_down[:,:,0]) + 0.5*np.abs(flux_down[:,:,2])
-
-                    q_in[0,1:-1] =   0.5*(np.maximum(flux_down[0,:-2,0],0.) \
-                                        - np.minimum(flux_down[0,2:,0],0.) \
-                                        + np.maximum(flux_down[0,2:,2],0.) \
-                                        - np.minimum(flux_down[0,:-2,2],0.))
-                else:
-                    # 2D interpretation
-                    flux_down[:,:,0][ix] = slope_diff[ix] * grad_h_down[:,:,0][ix] / grad_h[ix]
-                    flux_down[:,:,1][ix] = slope_diff[ix] * grad_h_down[:,:,1][ix] / grad_h[ix]
-                    flux_down[:,:,2][ix] = slope_diff[ix] * grad_h_down[:,:,2][ix] / grad_h[ix]
-                    flux_down[:,:,3][ix] = slope_diff[ix] * grad_h_down[:,:,3][ix] / grad_h[ix]
-
-                    # Calculation of change in bed level
-
-                    q_in = np.zeros((ny,nx))
-
-                    q_out = 0.5*np.abs(flux_down[:,:,0]) + 0.5* np.abs(flux_down[:,:,1]) + 0.5*np.abs(flux_down[:,:,2]) + 0.5* np.abs(flux_down[:,:,3])
-
-                    q_in[1:-1,1:-1] =   0.5*(np.maximum(flux_down[1:-1,:-2,0],0.) \
-                                        - np.minimum(flux_down[1:-1,2:,0],0.) \
-                                        + np.maximum(flux_down[:-2,1:-1,1],0.) \
-                                        - np.minimum(flux_down[2:,1:-1,1],0.) \
-
-                                        + np.maximum(flux_down[1:-1,2:,2],0.) \
-                                        - np.minimum(flux_down[1:-1,:-2,2],0.) \
-                                        + np.maximum(flux_down[2:,1:-1,3],0.) \
-                                        - np.minimum(flux_down[:-2,1:-1,3],0.))
-
-                s['zb'] += E * (q_in - q_out)
 
+        # call the njit-compiled loop for performance
+        zb, grad_h = avalanche_loop(
+            s['zb'].copy(), s['zne'], s['ds'], s['dn'], nx, ny, E, max_iter_ava, tan_dyn
+            )
+
         # Ensure water level is up-to-date with bed level
-        s['zs'] = s['SWL'].copy()
+        s['zb'] = zb
+        s['gradh'] = grad_h
+        s['zs'] = s['SWL']
         ix = (s['zb'] > s['zs'])
         s['zs'][ix] = s['zb'][ix]
-    return s	    
 
+    return s
 
-def calc_gradients(zb, nx, ny, ds, dn, zne):
-    '''Calculates the downslope gradients in the bed that are needed for
-    avalanching module
-
-    Parameters
-    ----------
-
+@njit(cache=True)
+def avalanche_loop(zb, zne, ds, dn, nx, ny, E, max_iter_ava, tan_dyn):
+    # Rewritten to use explicit loops and avoid numpy boolean indexing
+    # Allocate temporaries once
+    grad_h_down = np.zeros((ny, nx, 2))
+    flux_down = np.zeros((ny, nx, 2))
+    slope_diff = np.zeros((ny, nx))
+    grad_h = np.zeros((ny, nx))
+    for it in range(max_iter_ava):
+        # Reset temporaries to zero
+        grad_h_down.fill(0)
+        flux_down.fill(0)
+        slope_diff.fill(0)
+        grad_h.fill(0)
+
+        # first calculate the downslope gradients to see if there is avalanching
+        # Compute downslope gradients grad_h_down (ny,nx,2), grad_h (ny,nx), and max_grad_h
+        # Initialize
+        max_grad_h = 0.0
+
+        # Directions: 0 => +X, 1 => +Y
+        for i in range(ny):
+            for j in range(nx):
+                # disable avalanching where zne >= zb
+                if zne[i, j] >= zb[i, j]:
+                    continue
+                else:
+                    center = zb[i, j]
+                    # +X direction
+                    g0 = 0.0
+                    # Handle boundaries: set gradient to zero at edges
+                    if j == 0 or j == nx - 1:
+                        grad_h_down[i, j, 0] = 0.0
+                    else:
+                        right = zb[i, j + 1]
+                        left = zb[i, j - 1]
+                        if not ((right > center) and (left > center)):
+                            if right > left:
+                                g0 = left - center
+                            else:
+                                g0 = center - right
+                        grad_h_down[i, j, 0] = g0 / ds[i, j]
+
+                    # +Y direction
+                    g1 = 0.0
+                    if i == 0 or i == ny - 1:
+                        grad_h_down[i, j, 1] = 0.0
+                    else:
+                        down = zb[i + 1, j]
+                        up = zb[i - 1, j]
+                        if not ((down > center) and (up > center)):
+                            if down > up:
+                                g1 = up - center
+                            else:
+                                g1 = center - down
+                        grad_h_down[i, j, 1] = g1 / dn[i, j]
+
+                # gradient magnitude and maximum
+                gh2 = grad_h_down[i, j, 0] * grad_h_down[i, j, 0] + grad_h_down[i, j, 1] * grad_h_down[i, j, 1]
+                gh = np.sqrt(gh2)
+                grad_h[i, j] = gh
+                # derive maximum slope
+                if gh > max_grad_h:
+                    max_grad_h = gh
+
+        # ok now the gradients are calculated
+        # these are max_grad_h, grad_h, grad_h_down
+        # check for stopping criterion
+        if max_grad_h < tan_dyn:
+            break       
 
-    Returns
-    -------
-    np.ndarray
-        Downslope gradients in 4 different directions (nx*ny, 4)
-    '''
-
-    grad_h_down = np.zeros((ny,nx,4))
-
-    # Calculation of slope (positive x-direction)
-    grad_h_down[:,1:-1,0] = zb[:,1:-1] - zb[:,2:] 
-    ix = zb[:,2:] > zb[:,:-2]
-    grad_h_down[:,1:-1,0][ix] = - (zb[:,1:-1][ix] - zb[:,:-2][ix])    
-    ix = np.logical_and(zb[:,2:]>zb[:,1:-1], zb[:,:-2]>zb[:,1:-1])
-    grad_h_down[:,1:-1,0][ix] = 0.
-
-    # Calculation of slope (positive y-direction)
-    grad_h_down[1:-1,:,1] = zb[1:-1,:] - zb[2:,:]    
-    ix = zb[2:,:] > zb[:-2,:]
-    grad_h_down[1:-1,:,1][ix] = - (zb[1:-1,:][ix] - zb[:-2,:][ix])    
-    ix = np.logical_and(zb[2:,:]>zb[1:-1,:], zb[:-2,:]>zb[1:-1,:])
-    grad_h_down[1:-1,:,1][ix] = 0.
-
-    # Calculation of slope (negative x-direction)
-    grad_h_down[:,1:-1,2] = zb[:,1:-1] - zb[:,:-2]    
-    ix = zb[:,:-2] > zb[:,2:]
-    grad_h_down[:,1:-1,2][ix] = - (zb[:,1:-1][ix] - zb[:,2:][ix])    
-    ix = np.logical_and(zb[:,:-2]>zb[:,1:-1], zb[:,2:]>zb[:,1:-1])
-    grad_h_down[:,1:-1,2][ix] = 0.
-
-    # Calculation of slope (negative y-direction)
-    grad_h_down[1:-1,:,3] = zb[1:-1,:] - zb[:-2,:]    
-    ix = zb[:-2,:] > zb[2:,:]
-    grad_h_down[1:-1,:,3][ix] = - (zb[1:-1,:][ix] - zb[2:,:][ix])    
-    ix = np.logical_and(zb[:-2,:]>zb[1:-1,:], zb[2:,:]>zb[1:-1,:])
-    grad_h_down[1:-1,:,3][ix] = 0.
-
-    if ny == 1:
-        #1D interpretation
-        grad_h_down[:,0,:]  = 0
-        grad_h_down[:,-1,:] = 0
-
-    else:
-        # 2D interpretation
-        grad_h_down[:,0,:]  = 0
-        grad_h_down[:,-1,:] = 0
-        grad_h_down[0,:,:]  = 0
-        grad_h_down[-1,:,:] = 0
+        # we continue to compute fluxes and update zb
 
-    grad_h_down[:,:,0] /= ds
-    grad_h_down[:,:,1] /= dn
-    grad_h_down[:,:,2] /= ds
-    grad_h_down[:,:,3] /= dn
+        # compute grad_h_nonerod and slope_diff per cell using explicit loops
+        for i in range(ny):
+            for j in range(nx):
+                if grad_h[i, j] > tan_dyn: 
+                    slope_diff[i, j] = np.tanh(grad_h[i, j]) - np.tanh(0.9 * tan_dyn)
+                    flux_down[i, j, 0] = slope_diff[i, j] * grad_h_down[i, j, 0]# / grad_h[i, j]
+                    flux_down[i, j, 1] = slope_diff[i, j] * grad_h_down[i, j, 1]# / grad_h[i, j]
+
+        # Build q_in and q_out from 2-component flux representation
+        f_x = flux_down[:, :, 0]
+        f_y = flux_down[:, :, 1]
 
-    grad_h2 = 0.5*grad_h_down[:,:,0]**2 + 0.5*grad_h_down[:,:,1]**2 + 0.5*grad_h_down[:,:,2]**2 + 0.5*grad_h_down[:,:,3]**2
+        # preserve sign: compute positive outgoing components per direction
+        out_east = np.maximum(f_x, 0.0)
+        out_south = np.maximum(f_y, 0.0)
 
-    if 0: #Sierd_com; to be changed in future release
-        ix = zb < zne + 0.005
-        grad_h2[ix] = 0.
+        # average with neighbor contributions at faces (keeps sign info)
+        out_north = np.maximum(-f_y, 0.0)
+        out_west  = np.maximum(-f_x, 0.0)
 
-    grad_h = np.sqrt(grad_h2)
+        q_out = out_east + out_west + out_south + out_north
 
-    max_grad_h = np.max(grad_h)
+        inc_west = np.zeros_like(f_x)
+        # from west neighbor (positive f_x of west cell) with periodic wrap
+        inc_west[:, 1:] = np.maximum(f_x[:, :-1], 0.0)
+        inc_west[:, 0] = np.maximum(f_x[:, -1], 0.0)
 
-    return max_grad_h, grad_h, grad_h_down
-
+        inc_east = np.zeros_like(f_x)
+        # from east neighbor (negative f_x of east cell) with periodic wrap
+        inc_east[:, :-1] = np.maximum(-f_x[:, 1:], 0.0)
+        inc_east[:, -1] = np.maximum(-f_x[:, 0], 0.0)
+
+        inc_north = np.zeros_like(f_y)
+        # from north neighbor (positive f_y of north cell) with periodic wrap
+        inc_north[1:, :] = np.maximum(f_y[:-1, :], 0.0)
+        inc_north[0, :] = np.maximum(f_y[-1, :], 0.0)
+
+        inc_south = np.zeros_like(f_y)
+        # from south neighbor (negative f_y of south cell) with periodic wrap
+        inc_south[:-1, :] = np.maximum(-f_y[1:, :], 0.0)
+        inc_south[-1, :] = np.maximum(-f_y[0, :], 0.0)
+
+        q_in = (inc_west + inc_east + inc_north + inc_south)
+
+        # update bed level without non-erodible layer       
+        zb += E * (q_in - q_out)
 
+    return zb, grad_h

aeolis/constants.py

@@ -98,9 +98,8 @@
         'ustarn',                           # [m/s] Component of shear velocity in y-direction by wind
         'ustar0',                           # [m/s] Initial shear velocity (without perturbation)
         'zsep',                             # [m] Z level of polynomial that defines the separation bubble
-        'hsep',                             # [m] Height of separation bubbel = difference between z-level of zsep and of the bed level zb
-        'theta_stat',                       # [degrees] Updated, spatially varying static angle of repose
-        'theta_dyn',                        # [degrees] Updated, spatially varying dynamic angle of repose
+        'hsep',                             # [m] Height of separation bubble = difference between z-level of zsep and of the bed level zb
+        'theta_dyn',                        # [degrees] spatially varying dynamic angle of repose for avalanching
         'rhoveg',                           # [-] Vegetation cover
         'drhoveg',                          # Change in vegetation cover
         'hveg',                             # [m] height of vegetation

aeolis/inout.py

@@ -118,13 +118,6 @@ def read_configfile(configfile, parse_files=True, load_defaults=True):
     if 'nsavetimes' in p and not p['nsavetimes']:
         p['nsavetimes'] = int(p['dzb_interval']/p['dt'])
 
-    # catch some incompatible parameter combinations.
-    if (p['ne_file'] is None) and (p['process_avalanche'] == True):
-        print('Please provide a valid ne_file path in the configuration file when using Avalanching.')
-        print('Code does not proceed until this is provided.')
-        print('Hint: If you do not have a ne_file, you can create one with all zeros, with the same dimensions as your grid.')        
-        exit("Exiting due to error")       
-
     return p
 
 
@@ -227,6 +220,13 @@ def check_configuration(p):
         logger.warning('Warning: the used roughness method (constant) defines the z0 as '
                        'k (z0 = k), this was implemented to ensure backward compatibility '
                        'and does not follow the definition of Nikuradse (z0 = k / 30).')
+
+    # check if steadystate solver is used with multiple sediment fractions
+    if p['solver'].lower() in ['steadystate', 'steadystatepieter']:
+        if len(p['grain_size']) > 1:
+            logger.log_and_raise('The steadystate solver is not compatible with multiple sediment fractions. '
+                                 'Please use a single sediment fraction or switch to a different solver (e.g., trunk or pieter).', 
+                                 exc=ValueError)
 
 
 def parse_value(val, parse_files=True, force_list=False):