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"""