Initial code for wind field reconstruction from vorticity and divergence

docs/_templates/overrides/metpy.calc.rst

@@ -130,7 +130,8 @@ Dynamic/Kinematic
       wind_components
       wind_direction
       wind_speed
-
+      rotational_wind_from_inversion
+      divergent_wind_from_inversion
 
 Boundary Layer/Turbulence
 -------------------------
@@ -207,11 +208,14 @@ Other
 
       angle_to_direction
       azimuth_range_to_lat_lon
+      bounding_box_mask
       find_bounding_indices
+      find_bounding_box_indices
       find_intersections
       get_layer
       get_layer_heights
       get_perturbation
+      get_vectorized_array_indices
       isentropic_interpolation
       isentropic_interpolation_as_dataset
       nearest_intersection_idx

examples/calculations/Vorticity_Divergence_Inversion.py

@@ -0,0 +1,221 @@
+# Copyright (c) 2022 MetPy Developers.
+# Distributed under the terms of the BSD 3-Clause License.
+# SPDX-License-Identifier: BSD-3-Clause
+"""
+=========
+Vorticity
+=========
+
+Use `metpy.calc.vorticity`.
+
+This example demonstrates the calculation of reconstructed wind field for
+cyclone dora and plotting using matplotlib.
+"""
+import xarray as xr
+import numpy as np
+import metpy.calc as mpcalc
+import matplotlib.pyplot as plt
+from matplotlib.cm import get_cmap
+import matplotlib.ticker as mticker
+import cartopy.crs as crs
+from cartopy.feature import NaturalEarthFeature
+from cartopy.mpl.gridliner import LONGITUDE_FORMATTER, LATITUDE_FORMATTER
+
+u850 = xr.open_dataset('gfs.t12z.pgrb2.0p25.f000', engine='cfgrib',
+                       backend_kwargs={'filter_by_keys':
+                                       {'typeOfLevel': 'isobaricInhPa', 'shortName': 'u',
+                                        'level': 850}})
+u = u850.u
+
+v850 = xr.open_dataset('gfs.t12z.pgrb2.0p25.f000', engine='cfgrib',
+                       backend_kwargs={'filter_by_keys':
+                                       {'typeOfLevel': 'isobaricInhPa', 'shortName': 'v',
+                                        'level': 850}})
+v = v850.v
+
+
+# Compute the 850 hPa relative vorticity.
+
+
+vort850 = mpcalc.vorticity(u, v)
+fig = plt.figure(figsize=(12, 9), dpi=300.)
+# Create a set of axes for the figure and set
+# its map projection to that of the input data.
+ax = plt.axes(projection=crs.PlateCarree())
+
+# Add country borders and coastlines.
+countries = NaturalEarthFeature(category='cultural', scale='50m',
+                                facecolor='none',
+                                name='admin_0_countries')
+ax.add_feature(countries, linewidth=.5, edgecolor='black')
+ax.coastlines('50m', linewidth=0.8)
+
+plot = vort850.plot(levels=np.arange(-1.e-4, 1.e-4, 0.2e-5),
+                    cmap=get_cmap('PRGn'), transform=crs.PlateCarree(), cbar_kwargs={'label':
+                    'relative vorticity (x$10^{-5} s^{-1}$)', 'shrink': 0.98})
+
+# Set the map's extent to cover just Hurricane Dora.
+ax.set_extent([-180., -150., 0., 20.], crs=crs.PlateCarree())
+
+# Add latitude/longitude gridlines.
+gridlines = ax.gridlines(color='grey', linestyle='dotted', draw_labels=True)
+gridlines.xlabels_top = False
+gridlines.ylabels_right = False
+gridlines.xlocator = mticker.FixedLocator(np.arange(-180., 149., 5.))
+gridlines.ylocator = mticker.FixedLocator(np.arange(0., 21., 5.))
+gridlines.xlabel_style = {'size': 12, 'color': 'black'}
+gridlines.ylabel_style = {'size': 12, 'color': 'black'}
+gridlines.xformatter = LONGITUDE_FORMATTER
+gridlines.yformatter = LATITUDE_FORMATTER
+
+# Add a plot title, then show the image.
+plt.title('GFS 0-h 850 hPa relative vorticity (x$10^{-5} s^{-1}$) at 1200 UTC 9 August 2023')
+plt.savefig('vort.png')
+plt.show()
+
+# Compute the 850 hPa divergence.
+
+div850 = mpcalc.divergence(u, v)
+
+# Create a figure instance.
+fig = plt.figure(figsize=(12, 9), dpi=300.)
+
+# Create a set of axes for the figure and set
+# its map projection to that of the input data.
+ax = plt.axes(projection=crs.PlateCarree())
+
+# Add country borders and coastlines.
+countries = NaturalEarthFeature(category='cultural', scale='50m',
+                                facecolor='none',
+                                name='admin_0_countries')
+ax.add_feature(countries, linewidth=.5, edgecolor='black')
+ax.coastlines('50m', linewidth=0.8)
+
+# Plot the 850 hPa divergence using xarray's plot functionality.
+plot = div850.plot(levels=np.arange(-1.e-4, 1.e-4, 0.2e-5),
+                   cmap=get_cmap('PRGn'), transform=crs.PlateCarree(),
+                   cbar_kwargs={'label': 'relative vorticity (x$10^{-5} s^{-1}$)',
+                                'shrink': 0.98})
+
+# Set the map's extent to cover just Hurricane Dora.
+ax.set_extent([-180., -150., 0., 20.], crs=crs.PlateCarree())
+
+# Add latitude/longitude gridlines.
+gridlines = ax.gridlines(color='grey', linestyle='dotted', draw_labels=True)
+gridlines.xlabels_top = False
+gridlines.ylabels_right = False
+gridlines.xlocator = mticker.FixedLocator(np.arange(-180., 149., 5.))
+gridlines.ylocator = mticker.FixedLocator(np.arange(0., 21., 5.))
+gridlines.xlabel_style = {'size': 12, 'color': 'black'}
+gridlines.ylabel_style = {'size': 12, 'color': 'black'}
+gridlines.xformatter = LONGITUDE_FORMATTER
+gridlines.yformatter = LATITUDE_FORMATTER
+
+# Add a plot title, then show the image.
+plt.title('GFS 0-h 850 hPa divergence (x$10^{-5} s^{-1}$) at 1200 UTC 9 August 2023')
+plt.savefig('div.png')
+plt.show()
+
+umask = mpcalc.bounding_box_mask(u, 5., 13.5, 191., 202.)
+
+vmask = mpcalc.bounding_box_mask(v, 5., 13.5, 191., 202.)
+
+
+vortmask = mpcalc.bounding_box_mask(vort850, 5., 13.5, 191., 202.)
+
+
+divmask = mpcalc.bounding_box_mask(div850, 5., 13.5, 191., 202.)
+
+i_bb_indices = mpcalc.find_bounding_box_indices(vortmask, 5., 13.5, 191., 202.)
+
+
+o_bb_indices = mpcalc.find_bounding_box_indices(vortmask, 0., 30., 180., 220.)
+
+
+dx, dy = mpcalc.lat_lon_grid_deltas(vortmask.longitude, vortmask.latitude)
+
+upsi, vpsi = mpcalc.rotational_wind_from_inversion(umask, vmask, vortmask, dx, dy,
+                                                   o_bb_indices, i_bb_indices)
+
+# Create a figure instance.
+fig = plt.figure(figsize=(12, 9), dpi=300.)
+
+# Create a set of axes for the figure and set
+# its map projection to that of the input data.
+ax = plt.axes(projection=crs.PlateCarree())
+
+# Add country borders and coastlines.
+countries = NaturalEarthFeature(category='cultural', scale='50m',
+                                facecolor='none',
+                                name='admin_0_countries')
+ax.add_feature(countries, linewidth=.5, edgecolor='black')
+ax.coastlines('50m', linewidth=0.8)
+
+# Compute the magnitude of the non-divergent component of the 850 hPa wind.
+nd_spd = np.sqrt(upsi**2 + vpsi**2)
+
+# Plot this using xarray's plot functionality.
+plot = nd_spd.plot(levels=np.arange(0., 13., 1.),
+                   cmap=get_cmap('YlGnBu'), transform=crs.PlateCarree(),
+                   cbar_kwargs={'label': 'non-divergent wind ($m s^{-1}$)', 'shrink': 0.98})
+
+# Set the map's extent to match that over which we computed the non-divergent wind.
+ax.set_extent([-180., -140., 0., 30.], crs=crs.PlateCarree())
+
+# Add latitude/longitude gridlines.
+gridlines = ax.gridlines(color='grey', linestyle='dotted', draw_labels=True)
+gridlines.xlabels_top = False
+gridlines.ylabels_right = False
+gridlines.xlocator = mticker.FixedLocator(np.arange(-180., 139., 5.))
+gridlines.ylocator = mticker.FixedLocator(np.arange(0., 31., 5.))
+gridlines.xlabel_style = {'size': 12, 'color': 'black'}
+gridlines.ylabel_style = {'size': 12, 'color': 'black'}
+gridlines.xformatter = LONGITUDE_FORMATTER
+gridlines.yformatter = LATITUDE_FORMATTER
+
+# Add a plot title, then show the image.
+plt.title('850 hPa non-divergent wind magnitude  due to Dora at 1200 UTC 9 August 2023')
+plt.savefig('reconstructed_rotational_wind.png')
+plt.show()
+
+uchi, vchi = mpcalc.divergent_wind_from_inversion(umask, vmask, divmask, dx, dy,
+                                                  o_bb_indices, i_bb_indices)
+
+# Create a set of axes for the figure and set
+# its map projection to that of the input data.
+
+ax = plt.axes(projection=crs.PlateCarree())
+
+# Add country borders and coastlines.
+countries = NaturalEarthFeature(category='cultural', scale='50m',
+                                facecolor='none',
+                                name='admin_0_countries')
+ax.add_feature(countries, linewidth=.5, edgecolor='black')
+ax.coastlines('50m', linewidth=0.8)
+
+# Compute the magnitude of the non-divergent component of the 850 hPa wind.
+nd_spd = np.sqrt(uchi**2 + vchi**2)
+
+# Plot this using xarray's plot functionality.
+plot = nd_spd.plot(levels=np.arange(0., 13., 1.),
+                   cmap=get_cmap('YlGnBu'), transform=crs.PlateCarree(),
+                   cbar_kwargs={'label': 'non-divergent wind ($m s^{-1}$)', 'shrink': 0.98})
+
+# Set the map's extent to match that over which we computed the non-divergent wind.
+ax.set_extent([-180., -140., 0., 30.], crs=crs.PlateCarree())
+
+# Add latitude/longitude gridlines.
+gridlines = ax.gridlines(color='grey', linestyle='dotted', draw_labels=True)
+gridlines.top_labels = False
+gridlines.right_labels = False
+gridlines.xlocator = mticker.FixedLocator(np.arange(-180., 139., 5.))
+gridlines.ylocator = mticker.FixedLocator(np.arange(0., 31., 5.))
+gridlines.xlabel_style = {'size': 12, 'color': 'black'}
+gridlines.ylabel_style = {'size': 12, 'color': 'black'}
+gridlines.xformatter = LONGITUDE_FORMATTER
+gridlines.yformatter = LATITUDE_FORMATTER
+
+# Add a plot title, then show the image.
+plt.title('850 hPa divergent wind magnitude due to Dora at 1200 UTC 9 August 2023')
+plt.savefig('irrotational_winds.png')
+plt.show()

