Merge pull request #622 from mandli/add-track-plotting

Add Plotting Capabilities to Storm Object

src/python/geoclaw/surge/plot.py

@@ -13,6 +13,9 @@
 
 from __future__ import absolute_import
 from __future__ import print_function
+
+import warnings
+
 import numpy as np
 import matplotlib.pyplot as plt
 import matplotlib.colors as colors
@@ -23,6 +26,7 @@
 import clawpack.visclaw.gaugetools as gaugetools
 import clawpack.visclaw.geoplot as geoplot
 import clawpack.geoclaw.data as geodata
+# import clawpack.geoclaw.surge.storm
 
 # TODO:  Assign these based on data files
 bathy_index = 0
@@ -41,7 +45,9 @@
 
 surge_data = geodata.SurgeData()
 
-
+# ==============================
+#  Track Plotting Functionality
+# ==============================
 class track_data(object):
     """Read in storm track data from run output"""
 
@@ -68,8 +74,8 @@ def get_track(self, frame):
 
             # Check to make sure that this fixed the problem
             if self._data.shape[0] < frame + 1:
-                print(" *** WARNING *** Could not find track data for ",
-                      "frame %s." % frame)
+                warnings.warn(f" *** WARNING *** Could not find track data",
+                               " for frame {frame}.")
                 return None, None, None
 
         return self._data[frame, 1:]
@@ -165,8 +171,15 @@ def pressure(cd):
     # The division by 100.0 is to convert from Pa to millibars
     return cd.aux[pressure_field, :, :] / 100.0
 
-# def category(Storm, cd):
-#     return cd.aux[Storm.category, :, :]
+
+def storm_radius(cd, track):
+    """Distance from center of storm"""
+    track_data = track.get_track(cd.frameno)
+
+    if track_data[0] is not None and track_data[1] is not None:
+        return np.sqrt((cd.x - track_data[0])**2 + (cd.y - track_data[1])**2)
+    else:
+        return None
 
 
 # ========================================================================
@@ -435,6 +448,15 @@ def add_bathy_contours(plotaxes, contour_levels=None, color='k'):
     plotitem.patchedges_show = 0
 
 
+def add_storm_radii(plotaxes, track, radii=[100e3], color='r'):
+    """Add radii to plots based on storm position"""
+    plotitem = plotaxes.new_plotitem(name="storm radius", 
+                                     plot_type="2d_contour")
+    plotitem.plot_var = lambda cd: storm_radius(cd, track)
+    plotitem.contour_levels = radii
+    plotitem.contour_colors = color
+
+
 # ===== Storm related plotting =======
 def sec2days(seconds):
     """Converst seconds to days."""

src/python/geoclaw/surge/storm.py

@@ -281,18 +281,25 @@ def read_geoclaw(self, path, verbose=False):
         # Read header
         with open(path, 'r') as data_file:
             num_casts = int(data_file.readline())
-            self.time_offset = datetime.datetime.strptime(
-                data_file.readline()[:19],
-                '%Y-%m-%dT%H:%M:%S')
+            time = data_file.readline()[:19]
+            try:
+                self.time_offset = datetime.datetime.strptime(
+                    time, '%Y-%m-%dT%H:%M:%S')
+            except ValueError:
+                self.time_offset = float(time)
         # Read rest of data
         data = np.loadtxt(path, skiprows=3)
         num_forecasts = data.shape[0]
-        self.eye_location = np.empty((2, num_forecasts))
+        self.eye_location = np.empty((num_forecasts, 2))
         assert(num_casts == num_forecasts)
-        self.t = [self.time_offset + datetime.timedelta(seconds=data[i, 0])
-                  for i in range(num_forecasts)]
-        self.eye_location[0, :] = data[:, 1]
-        self.eye_location[1, :] = data[:, 2]
+        if isinstance(self.time_offset, datetime.datetime):
+            self.t = np.array([self.time_offset
+                               + datetime.timedelta(seconds=data[i, 0])
+                               for i in range(num_forecasts)])
+        else:
+            self.t = data[:, 0]
+        self.eye_location[:, 0] = data[:, 1]
+        self.eye_location[:, 1] = data[:, 2]
         self.max_wind_speed = data[:, 3]
         self.max_wind_radius = data[:, 4]
         self.central_pressure = data[:, 5]
@@ -1154,6 +1161,108 @@ def write_tcvitals(self, path, verbose=False):
                                    "implemented yet but is planned for a ",
                                    "future release."))
 
