Merge pull request #18 from firedrakeproject/TBendall/templates

Templates for case study and plotting

.github/pull_request_template.md

@@ -1,16 +1,26 @@
 # Add or update a Gusto case study
-Here is a checklist of things that must be done before a new case study
-will be included:
-- [ ] Readme with:
-    - [ ] Instructions for running the case study.
-    - [ ] Instructions for reproducing included checkpoint and image.
-    - [ ]
-- [ ] Case study as a function
-- [ ] Test added to the root directory
-- [ ] Output checkpoint (or link to GH-LFS of the checkpoint)
-- [ ] Image of the case study
-- [ ]
+Here is a checklist of things that should be done to add a new case study to
+the repository:
+- [ ] The case study has been prepared from the case studies template in `templates/template_case_study.py`. This ensures that the case study:
+  - [ ] begins with documentation of the case
+  - [ ] includes a dictionary of default argument values
+  - [ ] is run through a function
+  - [ ] follows the standard order of sections:
+    1. test case parameters
+    2. set up of model objects
+    3. initial conditions
+    4. run
+  - [ ] includes a `__main__` routine with arg-parsing of command line arguments
+- [ ] The case study has a quick-to-run test form in the relevant `test_*.py` file, so that it will be run as part of CI
+- [ ] A plotting script has been added to the relevant `plotting` directory, with a name that matches the case study script
+- [ ] Neat figures have been added to the relevant `figures` directory, with names that match the case study script
 
