Add plotting script for thermal forcing and surface mass balance

landice/output_processing_li/plot_forcing_time_series.py

@@ -0,0 +1,272 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Plot time series of thermal forcing and/or surface mass balance at specified
+locations and depths. 
+@author: trevorhillebrand
+"""
+
+
+import numpy as np
+import csv
+from netCDF4 import Dataset
+from optparse import OptionParser
+from scipy.interpolate import LinearNDInterpolator, interp1d
+import matplotlib.pyplot as plt
+from matplotlib.pyplot import cm
+
+
+parser = OptionParser(description='Plot time series of thermal forcing and/or surface mass balance')
+parser.add_option("-t", "--tf", dest="thermal_forcing_file",
+                  help="List of MPAS netCDF files that contains the ismip6shelfMelt_3dThermalForcing" \
+                       " field and zOcean variable. Comma-separated, no spaces.")
+parser.add_option("-s", "--smb", dest="smb_file",
+                  help="List of MPAS netCDF files that contains the sfcMassBal" \
+                       " field. Comma-separated, no spaces.")
+parser.add_option("-m", "--mesh", dest="mesh_file",
+                  help="the MPAS netCDF file that contains the mesh variables, as well as thickness and bedTopography")
+parser.add_option("-r", "--regions", dest="regions_file", default=None,
+                  help="the MPAS netCDF file that contains the region masks")
+parser.add_option("--start_time", dest="start_time", default="0",
+                  help="beginning of time range to plot")
+parser.add_option("--end_time", dest="end_time", default="-1",
+                  help="end of time range to plot")
+parser.add_option('-c', dest='coords_file', default=None,
+                  help='CSV file containing x in first column, y in second. No header.')
+parser.add_option('-x', dest='x_coords', default=None,
+                  help='List of x coordinates of transect if not \
+                        using a csv file. Comma-separated, no spaces.')
+parser.add_option('-y', dest='y_coords', default=None,
+                  help='List of y coordinates of transect if not \
+                        using a csv file. Comma-separated, no spaces.')
+parser.add_option('-d', dest='depth', default=None,
+                  help='Depth in meters at which to plot thermal forcing.' \
+                        ' If a single value, the script will use linear 1d' \
+                        ' interpolation to determine the thermal forcing' \
+                        ' at that depth. If two comma-delimited values are given, the script' \
+                        ' will provide the average over that depth range.')
+parser.add_option('-n', dest='region_number', default=None,
+                  help='Region number to plot. If None, use entire domain.')
+parser.add_option("--seafloor", dest="plot_seafloor_thermal_forcing",
+                  action="store_true",
+                  help="plot thermal forcing at the seafloor, instead of at specific depth")
+parser.add_option("--shelf_base", dest="plot_shelf_base_thermal_forcing",
+                  action="store_true",
+                  help="plot thermal forcing at the base of floating ice, instead of at specific depth")
+parser.add_option("--shelf_range", dest="plot_shelf_depth_range_thermal_forcing",
+                  action="store_true",
+                  help="plot average thermal forcing over the whole depth range spanned by the ice shelf base")
+parser.add_option("--average", dest="plot_average_thermal_forcing",
+                  action="store_true", help='Whether to plot average' \
+                  ' thermal forcing across all coordinates provided.')
+parser.add_option("--n_samples", dest="n_samples", default=None,
+                  help="Number of random samples to take from provided coordinates.")
+parser.add_option("--save", dest="save_filename", default=None,
+                   help="File to save figure to.")
+
+options, args = parser.parse_args()
+
+rhoi = 910.
+rhosw = 1028.
+start_year = 1995  # first year in TF forcings
+scyr = 60. * 60. * 24. * 365.
+forcing_interval_years = 1.
+
+times_list = [options.start_time, options.end_time]  # list of string times for plotting
+times = [int(i) for i in times_list]  # list of integer time indices
+
+if options.coords_file is not None:
+    x = []
+    y = []
+    with open(options.coords_file, newline='') as csvfile:
+        reader = csv.reader(csvfile, delimiter=',')
+
+        for row in reader:
+            x.append(float(row[0]))
+            y.append(float(row[1]))
+    if [options.x_coords, options.y_coords] != [None, None]:
+        print('-c and -x/-y options were both provided. Reading from ',
+              f'{options.coords_file} and ignoring -x and -y settings.')
+    x = np.asarray(x)
+    y = np.asarray(y)
+else:
+    x = np.array([float(i) for i in options.x_coords.split(',')])
+    y = np.array([float(i) for i in options.y_coords.split(',')])
+
+if options.n_samples is not None:
+   rand_idx = np.random.choice(x.shape[0], int(options.n_samples), replace=False)
+   print(f"Using {options.n_samples} random samples from the {x.shape[0]} points provided.")
+   x = x[rand_idx]
+   y = y[rand_idx]
+
+# Mesh and geometry fields
+mesh = Dataset(options.mesh_file, 'r')
+mesh.set_always_mask(False)
+bed = mesh.variables["bedTopography"][:]
+thk = mesh.variables["thickness"][:]
+xCell = mesh.variables["xCell"][:]
+yCell = mesh.variables["yCell"][:]
+nCells = mesh.dimensions['nCells'].size
+areaCell = mesh.variables["areaCell"][:]
+ice_mask = thk[0, :] > 1.
+mesh.close()
+
+fig, ax = plt.subplots(1,2, sharex=True, figsize=(8,3), layout='constrained')
+
+# Get region information, if desired
+if options.region_number is not None:
+    region_number = int(options.region_number)
+    regions = Dataset(options.regions_file, 'r')
+    regionCellMasks = regions.variables["regionCellMasks"][:, region_number]
+    # Update ice_mask to only include desired region
+    ice_mask = np.logical_and(ice_mask, regionCellMasks)
+    regions.close()
+
+def interp_and_plot_tf(tf_file, plot_ax):
+    # Thermal forcing fields
+    tf_data = Dataset(tf_file, 'r')
+    tf_data.set_always_mask(False)
+    tf = tf_data.variables["ismip6shelfMelt_3dThermalForcing"][:]
+    z = tf_data.variables["ismip6shelfMelt_zOcean"][:]
+    n_time_levs = tf_data.dimensions["Time"].size
+    tf_data.close()
+    if times[1] == -1:
+        times[1] = n_time_levs - 1
+    plot_times = np.arange(times[0], times[1], step=1)  # annual posting
+
+    # Find nearest cell to desired x,y locations
+    nearest_cell = []
+    for x_i, y_i in zip(x, y):
+        nearest_cell.append(np.argmin( np.sqrt((xCell - x_i)**2. + (yCell - y_i)**2) ))
+
+    # Find depth to seafloor or ice-shelf base
+    plot_depth_range = False
+    if options.depth is not None:
+        depth = np.array([float(i) for i in options.depth.split(',')])
+        if len(depth) == 2:
+           plot_depth_range = True
+
+    if options.plot_seafloor_thermal_forcing:
+        depth = bed[0, nearest_cell]
+    elif options.plot_shelf_base_thermal_forcing:
+        # Assume this is floating ice
+        depth = -1. * rhoi / rhosw * thk[0, nearest_cell]
+    elif options.plot_shelf_depth_range_thermal_forcing:
+        depth = np.zeros(2)
+        depth[0] = np.max( (-1. * rhoi / rhosw * thk[0, nearest_cell])
+                           [thk[0, nearest_cell] > 10.] )  # ignore very thin ice
+        depth[1] = np.min( (-1. * rhoi / rhosw * thk[0, nearest_cell])
+                           [thk[0, nearest_cell] > 10.] )
+        plot_depth_range = True
+
+    # Clip depth to within the range of the TF data
+    depth[depth > np.max(z)] = np.max(z)
+    depth[depth < np.min(z)] = np.min(z)
+
+    # Vertical interpolation of ocean forcing.
+    if plot_depth_range:
+        # We'll just use the nearest zOcean depths to the
+        # requested top and bottom depths for simplicity and speed
+        print(f"Averaging TF over the depth range {depth}")
+        top_depth_index = np.argmin(np.abs(z - np.max(depth)))
+        bottom_depth_index = np.argmin(np.abs(z - np.min(depth)))
+        tf_depth = np.mean(tf[plot_times[0]:plot_times[-1] + 1, :,
+                              top_depth_index:bottom_depth_index+1], axis=2)
+    else:
+        tf_depth = []
+        for time in plot_times:
+            tf_tmp = []
+            for cell, cell_depth in zip(nearest_cell, depth):
+                tf_interp = interp1d(z, tf[time, cell, :])
+                tf_tmp.append(tf_interp(cell_depth))
+            tf_depth.append(tf_tmp)
+
+    if "UKESM" in tf_file:
+        if "SSP126" in tf_file:
+            plot_color = 'tab:green'
+        else:
+            plot_color = 'tab:blue'
+    else:
+        plot_color = 'tab:grey'
+
+    if "CESM" in tf_file:
+        linestyle = "dashed"
+    elif "CCSM" in tf_file:
+        linestyle = "dotted"
+    else:
+        linestyle = "solid"
+
+    if options.plot_average_thermal_forcing:
+        tf_avg = np.mean(tf_depth, axis=1)
+        tf_std = np.std(tf_depth, axis=1)
+        plot_ax.plot(plot_times + start_year, tf_avg, c=plot_color, linestyle=linestyle)
+        plot_ax.fill_between(plot_times + start_year, tf_avg - tf_std,
+                             tf_avg + tf_std, fc=plot_color,
+                             alpha = 0.5)
+    else:
+        plot_ax.plot(plot_times + start_year, tf_depth, c=plot_color, linestyle=linestyle)
+
+
+def plot_smb(smb_file, plot_ax):
+    smb_data = Dataset(smb_file, 'r')
+    smb_data.set_always_mask(False)
+    smb = smb_data.variables["sfcMassBal"][:, ice_mask]
+    smb_tot = np.sum(smb * areaCell[ice_mask] * scyr / 1.e12, axis=1)  # Gt/yr
+
+    n_time_levs = smb_data.dimensions["Time"].size
+    smb_data.close()
+    if times[1] == -1:
+        times[1] = n_time_levs - 1
+    plot_times = np.arange(times[0], times[1], step=1)  # annual posting
+
+    # filter smb for plotting
+    filtered_smb = np.ones_like(smb_tot)
+    filtered_smb_std = np.ones_like(smb_tot) 
+    window_width_years = 10
+    for time in range(1, n_time_levs):
+        n_t = min(time, window_width_years)
+        filtered_smb[time] = np.mean(smb_tot[time-n_t:time])
+        filtered_smb_std[time] = np.std(smb_tot[time-n_t:time])
+
+    if "UKESM" in smb_file:
+        if "SSP126" in tf_file:
+            plot_color = 'tab:green'
+        else:
+            plot_color = 'tab:blue'
+    else:
+        plot_color = 'tab:grey'
+
+    if "CESM" in smb_file:
+        linestyle = "dashed"
+    elif "CCSM" in smb_file:
+        linestyle = "dotted"
+    else:
+        linestyle = "solid"
+
+    plot_smb = filtered_smb[plot_times[0]:plot_times[-1]+1]
+    plot_smb_std = filtered_smb_std[plot_times[0]:plot_times[-1]+1]
+
+    plot_ax.plot(plot_times + start_year, plot_smb, c=plot_color, linestyle=linestyle)
+    plot_ax.fill_between(plot_times + start_year, plot_smb - plot_smb_std,
+                         plot_smb + plot_smb_std, fc=plot_color,
+                         alpha = 0.5)
+
+
+tf_files = [i for i in options.thermal_forcing_file.split(',')]
+smb_files = [i for i in options.smb_file.split(',')]
+for tf_file, smb_file in zip(tf_files, smb_files):
+    print(f"Plotting from {tf_file}")
+    interp_and_plot_tf(tf_file, ax[0])
+    print(f"Plotting from {smb_file}")
+    plot_smb(smb_file, ax[1])
+
+ax[0].set_xlabel("Year")
+ax[0].set_ylabel("Thermal forcing (°C)")
+ax[0].grid('on')
+ax[1].set_xlabel("Year")
+ax[1].set_ylabel("Total surface mass balance (Gt yr$^{-1}$)")
+ax[1].grid('on')
+if options.save_filename is not None:
+    fig.savefig(options.save_filename, dpi=400, bbox_inches='tight')
+    fig.savefig(options.save_filename + ".pdf", format='pdf', bbox_inches='tight')
+plt.show()