from typing import Optional, Union
import numpy as np
__all__ = [
"calculate_gradient",
"calculate_meridional_gradient",
"calculate_zonal_gradient",
"calculate_vertical_gradient",
"EARTH_RADIUS",
]
EARTH_RADIUS = 6371e3 # Earth's mean radius (m)
[docs]
def calculate_gradient(
field: np.ndarray,
coordinates: np.ndarray,
axis: int = -1,
radius: Optional[float] = None,
) -> np.ndarray:
"""Calculate gradient of an arbitrary dimensional array along specified coordinates
Args:
field: Data field for gradient calculation, can be any dimensional array, e.g., (time, level, lat, lon)
coordinates: Coordinate array along which to calculate the gradient, such as latitude or longitude values
axis: Specifies the dimensional axis for gradient calculation, defaults to the last dimension (-1)
radius: Earth radius, default is None. If provided, coordinates are treated as latitude in degrees
and converted to distance in meters. If None, coordinates are used directly.
Returns:
Gradient field with the same shape as the input field
"""
# Check if input data dimensions match
if coordinates.size != field.shape[axis]:
raise ValueError(
f"Coordinate array size ({coordinates.size}) does not match field size ({field.shape[axis]}) on the specified axis"
)
# If radius is provided, treat coordinates as latitude and convert to distance
if radius is not None:
# Convert latitude to actual distance
distances = coordinates * np.pi / 180.0 * radius
else:
# Use coordinates directly
distances = coordinates
# Create output array with the same shape as input
gradient = np.zeros_like(field, dtype=float)
# Handle edge case: single point (no gradient can be computed)
if field.shape[axis] < 2:
return gradient # Return zero gradient
# To use numpy's advanced indexing, we need to create index arrays
ndim = field.ndim
idx_ranges = [slice(None)] * ndim
# Handle the simple case where distances are uniform
dx = np.diff(distances)
if len(dx) > 0 and np.allclose(dx, dx[0]):
# Uniform spacing - use simple central differences
h = dx[0]
# Central differences for interior points
for i in range(1, field.shape[axis] - 1):
idx = idx_ranges.copy()
idx[axis] = i
idx_forward = idx_ranges.copy()
idx_forward[axis] = i + 1
idx_backward = idx_ranges.copy()
idx_backward[axis] = i - 1
gradient[tuple(idx)] = (
field[tuple(idx_forward)] - field[tuple(idx_backward)]
) / (2.0 * h)
# Forward difference for left boundary
idx_left = idx_ranges.copy()
idx_left[axis] = 0
idx_left_plus = idx_ranges.copy()
idx_left_plus[axis] = 1
gradient[tuple(idx_left)] = (
field[tuple(idx_left_plus)] - field[tuple(idx_left)]
) / h
# Backward difference for right boundary
idx_right = idx_ranges.copy()
idx_right[axis] = -1
idx_right_minus = idx_ranges.copy()
idx_right_minus[axis] = -2
gradient[tuple(idx_right)] = (
field[tuple(idx_right)] - field[tuple(idx_right_minus)]
) / h
else:
# Non-uniform spacing - use variable spacing formulas
for i in range(field.shape[axis]):
idx = idx_ranges.copy()
idx[axis] = i
if i == 0:
# Forward difference for first point
idx_plus = idx_ranges.copy()
idx_plus[axis] = 1
gradient[tuple(idx)] = (field[tuple(idx_plus)] - field[tuple(idx)]) / (
distances[1] - distances[0]
)
elif i == field.shape[axis] - 1:
# Backward difference for last point
idx_minus = idx_ranges.copy()
idx_minus[axis] = i - 1
gradient[tuple(idx)] = (field[tuple(idx)] - field[tuple(idx_minus)]) / (
distances[i] - distances[i - 1]
)
else:
# Central difference for interior points
idx_forward = idx_ranges.copy()
idx_forward[axis] = i + 1
idx_backward = idx_ranges.copy()
idx_backward[axis] = i - 1
h1 = distances[i] - distances[i - 1]
h2 = distances[i + 1] - distances[i]
# Use weighted central difference for non-uniform spacing
gradient[tuple(idx)] = (
-h2 / (h1 * (h1 + h2)) * field[tuple(idx_backward)]
+ (h2 - h1) / (h1 * h2) * field[tuple(idx)]
+ h1 / (h2 * (h1 + h2)) * field[tuple(idx_forward)]
)
return gradient
[docs]
def calculate_meridional_gradient(
field: np.ndarray,
latitudes: np.ndarray,
lat_axis: int = -1,
radius: float = 6371000.0,
) -> np.ndarray:
"""Calculate meridional gradient (gradient along latitude direction)
Args:
field: Data field for gradient calculation, can be any dimensional array
latitudes: Latitude array (degrees)
lat_axis: Specifies the axis for latitude, defaults to the last dimension (-1)
radius: Earth radius, default is 6371000.0 meters
Returns:
Meridional gradient field
"""
return calculate_gradient(field, latitudes, axis=lat_axis, radius=radius)
[docs]
def calculate_vertical_gradient(
field: np.ndarray, pressure: np.ndarray, pressure_axis: int = -3
) -> np.ndarray:
"""Calculate vertical gradient (gradient along pressure direction)
Args:
field: Data field for gradient calculation
pressure: Pressure array (Pa), must be monotonically decreasing
pressure_axis: Specifies the axis for pressure, defaults to the third-to-last dimension (-3)
Returns:
Vertical gradient field
"""
return calculate_gradient(field, pressure, axis=pressure_axis, radius=None)
[docs]
def calculate_zonal_gradient(
field: np.ndarray,
longitudes: np.ndarray,
latitudes: np.ndarray,
lon_axis: int = -1,
lat_axis: int = -2,
radius: float = 6371000.0,
) -> np.ndarray:
"""Calculate zonal gradient (gradient along longitude direction)
Args:
field: Data field for gradient calculation, can be any dimensional array
longitudes: Longitude array (degrees)
latitudes: Latitude array (degrees), used to calculate actual distance between longitudes at different latitudes
lon_axis: Specifies the axis for longitude, defaults to the last dimension (-1)
lat_axis: Specifies the axis for latitude, defaults to the second-to-last dimension (-2)
radius: Earth radius, default is 6371000.0 meters
Returns:
Zonal gradient field
"""
# Get latitude factor to adjust actual distance between longitudes at different latitudes
cos_lat = np.cos(np.radians(latitudes))
# If field is 4D (time, level, lat, lon)
if field.ndim == 4 and lon_axis == -1 and lat_axis == -2:
# Create a latitude factor array with shape suitable for broadcasting
cos_lat_expanded = cos_lat.reshape(1, 1, -1, 1)
# Convert longitudes to actual distances considering latitude
effective_distances = np.radians(longitudes) * radius * cos_lat_expanded
# Calculate gradient
return calculate_gradient(field, effective_distances, axis=lon_axis, radius=1.0)
# If field is 3D (time, lat, lon)
elif field.ndim == 3 and lon_axis == -1 and lat_axis == -2:
cos_lat_expanded = cos_lat.reshape(1, -1, 1)
effective_distances = np.radians(longitudes) * radius * cos_lat_expanded
return calculate_gradient(field, effective_distances, axis=lon_axis, radius=1.0)
else:
# For other dimension combinations, create appropriate broadcasting shape
broadcast_shape = [1] * field.ndim
broadcast_shape[lat_axis] = len(latitudes)
cos_lat_expanded = cos_lat.reshape(broadcast_shape)
# Create effective distance array
effective_longitudes = np.radians(longitudes) * radius
# Calculate gradient for each latitude
result = np.zeros_like(field)
# Loop through each latitude (implementation depends on specific data structure, may need adjustment)
for i in range(len(latitudes)):
idx = [slice(None)] * field.ndim
idx[lat_axis] = i
# Adjust longitude distance for current latitude
current_effective_dist = effective_longitudes * cos_lat[i]
# Calculate gradient for current latitude
result[tuple(idx)] = calculate_gradient(
field[tuple(idx)], current_effective_dist, axis=lon_axis, radius=1.0
)
return result
