Dev icenet-ai#279: Lots of mp4 output related fixes

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:

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

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,