diff --git a/icenet/plotting/forecast.py b/icenet/plotting/forecast.py index d2cc2a4..714bf01 100644 --- a/icenet/plotting/forecast.py +++ b/icenet/plotting/forecast.py @@ -33,7 +33,7 @@ get_seas_forecast_init_dates, lat_lon_box, show_img, get_plot_axes, process_probes, - process_regions) + process_regions, get_custom_cmap) from icenet.plotting.video import xarray_to_video @@ -1395,7 +1395,8 @@ def __init__(self, *args, forecast_date: bool = True, **kwargs): default=None, type=region_arg, help="Region specified as lon and lat min/max: lon_min, lat_min, lon_max, lat_max") - self.add_argument("--north-facing", + self.add_argument("-nf", + "--north-facing", action="store_true", default=False) @@ -1673,7 +1674,7 @@ def plot_forecast(show_plot=False): pred_da = pred_da.drop("time").drop("leadtime").\ rename(leadtime="time", forecast_date="time").set_index(time="time") - anim_args = dict(figsize=5) + anim_args = dict(figsize=(10, 8)) output_filename = os.path.join( output_path, @@ -1736,6 +1737,10 @@ def plot_forecast(show_plot=False): # ax.set_global() else: cmap.set_bad("dimgrey", alpha=0) + # Hack since cartopy needs transparency for nan regions to wraparound + # correctly with pcolormesh. + data = np.where(np.isnan(pred_da), -9999, pred_da) + custom_cmap = get_custom_cmap(cmap) lon, lat = fc.lon.values, fc.lat.values transformed_coords = source_crs.transform_points(target_crs, lon, lat) @@ -1749,7 +1754,7 @@ def plot_forecast(show_plot=False): proj=target_crs, north_facing=True ) - im = ax.pcolormesh(x, y, pred_da, transform=source_crs, vmin=0, vmax=vmax, cmap=cmap) + im = ax.pcolormesh(x, y, data, transform=source_crs, vmin=0, vmax=vmax, cmap=custom_cmap) if not args.no_coastlines: ax.add_feature(cfeature.LAND, facecolor="dimgrey") # ax.add_feature(cfeature.COASTLINE) @@ -1768,7 +1773,7 @@ def plot_forecast(show_plot=False): output_path, "{}.{}_reference.{}{}".format( forecast_name, (args.forecast_date + dt.timedelta(days=leadtime)).strftime("%Y%m%d"), - "" if not args.stddev else "stddev.", args.format)) + "" if not args.stddev else "stddev.", "jpg")) plt.savefig(output_filename) for handle in region_plot: diff --git a/icenet/plotting/utils.py b/icenet/plotting/utils.py index 188af7f..e0a937e 100644 --- a/icenet/plotting/utils.py +++ b/icenet/plotting/utils.py @@ -5,6 +5,7 @@ import re import cartopy.crs as ccrs +import matplotlib as mpl import matplotlib.pyplot as plt import numpy as np import pandas as pd @@ -498,3 +499,15 @@ def lat_lon_box(lon_bounds: np.array, lat_bounds: np.array, segments: int=1): lons[i*segments:(i+1)*segments] = np.linspace(lon_bounds[lon_min], lon_bounds[lon_max], num=segments) return lats, lons + +def get_custom_cmap(cmap, vmin=0, vmax=1): + """Creates a new colormap for valid array with range 0-1, but with nan set to <0. + + Hack since cartopy needs transparency for nan regions to wraparound + correctly with pcolormesh. + """ + colors = cmap(np.linspace(vmin, vmax, cmap.N)) + custom_cmap = mpl.colors.ListedColormap(colors) + custom_cmap.set_bad("dimgrey", alpha=0) + custom_cmap.set_under("dimgrey") + return custom_cmap diff --git a/icenet/plotting/video.py b/icenet/plotting/video.py index 5503167..c538421 100644 --- a/icenet/plotting/video.py +++ b/icenet/plotting/video.py @@ -18,7 +18,7 @@ from icenet.process.predict import get_refcube from icenet.utils import setup_logging - +from icenet.plotting.utils import get_custom_cmap # TODO: This can be a plotting or analysis util function elsewhere def get_dataarray_from_files(files: object, numpy: bool = False) -> object: @@ -87,7 +87,7 @@ def xarray_to_video( data_type: str = 'abs', video_dates: object = None, cmap: object = "viridis", - figsize: int = 12, + figsize: int = (12, 12), dpi: int = 150, imshow_kwargs: dict = None, ax_init: object = None, @@ -124,9 +124,15 @@ def xarray_to_video( source_crs = ccrs.LambertAzimuthalEqualArea(central_latitude=pole*90, central_longitude=0) target_crs = ccrs.PlateCarree() + def update(date): logging.debug("Plotting {}".format(date.strftime("%D"))) - image.set_array(da.sel(time=date)) + data = da.sel(time=date) + # Hack since cartopy needs transparency for nan regions to wraparound + # correctly with pcolormesh. + if north_facing: + data = np.where(np.isnan(data), -9999, data) + image.set_array(data) image_title.set_text("{:04d}/{:02d}/{:02d}".format( date.year, date.month, date.day)) @@ -167,12 +173,15 @@ def update(date): projection = target_crs if method == "lat_lon" else source_crs if ax_init is None: if north_facing: - fig, ax = plt.subplots(figsize=(figsize, figsize), - subplot_kw={"projection": target_crs} + fig, ax = plt.subplots(figsize=figsize, + subplot_kw={"projection": target_crs}, + layout="tight", ) else: - fig, ax = plt.subplots(figsize=(figsize, figsize), - subplot_kw={"projection": source_crs}) + fig, ax = plt.subplots(figsize=figsize, + subplot_kw={"projection": source_crs}, + layout="tight", + ) fig.set_dpi(dpi) else: ax = ax_init @@ -228,9 +237,14 @@ def update(date): x = transformed_coords[:, :, 0] y = transformed_coords[:, :, 1] + + # Hack since cartopy needs transparency for nan regions to wraparound + # correctly with pcolormesh. + custom_cmap = get_custom_cmap(cmap) + image = ax.pcolormesh(x, y, data, transform=source_crs, - cmap=cmap, + cmap=custom_cmap, clim=(n_min, n_max), animated=True, zorder=1,