Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Derivatives & map factors #2743

Merged
merged 83 commits into from
Dec 23, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
83 commits
Select commit Hold shift + click to select a range
9473dfd
Mockup/draft implementation of using a projection-correct vector deri…
jthielen Nov 20, 2021
391a7c2
Minimum commit for vorticity+xarray testing
dcamron Oct 17, 2022
e17c2db
Provide some test drafts and fixtures
dcamron Oct 17, 2022
7c4ada2
Improve agreement with GEMPAK
dopplershift Oct 19, 2022
7fc8ae8
Fix logic error determining when to calculate factors
dopplershift Nov 10, 2022
14dd8c3
Fix generation of scale arrays to produce DataArrays
dopplershift Nov 10, 2022
4550661
Update test values for new computation
dopplershift Nov 10, 2022
c3b8bc5
Refactor eliminate the need for unused function arguments
dopplershift Nov 21, 2022
5d616d5
Adjust sign in correction
dopplershift Nov 22, 2022
538bc6f
Consolidate check on dimensions
dopplershift Nov 22, 2022
fc59a21
Add test for vorticity on geographic grid
dopplershift Nov 22, 2022
080341e
Dynamically update function signature and docstrings
dopplershift Nov 23, 2022
0e7781c
Fix up divergence test and add another
dopplershift Nov 23, 2022
5da1800
Update deformation calculations using _vector_derivative
dopplershift Nov 23, 2022
cdca1a9
Remove unused GFS global dataset fixture
dopplershift Nov 23, 2022
8ae745d
Clean up some flake8 warnings
dopplershift Nov 23, 2022
dbab1ff
Update absolute_vorticity for proper vector derivative handling
dopplershift Nov 30, 2022
e7bd0a1
Fix up some ordering in function docstrings
dopplershift Nov 30, 2022
f8f9890
Make sure all newly added function arguments are kw-only
dopplershift Nov 30, 2022
e22b8a8
Remove unneeded comment
dopplershift Nov 30, 2022
cac5334
Remove unneeded test fixtures
dopplershift Nov 30, 2022
0940bee
Fix spelling
dopplershift Nov 30, 2022
b55ba9d
Fix up lat/lon handling to allow proper xarray broadcasting
dopplershift Nov 30, 2022
8b28b9b
Fix numpy handling broken by previous commit
dopplershift Nov 30, 2022
848d4b7
Add geospatial gradient
dcamron Nov 30, 2022
75529ff
Update frontogenesis for spatial calc
dcamron Dec 1, 2022
21de52c
Update q_vector for spatial calc
dcamron Dec 1, 2022
5b71230
Update geostrophic_wind for spatial calc
dcamron Dec 1, 2022
36ae749
Update ageostrophic_wind for spatial
dcamron Dec 1, 2022
648ddbd
Update advection for spatial calc
dcamron Dec 1, 2022
38389dd
Add z-x "cross-sections" test for advection
dcamron Dec 1, 2022
83a364e
Update inertial_advective_wind for spatial calc
dcamron Dec 1, 2022
9fdba53
Update baroclinic pv for spatial calc
dcamron Dec 1, 2022
f11000d
Add spatial+4d test for baroclinic pv
dcamron Dec 1, 2022
1b9924e
Remove unused gradient import
dcamron Dec 1, 2022
4d267b6
Fix grid coordinates for non-lat/lon case
dopplershift Dec 2, 2022
b0fabed
Fix up coordinate assignment for p and m scale arrays
dopplershift Dec 2, 2022
d8488dd
Clean up
dopplershift Dec 2, 2022
0b2ffb6
Remove new-style type hint
dopplershift Dec 2, 2022
4ea11af
DOC: Add parallel/meridional scales to docstrings
dopplershift Dec 2, 2022
b0dace8
Update potential_vorticity_barotropic for new derivatives
dopplershift Dec 2, 2022
d2d71ab
Remove unused import of add_grid_arguments_from_xarray
dopplershift Dec 2, 2022
094ced1
Add vector_derivative and geospatial gradient to func table
dopplershift Dec 2, 2022
50bd892
Docstring improvements
dopplershift Dec 5, 2022
58e775a
Add parametrized geospatial_gradient test
dcamron Dec 5, 2022
c2d5623
Simplify casting scale to xarray
dopplershift Dec 6, 2022
7cc5c12
Assume a lat/lon PyPROJ CRS for 1D lat/lon coords
dopplershift Dec 6, 2022
8a770ae
Assume a CRS for 1D lat/lon when finding dx/dy
dopplershift Dec 6, 2022
5c58496
Add test for projected deltas at pole
dcamron Dec 6, 2022
bf78ebd
Update error message
dopplershift Dec 6, 2022
e9cd4b6
Fix reference to pyproj.crs.CRS
dopplershift Dec 6, 2022
358e0ee
Fix up some docstring errors
dopplershift Dec 7, 2022
129d7e8
Update test values based on new implementation
dopplershift Dec 7, 2022
a55bc8e
Bump minimum PyPROJ to 2.6 for get_factors()
dopplershift Dec 7, 2022
12908f6
Fix u-first behavior of geospatial_gradient
dcamron Dec 7, 2022
0bcc445
Fix up absolute_vorticity docstring
dopplershift Dec 7, 2022
041e0a6
Reduce code duplication using a fixture
dopplershift Dec 7, 2022
49a9376
Move geog_data fixture to conftest
dopplershift Dec 7, 2022
3323cd1
Bump PyProj minimum to 2.6.1
dopplershift Dec 7, 2022
306101f
Update see also
dcamron Dec 7, 2022
4288e7a
Flatten vector_derivative output
dcamron Dec 7, 2022
03cdccd
Add preprocessing decorators
dcamron Dec 7, 2022
bf1ae26
Add geospatial_laplacian
dcamron Dec 7, 2022
34aca8a
Simplify test to use fixture arrays
dcamron Dec 7, 2022
a61224e
Add test for geospatial_laplacian
dcamron Dec 7, 2022
3b125ce
Add geospatial_laplacian to doc table
dcamron Dec 7, 2022
c869409
Collapse _vector_derivative to public function
dcamron Dec 7, 2022
043775d
Remove add_grid_arguments_from_xarray
dopplershift Dec 12, 2022
2c13191
Add some more tests for nominal_lat_lon_grid_deltas
dopplershift Dec 13, 2022
e735fde
Expand testing for parse_grid_arguments
dopplershift Dec 13, 2022
a348461
Add test for parse_grid_arguments with unknown coordinates
dopplershift Dec 13, 2022
6b9b7aa
Remove unused line from test
dopplershift Dec 13, 2022
3741258
Use correct wrapping syntax
dcamron Dec 13, 2022
0c839c9
Add test for geospatial_gradient return_only
dcamron Dec 13, 2022
a585e4a
Add test for vector_derivative return_only
dcamron Dec 13, 2022
287b942
Expand test docstring placeholder
dcamron Dec 13, 2022
e31427a
Add test for unknown dimensions
dopplershift Dec 13, 2022
de5872e
Add xarray vertical advection test
dcamron Dec 13, 2022
79b258a
Remove unnecessary dict construction
dcamron Dec 20, 2022
f4bbd69
Rename units in vertical advection test
dcamron Dec 20, 2022
74c6008
Add advection escape in dz parsing
dcamron Dec 20, 2022
55e0132
Test 1-d vertical advection with xarray
dcamron Dec 20, 2022
f963218
Fix test for minimum xarray
dcamron Dec 20, 2022
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 20 additions & 0 deletions conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -182,3 +182,23 @@ def array_type(request):
return lambda d, u, *, mask=None: quantity(numpy.array(d), u)
else:
raise ValueError(f'Unsupported array_type option {request.param}')


