Update mov run

.pre-commit-config.yaml

@@ -28,6 +28,7 @@ repos:
  rev: 5.12.0
  hooks:
  - id: isort
+ args: ["--honor-noqa"]
 
  # =======================
  # Python linting

conda-env/dev.yml

@@ -11,7 +11,7 @@ dependencies:
  - python=3.10.10
  - pip=23.1.2
  - numpy=1.23.5
- - cartopy=0.21.1
+ - cartopy=0.22.0
  - matplotlib=3.7.1
  - cdat_info=8.2.1
  - cdms2=3.1.5

pcmdi_metrics/mjo/lib/post_process_plot_ensemble_mean.py

@@ -8,10 +8,10 @@
 
 
 def main():
- mip = "cmip5"
- # mip = 'cmip6'
+ # mip = "cmip5"
+ mip = "cmip6"
  exp = "historical"
- version = "v20200807"
+ version = "v20230924"
  period = "1985-2004"
  datadir = (
  "/p/user_pub/pmp/pmp_results/pmp_v1.1.2/diagnostic_results/mjo/"

pcmdi_metrics/variability_mode/lib/argparse_functions.py

@@ -150,13 +150,25 @@ def AddParserArgument(P):
  "--nc_out",
  type=bool,
  default=True,
- help="Option for generate netCDF file output: True (default) / False",
+ help="Option for generate netCDF file output for models: True (default) / False",
  )
  P.add_argument(
  "--plot",
  type=bool,
  default=True,
- help="Option for generate individual plots: True (default) / False",
+ help="Option for generate individual plots for models: True (default) / False",
+ )
+ P.add_argument(
+ "--nc_out_obs",
+ type=bool,
+ default=True,
+ help="Option for generate netCDF file output for obs: True (default) / False",
+ )
+ P.add_argument(
+ "--plot_obs",
+ type=bool,
+ default=True,
+ help="Option for generate individual plots for obs: True (default) / False",
  )
  P.add_argument(
  "--parallel",

pcmdi_metrics/variability_mode/lib/plot_map.py

@@ -1,16 +1,25 @@
+import faulthandler
 import sys
 
-import cartopy
 import cartopy.crs as ccrs
 import matplotlib.path as mpath
 import matplotlib.pyplot as plt
 import matplotlib.ticker as mticker
 import numpy as np
+import xarray as xr
+from cartopy.feature import LAND as cartopy_land
+from cartopy.feature import OCEAN as cartopy_ocean
 from cartopy.mpl.gridliner import LATITUDE_FORMATTER, LONGITUDE_FORMATTER
 from cartopy.mpl.ticker import LatitudeFormatter, LongitudeFormatter
 
+from pcmdi_metrics.variability_mode.lib import debug_print
 
-def plot_map(mode, model, syear, eyear, season, eof_Nth, frac_Nth, output_file_name):
+faulthandler.enable()
+
+
+def plot_map(
+ mode, model, syear, eyear, season, eof_Nth, frac_Nth, output_file_name, debug=False
+):
  """Plot dive down map and save
 
  Parameters
@@ -34,7 +43,6 @@ def plot_map(mode, model, syear, eyear, season, eof_Nth, frac_Nth, output_file_n
  """
  # Map Projection
  if "teleconnection" in mode:
- # projection = "PlateCarree"
  projection = "Robinson"
  elif mode in ["NAO", "PNA", "NPO", "PDO", "NPGO", "AMO"]:
  projection = "Lambert"
@@ -67,6 +75,10 @@ def plot_map(mode, model, syear, eyear, season, eof_Nth, frac_Nth, output_file_n
  + percentage
  )
 
+ debug_print(
+ "plot_map: projection, plot_title:" + projection + ", " + plot_title, debug
+ )
+
  gridline = True
 
  if mode in [
@@ -84,40 +96,53 @@ def plot_map(mode, model, syear, eyear, season, eof_Nth, frac_Nth, output_file_n
  maskout = None
 
  if mode in ["AMO", "AMO_teleconnection"]:
- center_lon_global = 0
+ central_longitude = 0
  else:
- center_lon_global = 180
+ central_longitude = 180
+
+ # Convert cdms variable to xarray
+ lons = eof_Nth.getLongitude()
+ lats = eof_Nth.getLatitude()
+ data = np.array(eof_Nth)
+ lon = np.array(lons)
+ lat = np.array(lats)
+ lon, lat = np.meshgrid(lon, lat)
+ data_array = xr.DataArray(
+ np.array(data), coords={"lon": lon[0, :], "lat": lat[:, 0]}, dims=("lat", "lon")
+ )
+ data_array = data_array.where(data_array != 1e20, np.nan)
 
  plot_map_cartopy(
- eof_Nth,
+ data_array,
  output_file_name,
  title=plot_title,
  proj=projection,
  gridline=gridline,
  levels=levels,
  maskout=maskout,
- center_lon_global=center_lon_global,
+ central_longitude=central_longitude,
+ debug=debug,
  )
 
 
 def plot_map_cartopy(
- data,
- filename,
+ data_array,
+ filename=None,
  title=None,
  gridline=True,
  levels=None,
  proj="PlateCarree",
  data_area="global",
  cmap="RdBu_r",
- center_lon_global=180,
+ central_longitude=180,
  maskout=None,
  debug=False,
 ):
  """
  Parameters
  ----------
- data : trainsisent variable
- 2D cdms2 TransientVariable with lat/lon coordinates attached.
+ data : data_array
+ 2D xarray DataArray with lat/lon coordinates attached.
  filename : str
  Output file name (it is okay to omit '.png')
  title : str, optional
@@ -138,54 +163,62 @@ def plot_map_cartopy(
  Switch for debugging print statements (default is False)
  """
 
- lons = data.getLongitude()
- lats = data.getLatitude()
+ debug_print("plot_map_cartopy starts", debug)
+
+ lon = data_array.lon
+ lat = data_array.lat
+
+ # Determine the extent based on the longitude range where data exists
+ lon_min = lon.min().item()
+ lon_max = lon.max().item()
+ lat_min = lat.min().item()
+ lat_max = lat.max().item()
 
- min_lon = min(lons)
- max_lon = max(lons)
- min_lat = min(lats)
- max_lat = max(lats)
  if debug:
- print(min_lon, max_lon, min_lat, max_lat)
+ print(lon_min, lon_max, lat_min, lat_max)
+
+ debug_print("Central longitude setup starts", debug)
+ debug_print("proj: " + proj, debug)
 
  # map types example:
  # https://github.com/SciTools/cartopy-tutorial/blob/master/tutorial/projections_crs_and_terms.ipynb
 
  if proj == "PlateCarree":
- projection = ccrs.PlateCarree(central_longitude=center_lon_global)
+ projection = ccrs.PlateCarree(central_longitude=central_longitude)
  elif proj == "Robinson":
- projection = ccrs.Robinson(central_longitude=center_lon_global)
+ projection = ccrs.Robinson(central_longitude=central_longitude)
  elif proj == "Stereo_north":
  projection = ccrs.NorthPolarStereo()
  elif proj == "Stereo_south":
  projection = ccrs.SouthPolarStereo()
  elif proj == "Lambert":
- max_lat = min(max_lat, 80)
+ lat_max = min(lat_max, 80)
  if debug:
- print("revised maxlat:", max_lat)
- central_longitude = (min_lon + max_lon) / 2.0
- central_latitude = (min_lat + max_lat) / 2.0
+ print("revised maxlat:", lat_max)
+ central_longitude = (lon_min + lon_max) / 2.0
+ central_latitude = (lat_min + lat_max) / 2.0
  projection = ccrs.AlbersEqualArea(
  central_longitude=central_longitude,
  central_latitude=central_latitude,
- standard_parallels=(20, max_lat),
+ standard_parallels=(20, lat_max),
  )
+ else:
+ print("Error: projection not defined!")
+
+ if debug:
+ debug_print("Central longitude setup completes", debug)
+ print("projection:", projection)
 
  # Generate plot
- fig = plt.figure(figsize=(8, 6))
- ax = plt.axes(projection=projection)
- im = ax.contourf(
- lons,
- lats,
- data,
- transform=ccrs.PlateCarree(),
- cmap=cmap,
- levels=levels,
- extend="both",
- )
+ fig, ax = plt.subplots(subplot_kw={"projection": projection}, figsize=(8, 6))
+ debug_print("fig, ax done", debug)
+
+ # Add coastlines
  ax.coastlines()
+ debug_print("Generate plot completed", debug)
 
  # Grid Lines and tick labels
+ debug_print("projection starts", debug)
  if proj == "PlateCarree":
  if data_area == "global":
  if gridline:
@@ -203,6 +236,7 @@ def plot_map_cartopy(
  if gridline:
  gl = ax.gridlines(alpha=0.5, linestyle="--")
  elif "Stereo" in proj:
+ debug_print(proj + " start", debug)
  if gridline:
  gl = ax.gridlines(draw_labels=True, alpha=0.5, linestyle="--")
  gl.xlocator = mticker.FixedLocator(
@@ -226,60 +260,82 @@ def plot_map_cartopy(
  verts = np.vstack([np.sin(theta), np.cos(theta)]).T
  circle = mpath.Path(verts * radius + center)
  ax.set_boundary(circle, transform=ax.transAxes)
+ debug_print(proj + " plotted", debug)
  elif proj == "Lambert":
  # Make a boundary path in PlateCarree projection, I choose to start in
  # the bottom left and go round anticlockwise, creating a boundary point
  # every 1 degree so that the result is smooth:
  # https://stackoverflow.com/questions/43463643/cartopy-albersequalarea-limit-region-using-lon-and-lat
+
  vertices = [
- (lon - 180, min_lat) for lon in range(int(min_lon), int(max_lon + 1), 1)
- ] + [(lon - 180, max_lat) for lon in range(int(max_lon), int(min_lon - 1), -1)]
+ (lon - 180, lat_min) for lon in range(int(lon_min), int(lon_max + 1), 1)
+ ] + [(lon - 180, lat_max) for lon in range(int(lon_max), int(lon_min - 1), -1)]
  boundary = mpath.Path(vertices)
- ax.set_boundary(boundary, transform=ccrs.PlateCarree(central_longitude=180))
- ax.set_extent([min_lon, max_lon, min_lat, max_lat], crs=ccrs.PlateCarree())
+ ax.set_boundary(
+ boundary, transform=ccrs.PlateCarree(central_longitude=180)
+ ) # Here, 180 should be hardcoded, otherwise AMO map will be at out of figure box
+
+ ax.set_extent([lon_min, lon_max, lat_min, lat_max], crs=ccrs.PlateCarree())
+
  if gridline:
  gl = ax.gridlines(
  draw_labels=True,
  alpha=0.8,
  linestyle="--",
- crs=cartopy.crs.PlateCarree(),
+ crs=ccrs.PlateCarree(),
  )
  gl.xformatter = LONGITUDE_FORMATTER
  gl.yformatter = LATITUDE_FORMATTER
  gl.ylocator = mticker.FixedLocator([30, 60])
  gl.xlocator = mticker.FixedLocator([120, 160, 200 - 360, 240 - 360])
  gl.top_labels = False # suppress top labels
  # suppress right labels
- # gl.right_labels = False
+ gl.right_labels = False
  for ea in gl.ylabel_artists:
  right_label = ea.get_position()[0] > 0
  if right_label:
  ea.set_visible(False)
 
- # Add title
- plt.title(title, pad=15, fontsize=15)
-
- # Add colorbar
- posn = ax.get_position()
- cbar_ax = fig.add_axes([0, 0, 0.1, 0.1])
- cbar_ax.set_position([posn.x0 + posn.width + 0.01, posn.y0, 0.01, posn.height])
- cbar = plt.colorbar(im, cax=cbar_ax)
- cbar.ax.tick_params(labelsize=10)
+ debug_print("projection completed", debug)
 
- if proj == "PlateCarree":
- ax.set_aspect("auto", adjustable=None)
+ # Plot contours from the data
+ contourf_plot = ax.contourf(
+ lon,
+ lat,
+ data_array,
+ levels=levels,
+ cmap=cmap,
+ extend="both",
+ transform=ccrs.PlateCarree(),
+ )
+ debug_print("contourf done", debug)
 
  # Maskout
  if maskout is not None:
  if maskout == "land":
  ax.add_feature(
- cartopy.feature.LAND, zorder=100, edgecolor="k", facecolor="lightgrey"
+ cartopy_land, zorder=100, edgecolor="k", facecolor="lightgrey"
  )
  if maskout == "ocean":
  ax.add_feature(
- cartopy.feature.OCEAN, zorder=100, edgecolor="k", facecolor="lightgrey"
+ cartopy_ocean, zorder=100, edgecolor="k", facecolor="lightgrey"
  )
+ if proj == "PlateCarree":
+ ax.set_aspect("auto", adjustable=None)
+
+ # Add title
+ ax.set_title(title, pad=15, fontsize=15)
+
+ # Add a colorbar
+ posn = ax.get_position()
+ cbar_ax = fig.add_axes([0, 0, 0.1, 0.1])
+ cbar_ax.set_position([posn.x0 + posn.width + 0.03, posn.y0, 0.01, posn.height])
+ cbar = plt.colorbar(contourf_plot, cax=cbar_ax)
+ cbar.ax.tick_params(labelsize=10)
 
  # Done, save figure
- fig.savefig(filename)
+ if filename is not None:
+ debug_print("plot done, save figure as " + filename, debug)
+ fig.savefig(filename)
+
  plt.close("all")