src/metpy/calc/kinematics.py

@@ -3,10 +3,10 @@
 # SPDX-License-Identifier: BSD-3-Clause
 """Contains calculation of kinematic parameters (e.g. divergence or vorticity)."""
 import numpy as np
-
+import xarray as xa
 from . import coriolis_parameter
-from .tools import (first_derivative, geospatial_gradient, get_layer_heights,
-                    parse_grid_arguments, vector_derivative)
+from .tools import (first_derivative, geospatial_gradient, get_vectorized_array_indices,
+                    get_layer_heights, parse_grid_arguments, vector_derivative)
 from .. import constants as mpconsts
 from ..package_tools import Exporter
 from ..units import check_units, units
@@ -1430,3 +1430,119 @@ def geospatial_laplacian(f, *, dx=None, dy=None, x_dim=-1, y_dim=-2,
                                          meridional_scale=meridional_scale)
     return divergence(grad_u, grad_y, dx=dx, dy=dy, x_dim=x_dim, y_dim=y_dim,
                       parallel_scale=parallel_scale, meridional_scale=meridional_scale)
+
+
+@exporter.export
+@parse_grid_arguments
+def rotational_wind_from_inversion(umask, vmask, vortmask, dx, dy, o_bb_indices, i_bb_indices):
+    r"""Calculate reconstructed rotational wind field from vorticity.
+
+    Parameters
+    ----------
+    vortmask : 'xarray DataArray'  subset of the original vorticity for the entire globe
+    dx : `pint.Quantity`,required
+        The grid spacing(s) in the x-direction. If an array, there should be one item less than
+        the size of `u` along the applicable axis. Optional if `xarray.DataArray` with
+        latitude/longitude coordinates used as input. Also optional if one-dimensional
+        longitude and latitude arguments are given for your data on a non-projected grid.
+        Keyword-only argument.
+    dy : `pint.Quantity`, required
+        The grid spacing(s) in the y-direction. If an array, there should be one item less than
+        the size of `u` along the applicable axis. Optional if `xarray.DataArray` with
+        latitude/longitude coordinates used as input. Also optional if one-dimensional
+        longitude and latitude arguments are given for your data on a non-projected grid.
+        Keyword-only argument.
+    o_bb_indices : contains the x and y lower left and upper right indices
+    i_bb_indices : contains the x and y lower left and upper right indices
+    """
+    upsi = xa.zeros_like(umask)
+    vpsi = xa.zeros_like(vmask)
+    dx1 = dx.magnitude
+    dy1 = dy.magnitude
+    [xindex, yindex] = get_vectorized_array_indices(i_bb_indices)
+    iindex = np.zeros_like(yindex)
+    jindex = np.zeros_like(yindex)
+    o_x_ll = o_bb_indices.x_ll
+    o_x_ur = o_bb_indices.x_ur
+    o_y_ll = o_bb_indices.y_ll
+    o_y_ur = o_bb_indices.y_ur
+    x_ll = i_bb_indices.x_ll
+    x_ur = i_bb_indices.x_ur
+    y_ll = i_bb_indices.y_ll
+    y_ur = i_bb_indices.y_ur
+    vortmask1 = vortmask.values
+    for i in range(o_x_ll, o_x_ur):
+        for j in range(o_y_ur, o_y_ll):
+            iindex[:, :] = i
+            jindex[:, :] = j
+            xdiff = (iindex - xindex) * dx1[y_ur:y_ll, x_ll:x_ur]
+            ydiff = (jindex - yindex) * dy1[y_ur:y_ll, x_ll:x_ur]
+            rsq = xdiff * xdiff + ydiff * ydiff
+            upsi[j, i] = np.where(rsq > 0., vortmask1[y_ur:y_ll, x_ll:x_ur] * -1.0 * (
+                ydiff / rsq) * dx1[y_ur:y_ll, x_ll:x_ur] * dy1[y_ur:y_ll,
+                                                               x_ll:x_ur], 0.0).sum()
+            vpsi[j, i] = np.where(rsq > 0., vortmask1[y_ur:y_ll, x_ll:x_ur] * (
+                xdiff / rsq) * dx1[y_ur:y_ll, x_ll:x_ur] * dy1[y_ur:y_ll,
+                                                               x_ll:x_ur], 0.0).sum()
+    upsi[:, :] = (1 / (2 * np.pi)) * upsi[:, :]
+    vpsi[:, :] = (1 / (2 * np.pi)) * vpsi[:, :]
+
+    return upsi, vpsi
+
+
+@exporter.export
+def divergent_wind_from_inversion(umask, vmask, divmask, dx, dy, o_bb_indices, i_bb_indices):
+    r"""Calculate reconstructed divergent wind field from divergence.
+
+    Parameters
+    ----------
+    divmask : 'xarray DataArray'  subset of the original vorticity for the entire globe
+    dx : `pint.Quantity`,required
+        The grid spacing(s) in the x-direction. If an array, there should be one item less than
+        the size of `u` along the applicable axis. Optional if `xarray.DataArray` with
+        latitude/longitude coordinates used as input. Also optional if one-dimensional
+        longitude and latitude arguments are given for your data on a non-projected grid.
+        Keyword-only argument.
+    dy : `pint.Quantity`, required
+        The grid spacing(s) in the y-direction. If an array, there should be one item less than
+        the size of `u` along the applicable axis. Optional if `xarray.DataArray` with
+        latitude/longitude coordinates used as input. Also optional if one-dimensional
+        longitude and latitude arguments are given for your data on a non-projected grid.
+        Keyword-only argument.
+    o_bb_indices : contains the x and y lower left and upper right indices
+    i_bb_indices : contains the x and y lower left and upper right indices
+    """
+    uchi = xa.zeros_like(umask)
+    vchi = xa.zeros_like(vmask)
+    dx1 = dx.magnitude
+    dy1 = dy.magnitude
+    divmask1 = divmask.values
+    [xindex, yindex] = get_vectorized_array_indices(i_bb_indices)
+    iindex = np.zeros_like(yindex)
+    jindex = np.zeros_like(yindex)
+    o_x_ll = o_bb_indices.x_ll
+    o_x_ur = o_bb_indices.x_ur
+    o_y_ll = o_bb_indices.y_ll
+    o_y_ur = o_bb_indices.y_ur
+    x_ll = i_bb_indices.x_ll
+    x_ur = i_bb_indices.x_ur
+    y_ll = i_bb_indices.y_ll
+    y_ur = i_bb_indices.y_ur
+    for i in range(o_x_ll, o_x_ur):
+        for j in range(o_y_ur, o_y_ll):
+            iindex[:, :] = i
+            jindex[:, :] = j
+            xdiff = (iindex - xindex) * dx1[y_ur:y_ll, x_ll:x_ur]
+            ydiff = (jindex - yindex) * dy1[y_ur:y_ll, x_ll:x_ur]
+            rsq = xdiff * xdiff + ydiff * ydiff
+            uchi[j, i] = np.where(rsq > 0., divmask1[y_ur:y_ll, x_ll:x_ur] * (
+                xdiff / rsq) * dx1[y_ur:y_ll, x_ll:x_ur] * dy1[y_ur:y_ll,
+                                                               x_ll:x_ur], 0.0).sum()
+            vchi[j, i] = np.where(rsq > 0., divmask1[y_ur:y_ll, x_ll:x_ur] * (
+                ydiff / rsq) * dx1[y_ur:y_ll, x_ll:x_ur] * dy1[y_ur:y_ll,
+                                                               x_ll:x_ur], 0.0).sum()
+
+    uchi[:, :] = (1 / (2 * np.pi)) * uchi[:, :]
+    vchi[:, :] = (1 / (2 * np.pi)) * vchi[:, :]
+
+    return uchi, vchi