+    # ================
+    #  Track Plotting
+    # ================
+    def plot(self, ax, radius=None, t_range=None, coordinate_system=2, track_style='ko--',
+             categorization="NHC", fill_alpha=0.25, fill_color='red'):
+        """TO DO:  Write doc-string"""
+
+        import matplotlib.pyplot as plt
+
+        if isinstance(track_style, str):
+            style = track_style
+
+        # Extract information for plotting the track/swath
+        t = self.t
+        x = self.eye_location[:, 0]
+        y = self.eye_location[:, 1]
+        if t_range is not None:
+            t = np.ma.masked_outside(t, t_range[0], t_range[1])
+            x = np.ma.array(x, mask=t.mask).compressed()
+            y = np.ma.array(y, mask=t.mask).compressed()
+            t = t.compressed()
+
+        # Plot track
+        if isinstance(track_style, str):
+            # Plot the track as a simple line with the given style
+            ax.plot(x, y, track_style)
+        elif isinstance(track_style, dict):
+            if self.max_wind_speed is None:
+                raise ValueError("Maximum wind speed not available so "
+                                 "plotting catgories is not available.")
+
+            # Plot the track using the colors provided in the dictionary
+            cat_color_defaults = {5: 'red', 4: 'yellow', 3: 'orange', 2: 'green',
+                                  1: 'blue', 0: 'gray', -1: 'lightgray'}
+            colors = [track_style.get(category, cat_color_defaults[category])
+                      for category in self.category(categorization=categorization)]
+            for i in range(t.shape[0] - 1):
+                ax.plot(x[i:i+2], y[i:i+2], color=colors[i], marker="o")
+
+        else:
+            raise ValueError("The `track_style` should be a string or dict.")
+
+        # Plot swath
+        if (isinstance(radius, float) or isinstance(radius, np.ndarray)
+                or radius is None):
+
+            if radius is None:
+                # Default behavior
+                if self.storm_radius is None:
+                    raise ValueError("Cannot use storm radius for plotting "
+                                     "the swath as the data is not available.")
+                else:
+                    if coordinate_system == 1:
+                        _radius = self.storm_radius
+                    elif coordinate_system == 2:
+                        _radius = units.convert(self.storm_radius,
+                                                'm', 'lat-long')
+                    else:
+                        raise ValueError(f"Unknown coordinate system "
+                                         f"{coordinate_system} provided.")
+
+            elif isinstance(radius, float):
+                # Only one value for the radius was given, replicate
+                _radius = np.ones(self.t.shape) * radius
+            elif isinstance(radius, np.ndarray):
+                # The array passed is the array to use
+                _radius = radius
+            else:
+                raise ValueError("Invalid input argument for radius.  Should "
+                                 "be a float or None")
+
+            # Draw first and last points
+            ax.add_patch(plt.Circle(
+                (x[0], y[0]), _radius[0], color=fill_color))
+            if t.shape[0] > 1:
+                ax.add_patch(plt.Circle((x[-1], y[-1]), _radius[-1],
+                                        color=fill_color))
+
+            # Draw path around inner points
+            if t.shape[0] > 2:
+                for i in range(t.shape[0] - 1):
+                    p = np.array([(x[i], y[i]), (x[i + 1], y[i + 1])])
+                    v = p[1] - p[0]
+                    if abs(v[1]) > 1e-16:
+                        n = np.array([1, -v[0] / v[1]], dtype=float)
+                    elif abs(v[0]) > 1e-16:
+                        n = np.array([-v[1] / v[0], 1], dtype=float)
+                    else:
+                        raise Exception("Zero-vector given")
+                    n /= np.linalg.norm(n)
+                    n *= _radius[i]
+
+                    ax.fill((p[0, 0] + n[0], p[0, 0] - n[0],
+                             p[1, 0] - n[0],
+                             p[1, 0] + n[0]),
+                            (p[0, 1] + n[1], p[0, 1] - n[1],
+                             p[1, 1] - n[1],
+                             p[1, 1] + n[1]),
+                            facecolor=fill_color, alpha=fill_alpha)
+                    ax.add_patch(plt.Circle((p[1][0], p[1, 1]), _radius[i],
+                                            color=fill_color, alpha=fill_alpha))
+
     # =========================================================================
     # Other Useful Routines
 

src/python/geoclaw/topotools.py

@@ -414,7 +414,7 @@ def delta(self):
 
 
     def __init__(self, path=None, topo_type=None, topo_func=None, 
-                       unstructured=False):
+                       unstructured=False, **kwargs):
         r"""Topography initialization routine.
         
         See :class:`Topography` for more info.
@@ -444,7 +444,7 @@ def __init__(self, path=None, topo_type=None, topo_func=None,
 
         if path:
             self.read(path=path, topo_type=topo_type,
-                      unstructured=unstructured)
+                      unstructured=unstructured, **kwargs)
 
     def set_xyZ(self, X, Y, Z):
         r"""