diff --git a/plotting/compressible_euler/plot_dcmip_3_1_gravity_wave.py b/plotting/compressible_euler/plot_dcmip_3_1_gravity_wave.py new file mode 100644 index 00000000..1575bfab --- /dev/null +++ b/plotting/compressible_euler/plot_dcmip_3_1_gravity_wave.py @@ -0,0 +1,277 @@ +""" +Plots the DCMIP 3-1 gravity wave test case. + +This plots the initial conditions @ t = 0 s, with +(a) zonal wind, (b) theta (c) theta perturbation: all on a lon-lat slice, +(d) zonal wind, (e) theta (f) theta perturbation: on a lat-z slice, + +and the final state @ t = 3600 s, with +(a) theta perturbation on a lon-lat slice, +(b) theta perturbation on a lon-z slice. +""" +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_ax, tomplot_field_title, extract_gusto_coords, + extract_gusto_field, reshape_gusto_data, extract_gusto_vertical_slice, + regrid_vertical_slice +) + +test = 'dcmip_3_1_gravity_wave' + +# ---------------------------------------------------------------------------- # +# 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/compressible_euler/{test}' + +# ---------------------------------------------------------------------------- # +# Initial plot details +# ---------------------------------------------------------------------------- # +init_field_names = ['u_zonal', 'theta', 'theta_perturbation', + 'u_zonal', 'theta', 'theta_perturbation'] +init_colour_schemes = ['YlOrBr', 'PuRd', 'OrRd', + 'YlOrBr', 'PuRd', 'OrRd',] +init_field_labels = [r'$u$ (m s$^{-1}$)', r'$\theta$ (K)', r'$\Delta\theta$ (K)', + r'$u$ (m s$^{-1}$)', r'$\theta$ (K)', r'$\Delta\theta$ (K)'] +init_contours = [np.linspace(0, 25, 11), + np.linspace(300, 335, 13), + np.linspace(0.0, 1.0, 11), + np.linspace(0, 25, 11), + np.linspace(300, 335, 13), + np.linspace(0.0, 1.0, 11)] +init_contours_to_remove = [None, None, None, None, None, None] +init_slice_along = ['z', 'z', 'z', 'lon', 'lon', 'lon'] + +# ---------------------------------------------------------------------------- # +# Final plot details +# ---------------------------------------------------------------------------- # +final_field_names = ['theta_perturbation', 'theta_perturbation'] +final_colour_schemes = ['RdBu_r', 'RdBu_r'] +final_field_labels = [r'$\Delta\theta$ (K)', r'$\Delta\theta$ (K)'] +final_contours = [np.linspace(-0.1, 0.1, 21), + np.linspace(-0.1, 0.1, 21)] +final_contours_to_remove = [0.0, 0.0] +final_slice_along = ['z', 'lat'] + +# ---------------------------------------------------------------------------- # +# General options +# ---------------------------------------------------------------------------- # +contour_method = 'tricontour' +lon_lims = [-180, 180] +lat_lims = [-90, 90] +z_lims = [0, 10] +level = 5 +slice_at_lon = 120.0 +slice_at_lat = 0.0 + +# 1D grids for vertical regridding +coords_lon_1d = np.linspace(-180, 180, 50) +coords_lat_1d = np.linspace(-90, 90, 50) +# Dictionary to hold plotting grids -- keys are "slice_along" values +plotting_grid = {'lat': coords_lon_1d, 'lon': coords_lat_1d} + +cbar_format = {'u_zonal': '1.0f', + 'theta': '1.0f', + 'theta_perturbation': '1.1f'} + +# Things that are likely the same for all plots -------------------------------- +set_tomplot_style() +data_file = Dataset(results_file_name, 'r') + +# ---------------------------------------------------------------------------- # +# INITIAL PLOTTING +# ---------------------------------------------------------------------------- # +fig, axarray = plt.subplots(2, 3, figsize=(18, 12), sharex='row', sharey='row') +time_idx = 0 + +for i, (ax, field_name, field_label, colour_scheme, contours, + to_remove, slice_along) in \ + enumerate(zip(axarray.flatten(), init_field_names, init_field_labels, + init_colour_schemes, init_contours, + init_contours_to_remove, init_slice_along)): + + # Data extraction ---------------------------------------------------------- + time = data_file['time'][time_idx] + + if slice_along == 'z': + field_full = extract_gusto_field(data_file, field_name, time_idx) + coords_X_full, coords_Y_full, coords_Z_full = \ + extract_gusto_coords(data_file, field_name) + + # Reshape + field_full, coords_X_full, coords_Y_full, _ = \ + reshape_gusto_data(field_full, coords_X_full, + coords_Y_full, coords_Z_full) + + # Take level for a horizontal slice + field_data = field_full[:, level] + # Abuse of names for coord variables but simplifies code below + coords_X = coords_X_full[:, level] + coords_Y = coords_Y_full[:, level] + + else: + orig_field_data, orig_coords_X, orig_coords_Y, orig_coords_Z = \ + extract_gusto_vertical_slice( + data_file, field_name, time_idx, + slice_along=slice_along, slice_at=slice_at_lon + ) + + # Slices need regridding as points don't cleanly live along lon or lat = 0.0 + field_data, coords_X, coords_Y = \ + regrid_vertical_slice( + plotting_grid[slice_along], slice_along, slice_at_lon, + orig_coords_X, orig_coords_Y, orig_coords_Z, orig_field_data + ) + # Scale coordinates + coords_Y /= 1000. + + # Plot data ---------------------------------------------------------------- + cmap, lines = tomplot_cmap(contours, colour_scheme, remove_contour=to_remove) + cf, lines = plot_contoured_field( + ax, coords_X, coords_Y, field_data, contour_method, contours, + cmap=cmap, line_contours=lines + ) + + add_colorbar_ax( + fig, cf, field_label, location='bottom', cbar_labelpad=-10, + cbar_format=cbar_format[field_name] + ) + if slice_along == 'z': + tomplot_field_title( + ax, '$z = $ 5 km', minmax=True, field_data=field_data + ) + elif slice_along == 'lon': + tomplot_field_title( + ax, r'$\lambda = $ 120 deg', minmax=True, field_data=field_data + ) + + # Labels ------------------------------------------------------------------- + if i == 0: + ax.set_ylabel(r'$\vartheta$ (deg)', labelpad=-20) + ax.set_ylim(lat_lims) + ax.set_yticks(lat_lims) + ax.set_yticklabels(lat_lims) + elif i == 3: + ax.set_ylabel(r'$z$ (km)', labelpad=-20) + ax.set_ylim(z_lims) + ax.set_yticks(z_lims) + ax.set_yticklabels(z_lims) + + if i < 3: + ax.set_xlabel(r'$\lambda$ (deg)', labelpad=-10) + ax.set_xlim(lon_lims) + ax.set_xticks(lon_lims) + ax.set_xticklabels(lon_lims) + else: + ax.set_xlabel(r'$\vartheta$ (deg)', labelpad=-10) + ax.set_xlim(lat_lims) + ax.set_xticks(lat_lims) + ax.set_xticklabels(lat_lims) + +# Save figure ------------------------------------------------------------------ +plt.suptitle(f't = {time:.1f} s', y=0.95) +fig.subplots_adjust(wspace=0.25, hspace=0.1) +plot_name = f'{plot_stem}_initial.png' +print(f'Saving figure to {plot_name}') +fig.savefig(plot_name, bbox_inches='tight') +plt.close() + +# ---------------------------------------------------------------------------- # +# FINAL PLOTTING +# ---------------------------------------------------------------------------- # +fig, axarray = plt.subplots(1, 2, figsize=(12, 6)) +time_idx = -1 +time = data_file['time'][time_idx] + +for i, (ax, field_name, field_label, colour_scheme, contours, + to_remove, slice_along) in \ + enumerate(zip(axarray.flatten(), final_field_names, final_field_labels, + final_colour_schemes, final_contours, + final_contours_to_remove, final_slice_along)): + + # Data extraction ---------------------------------------------------------- + if slice_along == 'z': + field_full = extract_gusto_field(data_file, field_name, time_idx) + coords_X_full, coords_Y_full, coords_Z_full = \ + extract_gusto_coords(data_file, field_name) + + # Reshape + field_full, coords_X_full, coords_Y_full, _ = \ + reshape_gusto_data(field_full, coords_X_full, + coords_Y_full, coords_Z_full) + + # Take level for a horizontal slice + field_data = field_full[:, level] + # Abuse of names for coord variables but simplifies code below + coords_X = coords_X_full[:, level] + coords_Y = coords_Y_full[:, level] + + else: + orig_field_data, orig_coords_X, orig_coords_Y, orig_coords_Z = \ + extract_gusto_vertical_slice( + data_file, field_name, time_idx, + slice_along=slice_along, slice_at=slice_at_lat + ) + + # Slices need regridding as points don't cleanly live along lon or lat = 0.0 + field_data, coords_X, coords_Y = \ + regrid_vertical_slice( + plotting_grid[slice_along], slice_along, slice_at_lat, + orig_coords_X, orig_coords_Y, orig_coords_Z, orig_field_data + ) + # Scale coordinates + coords_Y /= 1000. + + # Plot data ---------------------------------------------------------------- + cmap, lines = tomplot_cmap(contours, colour_scheme, remove_contour=to_remove) + cf, lines = plot_contoured_field( + ax, coords_X, coords_Y, field_data, contour_method, contours, + cmap=cmap, line_contours=lines + ) + + add_colorbar_ax( + fig, cf, field_label, location='bottom', cbar_labelpad=-10, + cbar_format=cbar_format[field_name] + ) + if slice_along == 'z': + tomplot_field_title( + ax, r'$z = $ 5 km', minmax=True, field_data=field_data + ) + elif slice_along == 'lat': + tomplot_field_title( + ax, r'$\vartheta = $ 0 deg', minmax=True, field_data=field_data + ) + + # Labels ------------------------------------------------------------------- + if i == 0: + ax.set_xlabel(r'$\lambda$ (deg)', labelpad=-10) + ax.set_xlim(lon_lims) + ax.set_xticks(lon_lims) + ax.set_xticklabels(lon_lims) + ax.set_ylabel(r'$\vartheta$ (deg)', labelpad=-20) + ax.set_ylim(lat_lims) + ax.set_yticks(lat_lims) + ax.set_yticklabels(lat_lims) + else: + ax.set_xlabel(r'$\lambda$ (deg)', labelpad=-10) + ax.set_xlim(lon_lims) + ax.set_xticks(lon_lims) + ax.set_xticklabels(lon_lims) + ax.set_ylabel(r'$z$ (km)', labelpad=-20) + ax.set_ylim(z_lims) + ax.set_yticks(z_lims) + ax.set_yticklabels(z_lims) + +# Save figure ------------------------------------------------------------------ +plt.suptitle(f't = {time:.0f} s') +fig.subplots_adjust(wspace=0.18) +plot_name = f'{plot_stem}_final.png' +print(f'Saving figure to {plot_name}') +fig.savefig(plot_name, bbox_inches='tight') +plt.close() diff --git a/plotting/shallow_water/plot_linear_williamson_2.py b/plotting/shallow_water/plot_linear_williamson_2.py new file mode 100644 index 00000000..4851e018 --- /dev/null +++ b/plotting/shallow_water/plot_linear_williamson_2.py @@ -0,0 +1,172 @@ +""" +Plots the linear Williamson 2 test case. + +The initial conditions are plotted: +(a) the velocity field, (b) the depth field. + +And after 5 days, this plots: +(a) the relative vorticity field, (b) the error in the depth field. +""" +from os.path import abspath, dirname +import cartopy.crs as ccrs +import matplotlib.pyplot as plt +import numpy as np +from netCDF4 import Dataset +from tomplot import ( + set_tomplot_style, tomplot_contours, tomplot_cmap, plot_contoured_field, + add_colorbar_ax, plot_field_quivers, tomplot_field_title, + extract_gusto_coords, extract_gusto_field, regrid_horizontal_slice +) + +test_name = 'linear_williamson_2' + +# ---------------------------------------------------------------------------- # +# 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_name}/field_output.nc' +plot_stem = f'{abspath(dirname(__file__))}/../../figures/{test_name}' + +# ---------------------------------------------------------------------------- # +# Initial plot details +# ---------------------------------------------------------------------------- # +init_field_names = ['u', 'D'] +init_colour_schemes = ['Oranges', 'YlGnBu'] +init_field_labels = [r'$|u|$ (m s$^{-1}$)', r'$D$ (m)'] +init_contours_to_remove = [None, None, None] +init_contours = [np.linspace(0, 40, 9), + np.linspace(-2000, 0, 11)] + +# ---------------------------------------------------------------------------- # +# Final plot details +# ---------------------------------------------------------------------------- # +final_field_names = ['RelativeVorticity', 'D_error'] +final_colour_schemes = ['RdBu_r', 'PiYG'] +final_field_labels = [r'$\zeta \ / $ s$^{-1}$', r'$\Delta D$ (m)'] +final_contours_to_remove = [0.0, 0.0] + +# ---------------------------------------------------------------------------- # +# General options +# ---------------------------------------------------------------------------- # +projection = ccrs.Robinson() +contour_method = 'contour' +xlims = [-180, 180] +ylims = [-90, 90] + +# We need to regrid onto lon-lat grid -- specify that here +lon_1d = np.linspace(-180.0, 180.0, 120) +lat_1d = np.linspace(-90, 90, 120) +lon_2d, lat_2d = np.meshgrid(lon_1d, lat_1d, indexing='ij') + +# Things that are likely the same for all plots -------------------------------- +set_tomplot_style() +data_file = Dataset(results_file_name, 'r') + +# ---------------------------------------------------------------------------- # +# INITIAL PLOTTING +# ---------------------------------------------------------------------------- # +fig = plt.figure(figsize=(15, 5)) +time_idx = 0 + +for i, (field_name, colour_scheme, field_label, contour_to_remove, contours) in \ + enumerate(zip( + init_field_names, init_colour_schemes, + init_field_labels, init_contours_to_remove, init_contours)): + + # Make axes + ax = fig.add_subplot(1, 2, 1+i, projection=projection) + + # Data extraction ---------------------------------------------------------- + if field_name == 'u': + zonal_data = extract_gusto_field(data_file, 'u_zonal', time_idx=time_idx) + meridional_data = extract_gusto_field(data_file, 'u_meridional', time_idx=time_idx) + field_data = np.sqrt(zonal_data**2 + meridional_data**2) + coords_X, coords_Y = extract_gusto_coords(data_file, 'u_zonal') + + else: + 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] / (24.*60.*60.) + + # Regrid onto lon-lat grid + field_data = regrid_horizontal_slice( + lon_2d, lat_2d, coords_X, coords_Y, field_data, periodic_fix='sphere' + ) + + # Plot data ---------------------------------------------------------------- + cmap, lines = tomplot_cmap(contours, colour_scheme, remove_contour=contour_to_remove) + cf, _ = plot_contoured_field( + ax, lon_2d, lat_2d, field_data, contour_method, contours, + cmap=cmap, line_contours=lines, projection=projection + ) + + add_colorbar_ax(ax, cf, field_label, location='bottom', cbar_labelpad=-10) + tomplot_field_title(ax, None, minmax=True, field_data=field_data) + + # Add quivers -------------------------------------------------------------- + if field_name == 'u': + # Need to re-grid to lat-lon grid to get sensible looking quivers + regrid_zonal_data = regrid_horizontal_slice( + lon_2d, lat_2d, coords_X, coords_Y, zonal_data, + periodic_fix='sphere' + ) + regrid_meridional_data = regrid_horizontal_slice( + lon_2d, lat_2d, coords_X, coords_Y, meridional_data, + periodic_fix='sphere' + ) + plot_field_quivers( + ax, lon_2d, lat_2d, regrid_zonal_data, regrid_meridional_data, + spatial_filter_step=6, magnitude_filter=1.0, + projection=ccrs.PlateCarree() + ) + +# Save figure ------------------------------------------------------------------ +fig.subplots_adjust(wspace=0.25) +plt.suptitle(f't = {time:.1f} days') +plot_name = f'{plot_stem}_initial.png' +print(f'Saving figure to {plot_name}') +fig.savefig(plot_name, bbox_inches='tight') +plt.close() + +# ---------------------------------------------------------------------------- # +# FINAL PLOTTING +# ---------------------------------------------------------------------------- # +fig = plt.figure(figsize=(15, 5)) +time_idx = -1 + +for i, (field_name, colour_scheme, field_label, contour_to_remove) in \ + enumerate(zip( + final_field_names, final_colour_schemes, + final_field_labels, final_contours_to_remove)): + + # Make axes + ax = fig.add_subplot(1, 2, 1+i, projection=projection) + + # 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] / (24.*60.*60.) + + # Regrid onto lon-lat grid + field_data = regrid_horizontal_slice( + lon_2d, lat_2d, coords_X, coords_Y, field_data, periodic_fix='sphere' + ) + + # Plot data ---------------------------------------------------------------- + contours = tomplot_contours(field_data) + cmap, lines = tomplot_cmap(contours, colour_scheme, remove_contour=contour_to_remove) + cf, _ = plot_contoured_field( + ax, lon_2d, lat_2d, field_data, contour_method, contours, + cmap=cmap, line_contours=lines, projection=projection + ) + + add_colorbar_ax(ax, cf, field_label, location='bottom', cbar_labelpad=-10) + tomplot_field_title(ax, None, minmax=True, field_data=field_data) + +# Save figure ------------------------------------------------------------------ +plt.suptitle(f't = {time:.1f} days') +plot_name = f'{plot_stem}_final.png' +print(f'Saving figure to {plot_name}') +fig.savefig(plot_name, bbox_inches='tight') +plt.close() diff --git a/plotting/shallow_water/plot_moist_convective_williamson_2.py b/plotting/shallow_water/plot_moist_convective_williamson_2.py new file mode 100644 index 00000000..3c0a4254 --- /dev/null +++ b/plotting/shallow_water/plot_moist_convective_williamson_2.py @@ -0,0 +1,186 @@ +""" +Plots the moist convective Williamson 2 test case. + +The initial conditions are plotted: +(a) the velocity field, (b) the depth field, (c) water vapour field. + +And after 5 days, this plots: +(a) relative vorticity, (b) depth error (c) cloud. +""" +from os.path import abspath, dirname +import cartopy.crs as ccrs +import matplotlib.pyplot as plt +import numpy as np +from netCDF4 import Dataset +from tomplot import ( + set_tomplot_style, tomplot_contours, tomplot_cmap, plot_contoured_field, + add_colorbar_ax, plot_field_quivers, tomplot_field_title, + extract_gusto_coords, extract_gusto_field, regrid_horizontal_slice +) + +test_name = 'moist_convective_williamson_2' + +# ---------------------------------------------------------------------------- # +# 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_name}/field_output.nc' +plot_stem = f'{abspath(dirname(__file__))}/../../figures/shallow_water/{test_name}' + +# ---------------------------------------------------------------------------- # +# Initial plot details +# ---------------------------------------------------------------------------- # +init_field_names = ['u', 'D', 'water_vapour'] +init_colour_schemes = ['Oranges', 'YlGnBu', 'Purples'] +init_field_labels = [r'$|u|$ (m s$^{-1}$)', r'$D$ (m)', r'$m_v$ (kg kg$^{-1}$)'] +init_contours_to_remove = [None, None, None] +init_contours = [np.linspace(0, 20, 9), + np.linspace(1900, 3100, 13), + np.linspace(0, 0.012, 13)] +init_data_formats = ['1.0f', '1.0f', '1.2e'] + +# ---------------------------------------------------------------------------- # +# Final plot details +# ---------------------------------------------------------------------------- # +final_field_names = ['RelativeVorticity', 'D_error', 'cloud_water'] +final_colour_schemes = ['RdBu_r', 'PiYG', 'Blues'] +final_field_labels = [r'$\zeta \ / $ s$^{-1}$', r'$\Delta D$ (m)', r'$m_{cl}$ (kg kg$^{-1}$)'] +final_contours_to_remove = [0.0, 0.0, None] +final_data_formats = ['1.1e', '1.1f', '1.2e'] + +# ---------------------------------------------------------------------------- # +# General options +# ---------------------------------------------------------------------------- # +projection = ccrs.Robinson() +contour_method = 'contour' +xlims = [-180, 180] +ylims = [-90, 90] + +# We need to regrid onto lon-lat grid -- specify that here +lon_1d = np.linspace(-180.0, 180.0, 120) +lat_1d = np.linspace(-90, 90, 120) +lon_2d, lat_2d = np.meshgrid(lon_1d, lat_1d, indexing='ij') + +# Things that are likely the same for all plots -------------------------------- +set_tomplot_style() +data_file = Dataset(results_file_name, 'r') + +# ---------------------------------------------------------------------------- # +# INITIAL PLOTTING +# ---------------------------------------------------------------------------- # +fig = plt.figure(figsize=(21, 5)) +time_idx = 0 + +for i, (field_name, colour_scheme, field_label, + contour_to_remove, contours, data_format) in \ + enumerate(zip( + init_field_names, init_colour_schemes, init_field_labels, + init_contours_to_remove, init_contours, init_data_formats)): + + # Make axes + ax = fig.add_subplot(1, 3, 1+i, projection=projection) + + # Data extraction ---------------------------------------------------------- + if field_name == 'u': + zonal_data = extract_gusto_field(data_file, 'u_zonal', time_idx=time_idx) + meridional_data = extract_gusto_field(data_file, 'u_meridional', time_idx=time_idx) + field_data = np.sqrt(zonal_data**2 + meridional_data**2) + coords_X, coords_Y = extract_gusto_coords(data_file, 'u_zonal') + + else: + 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] / (24.*60.*60.) + + # Regrid onto lon-lat grid + field_data = regrid_horizontal_slice( + lon_2d, lat_2d, coords_X, coords_Y, field_data, periodic_fix='sphere' + ) + + # Plot data ---------------------------------------------------------------- + cmap, lines = tomplot_cmap(contours, colour_scheme, remove_contour=contour_to_remove) + cf, _ = plot_contoured_field( + ax, lon_2d, lat_2d, field_data, contour_method, contours, + cmap=cmap, line_contours=lines, projection=projection + ) + + add_colorbar_ax( + ax, cf, field_label, location='bottom', cbar_labelpad=-10, + cbar_format=data_format + ) + tomplot_field_title( + ax, None, minmax=True, field_data=field_data, minmax_format=data_format + ) + + # Add quivers -------------------------------------------------------------- + if field_name == 'u': + # Need to re-grid to lat-lon grid to get sensible looking quivers + regrid_zonal_data = regrid_horizontal_slice( + lon_2d, lat_2d, coords_X, coords_Y, zonal_data, + periodic_fix='sphere' + ) + regrid_meridional_data = regrid_horizontal_slice( + lon_2d, lat_2d, coords_X, coords_Y, meridional_data, + periodic_fix='sphere' + ) + plot_field_quivers( + ax, lon_2d, lat_2d, regrid_zonal_data, regrid_meridional_data, + spatial_filter_step=6, magnitude_filter=1.0, + projection=ccrs.PlateCarree() + ) + +# Save figure ------------------------------------------------------------------ +fig.subplots_adjust(wspace=0.25) +plt.suptitle(f't = {time:.1f} days') +plot_name = f'{plot_stem}_initial.png' +print(f'Saving figure to {plot_name}') +fig.savefig(plot_name, bbox_inches='tight') +plt.close() + +# ---------------------------------------------------------------------------- # +# FINAL PLOTTING +# ---------------------------------------------------------------------------- # +fig = plt.figure(figsize=(21, 5)) +time_idx = -1 + +for i, (field_name, colour_scheme, field_label, contour_to_remove, data_format) in \ + enumerate(zip( + final_field_names, final_colour_schemes, + final_field_labels, final_contours_to_remove, final_data_formats)): + + # Make axes + ax = fig.add_subplot(1, 3, 1+i, projection=projection) + + # 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] / (24.*60.*60.) + + # Regrid onto lon-lat grid + field_data = regrid_horizontal_slice( + lon_2d, lat_2d, coords_X, coords_Y, field_data, periodic_fix='sphere' + ) + + # Plot data ---------------------------------------------------------------- + contours = tomplot_contours(field_data) + cmap, lines = tomplot_cmap(contours, colour_scheme, remove_contour=contour_to_remove) + cf, _ = plot_contoured_field( + ax, lon_2d, lat_2d, field_data, contour_method, contours, + cmap=cmap, line_contours=lines, projection=projection + ) + + add_colorbar_ax( + ax, cf, field_label, location='bottom', cbar_labelpad=-10, + cbar_format=data_format + ) + tomplot_field_title( + ax, None, minmax=True, field_data=field_data, minmax_format=data_format + ) + +# Save figure ------------------------------------------------------------------ +plt.suptitle(f't = {time:.1f} days') +plot_name = f'{plot_stem}_final.png' +print(f'Saving figure to {plot_name}') +fig.savefig(plot_name, bbox_inches='tight') +plt.close() diff --git a/plotting/shallow_water/plot_moist_thermal_williamson_5.py b/plotting/shallow_water/plot_moist_thermal_williamson_5.py new file mode 100644 index 00000000..9c5bf8ec --- /dev/null +++ b/plotting/shallow_water/plot_moist_thermal_williamson_5.py @@ -0,0 +1,195 @@ +""" +Plots the moist thermal Williamson 5 test case. + +The initial conditions are plotted: +(a) velocity, (b) depth field, +(c) buoyancy, (d) water vapour. + +And after 50 days, this plots: +(a) relative vorticity, (b) free-surface height, +(c) buoyancy, (d) cloud. +""" +from os.path import abspath, dirname +import cartopy.crs as ccrs +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_ax, plot_field_quivers, tomplot_field_title, + extract_gusto_coords, extract_gusto_field, regrid_horizontal_slice +) + +test_name = 'moist_thermal_williamson_5' + +# ---------------------------------------------------------------------------- # +# 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_name}/field_output.nc' +plot_stem = f'{abspath(dirname(__file__))}/../../figures/shallow_water/{test_name}' + +# ---------------------------------------------------------------------------- # +# Initial plot details +# ---------------------------------------------------------------------------- # +init_field_names = ['u', 'D', 'b', 'water_vapour'] +init_colour_schemes = ['Oranges', 'YlGnBu', 'PuRd_r', 'Purples'] +init_field_labels = [r'$|u|$ (m s$^{-1}$)', r'$D$ (m)', + r'$b$ (m s$^{-2}$)', r'$m_v$ (kg kg$^{-1}$)'] +init_contours_to_remove = [None, None, None, None] +init_contours = [np.linspace(0, 20, 9), + np.linspace(3800, 6000, 13), + np.linspace(8.8, 11.2, 13), + np.linspace(0.0, 0.02, 11)] + +# ---------------------------------------------------------------------------- # +# Final plot details +# ---------------------------------------------------------------------------- # +final_field_names = ['RelativeVorticity', 'D_plus_topography', 'b', 'cloud_water'] +final_colour_schemes = ['RdBu_r', 'YlGnBu', 'PuRd_r', 'Blues'] +final_field_labels = [r'$\zeta \ / $ s$^{-1}$', r'$D+B$ (m)', + r'$b$ (m s$^{-2}$)', r'$m_{cl}$ (kg kg$^{-1}$)'] +final_contours_to_remove = [0.0, None, None, 0.0] +final_contours = [np.linspace(-5e-5, 5e-5, 11), + np.linspace(4800, 6000, 13), + np.linspace(8.8, 11.2, 13), + np.linspace(-5e-5, 5e-4, 12)] + +# ---------------------------------------------------------------------------- # +# General options +# ---------------------------------------------------------------------------- # +projection = ccrs.Robinson() +contour_method = 'contour' +xlims = [-180, 180] +ylims = [-90, 90] + +# We need to regrid onto lon-lat grid -- specify that here +lon_1d = np.linspace(-180.0, 180.0, 120) +lat_1d = np.linspace(-90, 90, 120) +lon_2d, lat_2d = np.meshgrid(lon_1d, lat_1d, indexing='ij') + +cbar_format = {'RelativeVorticity': '1.1e', + 'u': '1.0f', + 'D': '1.0f', + 'D_plus_topography': '1.0f', + 'b': '1.1f', + 'water_vapour': '1.2f', + 'cloud_water': '1.1e'} + +# Things that are likely the same for all plots -------------------------------- +set_tomplot_style() +data_file = Dataset(results_file_name, 'r') + +# ---------------------------------------------------------------------------- # +# INITIAL PLOTTING +# ---------------------------------------------------------------------------- # +fig = plt.figure(figsize=(15, 10)) +time_idx = 0 + +for i, (field_name, colour_scheme, field_label, contour_to_remove, contours) in \ + enumerate(zip( + init_field_names, init_colour_schemes, + init_field_labels, init_contours_to_remove, init_contours)): + + # Make axes + ax = fig.add_subplot(2, 2, 1+i, projection=projection) + + # Data extraction ---------------------------------------------------------- + if field_name == 'u': + zonal_data = extract_gusto_field(data_file, 'u_zonal', time_idx=time_idx) + meridional_data = extract_gusto_field(data_file, 'u_meridional', time_idx=time_idx) + field_data = np.sqrt(zonal_data**2 + meridional_data**2) + coords_X, coords_Y = extract_gusto_coords(data_file, 'u_zonal') + + else: + 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] / (24.*60.*60.) + + # Regrid onto lon-lat grid + field_data = regrid_horizontal_slice( + lon_2d, lat_2d, coords_X, coords_Y, field_data, periodic_fix='sphere' + ) + + # Plot data ---------------------------------------------------------------- + cmap, lines = tomplot_cmap(contours, colour_scheme, remove_contour=contour_to_remove) + cf, _ = plot_contoured_field( + ax, lon_2d, lat_2d, field_data, contour_method, contours, + cmap=cmap, line_contours=lines, projection=projection + ) + + add_colorbar_ax( + ax, cf, field_label, location='bottom', cbar_labelpad=-10, + cbar_format=cbar_format[field_name] + ) + tomplot_field_title(ax, None, minmax=True, field_data=field_data) + + # Add quivers -------------------------------------------------------------- + if field_name == 'u': + # Need to re-grid to lat-lon grid to get sensible looking quivers + regrid_zonal_data = regrid_horizontal_slice( + lon_2d, lat_2d, coords_X, coords_Y, zonal_data, + periodic_fix='sphere' + ) + regrid_meridional_data = regrid_horizontal_slice( + lon_2d, lat_2d, coords_X, coords_Y, meridional_data, + periodic_fix='sphere' + ) + plot_field_quivers( + ax, lon_2d, lat_2d, regrid_zonal_data, regrid_meridional_data, + spatial_filter_step=6, magnitude_filter=1.0, + projection=ccrs.PlateCarree() + ) + +# Save figure ------------------------------------------------------------------ +fig.subplots_adjust(wspace=0.25) +plt.suptitle(f't = {time:.1f} days') +plot_name = f'{plot_stem}_initial.png' +print(f'Saving figure to {plot_name}') +fig.savefig(plot_name, bbox_inches='tight') +plt.close() + +# ---------------------------------------------------------------------------- # +# FINAL PLOTTING +# ---------------------------------------------------------------------------- # +fig = plt.figure(figsize=(15, 10)) +time_idx = -1 + +for i, (field_name, colour_scheme, field_label, contour_to_remove, contours) in \ + enumerate(zip( + final_field_names, final_colour_schemes, + final_field_labels, final_contours_to_remove, final_contours)): + + # Make axes + ax = fig.add_subplot(2, 2, 1+i, projection=projection) + + # 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] / (24.*60.*60.) + + # Regrid onto lon-lat grid + field_data = regrid_horizontal_slice( + lon_2d, lat_2d, coords_X, coords_Y, field_data, periodic_fix='sphere' + ) + + # Plot data ---------------------------------------------------------------- + cmap, lines = tomplot_cmap(contours, colour_scheme, remove_contour=contour_to_remove) + cf, _ = plot_contoured_field( + ax, lon_2d, lat_2d, field_data, contour_method, contours, + cmap=cmap, line_contours=lines, projection=projection + ) + + add_colorbar_ax( + ax, cf, field_label, location='bottom', cbar_labelpad=-10, + cbar_format=cbar_format[field_name] + ) + tomplot_field_title(ax, None, minmax=True, field_data=field_data) + +# Save figure ------------------------------------------------------------------ +plt.suptitle(f't = {time:.1f} days') +plot_name = f'{plot_stem}_final.png' +print(f'Saving figure to {plot_name}') +fig.savefig(plot_name, bbox_inches='tight') +plt.close() diff --git a/plotting/shallow_water/plot_shallow_water_1d_wave.py b/plotting/shallow_water/plot_shallow_water_1d_wave.py new file mode 100644 index 00000000..256ac1c1 --- /dev/null +++ b/plotting/shallow_water/plot_shallow_water_1d_wave.py @@ -0,0 +1,102 @@ +""" +Plots the fields from the 1D shallow water wave. + +This plots: +(a) u @ t = 0 s, (b) v @ t = 0 s, (c) D @ t = 0 s +(d) u @ t = 1 s, (e) v @ t = 1 s, (f) D @ t = 1 s +""" +from os.path import abspath, dirname +import matplotlib.pyplot as plt +import numpy as np +from netCDF4 import Dataset +import pandas as pd +from tomplot import ( + set_tomplot_style, tomplot_field_title, extract_gusto_coords, + extract_gusto_field +) + +test = 'shallow_water_1d_wave' + +# ---------------------------------------------------------------------------- # +# 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/shallow_water/{test}' + +# ---------------------------------------------------------------------------- # +# Plot details +# ---------------------------------------------------------------------------- # +field_names = ['u', 'v', 'D', 'u', 'v', 'D'] +time_idxs = [0, 0, 0, -1, -1, -1] + +# ---------------------------------------------------------------------------- # +# General options +# ---------------------------------------------------------------------------- # +xlims = [0, 2*np.pi] +xlims_labels = [0, r'$2\pi$'] + +ylims = { + 'u': [-0.5, 0.5], + 'v': [-0.5, 0.5], + 'D': [8, 12] +} +field_labels = { + 'u': r'$u$ (m s$^{-1}$)', + 'v': r'$v$ (m s$^{-1}$)', + 'D': r'$D$ (m)' +} + +# Things that are likely the same for all plots -------------------------------- +set_tomplot_style() +data_file = Dataset(results_file_name, 'r') + +# ---------------------------------------------------------------------------- # +# PLOTTING +# ---------------------------------------------------------------------------- # +fig, axarray = plt.subplots(2, 3, figsize=(16, 6), sharex='all', sharey='col') + +for i, (ax, time_idx, field_name) in \ + enumerate(zip(axarray.flatten(), time_idxs, field_names)): + + # Data extraction ---------------------------------------------------------- + field_data = extract_gusto_field(data_file, field_name, time_idx=time_idx) + coords_X = extract_gusto_coords(data_file, field_name) + time = data_file['time'][time_idx] + + # Get coordinates in order + data_frame = pd.DataFrame({'coords': coords_X, 'field': field_data}) + data_frame = data_frame.sort_values(by=['coords']) + coords_X = data_frame['coords'] + field_data = data_frame['field'] + + # Convert coordinates to m + coords_X *= 1000. + + # Plot data ---------------------------------------------------------------- + ax.plot(coords_X, field_data, color='black', linestyle='-', marker='') + + tomplot_field_title( + ax, f't = {time:.1f} s', minmax=True, field_data=field_data, + minmax_format='1.2f' + ) + + # Labels ------------------------------------------------------------------- + ax.set_ylabel(field_labels[field_name], labelpad=-15) + ax.set_ylim(ylims[field_name]) + ax.set_yticks(ylims[field_name]) + ax.set_yticklabels(ylims[field_name]) + + if i > 2: + ax.set_xlabel(r'$x$ (m)', labelpad=-10) + ax.set_xlim(xlims) + ax.set_xticks(xlims) + ax.set_xticklabels(xlims_labels) + +# Save figure ------------------------------------------------------------------ +fig.subplots_adjust(wspace=0.2) +plot_name = f'{plot_stem}.png' +print(f'Saving figure to {plot_name}') +fig.savefig(plot_name, bbox_inches='tight') +plt.close() diff --git a/plotting/shallow_water/plot_thermal_williamson_2.py b/plotting/shallow_water/plot_thermal_williamson_2.py new file mode 100644 index 00000000..32ab1483 --- /dev/null +++ b/plotting/shallow_water/plot_thermal_williamson_2.py @@ -0,0 +1,173 @@ +""" +Plots the thermal Williamson 2 test case. + +The initial conditions are plotted: +(a) the velocity field, (b) the depth field, (c) buoyancy field. + +And after 5 days, this plots: +(a) relative vorticity, (b) depth error (c) buoyancy error. +""" +from os.path import abspath, dirname +import cartopy.crs as ccrs +import matplotlib.pyplot as plt +import numpy as np +from netCDF4 import Dataset +from tomplot import ( + set_tomplot_style, tomplot_contours, tomplot_cmap, plot_contoured_field, + add_colorbar_ax, plot_field_quivers, tomplot_field_title, + extract_gusto_coords, extract_gusto_field, regrid_horizontal_slice +) + +test_name = 'thermal_williamson_2' + +# ---------------------------------------------------------------------------- # +# 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_name}/field_output.nc' +plot_stem = f'{abspath(dirname(__file__))}/../../figures/{test_name}' + +# ---------------------------------------------------------------------------- # +# Initial plot details +# ---------------------------------------------------------------------------- # +init_field_names = ['u', 'D', 'b'] +init_colour_schemes = ['Oranges', 'YlGnBu', 'PuRd_r'] +init_field_labels = [r'$|u|$ (m s$^{-1}$)', r'$D$ (m)', r'$b$ (m s$^{-2}$)'] +init_contours_to_remove = [None, None, None] +init_contours = [np.linspace(0, 20, 9), + np.linspace(1900, 3100, 13), + np.linspace(9, 10, 11)] + +# ---------------------------------------------------------------------------- # +# Final plot details +# ---------------------------------------------------------------------------- # +final_field_names = ['RelativeVorticity', 'D_error', 'b_error'] +final_colour_schemes = ['RdBu_r', 'PiYG', 'PuOr'] +final_field_labels = [r'$\zeta \ / $ s$^{-1}$', r'$\Delta D$ (m)', r'$\Delta b$ (m s$^{-2}$)'] +final_contours_to_remove = [0.0, 0.0, 0.0] + +# ---------------------------------------------------------------------------- # +# General options +# ---------------------------------------------------------------------------- # +projection = ccrs.Robinson() +contour_method = 'contour' +xlims = [-180, 180] +ylims = [-90, 90] + +# We need to regrid onto lon-lat grid -- specify that here +lon_1d = np.linspace(-180.0, 180.0, 120) +lat_1d = np.linspace(-90, 90, 120) +lon_2d, lat_2d = np.meshgrid(lon_1d, lat_1d, indexing='ij') + +# Things that are likely the same for all plots -------------------------------- +set_tomplot_style() +data_file = Dataset(results_file_name, 'r') + +# ---------------------------------------------------------------------------- # +# INITIAL PLOTTING +# ---------------------------------------------------------------------------- # +fig = plt.figure(figsize=(21, 5)) +time_idx = 0 + +for i, (field_name, colour_scheme, field_label, contour_to_remove, contours) in \ + enumerate(zip( + init_field_names, init_colour_schemes, + init_field_labels, init_contours_to_remove, init_contours)): + + # Make axes + ax = fig.add_subplot(1, 3, 1+i, projection=projection) + + # Data extraction ---------------------------------------------------------- + if field_name == 'u': + zonal_data = extract_gusto_field(data_file, 'u_zonal', time_idx=time_idx) + meridional_data = extract_gusto_field(data_file, 'u_meridional', time_idx=time_idx) + field_data = np.sqrt(zonal_data**2 + meridional_data**2) + coords_X, coords_Y = extract_gusto_coords(data_file, 'u_zonal') + + else: + 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] / (24.*60.*60.) + + # Regrid onto lon-lat grid + field_data = regrid_horizontal_slice( + lon_2d, lat_2d, coords_X, coords_Y, field_data, periodic_fix='sphere' + ) + + # Plot data ---------------------------------------------------------------- + cmap, lines = tomplot_cmap(contours, colour_scheme, remove_contour=contour_to_remove) + cf, _ = plot_contoured_field( + ax, lon_2d, lat_2d, field_data, contour_method, contours, + cmap=cmap, line_contours=lines, projection=projection + ) + + add_colorbar_ax(ax, cf, field_label, location='bottom', cbar_labelpad=-10) + tomplot_field_title(ax, None, minmax=True, field_data=field_data) + + # Add quivers -------------------------------------------------------------- + if field_name == 'u': + # Need to re-grid to lat-lon grid to get sensible looking quivers + regrid_zonal_data = regrid_horizontal_slice( + lon_2d, lat_2d, coords_X, coords_Y, zonal_data, + periodic_fix='sphere' + ) + regrid_meridional_data = regrid_horizontal_slice( + lon_2d, lat_2d, coords_X, coords_Y, meridional_data, + periodic_fix='sphere' + ) + plot_field_quivers( + ax, lon_2d, lat_2d, regrid_zonal_data, regrid_meridional_data, + spatial_filter_step=6, magnitude_filter=1.0, + projection=ccrs.PlateCarree() + ) + +# Save figure ------------------------------------------------------------------ +fig.subplots_adjust(wspace=0.25) +plt.suptitle(f't = {time:.1f} days') +plot_name = f'{plot_stem}_initial.png' +print(f'Saving figure to {plot_name}') +fig.savefig(plot_name, bbox_inches='tight') +plt.close() + +# ---------------------------------------------------------------------------- # +# FINAL PLOTTING +# ---------------------------------------------------------------------------- # +fig = plt.figure(figsize=(21, 5)) +time_idx = -1 + +for i, (field_name, colour_scheme, field_label, contour_to_remove) in \ + enumerate(zip( + final_field_names, final_colour_schemes, + final_field_labels, final_contours_to_remove)): + + # Make axes + ax = fig.add_subplot(1, 3, 1+i, projection=projection) + + # 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] / (24.*60.*60.) + + # Regrid onto lon-lat grid + field_data = regrid_horizontal_slice( + lon_2d, lat_2d, coords_X, coords_Y, field_data, periodic_fix='sphere' + ) + + # Plot data ---------------------------------------------------------------- + contours = tomplot_contours(field_data) + cmap, lines = tomplot_cmap(contours, colour_scheme, remove_contour=contour_to_remove) + cf, _ = plot_contoured_field( + ax, lon_2d, lat_2d, field_data, contour_method, contours, + cmap=cmap, line_contours=lines, projection=projection + ) + + add_colorbar_ax(ax, cf, field_label, location='bottom', cbar_labelpad=-10) + tomplot_field_title(ax, None, minmax=True, field_data=field_data) + +# Save figure ------------------------------------------------------------------ +plt.suptitle(f't = {time:.1f} days') +plot_name = f'{plot_stem}_final.png' +print(f'Saving figure to {plot_name}') +fig.savefig(plot_name, bbox_inches='tight') +plt.close() diff --git a/plotting/shallow_water/plot_williamson_2.py b/plotting/shallow_water/plot_williamson_2.py new file mode 100644 index 00000000..1f8eaef5 --- /dev/null +++ b/plotting/shallow_water/plot_williamson_2.py @@ -0,0 +1,172 @@ +""" +Plots the Williamson 2 test case. + +The initial conditions are plotted: +(a) the velocity field, (b) the depth field. + +And after 5 days, this plots: +(a) the relative vorticity field, (b) the error in the depth field. +""" +from os.path import abspath, dirname +import cartopy.crs as ccrs +import matplotlib.pyplot as plt +import numpy as np +from netCDF4 import Dataset +from tomplot import ( + set_tomplot_style, tomplot_contours, tomplot_cmap, plot_contoured_field, + add_colorbar_ax, plot_field_quivers, tomplot_field_title, + extract_gusto_coords, extract_gusto_field, regrid_horizontal_slice +) + +test_name = 'williamson_2' + +# ---------------------------------------------------------------------------- # +# 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_name}/field_output.nc' +plot_stem = f'{abspath(dirname(__file__))}/../../figures/{test_name}' + +# ---------------------------------------------------------------------------- # +# Initial plot details +# ---------------------------------------------------------------------------- # +init_field_names = ['u', 'D'] +init_colour_schemes = ['Oranges', 'YlGnBu'] +init_field_labels = [r'$|u|$ (m s$^{-1}$)', r'$D$ (m)'] +init_contours_to_remove = [None, None, None] +init_contours = [np.linspace(0, 40, 9), + np.linspace(4000, 6000, 9)] + +# ---------------------------------------------------------------------------- # +# Final plot details +# ---------------------------------------------------------------------------- # +final_field_names = ['RelativeVorticity', 'D_error'] +final_colour_schemes = ['RdBu_r', 'PiYG'] +final_field_labels = [r'$\zeta \ / $ s$^{-1}$', r'$\Delta D$ (m)'] +final_contours_to_remove = [0.0, 0.0] + +# ---------------------------------------------------------------------------- # +# General options +# ---------------------------------------------------------------------------- # +projection = ccrs.Robinson() +contour_method = 'contour' +xlims = [-180, 180] +ylims = [-90, 90] + +# We need to regrid onto lon-lat grid -- specify that here +lon_1d = np.linspace(-180.0, 180.0, 120) +lat_1d = np.linspace(-90, 90, 120) +lon_2d, lat_2d = np.meshgrid(lon_1d, lat_1d, indexing='ij') + +# Things that are likely the same for all plots -------------------------------- +set_tomplot_style() +data_file = Dataset(results_file_name, 'r') + +# ---------------------------------------------------------------------------- # +# INITIAL PLOTTING +# ---------------------------------------------------------------------------- # +fig = plt.figure(figsize=(15, 5)) +time_idx = 0 + +for i, (field_name, colour_scheme, field_label, contour_to_remove, contours) in \ + enumerate(zip( + init_field_names, init_colour_schemes, + init_field_labels, init_contours_to_remove, init_contours)): + + # Make axes + ax = fig.add_subplot(1, 2, 1+i, projection=projection) + + # Data extraction ---------------------------------------------------------- + if field_name == 'u': + zonal_data = extract_gusto_field(data_file, 'u_zonal', time_idx=time_idx) + meridional_data = extract_gusto_field(data_file, 'u_meridional', time_idx=time_idx) + field_data = np.sqrt(zonal_data**2 + meridional_data**2) + coords_X, coords_Y = extract_gusto_coords(data_file, 'u_zonal') + + else: + 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] / (24.*60.*60.) + + # Regrid onto lon-lat grid + field_data = regrid_horizontal_slice( + lon_2d, lat_2d, coords_X, coords_Y, field_data, periodic_fix='sphere' + ) + + # Plot data ---------------------------------------------------------------- + cmap, lines = tomplot_cmap(contours, colour_scheme, remove_contour=contour_to_remove) + cf, _ = plot_contoured_field( + ax, lon_2d, lat_2d, field_data, contour_method, contours, + cmap=cmap, line_contours=lines, projection=projection + ) + + add_colorbar_ax(ax, cf, field_label, location='bottom', cbar_labelpad=-10) + tomplot_field_title(ax, None, minmax=True, field_data=field_data) + + # Add quivers -------------------------------------------------------------- + if field_name == 'u': + # Need to re-grid to lat-lon grid to get sensible looking quivers + regrid_zonal_data = regrid_horizontal_slice( + lon_2d, lat_2d, coords_X, coords_Y, zonal_data, + periodic_fix='sphere' + ) + regrid_meridional_data = regrid_horizontal_slice( + lon_2d, lat_2d, coords_X, coords_Y, meridional_data, + periodic_fix='sphere' + ) + plot_field_quivers( + ax, lon_2d, lat_2d, regrid_zonal_data, regrid_meridional_data, + spatial_filter_step=6, magnitude_filter=1.0, + projection=ccrs.PlateCarree() + ) + +# Save figure ------------------------------------------------------------------ +fig.subplots_adjust(wspace=0.25) +plt.suptitle(f't = {time:.1f} days') +plot_name = f'{plot_stem}_initial.png' +print(f'Saving figure to {plot_name}') +fig.savefig(plot_name, bbox_inches='tight') +plt.close() + +# ---------------------------------------------------------------------------- # +# FINAL PLOTTING +# ---------------------------------------------------------------------------- # +fig = plt.figure(figsize=(15, 5)) +time_idx = -1 + +for i, (field_name, colour_scheme, field_label, contour_to_remove) in \ + enumerate(zip( + final_field_names, final_colour_schemes, + final_field_labels, final_contours_to_remove)): + + # Make axes + ax = fig.add_subplot(1, 2, 1+i, projection=projection) + + # 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] / (24.*60.*60.) + + # Regrid onto lon-lat grid + field_data = regrid_horizontal_slice( + lon_2d, lat_2d, coords_X, coords_Y, field_data, periodic_fix='sphere' + ) + + # Plot data ---------------------------------------------------------------- + contours = tomplot_contours(field_data) + cmap, lines = tomplot_cmap(contours, colour_scheme, remove_contour=contour_to_remove) + cf, _ = plot_contoured_field( + ax, lon_2d, lat_2d, field_data, contour_method, contours, + cmap=cmap, line_contours=lines, projection=projection + ) + + add_colorbar_ax(ax, cf, field_label, location='bottom', cbar_labelpad=-10) + tomplot_field_title(ax, None, minmax=True, field_data=field_data) + +# Save figure ------------------------------------------------------------------ +plt.suptitle(f't = {time:.1f} days') +plot_name = f'{plot_stem}_final.png' +print(f'Saving figure to {plot_name}') +fig.savefig(plot_name, bbox_inches='tight') +plt.close() diff --git a/plotting/shallow_water/plot_williamson_5.py b/plotting/shallow_water/plot_williamson_5.py new file mode 100644 index 00000000..51ce0e80 --- /dev/null +++ b/plotting/shallow_water/plot_williamson_5.py @@ -0,0 +1,184 @@ +""" +Plots the Williamson 5 test case. + +The initial conditions are plotted: +(a) the velocity field, (b) the depth field. + +And after 50 days, this plots: +(a) the relative vorticity field, (b) free-surface height. +""" +from os.path import abspath, dirname +import cartopy.crs as ccrs +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_ax, plot_field_quivers, tomplot_field_title, + extract_gusto_coords, extract_gusto_field, regrid_horizontal_slice +) + +test_name = 'williamson_5' + +# ---------------------------------------------------------------------------- # +# 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_name}/field_output.nc' +plot_stem = f'{abspath(dirname(__file__))}/../../figures/shallow_water/{test_name}' + +# ---------------------------------------------------------------------------- # +# Initial plot details +# ---------------------------------------------------------------------------- # +init_field_names = ['u', 'D'] +init_colour_schemes = ['Oranges', 'YlGnBu'] +init_field_labels = [r'$|u|$ (m s$^{-1}$)', r'$D$ (m)'] +init_contours_to_remove = [None, None, None] +init_contours = [np.linspace(0, 20, 9), + np.linspace(3800, 6000, 12)] + +# ---------------------------------------------------------------------------- # +# Final plot details +# ---------------------------------------------------------------------------- # +final_field_names = ['RelativeVorticity', 'D_plus_topography'] +final_colour_schemes = ['RdBu_r', 'PiYG'] +final_field_labels = [r'$\zeta \ / $ s$^{-1}$', r'$D+B$ (m)'] +final_contours = [np.linspace(-1e-4, 1e-4, 21), + np.linspace(5000, 6000, 11)] +final_contours_to_remove = [0.0, None] + +# ---------------------------------------------------------------------------- # +# General options +# ---------------------------------------------------------------------------- # +projection = ccrs.Robinson() +contour_method = 'contour' +xlims = [-180, 180] +ylims = [-90, 90] + +cbar_format = {'RelativeVorticity': '1.1e', + 'u': '1.0f', + 'D': '1.0f', + 'D_plus_topography': '1.0f'} + +# We need to regrid onto lon-lat grid -- specify that here +lon_1d = np.linspace(-180.0, 180.0, 120) +lat_1d = np.linspace(-90, 90, 120) +lon_2d, lat_2d = np.meshgrid(lon_1d, lat_1d, indexing='ij') + +# Things that are likely the same for all plots -------------------------------- +set_tomplot_style() +data_file = Dataset(results_file_name, 'r') + +# ---------------------------------------------------------------------------- # +# INITIAL PLOTTING +# ---------------------------------------------------------------------------- # +fig = plt.figure(figsize=(15, 5)) +time_idx = 0 + +for i, (field_name, colour_scheme, field_label, contour_to_remove, contours) in \ + enumerate(zip( + init_field_names, init_colour_schemes, + init_field_labels, init_contours_to_remove, init_contours)): + + # Make axes + ax = fig.add_subplot(1, 2, 1+i, projection=projection) + + # Data extraction ---------------------------------------------------------- + if field_name == 'u': + zonal_data = extract_gusto_field(data_file, 'u_zonal', time_idx=time_idx) + meridional_data = extract_gusto_field(data_file, 'u_meridional', time_idx=time_idx) + field_data = np.sqrt(zonal_data**2 + meridional_data**2) + coords_X, coords_Y = extract_gusto_coords(data_file, 'u_zonal') + + else: + 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] / (24.*60.*60.) + + # Regrid onto lon-lat grid + field_data = regrid_horizontal_slice( + lon_2d, lat_2d, coords_X, coords_Y, field_data, periodic_fix='sphere' + ) + + # Plot data ---------------------------------------------------------------- + cmap, lines = tomplot_cmap(contours, colour_scheme, remove_contour=contour_to_remove) + cf, _ = plot_contoured_field( + ax, lon_2d, lat_2d, field_data, contour_method, contours, + cmap=cmap, line_contours=lines, projection=projection + ) + + add_colorbar_ax( + ax, cf, field_label, location='bottom', cbar_labelpad=-10, + cbar_format=cbar_format[field_name] + ) + tomplot_field_title(ax, None, minmax=True, field_data=field_data) + + # Add quivers -------------------------------------------------------------- + if field_name == 'u': + # Need to re-grid to lat-lon grid to get sensible looking quivers + regrid_zonal_data = regrid_horizontal_slice( + lon_2d, lat_2d, coords_X, coords_Y, zonal_data, + periodic_fix='sphere' + ) + regrid_meridional_data = regrid_horizontal_slice( + lon_2d, lat_2d, coords_X, coords_Y, meridional_data, + periodic_fix='sphere' + ) + plot_field_quivers( + ax, lon_2d, lat_2d, regrid_zonal_data, regrid_meridional_data, + spatial_filter_step=6, magnitude_filter=1.0, + projection=ccrs.PlateCarree() + ) + +# Save figure ------------------------------------------------------------------ +fig.subplots_adjust(wspace=0.25) +plt.suptitle(f't = {time:.1f} days') +plot_name = f'{plot_stem}_initial.png' +print(f'Saving figure to {plot_name}') +fig.savefig(plot_name, bbox_inches='tight') +plt.close() + +# ---------------------------------------------------------------------------- # +# FINAL PLOTTING +# ---------------------------------------------------------------------------- # +fig = plt.figure(figsize=(15, 5)) +time_idx = -1 + +for i, (field_name, colour_scheme, field_label, contours, contour_to_remove) in \ + enumerate(zip( + final_field_names, final_colour_schemes, + final_field_labels, final_contours, final_contours_to_remove)): + + # Make axes + ax = fig.add_subplot(1, 2, 1+i, projection=projection) + + # 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] / (24.*60.*60.) + + # Regrid onto lon-lat grid + field_data = regrid_horizontal_slice( + lon_2d, lat_2d, coords_X, coords_Y, field_data, periodic_fix='sphere' + ) + + # Plot data ---------------------------------------------------------------- + cmap, lines = tomplot_cmap(contours, colour_scheme, remove_contour=contour_to_remove) + cf, _ = plot_contoured_field( + ax, lon_2d, lat_2d, field_data, contour_method, contours, + cmap=cmap, line_contours=lines, projection=projection + ) + + add_colorbar_ax( + ax, cf, field_label, location='bottom', cbar_labelpad=-10, + cbar_format=cbar_format[field_name] + ) + tomplot_field_title(ax, None, minmax=True, field_data=field_data) + +# Save figure ------------------------------------------------------------------ +plt.suptitle(f't = {time:.1f} days') +plot_name = f'{plot_stem}_final.png' +print(f'Saving figure to {plot_name}') +fig.savefig(plot_name, bbox_inches='tight') +plt.close()