@pytest.fixture
def geog_data(request):
"""Create data to use for testing calculations on geographic coordinates."""
# Generate a field of u and v on a lat/lon grid
crs = pyproj.CRS(request.param)
proj = pyproj.Proj(crs)
a = numpy.arange(4)[None, :]
arr = numpy.r_[a, a, a] * metpy.units.units('m/s')
lons = numpy.array([-100, -90, -80, -70]) * metpy.units.units.degree
lats = numpy.array([45, 55, 65]) * metpy.units.units.degree
lon_arr, lat_arr = numpy.meshgrid(lons.m_as('degree'), lats.m_as('degree'))
factors = proj.get_factors(lon_arr, lat_arr)

return (crs, lons, lats, arr, arr, factors.parallel_scale, factors.meridional_scale,
metpy.calc.lat_lon_grid_deltas(lons.m, numpy.zeros_like(lons.m),
geod=crs.get_geod())[0][0],
metpy.calc.lat_lon_grid_deltas(numpy.zeros_like(lats.m), lats.m,
geod=crs.get_geod())[1][:, 0])
3 changes: 3 additions & 0 deletions docs/_templates/overrides/metpy.calc.rst
Original file line number Diff line number Diff line change
Expand Up @@ -149,13 +149,16 @@ Mathematical Functions

cross_section_components
first_derivative
geospatial_gradient
geospatial_laplacian
gradient
laplacian
lat_lon_grid_deltas
normal_component
second_derivative
tangential_component
unit_vectors_from_cross_section
vector_derivative


