如何在维度变化时使用xr.apply_ufunc？附Xarray气候数据集示例

Great question! Working with dimension transformations using xr.apply_ufunc is a staple for climate data analysis—let’s walk through practical examples tailored to your dataset:

import xarray as xr
import numpy as np
from scipy.stats import linregress

# Your loaded dataset
climate = xr.open_dataset(data_file)

First, let’s recap your dataset structure for context:

• Dimensions: time (424), lat (621), lon (1405)
• Data variables: tmean (shape: (time, lat, lon)), status (shape: (time))

Scenario 1: Collapse Spatial Dimensions to Time #

Let’s say you want to compute a daily spatial average of tmean—this reduces the (time, lat, lon) array to a (time) array.

The key here is telling apply_ufunc which dimensions to operate on using input_core_dims, and defining the output dimensions with output_core_dims:

def spatial_mean(arr):
    # arr is a 2D (lat, lon) array for each time step
    return np.nanmean(arr, axis=(-2, -1))  # Handle NaNs in your climate data

# Apply the function
tmean_spatial_avg = xr.apply_ufunc(
    spatial_mean,
    climate['tmean'],
    input_core_dims=[['lat', 'lon']],  # We operate on lat/lon for each time slice
    output_core_dims=[[]],  # Output has no core dimensions (scalar per time step)
    vectorize=True,  # Auto-loop over non-core dimensions (time)
    dtypes=[float]
)

# Result will be a DataArray with only the `time` dimension
tmean_spatial_avg

Scenario 2: Compute a Time Trend for Each Spatial Grid Point #

Next, let’s calculate the linear trend of tmean over time for every (lat, lon) grid point. This transforms the (time, lat, lon) array to a (lat, lon) array.

Here, we’ll use scipy.stats.linregress and target the time dimension:

def compute_time_trend(arr):
    # arr is a 1D (time) array for each (lat, lon) grid point
    x = np.arange(len(arr))  # Time index for regression
    slope, _, _, _, _ = linregress(x, arr)
    return slope

tmean_trend = xr.apply_ufunc(
    compute_time_trend,
    climate['tmean'],
    input_core_dims=[['time']],  # Operate on the time dimension for each grid point
    output_core_dims=[[]],  # Output is a scalar per (lat, lon)
    vectorize=True,
    dtypes=[float]
)

# Result will be a DataArray with `lat` and `lon` dimensions only
tmean_trend

Scenario 3: Combine Multiple Variables with Different Dimensions #

Suppose you want to filter tmean based on the status variable (which only has a time dimension) and compute a spatial average for valid time steps.

Here, we pass two inputs to apply_ufunc and define core dimensions for each:

def filtered_spatial_avg(tmean_slice, status_val):
    # tmean_slice: 2D (lat, lon) array; status_val: scalar for the time step
    if status_val == 1:  # Example: only compute average if status is 1
        return np.nanmean(tmean_slice)
    else:
        return np.nan

filtered_avg = xr.apply_ufunc(
    filtered_spatial_avg,
    climate['tmean'],
    climate['status'],
    input_core_dims=[['lat', 'lon'], []],  # tmean uses lat/lon; status has no core dims
    output_core_dims=[[]],
    vectorize=True,
    dtypes=[float]
)

Key Parameters to Remember #

• input_core_dims: List of lists specifying which dimensions each input function operates on. Non-core dimensions are looped over automatically.
• output_core_dims: List of lists defining the dimensions of the output. Use [] for scalar outputs per loop iteration.
• vectorize=True: Critical for functions that don’t handle multi-dimensional broadcasting (like linregress). It tells Xarray to loop over non-core dimensions.
• dtypes: Specify the output data type to avoid type mismatches.

内容的提问来源于stack exchange，提问作者Shawn