Merge pull request #21 from firedrakeproject/TBendall/transport_plots

Transport and vertical slice plotting scripts

compressible_euler/moist_baroclinic_channel.py

@@ -100,7 +100,7 @@ def moist_baroclinic_channel(
     # I/O
     dirname = 'moist_baroclinic_channel'
     output = OutputParameters(
-        dirname=dirname, dumpfreq=dumpfreq, dump_nc=True,
+        dirname=dirname, dumpfreq=dumpfreq, dump_nc=True, dump_vtus=False,
         dumplist=['cloud_water']
     )
     diagnostic_fields = [Perturbation('theta')]

compressible_euler/moist_bryan_fritsch.py

@@ -51,6 +51,8 @@ def moist_bryan_fritsch(
     zc = 2000.                # vertical centre of bubble, in m
     rc = 2000.                # radius of bubble, in m
     Tdash = 2.0               # strength of temperature perturbation, in K
+    Tsurf = 320.0             # background theta_e value, in K
+    total_water = 0.02        # total moisture mixing ratio, in kg/kg
 
     # ------------------------------------------------------------------------ #
     # Our settings for this set up
@@ -77,7 +79,7 @@ def moist_bryan_fritsch(
 
     # I/O
     output = OutputParameters(
-        dirname=dirname, dumpfreq=dumpfreq, dumplist=['u']
+        dirname=dirname, dumpfreq=dumpfreq, dump_vtus=False, dump_nc=True
     )
     diagnostic_fields = [Theta_e(eqns)]
     io = IO(domain, output, diagnostic_fields=diagnostic_fields)
@@ -126,8 +128,6 @@ def moist_bryan_fritsch(
     dxp = dx(degree=(quadrature_degree))
 
     # Define constant theta_e and water_t
-    Tsurf = 320.0
-    total_water = 0.02
     theta_e = Function(Vt).assign(Tsurf)
     water_t = Function(Vt).assign(total_water)
 

compressible_euler/mountain_nonhydrostatic.py

@@ -11,14 +11,14 @@
 
 from firedrake import (
     as_vector, VectorFunctionSpace, PeriodicIntervalMesh, ExtrudedMesh,
-    SpatialCoordinate, exp, pi, cos, Function, conditional, Mesh, op2
+    SpatialCoordinate, exp, pi, cos, Function, conditional, Mesh, Constant
 )
 from gusto import (
-    Domain, CompressibleParameters, CompressibleSolver,
+    Domain, CompressibleParameters, CompressibleSolver, logger,
     CompressibleEulerEquations, OutputParameters, IO, SSPRK3,
     DGUpwind, SemiImplicitQuasiNewton, compressible_hydrostatic_balance,
-    SpongeLayerParameters, CourantNumber, ZComponent, Perturbation,
-    SUPGOptions, TrapeziumRule, remove_initial_w
+    SpongeLayerParameters, Exner, ZComponent, Perturbation,
+    SUPGOptions, TrapeziumRule, remove_initial_w, MaxKernel, MinKernel
 )
 
 mountain_nonhydrostatic_defaults = {
@@ -55,6 +55,9 @@ def mountain_nonhydrostatic(
     g = 9.80665              # acceleration due to gravity, in m/s^2
     cp = 1004.               # specific heat capacity at constant pressure
     sponge_mu = 0.15         # parameter for strength of sponge layer, in J/kg/K
+    exner_surf = 1.0         # maximum value of Exner pressure at surface
+    max_iterations = 10      # maximum number of hydrostatic balance iterations
+    tolerance = 1e-7         # tolerance for hydrostatic balance iteration
 
     # ------------------------------------------------------------------------ #
     # Our settings for this set up
@@ -98,12 +101,11 @@ def mountain_nonhydrostatic(
     )
 
     # I/O
-    dirname = 'nonhydrostatic_mountain'
     output = OutputParameters(
-        dirname=dirname, dumpfreq=dumpfreq, dumplist=['u']
+        dirname=dirname, dumpfreq=dumpfreq, dump_vtus=False, dump_nc=True
     )
     diagnostic_fields = [
-        CourantNumber(), ZComponent('u'), Perturbation('theta'),
+        Exner(parameters), ZComponent('u'), Perturbation('theta'),
         Perturbation('rho')
     ]
     io = IO(domain, output, diagnostic_fields=diagnostic_fields)
@@ -155,42 +157,65 @@ def mountain_nonhydrostatic(
     exner = Function(Vr)
     rho_b = Function(Vr)
 
-    exner_params = {'ksp_type': 'gmres',
-                    'ksp_monitor_true_residual': None,
-                    'pc_type': 'python',
-                    'mat_type': 'matfree',
-                    'pc_python_type': 'gusto.VerticalHybridizationPC',
-                    # Vertical trace system is only coupled vertically in columns
-                    # block ILU is a direct solver!
-                    'vert_hybridization': {'ksp_type': 'preonly',
-                                           'pc_type': 'bjacobi',
-                                           'sub_pc_type': 'ilu'}}
+    # Set up kernels to evaluate global minima and maxima of fields
+    min_kernel = MinKernel()
+    max_kernel = MaxKernel()
 
+    # First solve hydrostatic balance that gives Exner = 1 at bottom boundary
+    # This gives us a guess for the top boundary condition
+    bottom_boundary = Constant(exner_surf, domain=mesh)
+    logger.info(f'Solving hydrostatic with bottom Exner of {exner_surf}')
     compressible_hydrostatic_balance(
-        eqns, theta_b, rho_b, exner, top=True, exner_boundary=0.5,
-        params=exner_params
+        eqns, theta_b, rho_b, exner, top=False, exner_boundary=bottom_boundary
     )
 
-    def minimum(f):
-        fmin = op2.Global(1, [1000], dtype=float)
-        op2.par_loop(op2.Kernel("""
-    static void minify(double *a, double *b) {
-        a[0] = a[0] > fabs(b[0]) ? fabs(b[0]) : a[0];
-    }
-    """, "minify"), f.dof_dset.set, fmin(op2.MIN), f.dat(op2.READ))
-        return fmin.data[0]
-
-    p0 = minimum(exner)
+    # Solve hydrostatic balance again, but now use minimum value from first
+    # solve as the *top* boundary condition for Exner
+    top_value = min_kernel.apply(exner)
+    top_boundary = Constant(top_value, domain=mesh)
+    logger.info(f'Solving hydrostatic with top Exner of {top_value}')
     compressible_hydrostatic_balance(
-        eqns, theta_b, rho_b, exner, top=True, params=exner_params
+        eqns, theta_b, rho_b, exner, top=True, exner_boundary=top_boundary
     )
-    p1 = minimum(exner)
-    alpha = 2.*(p1 - p0)
-    beta = p1 - alpha
-    exner_top = (1. - beta) / alpha
+
+    max_bottom_value = max_kernel.apply(exner)
+
+    # Now we iterate, adjusting the top boundary condition, until this gives
+    # a maximum value of 1.0 at the surface
+    lower_top_guess = 0.9*top_value
+    upper_top_guess = 1.2*top_value
+    for i in range(max_iterations):
+        # If max bottom Exner value is equal to desired value, stop iteration
+        if abs(max_bottom_value - exner_surf) < tolerance:
+            break
+
+        # Make new guess by average of previous guesses
+        top_guess = 0.5*(lower_top_guess + upper_top_guess)
+        top_boundary.assign(top_guess)
+
+        logger.info(
+            f'Solving hydrostatic balance iteration {i}, with top Exner value '
+            + f'of {top_guess}'
+        )
+
+        compressible_hydrostatic_balance(
+            eqns, theta_b, rho_b, exner, top=True, exner_boundary=top_boundary
+        )
+
+        max_bottom_value = max_kernel.apply(exner)
+
+        # Adjust guesses based on new value
+        if max_bottom_value < exner_surf:
+            lower_top_guess = top_guess
+        else:
+            upper_top_guess = top_guess
+
+    logger.info(f'Final max bottom Exner value of {max_bottom_value}')
+
+    # Perform a final solve to obtain hydrostatically balanced rho
     compressible_hydrostatic_balance(
-        eqns, theta_b, rho_b, exner, top=True, exner_boundary=exner_top,
-        solve_for_rho=True, params=exner_params
+        eqns, theta_b, rho_b, exner, top=True, exner_boundary=top_boundary,
+        solve_for_rho=True
     )
 
     theta0.assign(theta_b)

compressible_euler/plotting/plot_moist_bryan_fritsch.py

@@ -0,0 +1,95 @@
+"""
+Plots the moist Bryan & Fritsch bubble test case.
+
+This plots:
+(a) theta_e @ t = 0 s, (b) theta_e @ t = 1000 s
+"""
+from os.path import abspath, dirname
+import matplotlib.pyplot as plt
+import numpy as np
+from netCDF4 import Dataset
+from tomplot import (
+    set_tomplot_style, tomplot_cmap, plot_contoured_field,
+    add_colorbar_fig, tomplot_field_title, extract_gusto_coords,
+    extract_gusto_field
+)
+
+test = 'moist_bryan_fritsch'
+
+# ---------------------------------------------------------------------------- #
+# Directory for results and plots
+# ---------------------------------------------------------------------------- #
+# When copying this example these paths need editing, which will usually involve
+# removing the abspath part to set directory paths relative to this file
+results_file_name = f'{abspath(dirname(__file__))}/../../results/{test}/field_output.nc'
+plot_stem = f'{abspath(dirname(__file__))}/../figures/{test}'
+
+# ---------------------------------------------------------------------------- #
+# Plot details
+# ---------------------------------------------------------------------------- #
+field_names = ['Theta_e', 'Theta_e']
+time_idxs = [0, -1]
+cbars = [False, True]
+
+# ---------------------------------------------------------------------------- #
+# General options
+# ---------------------------------------------------------------------------- #
+contours = np.linspace(315.0, 325.0, 21)
+remove_contour = 320.0
+colour_scheme = 'RdBu_r'
+field_label = r'$\theta_e$ (K)'
+contour_method = 'tricontour'
+xlims = [0., 10.]
+ylims = [0., 10.]
+
+# Things that are likely the same for all plots --------------------------------
+set_tomplot_style()
+data_file = Dataset(results_file_name, 'r')
+
+# ---------------------------------------------------------------------------- #
+# PLOTTING
+# ---------------------------------------------------------------------------- #
+fig, axarray = plt.subplots(1, 2, figsize=(15, 6), sharex='all', sharey='all')
+time_idx = 0
+
+for i, (ax, time_idx, field_name, cbar) in \
+        enumerate(zip(axarray.flatten(), time_idxs, field_names, cbars)):
+
+    # Data extraction ----------------------------------------------------------
+    field_data = extract_gusto_field(data_file, field_name, time_idx=time_idx)
+    coords_X, coords_Y = extract_gusto_coords(data_file, field_name)
+    time = data_file['time'][time_idx]
+
+    # Plot data ----------------------------------------------------------------
+    cmap, lines = tomplot_cmap(contours, colour_scheme, remove_contour=remove_contour)
+    cf, lines = plot_contoured_field(
+        ax, coords_X, coords_Y, field_data, contour_method, contours,
+        cmap=cmap, line_contours=lines
+    )
+
+    if cbar:
+        add_colorbar_fig(
+            fig, cf, field_label, ax_idxs=[i], location='right'
+        )
+    tomplot_field_title(
+        ax, f't = {time:.1f} s', minmax=True, field_data=field_data
+    )
+
+    # Labels -------------------------------------------------------------------
+    if i == 0:
+        ax.set_ylabel(r'$z$ (km)', labelpad=-20)
+        ax.set_ylim(ylims)
+        ax.set_yticks(ylims)
+        ax.set_yticklabels(ylims)
+
+    ax.set_xlabel(r'$x$ (km)', labelpad=-10)
+    ax.set_xlim(xlims)
+    ax.set_xticks(xlims)
+    ax.set_xticklabels(xlims)
+
+# Save figure ------------------------------------------------------------------
+fig.subplots_adjust(wspace=0.15)
+plot_name = f'{plot_stem}.png'
+print(f'Saving figure to {plot_name}')
+fig.savefig(plot_name, bbox_inches='tight')
+plt.close()