Add config file for diagnostic scripts and use it to generalize sst_eval.py (#74)

* Added config file for diagnostic scripts, modified sst_eval to use config file

* Added options to set location of annotations and fig size, used rename_map keys instead of separate list

* Added option to choose projection

* Added transforms and  changed x/ylim to setextent to handle different projections

* Added else statement to try/except

* Abstracted glorys/oisst processing to function, added logging, added option to select data projection

* Align model and model grid before regridding, remove else statement from process_glorys

* Added support for plotting along lines of latitude, removed corners from sst_eval import list

* Dropped lat points that caused regridding error, factored out common line from annotate function, addressed label_mode error

* Made command line flags required

---------

Co-authored-by: Utheri Wagura <Utheri.Wagura@an105.princeton.rdhpcs.noaa.gov>
Co-authored-by: Utheri Wagura <Utheri.Wagura@an002.princeton.rdhpcs.noaa.gov>
Co-authored-by: Utheri Wagura <Utheri.Wagura@an202.princeton.rdhpcs.noaa.gov>
Co-authored-by: Utheri Wagura <Utheri.Wagura@an200.princeton.rdhpcs.noaa.gov>
Co-authored-by: Utheri Wagura <Utheri.Wagura@an106.princeton.rdhpcs.noaa.gov>
Co-authored-by: Utheri Wagura <Utheri.Wagura@an102.princeton.rdhpcs.noaa.gov>
Co-authored-by: Utheri Wagura <Utheri.Wagura@an104.princeton.rdhpcs.noaa.gov>

Author: 8 people

diagnostics/physics/config.yaml

@@ -0,0 +1,61 @@
+figures_dir: 'figures/'
+glorys: '/work/acr/mom6/diagnostics/glorys/glorys_sfc.nc'
+model_grid: '../data/geography/ocean_static.nc'
+
+# Variables to rename
+rename_map:
+    geolon: lon
+    geolat: lat
+    geolon_c: lon_b
+    geolat_c: lat_b
+
+# Domain and plotting details
+domain: ocean_monthly
+
+# Lat/lon range 
+lat: 
+    south: 0
+    north: 60
+
+lon:
+    west: 260
+    east: 330
+
+# Xlim/ylim to plot
+x: 
+    min: -99
+    max: -35
+
+y:
+    min: 4
+    max: 59
+
+# Projection ( current options are either NorthPolarStereo for the arctic, or PlateCarree for all other domains
+projection_grid: 'PlateCarree'
+projection_data: 'PlateCarree'
+
+# Figure size
+fig_width: 11
+fig_height: 14
+
+# Location of skill score annotations in plot
+text_x: -98.5
+text_y: 54
+
+# Space between skill score text
+text_xint: 4 # This is unused if col=1, so it is only included for completeness
+text_yint: 4
+plot_lat: False
+
+# SST_eval settings
+oisst: '/work/acr/oisstv2/' 
+
+# Colorbar for sst plots
+levels_min: 2
+levels_max: 31
+levels_step: 2
+
+# Colorbar for sst difference plots
+bias_min: -2
+bias_max: 2.1
+bias_step: 0.25

diagnostics/physics/plot_common.py

@@ -7,6 +7,12 @@
 import os
 import xarray
 import xskillscore
+import logging
+import yaml
+import xesmf
+
+# Configure logging for plot_common
+logger = logging.getLogger(__name__)
 
 def center_to_outer(center, left=None, right=None):
     """
@@ -46,25 +52,39 @@ def get_map_norm(cmap, levels, no_offset=True):
     norm = BoundaryNorm(levels, ncolors=nlev, clip=False)
     return cmap, norm
 
-def annotate_skill(model, obs, ax, dim=['yh', 'xh'], x0=-98.5, y0=54, yint=4, xint=4, weights=None, cols=1, **kwargs):
+def annotate_skill(model, obs, ax, dim=['yh', 'xh'], x0=-98.5, y0=54, yint=4, xint=4, weights=None, cols=1, proj = ccrs.PlateCarree(), plot_lat=False,**kwargs):
     """
     Annotate an axis with model vs obs skill metrics
     """
     bias = xskillscore.me(model, obs, dim=dim, skipna=True, weights=weights)
     rmse = xskillscore.rmse(model, obs, dim=dim, skipna=True, weights=weights)
     corr = xskillscore.pearson_r(model, obs, dim=dim, skipna=True, weights=weights)
     medae = xskillscore.median_absolute_error(model, obs, dim=dim, skipna=True)
-    ax.text(x0, y0, f'Bias: {float(bias):2.2f}', **kwargs)
-    ax.text(x0, y0-yint, f'RMSE: {float(rmse):2.2f}', **kwargs)
-    if cols == 1:
-        ax.text(x0, y0-yint*2, f'MedAE: {float(medae):2.2f}', **kwargs)
-        ax.text(x0, y0-yint*3, f'Corr: {float(corr):2.2f}', **kwargs)
-    elif cols == 2:
-        ax.text(x0+xint, y0, f'MedAE: {float(medae):2.2f}', **kwargs)
-        ax.text(x0+xint, y0-yint, f'Corr: {float(corr):2.2f}', **kwargs)
+
+    ax.text(x0, y0, f'Bias: {float(bias):2.2f}', transform=proj, **kwargs)
+    # Set plot_lat=True in order to plot skill along a line of latitude. Otherwise, plot along longitude
+    if plot_lat:
+        ax.text(x0-xint, y0, f'RMSE: {float(rmse):2.2f}', transform=proj, **kwargs)
+        if cols == 1:
+            ax.text(x0-xint*2, y0, f'MedAE: {float(medae):2.2f}', transform=proj, **kwargs)
+            ax.text(x0-xint*3, y0, f'Corr: {float(corr):2.2f}', transform=proj, **kwargs)
+        elif cols == 2:
+            ax.text(x0, y0+yint, f'MedAE: {float(medae):2.2f}', transform=proj, **kwargs)
+            ax.text(x0-xint, y0+yint, f'Corr: {float(corr):2.2f}', transform=proj, **kwargs)
+        else:
+            raise ValueError(f'Unsupported number of columns: {cols}')
+
     else:
-        raise ValueError(f'Unsupported number of columns: {cols}')
-    
+        ax.text(x0, y0-yint, f'RMSE: {float(rmse):2.2f}', transform=proj, **kwargs)
+        if cols == 1:
+            ax.text(x0, y0-yint*2, f'MedAE: {float(medae):2.2f}', transform=proj, **kwargs)
+            ax.text(x0, y0-yint*3, f'Corr: {float(corr):2.2f}', transform=proj, **kwargs)
+        elif cols == 2:
+            ax.text(x0+xint, y0, f'MedAE: {float(medae):2.2f}', transform=proj, **kwargs)
+            ax.text(x0+xint, y0-yint, f'Corr: {float(corr):2.2f}', transform=proj, **kwargs)
+        else:
+            raise ValueError(f'Unsupported number of columns: {cols}')
+
 def autoextend_colorbar(ax, plot, plot_array=None, **kwargs):
     """
     Add a colorbar, setting the extend metric based on 
@@ -141,3 +161,67 @@ def save_figure(fname, label='', pdf=False, output_dir='figures'):
         plt.savefig(os.path.join(output_dir, f'{fname}_{label}.png'), dpi=200, bbox_inches='tight')
         if pdf:
             plt.savefig(os.path.join(output_dir, f'{fname}_{label}.pdf'), bbox_inches='tight')
+
+def load_config(config_path: str):
+    """Load the configuration file."""
+    try:
+        with open(config_path, 'r') as f:
+            config = yaml.safe_load(f)
+            logger.info(f"Loaded configuration from {config_path}")
+            return config
+    except Exception as e:
+        logger.error(f"Error loading configuration from {config_path}: {e}")
+        raise
+
+def process_oisst(config, target_grid, model_ave):
+    """Open and regrid OISST dataset, return relevant vars from dataset."""
+    try:
+        oisst = (
+            xarray.open_mfdataset([config['oisst'] + f'sst.month.mean.{y}.nc' for y in range(1993, 2020)])
+            .sst
+            .sel(lat=slice(config['lat']['south'], config['lat']['north']), lon=slice(config['lon']['west'], config['lon']['east']))
+        )
+    except Exception as e:
+        logger.error(f"Error processing OISST data: {e}")
+        raise e("Could not open OISST dataset")
+
+    # Drop any latitude points greater than or equal to 90 to prevent errors when regridding
+    oisst = oisst.where( oisst.lat < 90, drop = True)
+
+    oisst_lonc, oisst_latc = corners(oisst.lon, oisst.lat)
+    oisst_lonc -= 360
+    mom_to_oisst = xesmf.Regridder(
+        target_grid,
+        {'lat': oisst.lat, 'lon': oisst.lon, 'lat_b': oisst_latc, 'lon_b': oisst_lonc},
+        method='conservative_normed',
+        unmapped_to_nan=True
+    )
+
+    oisst_ave = oisst.mean('time').load()
+    mom_rg = mom_to_oisst(model_ave)
+    logger.info("OISST data processed successfully.")
+    return mom_rg, oisst_ave, oisst_lonc, oisst_latc
+
+def process_glorys(config, target_grid):
+    """ Open and regrid glorys data, return regridded glorys data """
+    glorys = xarray.open_dataset( config['glorys'] ).squeeze(drop=True)['thetao'] #.rename({'longitude': 'lon', 'latitude': 'lat'})
+
+    # Drop any latitude points greater than or equal to 90 to prevent errors when regridding, then get corner points
+    try:
+        glorys = glorys.where( glorys.lat < 90, drop = True)
+        glorys_lonc, glorys_latc = corners(glorys.lon, glorys.lat)
+        logger.info("Glorys data is using lon/lat")
+    except AttributeError:
+        glorys = glorys.where( glorys.latitude < 90, drop = True)
+        glorys_lonc, glorys_latc = corners(glorys.longitude, glorys.latitude)
+        logger.info("Glorys data is using longitude/latitude")
+    except:
+        logger.error("Name of longitude and latitude variables is unknown")
+        raise Exception("Error: Lat/Latitude, Lon/Longitdue not found in glorys data")
+
+    glorys_ave = glorys.mean('time').load()
+    glorys_to_mom = xesmf.Regridder(glorys_ave, target_grid, method='bilinear', unmapped_to_nan=True)
+    glorys_rg = glorys_to_mom(glorys_ave)
+
+    logger.info("Glorys data processed successfully.")
+    return glorys_rg, glorys_ave, glorys_lonc, glorys_latc

diagnostics/physics/sst_eval.py

@@ -3,7 +3,7 @@
 Uses whatever model data can be found within the directory pp_root,
 and does not try to match the model and observed time periods.
 How to use:
-python sst_eval.py /archive/acr/fre/NWA/2023_04/NWA12_COBALT_2023_04_kpo4-coastatten-physics/gfdl.ncrc5-intel22-prod 
+python sst_eval.py -p /archive/acr/fre/NWA/2023_04/NWA12_COBALT_2023_04_kpo4-coastatten-physics/gfdl.ncrc5-intel22-prod -c config.yaml
 """
 
 import cartopy.crs as ccrs
@@ -12,55 +12,61 @@
 from mpl_toolkits.axes_grid1 import AxesGrid
 import numpy as np
 import xarray
-import xesmf
+import logging
 
-from plot_common import annotate_skill, autoextend_colorbar, corners, get_map_norm, open_var
+from plot_common import annotate_skill, autoextend_colorbar, get_map_norm, open_var, load_config, process_oisst, process_glorys
 
+# Configure logging for sst_eval
+logger = logging.getLogger(__name__)
+logging.basicConfig(filename="sst_eval.log", format='%(asctime)s %(levelname)s:%(name)s: %(message)s',level=logging.INFO)
 
-def plot_sst_eval(pp_root):
-    model = open_var(pp_root, 'ocean_monthly', 'tos')
-    model_grid = xarray.open_dataset('../data/geography/ocean_static.nc')
-    target_grid = model_grid[['geolon', 'geolat', 'geolon_c', 'geolat_c']].rename({'geolon': 'lon', 'geolat': 'lat', 'geolon_c': 'lon_b', 'geolat_c': 'lat_b'})
+def plot_sst_eval(pp_root, config):
+    
+    model = open_var(pp_root, config['domain'], 'tos')
+    model_grid = xarray.open_dataset( config['model_grid'] )
+    target_grid = model_grid[config['rename_map'].keys()].rename(config['rename_map'])
     model_ave = model.mean('time').load()
+    logger.info("MODEL_GRID: %s",model_grid)
+    logger.info("MODEL_AVE: %s",model_ave)
+    logger.info("Successfully opened model grid and took mean over time")
 
-    glorys = xarray.open_dataset('/work/acr/mom6/diagnostics/glorys/glorys_sfc.nc')['thetao'] #.rename({'longitude': 'lon', 'latitude': 'lat'})
-    glorys_lonc, glorys_latc = corners(glorys.lon, glorys.lat)
-    glorys_ave = glorys.mean('time').load()
-    glorys_to_mom = xesmf.Regridder(glorys_ave, target_grid, method='bilinear', unmapped_to_nan=True)
-    glorys_rg = glorys_to_mom(glorys_ave)
+    # Verify that xh/yh are set as coordinates, then make sure model coordinates match grid data
+    model_grid = model_grid.assign_coords( {'xh':model_grid.xh, 'yh':model_grid.yh } )
+    model_ave = xarray.align(model_grid, model_ave,join='override')[1]
+
+    glorys_rg, glorys_ave, glorys_lonc, glorys_latc = process_glorys(config, target_grid)
     delta_glorys = model_ave - glorys_rg
+    logger.info("GLORYS_RG: %s",glorys_rg)
+    logger.info("DELTA_GLORYS: %s",delta_glorys)
 
-    oisst = (
-        xarray.open_mfdataset([f'/work/acr/oisstv2/sst.month.mean.{y}.nc' for y in range(1993, 2020)])
-        .sst
-        .sel(lat=slice(0, 60), lon=slice(360-100, 360-30))
-    )
-    oisst_lonc, oisst_latc = corners(oisst.lon, oisst.lat)
-    oisst_lonc -= 360
-    mom_to_oisst = xesmf.Regridder(
-        target_grid, 
-        {'lat': oisst.lat, 'lon': oisst.lon, 'lat_b': oisst_latc, 'lon_b': oisst_lonc}, 
-        method='conservative_normed', 
-        unmapped_to_nan=True
-    )
-    oisst_ave = oisst.mean('time').load()
-    mom_rg = mom_to_oisst(model_ave)
+    mom_rg, oisst_ave, oisst_lonc, oisst_latc = process_oisst(config, target_grid, model_ave)
     delta_oisst = mom_rg - oisst_ave
+    logger.info("OISST_AVE: %s",oisst_ave)
+    logger.info("DELTA_OISST: %s",delta_oisst)
+
+    fig = plt.figure(figsize=(config['fig_width'], config['fig_height']))
+
+    # Set projection of each grid in the plot
+    # For now, sst_eval.py will only support a projection for the arctic and a projection for all other domains
+    if config['projection_grid'] == 'NorthPolarStereo':
+        p = ccrs.NorthPolarStereo()
+    else:
+        p = ccrs.PlateCarree()
 
-    fig = plt.figure(figsize=(11, 14))
     grid = AxesGrid(fig, 111,
-        axes_class=(GeoAxes, dict(projection=ccrs.PlateCarree())),
+        axes_class=(GeoAxes, dict( projection = p )),
         nrows_ncols=(2, 3),
         axes_pad=0.3,
         cbar_location='bottom',
         cbar_mode='edge',
         cbar_pad=0.2,
         cbar_size='15%',
-        label_mode=''
+        label_mode='keep'
     )
+    logger.info("Successfully created grid")
 
     # Discrete levels and colorbar for SST plots
-    levels = np.arange(2, 31, 2)
+    levels = np.arange(config['levels_min'], config['levels_max'], config['levels_step'])
     try:
         import cmcrameri.cm as cmc
     except ModuleNotFoundError:
@@ -70,55 +76,70 @@ def plot_sst_eval(pp_root):
     cmap, norm = get_map_norm(cmap, levels=levels)
 
     # Discrete levels and colorbar for difference plots
-    bias_levels = np.arange(-2, 2.1, 0.25)
+    bias_levels = np.arange(config['bias_min'], config['bias_max'], config['bias_step'])
     bias_cmap, bias_norm = get_map_norm('coolwarm', levels=bias_levels)
     bias_common = dict(cmap=bias_cmap, norm=bias_norm)
 
+    # Set projection of input data files so that data is correctly tranformed when plotting
+    # For now, sst_eval.py will only support a projection for the arctic and a projection for all other domains
+    if config['projection_data'] == 'NorthPolarStereo':
+        proj = ccrs.NorthPolarStereo()
+    else:
+        proj = ccrs.PlateCarree()
+
     # Model 
-    p0 = grid[0].pcolormesh(model_grid.geolon_c, model_grid.geolat_c, model_ave, cmap=cmap, norm=norm)
+    p0 = grid[0].pcolormesh(model_grid.geolon_c, model_grid.geolat_c, model_ave, cmap=cmap, norm=norm, transform=proj)
     grid[0].set_title('(a) Model')
     cbar1 = grid.cbar_axes[0].colorbar(p0)
     cbar1.ax.set_xlabel('Mean SST (°C)')
+    logger.info("Successfully plotted model data")
 
     # OISST
-    grid[1].pcolormesh(oisst_lonc, oisst_latc, oisst_ave, cmap=cmap, norm=norm)
+    grid[1].pcolormesh(oisst_lonc, oisst_latc, oisst_ave, cmap=cmap, norm=norm, transform=proj)
     grid[1].set_title('(b) OISST')
+    logger.info("Successfully plotted oisst")
 
     # Model - OISST
-    grid[2].pcolormesh(oisst_lonc, oisst_latc, delta_oisst, **bias_common)
+    grid[2].pcolormesh(oisst_lonc, oisst_latc, delta_oisst, transform=proj,**bias_common)
     grid[2].set_title('(c) Model - OISST')
-    annotate_skill(mom_rg, oisst_ave, grid[2], dim=['lat', 'lon'])
+    annotate_skill(mom_rg, oisst_ave, grid[2], dim=['lat', 'lon'], x0=config['text_x'], y0=config['text_y'], xint=config['text_xint'], plot_lat=config['plot_lat'])
+    logger.info("Successfully plotted difference between model and oisst")
 
     # GLORYS
-    p1 = grid[4].pcolormesh(glorys_lonc, glorys_latc, glorys_ave, cmap=cmap, norm=norm)
+    p1 = grid[4].pcolormesh(glorys_lonc, glorys_latc, glorys_ave, cmap=cmap, norm=norm, transform=proj)
     grid[4].set_title('(d) GLORYS12')
     cbar1 = autoextend_colorbar(grid.cbar_axes[1], p1)
     cbar1.ax.set_xlabel('Mean SST (°C)')
+    logger.info("Successfully plotted glorys")
 
     # Model - GLORYS
-    p2 = grid[5].pcolormesh(model_grid.geolon_c, model_grid.geolat_c, delta_glorys, **bias_common)
+    p2 = grid[5].pcolormesh(model_grid.geolon_c, model_grid.geolat_c, delta_glorys, transform=proj, **bias_common)
     cbar2 = autoextend_colorbar(grid.cbar_axes[2], p2)
     cbar2.ax.set_xlabel('SST difference (°C)')
     cbar2.ax.set_xticks([-2, -1, 0, 1, 2])
     grid[5].set_title('(e) Model - GLORYS12')
-    annotate_skill(model_ave, glorys_rg, grid[5], weights=model_grid.areacello)
+    annotate_skill(model_ave, glorys_rg, grid[5], weights=model_grid.areacello, x0=config['text_x'], y0=config['text_y'], xint=config['text_xint'], plot_lat=config['plot_lat'])
+    logger.info("Successfully plotted difference between glorys and model")
 
     for ax in grid:
-        ax.set_xlim(-99, -35)
-        ax.set_ylim(4, 59)
+        ax.set_extent([ config['x']['min'], config['x']['max'], config['y']['min'], config['y']['max'] ], crs=proj )
         ax.set_xticks([])
         ax.set_yticks([])
         ax.set_xticklabels([])
         ax.set_yticklabels([])
         for s in ax.spines.values():
             s.set_visible(False)
-        
-    plt.savefig('figures/sst_eval.png', dpi=300, bbox_inches='tight')
+    logger.info("Successfully set extent of each axis")
+
+    plt.savefig(config['figures_dir']+'sst_eval.png', dpi=300, bbox_inches='tight')
+    logger.info("Successfully saved figure")
 
 
 if __name__ == '__main__':
     from argparse import ArgumentParser
     parser = ArgumentParser()
-    parser.add_argument('pp_root', help='Path to postprocessed data (up to but not including /pp/)')
+    parser.add_argument('-p','--pp_root', type=str, help='Path to postprocessed data (up to but not including /pp/)', required=True)
+    parser.add_argument('-c','--config', type=str, help='Path to config.yaml file containing relevant paths for diagnostic scripts', required=True)
     args = parser.parse_args()
-    plot_sst_eval(args.pp_root)
+    config = load_config(args.config)
+    plot_sst_eval(args.pp_root, config)