Apparent Temperature
Expand Down
2 changes: 1 addition & 1 deletion setup.cfg
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,7 @@ install_requires =
pandas>=1.0.0
pint>=0.15
pooch>=1.2.0
pyproj>=2.5.0
pyproj>=2.6.1
scipy>=1.4.0
traitlets>=5.0.5
xarray>=0.18.0
Expand Down
660 changes: 504 additions & 156 deletions src/metpy/calc/kinematics.py

Large diffs are not rendered by default.

442 changes: 439 additions & 3 deletions src/metpy/calc/tools.py

Large diffs are not rendered by default.

175 changes: 78 additions & 97 deletions src/metpy/xarray.py
Original file line number Diff line number Diff line change
Expand Up @@ -280,6 +280,11 @@ def pyproj_crs(self):
"""Return the coordinate reference system (CRS) as a pyproj object."""
return self.crs.to_pyproj()

@property
def pyproj_proj(self):
"""Return the Proj object corresponding to the coordinate reference system (CRS)."""
return Proj(self.pyproj_crs)

def _fixup_coordinate_map(self, coord_map):
"""Ensure sure we have coordinate variables in map, not coordinate names."""
new_coord_map = {}
Expand Down Expand Up @@ -430,9 +435,46 @@ def coordinates(self, *args):
access a single coordinate, use the appropriate attribute on the accessor, or use tuple
unpacking.

If latitude and/or longitude are requested here, and yet are not present on the
DataArray, an on-the-fly computation from the CRS and y/x dimension coordinates is
attempted.

"""
latitude = None
longitude = None
for arg in args:
yield self._axis(arg)
try:
yield self._axis(arg)
except AttributeError as exc:
if (
(arg == 'latitude' and latitude is None)
or (arg == 'longitude' and longitude is None)
):
# Try to compute on the fly
try:
latitude, longitude = _build_latitude_longitude(self._data_array)
except Exception:
# Attempt failed, re-raise original error
raise exc from None
# Otherwise, warn and yield result
warnings.warn(
'Latitude and longitude computed on-demand, which may be an '
'expensive operation. To avoid repeating this computation, assign '
'these coordinates ahead of time with '
'.metpy.assign_latitude_longitude().'
)
if arg == 'latitude':
yield latitude
else:
yield longitude
elif arg == 'latitude' and latitude is not None:
# We have this from previous computation
yield latitude
elif arg == 'longitude' and longitude is not None:
# We have this from previous computation
yield longitude
else:
raise exc

@property
def time(self):
Expand Down Expand Up @@ -465,7 +507,7 @@ def longitude(self):
return self._axis('longitude')

def coordinates_identical(self, other):
"""Return whether or not the coordinates of other match this DataArray's."""
"""Return whether the coordinates of other match this DataArray's."""
return (len(self._data_array.coords) == len(other.coords)
and all(coord_name in other.coords and other[coord_name].identical(coord_var)
for coord_name, coord_var in self._data_array.coords.items()))
Expand All @@ -476,6 +518,36 @@ def time_deltas(self):
us_diffs = np.diff(self._data_array.values).astype('timedelta64[us]').astype('int64')
return units.Quantity(us_diffs / 1e6, 's')

@property
def grid_deltas(self):
"""Return the horizontal dimensional grid deltas suitable for vector derivatives."""
if (
(hasattr(self, 'crs')
and self.crs._attrs['grid_mapping_name'] == 'latitude_longitude')
or (hasattr(self, 'longitude') and self.longitude.squeeze().ndim == 1
and hasattr(self, 'latitude') and self.latitude.squeeze().ndim == 1)
):
# Calculate dx and dy on ellipsoid (on equator and 0 deg meridian, respectively)
from .calc.tools import nominal_lat_lon_grid_deltas
crs = getattr(self, 'pyproj_crs', CRS('+proj=latlon'))

Check notice

Code scanning / CodeQL

Cyclic import

Import of module [metpy.calc.tools](1) begins an import cycle.
dx, dy = nominal_lat_lon_grid_deltas(
Fixed Show fixed Hide fixed
self.longitude.metpy.unit_array,
self.latitude.metpy.unit_array,
crs.get_geod()
)
else:
# Calculate dx and dy in projection space
try:
dx = np.diff(self.x.metpy.unit_array)
dy = np.diff(self.y.metpy.unit_array)
except AttributeError:
raise AttributeError(
'Grid deltas cannot be calculated since horizontal dimension coordinates '
'cannot be found.'
)

return {'dx': dx, 'dy': dy}