+# Add or update a plotting script
+Here is a checklist of things that should be done to add a new plotting script to the repository:
+- [ ] The plotting script has been prepared from the template in `templates/template_plotting_script.py` or another acceptable plotting script
+- [ ] The plot follows the Good Plot Guide in [`tomplot/good_plot_guide.md`](https://github.com/tommbendall/tomplot/blob/main/good_plot_guide.md)
+- [ ] Relevant initial and final fields are plotted
+- [ ] The figures produced have been added to the repository
 
 <!--
 Here is a comment that can include verbose instructions that will not

README.md

@@ -1 +1,18 @@
-A repository for holding scripts setting up working case studies for Gusto.
+# Gusto Case Studies
+
+Welcome to the Gusto Case Studies repository!
+
+This stores a collection of Python scripts which implement standard (and more exotic) test cases for geophysical fluids, particularly those used in the development of the dynamical cores used in numerical weather prediction models.
+
+The test cases are run using the [Gusto](https://www.firedrakeproject.org/gusto/) code library, which uses compatible finite element methods to discretise the fluid's equations of motion.
+Gusto is built upon the [Firedrake](https://www.firedrakeproject.org/) software, which provides the finite element infrastructure used by Gusto.
+
+Case studies are organised by governing equation. Each case study is accompanied by a plotting script, and reference figures.
+
+Our continuous integration is intended to ensure that this repository runs at the Gusto head-of-trunk.
+
+## Visualisation
+
+Gusto can produce output in two formats:
+- VTU files, which can be viewed with the [Paraview](https://www.paraview.org/) software
+- netCDF files, which has data that can be plotted using standard python packages such as matplotlib. We suggest using the [tomplot](https://github.com/tommbendall/tomplot) python library, which contains several routines to simplify the plotting of Gusto output.

templates/template_case_study.py

@@ -0,0 +1,178 @@
+"""
+The <test_name> test of <authors>, <year>:
+``<paper_title>'', <journal>.
+
+<Description of test>
+
+The setup here uses <details of our configuration>.
+"""
+from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
+
+from firedrake import (
+    SpatialCoordinate, Constant, pi, cos, exp
+)
+from gusto import (
+    Domain, IO, OutputParameters,
+    GeneralCubedSphereMesh, ZonalComponent, MeridionalComponent,
+    lonlatr_from_xyz, xyz_vector_from_lonlatr, great_arc_angle,
+    AdvectionEquation, PrescribedTransport,
+    SSPRK3, DGUpwind
+)
+
+# Dictionary containing default values of arguments to case study
+""" TO REMOVE
+Should include resolution, dt, tmax, dumpfreq and dirname
+May also include other important configuration options
+"""
+test_name_defaults = {
+    'ncells_1d': 8,
+    'dt': 1.0,
+    'tmax': 10.0,
+    'dumpfreq': 10,
+    'dirname': 'test_name'
+}
+
+
+def test_name(
+        ncells_1d=test_name_defaults['ncells_1d'],
+        dt=test_name_defaults['dt'],
+        tmax=test_name_defaults['tmax'],
+        dumpfreq=test_name_defaults['dumpfreq'],
+        dirname=test_name_defaults['dirname']
+):
+
+    # ------------------------------------------------------------------------ #
+    # Parameters for test case
+    # ------------------------------------------------------------------------ #
+
+    """ TO REMOVE
+    These should contain inline comments to make the parameters clear
+    """
+    tau = 200.         # time period relating to reversible flow, in s
+    radius = 6371220.  # radius of sphere, in m
+    lamda_c = 0.       # central longtiude of transported blob, in rad
+    theta_c = -pi/6.   # central latitude of transported blob, in rad
+    F0 = 3.            # max magnitude of transported blob, in kg/kg
+    r0 = 0.25          # width parameter of transported blob, dimensionless
+
+    # ------------------------------------------------------------------------ #
+    # Our settings for this set up
+    # ------------------------------------------------------------------------ #
+
+    degree = 1
+    hdiv_family = 'BDM'
+
+    # ------------------------------------------------------------------------ #
+    # Set up model objects
+    # ------------------------------------------------------------------------ #
+
+    # Domain
+    mesh = GeneralCubedSphereMesh(radius, ncells_1d, degree=2)
+    xyz = SpatialCoordinate(mesh)
+    domain = Domain(mesh, dt, hdiv_family, degree)
+
+    # Equation
+    V = domain.spaces("DG")
+    eqn = AdvectionEquation(domain, V, "F")
+
+    # I/O and diagnostics
+    output = OutputParameters(
+        dirname=dirname, dumpfreq=dumpfreq, dump_nc=True, dump_vtus=False
+    )
+
+    diagnostic_fields = [ZonalComponent('u'), MeridionalComponent('u')]
+
+    io = IO(domain, output, diagnostic_fields=diagnostic_fields)
+
+    # Details of transport
+    transport_scheme = SSPRK3(domain)
+    transport_method = DGUpwind(eqn, "F")
+
+    # Transporting wind ------------------------------------------------------ #
+    """ TO REMOVE
+    Any prescribed wind for transport tests should be added here
+    """
+    def u_t(t):
+        _, theta, _ = lonlatr_from_xyz(xyz[0], xyz[1], xyz[2])
+        umax = 2*pi*radius / tau
+        u_zonal = umax*cos(theta)
+
+        return xyz_vector_from_lonlatr(u_zonal, Constant(0.0), Constant(0.0), xyz)
+
+    # Physics parametrisation ------------------------------------------------ #
+    """ TO REMOVE
+    Any physics parametrisations need defining here
+    """
+
+    # Time stepper
+    stepper = PrescribedTransport(
+        eqn, transport_scheme, io, transport_method,
+        prescribed_transporting_velocity=u_t
+    )
+
+    # ------------------------------------------------------------------------ #
+    # Initial conditions
+    # ------------------------------------------------------------------------ #
+
+    # Initialise the field to be transported
+    lamda, theta, _ = lonlatr_from_xyz(xyz[0], xyz[1], xyz[2])
+    dist = great_arc_angle(lamda, theta, lamda_c, theta_c)
+    F_init_expr = F0*exp(-(dist/r0)**2)
+
+    # Set fields
+    u0 = stepper.fields("u")
+    F0 = stepper.fields("F")
+    u0.project(u_t(0))
+    F0.interpolate(F_init_expr)
+
+    # ------------------------------------------------------------------------ #
+    # Run
+    # ------------------------------------------------------------------------ #
+
+    stepper.run(t=0, tmax=tmax)
+
+
+# ---------------------------------------------------------------------------- #
+# MAIN
+# ---------------------------------------------------------------------------- #
+
+
+if __name__ == "__main__":
+
+    parser = ArgumentParser(
+        description=__doc__,
+        formatter_class=ArgumentDefaultsHelpFormatter
+    )
+    parser.add_argument(
+        '--ncells_1d',
+        help="The number of cells in one dimension",
+        type=int,
+        default=test_name_defaults['ncells_1d']
+    )
+    parser.add_argument(
+        '--dt',
+        help="The time step in seconds.",
+        type=float,
+        default=test_name_defaults['dt']
+    )
+    parser.add_argument(
+        "--tmax",
+        help="The end time for the simulation in seconds.",
+        type=float,
+        default=test_name_defaults['tmax']
+    )
+    parser.add_argument(
+        '--dumpfreq',
+        help="The frequency at which to dump field output.",
+        type=int,
+        default=test_name_defaults['dumpfreq']
+    )
+    parser.add_argument(
+        '--dirname',
+        help="The name of the directory to write to.",
+        type=str,
+        default=test_name_defaults['dirname']
+    )
+    args, unknown = parser.parse_known_args()
+
+    test_name(**vars(args))

templates/template_plotting_script.py

@@ -0,0 +1,130 @@
+"""
+Plots the <my_test_name> test case
+
+This plots:
+(a) <field1> @ t = 0 s, (b) <field1> @ t = <final_time> s
+(d) <field2> @ t = 0 s, (e) <field2> @ t = <final_time> 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 = 'my_test_name'
+
+# ---------------------------------------------------------------------------- #
+# 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
+# ---------------------------------------------------------------------------- #
+subplots_x = 2
+subplots_y = 2
+field_names = ['field1', 'field1',
+               'field2', 'field2']
+time_idxs = [0, -1,
+             0, -1]
+cbars = [False, True,
+         False, True]
+
+# ---------------------------------------------------------------------------- #
+# General options
+# ---------------------------------------------------------------------------- #
+# First field
+f1_contours = np.linspace(0.0, 3.6, 37)
+f1_colour_scheme = 'RdBu_r'
+f1_contour_to_remove = None
+f1_field_label = r'$f_1$ (kg m$^{-3}$)'
+
+# Second field
+f2_contours = np.linspace(0.01, 0.07, 13)
+f2_colour_scheme = 'Purples'
+f2_contour_to_remove = 0.02
+f2_field_label = r'$f_2$ (kg kg$^{-1}$)'
+
+# General
+contour_method = 'tricontour'
+xmax = 2.0
+ymax = 2.0
+xlims = [0., xmax]
+ylims = [0., ymax]
+
+# Things that are likely the same for all plots --------------------------------
+set_tomplot_style()
+assert len(field_names) == subplots_x*subplots_y
+data_file = Dataset(results_file_name, 'r')
+
+# ---------------------------------------------------------------------------- #
+# PLOTTING
+# ---------------------------------------------------------------------------- #
+fig, axarray = plt.subplots(
+    subplots_y, subplots_x, figsize=(16, 12), sharex='all', sharey='all'
+)
+
+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]
+
+    # Select options for each field --------------------------------------------
+    if field_name == 'field1':
+        contours = f1_contours
+        colour_scheme = f1_colour_scheme
+        field_label = f1_field_label
+        contour_to_remove = f1_contour_to_remove
+
+    elif field_name == 'field2':
+        contours = f2_contours
+        colour_scheme = f2_colour_scheme
+        field_label = f2_field_label
+        contour_to_remove = f2_contour_to_remove
+
+    cmap, lines = tomplot_cmap(
+        contours, colour_scheme, remove_contour=contour_to_remove
+    )
+
+    # Plot data ----------------------------------------------------------------
+    cf, _ = 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 % subplots_x == 0:
+        ax.set_ylabel(r'$z$ (km)', labelpad=-20)
+        ax.set_ylim(ylims)
+        ax.set_yticks(ylims)
+        ax.set_yticklabels(ylims)
+
+    if i > (subplots_y - 1)*subplots_x - 1:
+        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.25)
+plot_name = f'{plot_stem}.png'
+print(f'Saving figure to {plot_name}')
+fig.savefig(plot_name, bbox_inches='tight')
+plt.close()