diff --git a/docs/_templates/overrides/metpy.calc.rst b/docs/_templates/overrides/metpy.calc.rst index e5e7a4f787f..eb73243699f 100644 --- a/docs/_templates/overrides/metpy.calc.rst +++ b/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 diff --git a/examples/calculations/Vorticity_Divergence_Inversion.py b/examples/calculations/Vorticity_Divergence_Inversion.py new file mode 100644 index 00000000000..1b0dd7ee62f --- /dev/null +++ b/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() diff --git a/src/metpy/calc/kinematics.py b/src/metpy/calc/kinematics.py index 73f8d7a64c4..20ade04da85 100644 --- a/src/metpy/calc/kinematics.py +++ b/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 diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index 091545d8856..ab0b3488918 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -7,6 +7,7 @@ from inspect import Parameter, signature from operator import itemgetter +from dataclasses import dataclass import numpy as np # Can drop fallback once we rely on numpy>=2 @@ -1942,3 +1943,53 @@ def _remove_nans(*variables): for v in variables: ret.append(v[~mask]) return ret + + +@exporter.export +def bounding_box_mask(data_array, min_lat, max_lat, min_lon, max_lon): + """Construct a bounding box mask within the specified coordinates. + + All values that satisfy the condition below will have 1 in the mask + and others zero. + """ + mask = ((data_array.latitude <= max_lat) & (data_array.latitude >= min_lat) + & (data_array.longitude <= max_lon) & (data_array.longitude >= min_lon)) + data_mask = data_array.where(mask) + data_mask = data_mask.fillna(0.0) + return data_mask + + +@exporter.export +def find_bounding_box_indices(data_mask, min_lat, max_lat, min_lon, max_lon): + """Returns the array indices of a bounding box.""" + + @dataclass + class BoundingBoxIndices: + x_ll: int = None + x_ur: int = None + y_ll: int = None + y_ur: int = None + x_ll = list(data_mask.longitude.values).index(min_lon) + x_ur = list(data_mask.longitude.values).index(max_lon) + y_ll = list(data_mask.latitude.values).index(min_lat) + y_ur = list(data_mask.latitude.values).index(max_lat) + bounding_box_indices = BoundingBoxIndices(x_ll, x_ur, y_ll, y_ur) + return bounding_box_indices + + +@exporter.export +def get_vectorized_array_indices(i_bb_indices): + """Returns the vectorization indices for inner for loop in the + wind field reconstruction method. + """ + i_x_ll = i_bb_indices.x_ll + i_x_ur = i_bb_indices.x_ur + i_y_ll = i_bb_indices.y_ll + i_y_ur = i_bb_indices.y_ur + x = np.abs(i_x_ll - i_x_ur) + y = np.abs(i_y_ll - i_y_ur) + xindex = np.linspace(i_x_ll, i_x_ur, num=x * y, endpoint=False, dtype=np.int32) + yindex = np.linspace(i_y_ur, i_y_ll, num=y * x, endpoint=False, dtype=np.int32) + xindex = xindex.reshape((y, x), order='F') + yindex = yindex.reshape((y, x), order='C') + return [xindex, yindex] diff --git a/tests/calc/test_calc_tools.py b/tests/calc/test_calc_tools.py index c54acb21a0b..5116ee6c3ad 100644 --- a/tests/calc/test_calc_tools.py +++ b/tests/calc/test_calc_tools.py @@ -11,6 +11,8 @@ from pyproj import CRS, Geod import pytest import xarray as xr +import datetime + from metpy.calc import (angle_to_direction, find_bounding_indices, find_intersections, first_derivative, geospatial_gradient, get_layer, get_layer_heights, @@ -20,7 +22,9 @@ from metpy.calc.tools import (_delete_masked_points, _get_bound_pressure_height, _greater_or_close, _less_or_close, _next_non_masked_element, _remove_nans, azimuth_range_to_lat_lon, BASE_DEGREE_MULTIPLIER, - DIR_STRS, nominal_lat_lon_grid_deltas, parse_grid_arguments, UND) + DIR_STRS, nominal_lat_lon_grid_deltas, parse_grid_arguments, UND, + bounding_box_mask, find_bounding_box_indices, + get_vectorized_array_indices) from metpy.testing import (assert_almost_equal, assert_array_almost_equal, assert_array_equal, get_test_data) from metpy.units import units @@ -1557,3 +1561,97 @@ def test_vector_derivative_return_subset(return_only, length): u, v, longitude=lons, latitude=lats, crs=crs, return_only=return_only) assert len(ddx) == length + + +def test_bounding_box_mask(): + """ Test the mask of a bounding box. """ + rng = np.random.default_rng() + temperature = 273 + 20 * rng.random(size=(4, 17, 73, 144)) + latitude = np.linspace(-90., 90., 73) + longitude = np.linspace(0., 360., 144, endpoint=False) + press = ['1000', '925', '850', '700', '600', '500', '400', '300', '250', + '200', '150', '100', '70', '50', '30', '20', '10'] + level = np.array(press) + level = level.astype(float) * 100 + time = np.empty(4) + time[0] = np.datetime64(datetime.datetime(2023, 1, 14, 0)) + time[1] = np.datetime64(datetime.datetime(2023, 1, 14, 6)) + time[2] = np.datetime64(datetime.datetime(2023, 1, 14, 12)) + time[3] = np.datetime64(datetime.datetime(2023, 1, 14, 18)) + temp = xr.DataArray(temperature, coords=[time, level, latitude, + longitude], dims=['time', + 'level', + 'latitude', + 'longitude']) + min_lat = 0. + max_lat = 30. + min_lon = 30. + max_lon = 100. + temp_mask1 = bounding_box_mask(temp, min_lat, max_lat, min_lon, max_lon) + temp_mask2 = bounding_box_mask(temp, 0., 30., 30., 100.) + assert_array_equal(temp_mask1.values, temp_mask2.values) + + +def test_find_bounding_box_indices(): + """Tests the bounding box indices for 2 different cases. """ + rng = np.random.default_rng() + temperature = 273 + 20 * rng.random(size=(4, 17, 73, 144)) + latitude = np.linspace(-90., 90., 73) + longitude = np.linspace(0., 360., 144, endpoint=False) + press = ['1000', '925', '850', '700', '600', '500', '400', '300', '250', + '200', '150', '100', '70', '50', '30', '20', '10'] + level = np.array(press) + level = level.astype(float) * 100 + time = np.empty(4) + time[0] = np.datetime64(datetime.datetime(2023, 1, 14, 0)) + time[1] = np.datetime64(datetime.datetime(2023, 1, 14, 6)) + time[2] = np.datetime64(datetime.datetime(2023, 1, 14, 12)) + time[3] = np.datetime64(datetime.datetime(2023, 1, 14, 18)) + temp = xr.DataArray(temperature, coords=[time, level, latitude, + longitude], dims=['time', + 'level', + 'latitude', + 'longitude']) + min_lat = 0. + max_lat = 30. + min_lon = 30. + max_lon = 100. + temp_mask1 = bounding_box_mask(temp, min_lat, max_lat, min_lon, max_lon) + temp_mask2 = bounding_box_mask(temp, 0., 30., 30., 100.) + data_class1 = find_bounding_box_indices(temp_mask1, min_lat, max_lat, min_lon, max_lon) + data_class2 = find_bounding_box_indices(temp_mask2, 0., 30., 30., 100.) + assert data_class1 != data_class2 + + +def test_get_vectorized_array_indices(): + """Tests the vectorized array indices for two different bounding boxes. """ + rng = np.random.default_rng() + temperature = 273 + 20 * rng.random(size=(4, 17, 73, 144)) + latitude = np.linspace(-90., 90., 73) + longitude = np.linspace(0., 360., 144, endpoint=False) + press = ['1000', '925', '850', '700', '600', '500', '400', '300', '250', + '200', '150', '100', '70', '50', '30', '20', '10'] + level = np.array(press) + level = level.astype(float) * 100 + time = np.empty(4) + time[0] = np.datetime64(datetime.datetime(2023, 1, 14, 0)) + time[1] = np.datetime64(datetime.datetime(2023, 1, 14, 6)) + time[2] = np.datetime64(datetime.datetime(2023, 1, 14, 12)) + time[3] = np.datetime64(datetime.datetime(2023, 1, 14, 18)) + temp = xr.DataArray(temperature, coords=[time, level, latitude, + longitude], dims=['time', + 'level', + 'latitude', + 'longitude']) + min_lat = 0. + max_lat = 30. + min_lon = 30. + max_lon = 100. + temp_mask1 = bounding_box_mask(temp, min_lat, max_lat, min_lon, max_lon) + temp_mask2 = bounding_box_mask(temp, 0., 30., 30., 100.) + data_class1 = find_bounding_box_indices(temp_mask1, min_lat, max_lat, min_lon, max_lon) + data_class2 = find_bounding_box_indices(temp_mask2, 0., 30., 30., 100.) + [xindices1, yindices1] = get_vectorized_array_indices(data_class1) + [xindices2, yindices2] = get_vectorized_array_indices(data_class2) + assert_array_equal(xindices1, xindices2) + assert_array_equal(yindices1, yindices2) diff --git a/tests/calc/test_kinematics.py b/tests/calc/test_kinematics.py index f71a75bc530..7e0ab0b3b1b 100644 --- a/tests/calc/test_kinematics.py +++ b/tests/calc/test_kinematics.py @@ -6,7 +6,7 @@ import numpy as np import pytest import xarray as xr - +import sys from metpy.calc import (absolute_vorticity, advection, ageostrophic_wind, coriolis_parameter, divergence, first_derivative, frontogenesis, geospatial_laplacian, geostrophic_wind, inertial_advective_wind, lat_lon_grid_deltas, @@ -14,7 +14,8 @@ potential_vorticity_baroclinic, potential_vorticity_barotropic, q_vector, shearing_deformation, static_stability, storm_relative_helicity, stretching_deformation, total_deformation, - vorticity, wind_components) + vorticity, wind_components,rotational_wind_from_inversion, + divergent_wind_from_inversion) from metpy.constants import g, Re from metpy.testing import (assert_almost_equal, assert_array_almost_equal, assert_array_equal, get_test_data) @@ -70,8 +71,9 @@ def test_vorticity(): def test_vorticity_geographic(geog_data): """Test vorticity for simple case on geographic coordinates.""" crs, lons, lats, u, v, mx, my, dx, dy = geog_data + print(crs) vort = vorticity(u, v, longitude=lons, latitude=lats, crs=crs) - + sys.exit() # Calculate the true field using known map-correct approach truth = (mx * first_derivative(v, delta=dx, axis=1) - my * first_derivative(u, delta=dy, axis=0)