Output visual output (#197)

* Clean-up of the Visualization function in inout.py

* Added the variable "masstop"

Added the option to export the mass of the toplayer exclusively to the *.nc file

* Add SedTRAILS output

Added the option the output SedTRAILS-required output for each timestep (MAT-files)

aeolis/bed.py

@@ -323,6 +323,9 @@ def update(s, p):
     # reshape mass matrix
     s['mass'] = m.reshape((ny+1,nx+1,nl,nf))
 
+    # Store toplayer of 'mass' variable (ilayer = 0)
+    s['masstop'][:,:,:] = s['mass'][:,:,0,:].copy()
+
     # update bathy
     if p['process_bedupdate']:
 

aeolis/constants.py

@@ -138,7 +138,10 @@
         'u',                                # [m/s] Mean horizontal saltation velocity in saturated state
         'us',                               # [m/s] Component of the saltation velocity in x-direction
         'un',                               # [m/s] Component of the saltation velocity in y-direction
+        'usST',                            # [NEW] [m/s] Component of the saltation velocity in x-direction for SedTRAILS
+        'unST',                            # [NEW] [m/s] Component of the saltation velocity in y-direction for SedTRAILS
         'u0',
+        'masstop',                          # [kg/m^2] Sediment mass in bed toplayer, only stored for output
     ),
     ('ny','nx','nlayers') : (
         'thlyr',                            # [m] Bed composition layer thickness
@@ -184,6 +187,8 @@
     'process_dune_erosion'          : False,              # Enable the process of wave-driven dune erosion
     'process_seepage_face'          : False,              # Enable the process of groundwater seepage (NB. only applicable to positive beach slopes)
     'visualization'                 : False,              # Boolean for visualization of model interpretation before and just after initialization
+    'output_sedtrails'              : False,              # NEW! [T/F] Boolean to see whether additional output for SedTRAILS should be generated
+    'nfraction_sedtrails'           : 0,                  # [-] Index of selected fraction for SedTRAILS (0 if only one fraction)
     'xgrid_file'                    : None,               # Filename of ASCII file with x-coordinates of grid cells
     'ygrid_file'                    : None,               # Filename of ASCII file with y-coordinates of grid cells
     'bed_file'                      : None,               # Filename of ASCII file with bed level heights of grid cells

aeolis/inout.py

@@ -35,6 +35,7 @@
 from webbrowser import UnixBrowser
 import numpy as np
 from matplotlib import pyplot as plt
+from scipy.io import savemat
 
 # package modules
 from aeolis.utils import *
@@ -493,14 +494,17 @@ def visualize_spatial(s, p):
     fig, axs = plt.subplots(5, 3)
     pcs = [[None for _ in range(3)] for _ in range(5)]
 
+    # In the plotting below, prevent the UserWarning: The input coordinates to pcolormesh are interpreted as cell centers, but are not monotonically increasing or decreasing (...)
+    import warnings 
+    warnings.filterwarnings("ignore", category=UserWarning)
+
     # Plotting colormeshes
     if p['ny'] > 0:
         pcs[0][0] = axs[0,0].pcolormesh(x, y, s['zb'], cmap='viridis')
         pcs[0][1] = axs[0,1].pcolormesh(x, y, s['zne'], cmap='viridis')
         pcs[0][2] = axs[0,2].pcolormesh(x, y, s['rhoveg'], cmap='Greens', clim= [0, 1])
         pcs[1][0] = axs[1,0].pcolormesh(x, y, s['uw'], cmap='plasma')
         pcs[1][1] = axs[1,1].pcolormesh(x, y, s['ustar'], cmap='plasma')
-        # pcs[1][2] = axs[1,2].pcolormesh(x, y, s['tau'], cmap='plasma')
         pcs[1][2] = axs[1,2].pcolormesh(x, y, s['u'][:, :, 0], cmap='plasma')
         pcs[2][0] = axs[2,0].pcolormesh(x, y, s['moist'], cmap='Blues', clim= [0, 0.4])
         pcs[2][1] = axs[2,1].pcolormesh(x, y, s['gw'], cmap='viridis')
@@ -528,11 +532,13 @@ def visualize_spatial(s, p):
         pcs[4][1] = axs[4,1].scatter(x, y, c=tide_mask_add, cmap='binary', clim= [0, 1])
         pcs[4][2] = axs[4,2].scatter(x, y, c=wave_mask_add, cmap='binary', clim= [0, 1])
 
+    # Re-allow the UserWarning
+    warnings.filterwarnings("default", category=UserWarning)
+
     # Quiver for vectors
     skip = 10
     axs[1,0].quiver(x[::skip, ::skip], y[::skip, ::skip], s['uws'][::skip, ::skip], s['uwn'][::skip, ::skip])
     axs[1,1].quiver(x[::skip, ::skip], y[::skip, ::skip], s['ustars'][::skip, ::skip], s['ustarn'][::skip, ::skip])
-    # axs[1,2].quiver(x[::skip, ::skip], y[::skip, ::skip], s['taus'][::skip, ::skip], s['taun'][::skip, ::skip])
     axs[1,2].quiver(x[::skip, ::skip], y[::skip, ::skip], s['us'][::skip, ::skip, 0], s['un'][::skip, ::skip, 0])
 
     # Adding titles to the plots
@@ -541,7 +547,6 @@ def visualize_spatial(s, p):
     axs[0,2].set_title('Vegetation density, rhoveg (-)')
     axs[1,0].set_title('Wind velocity, uw (m/s)')
     axs[1,1].set_title('Shear velocity, ustar (m/s)')
-    # axs[1,2].set_title('Shear stress, tau (N/m2)')
     axs[1,2].set_title('Grain velocity, u (m/s)')
     axs[2,0].set_title('Soil moisture content, (-)')
     axs[2,1].set_title('Ground water level, gw (m)')
@@ -573,4 +578,29 @@ def visualize_spatial(s, p):
     fig.savefig('figure_params_initialization.png', dpi=300)
     plt.close()
 
-    return 
+    return 
+
+def output_sedtrails(s, p):
+    '''Create additional output for SedTRAILS and save as mat-files.
+    Chosen for seperate files, such that only relevant (Ct > 0) cells 
+    are exported for memory and speed efficiency''' 
+
+    nf = p['nfraction_sedtrails']
+
+    # Speed and concetration: Only for the first fraction now
+    x = s['x'].flatten()
+    y = s['y'].flatten()
+    us = s['usST'][:,:,nf].flatten()
+    un = s['unST'][:,:,nf].flatten()
+    pickup = s['pickup'][:,:,nf].flatten()
+    dzb = s['dzb'].flatten() # Store the bed level change (AEOLIAN ONLY) for every timestep
+
+    os.makedirs('sedtrails_output', exist_ok=True) 
+
+    time = p['_time']
+    if time == 0: # Save the x and y coordinates only once to save memory
+        mdic = {'x': x, 'y': y, 'us': us, 'un': un, 'dzb': dzb, 'pickup': pickup}
+        savemat(os.path.join('sedtrails_output', str(int(time)).zfill(12) + '.mat'), mdic)
+    else:
+        mdic = {'us': us, 'un': un, 'dzb': dzb, 'pickup': pickup}
+        savemat(os.path.join('sedtrails_output', str(int(time)).zfill(12) + '.mat'), mdic)

aeolis/transport.py

@@ -73,6 +73,7 @@ def duran_grainspeed(s, p):
     ustar0 = s['ustar0']
     uth = s['uth0'] # uth0 or uth???
     uth0 = s['uth0'] 
+    uthST = s['uth'] 
 
     # Wind input for filling up ets/ets (udir), where ustar == 0
     uw = s['uw']
@@ -97,6 +98,8 @@ def duran_grainspeed(s, p):
     u0 = np.zeros(uth.shape)
     us = np.zeros(uth.shape)
     un = np.zeros(uth.shape)
+    usST = np.zeros(uth.shape)
+    unST = np.zeros(uth.shape)
     u_approx = np.zeros(uth.shape)
     us_approx = np.zeros(uth.shape)
     un_approx = np.zeros(uth.shape)
@@ -209,7 +212,15 @@ def solve_u(u_i: complex, veff_i: complex, uf_i: float, alpha_i: float, dh_i: co
         else:
             logger.error('Grainspeed method not found!')
 
-    return u0, us, un, u
+        # For SedTRAILS: Set grainspeed to 0 whenever uth > ustar
+        usST[:,:,i] = us[:,:,i]
+        unST[:,:,i] = un[:,:,i]
+
+        ix_no_speed = (uthST[:,:,i] > ustar[:,:,i])
+        usST[ix_no_speed, i] *= 0.
+        unST[ix_no_speed, i] *= 0.
+
+    return u0, us, un, u, usST, unST
 
 
 
@@ -292,11 +303,13 @@ def equilibrium(s, p):
 
         if p['method_grainspeed']=='duran' or p['method_grainspeed']=='duran_full':
             #the syntax inside grainspeed needs to be cleaned up
-            u0, us, un, u = duran_grainspeed(s,p)
+            u0, us, un, u, usST, unST = duran_grainspeed(s,p)
             s['u0'] = u0
             s['us'] = us
             s['un'] = un
             s['u']  = u
+            s['usST'][:] = usST
+            s['unST'][:] = unST
 
         elif p['method_grainspeed']=='windspeed':
             s['u0'] = s['uw'][:,:,np.newaxis].repeat(nf, axis=2)