def find_axis_name(self, axis):
"""Return the name of the axis corresponding to the given identifier.

Expand Down Expand Up @@ -1115,7 +1187,7 @@ def _build_latitude_longitude(da):
"""Build latitude/longitude coordinates from DataArray's y/x coordinates."""
y, x = da.metpy.coordinates('y', 'x')
xx, yy = np.meshgrid(x.values, y.values)
lonlats = np.stack(Proj(da.metpy.pyproj_crs)(xx, yy, inverse=True, radians=False), axis=-1)
lonlats = np.stack(da.metpy.pyproj_proj(xx, yy, inverse=True, radians=False), axis=-1)
longitude = xr.DataArray(lonlats[..., 0], dims=(y.name, x.name),
coords={y.name: y, x.name: x},
attrs={'units': 'degrees_east', 'standard_name': 'longitude'})
Expand All @@ -1136,7 +1208,7 @@ def _build_y_x(da, tolerance):
'must be 2D')

# Convert to projected y/x
xxyy = np.stack(Proj(da.metpy.pyproj_crs)(
xxyy = np.stack(da.metpy.pyproj_proj(
longitude.values,
latitude.values,
inverse=False,
Expand Down Expand Up @@ -1305,7 +1377,8 @@ def _wrap_output_like_not_matching_units(result, match):
):
result = units.Quantity(result)
return (
xr.DataArray(result, coords=match.coords, dims=match.dims) if output_xarray
xr.DataArray(result, coords=match.coords, dims=match.dims)
if output_xarray and result is not None
else result
)

Expand Down Expand Up @@ -1450,98 +1523,6 @@ def dataarray_arguments(bound_args):
yield value


def add_grid_arguments_from_xarray(func):
"""Fill in optional arguments like dx/dy from DataArray arguments."""
@functools.wraps(func)
def wrapper(*args, **kwargs):
bound_args = signature(func).bind(*args, **kwargs)
bound_args.apply_defaults()

# Search for DataArray with valid latitude and longitude coordinates to find grid
# deltas and any other needed parameter
grid_prototype = None
for da in dataarray_arguments(bound_args):
if hasattr(da.metpy, 'latitude') and hasattr(da.metpy, 'longitude'):
grid_prototype = da
break

# Fill in x_dim/y_dim
if (
grid_prototype is not None
and 'x_dim' in bound_args.arguments
and 'y_dim' in bound_args.arguments
):
try:
bound_args.arguments['x_dim'] = grid_prototype.metpy.find_axis_number('x')
bound_args.arguments['y_dim'] = grid_prototype.metpy.find_axis_number('y')
except AttributeError:
# If axis number not found, fall back to default but warn.
warnings.warn('Horizontal dimension numbers not found. Defaulting to '
'(..., Y, X) order.')

# Fill in vertical_dim
if (
grid_prototype is not None
and 'vertical_dim' in bound_args.arguments
):
try:
bound_args.arguments['vertical_dim'] = (
grid_prototype.metpy.find_axis_number('vertical')
)
except AttributeError:
# If axis number not found, fall back to default but warn.
warnings.warn(
'Vertical dimension number not found. Defaulting to (..., Z, Y, X) order.'
)

# Fill in dz
if (
grid_prototype is not None
and 'dz' in bound_args.arguments
and bound_args.arguments['dz'] is None
):
try:
vertical_coord = grid_prototype.metpy.vertical
bound_args.arguments['dz'] = np.diff(vertical_coord.metpy.unit_array)
except (AttributeError, ValueError):
# Skip, since this only comes up in advection, where dz is optional (may not
# need vertical at all)
pass

# Fill in dx/dy
if (
'dx' in bound_args.arguments and bound_args.arguments['dx'] is None
and 'dy' in bound_args.arguments and bound_args.arguments['dy'] is None
):
if grid_prototype is not None:
bound_args.arguments['dx'], bound_args.arguments['dy'] = (
grid_deltas_from_dataarray(grid_prototype, kind='actual')
)
elif 'dz' in bound_args.arguments:
# Handle advection case, allowing dx/dy to be None but dz to not be None
if bound_args.arguments['dz'] is None:
raise ValueError(
'Must provide dx, dy, and/or dz arguments or input DataArray with '
'proper coordinates.'
)
else:
raise ValueError('Must provide dx/dy arguments or input DataArray with '
'latitude/longitude coordinates.')

# Fill in latitude
if 'latitude' in bound_args.arguments and bound_args.arguments['latitude'] is None:
if grid_prototype is not None:
bound_args.arguments['latitude'] = (
grid_prototype.metpy.latitude
)
else:
raise ValueError('Must provide latitude argument or input DataArray with '
'latitude/longitude coordinates.')

return func(*bound_args.args, **bound_args.kwargs)
return wrapper


def add_vertical_dim_from_xarray(func):
"""Fill in optional vertical_dim from DataArray argument."""
@functools.wraps(func)
Expand Down
Loading