基于xarray的7天滚动窗口温度分箱计算优化求助

Great question! Your current implementation gets the job done, but looping through each rolling window manually can be slow—especially with large datasets. Let's optimize this using xarray.apply_ufunc paired with Dask, which will leverage vectorization and parallel computing to speed things up, perfect for your dask.distributed.Client setup.

Here's the Optimized Approach #

First, we'll define a pure NumPy function to handle the per-window binning logic (this lets us leverage fast array operations without Python-level loops). Then we'll wrap this function with xarray.apply_ufunc to integrate it seamlessly with xarray's rolling windows and Dask parallelism.

import xarray as xr
import numpy as np
from dask.distributed import Client

# Initialize your Dask client (adjust settings as needed for your cluster)
client = Client()

def compute_temp_buckets(window):
    """
    Takes a 3D window array (time, lat, lon) and returns 2D bucket indices (lat, lon)
    for the last time step in the window, binned using min/max of the window.
    """
    # Calculate min/max temp across the window for each (lat, lon) point
    max_temp = np.nanmax(window, axis=0)
    min_temp = np.nanmin(window, axis=0)
    
    # Handle edge case where min == max (avoid empty bins)
    bins = np.where(
        min_temp != max_temp,
        np.arange(min_temp, max_temp, 2),  # Regular bins if min != max
        np.stack([min_temp - 1, min_temp + 1], axis=-1)  # Dummy bins if all values are the same
    )
    
    # Extract the last time step in the window (the one we want to bin)
    last_timestep = window[-1, :, :]
    
    # Compute bucket indices (np.digitize returns 1-based indices)
    buckets = np.digitize(last_timestep, bins=bins)
    
    # Fix edge cases: values below the first bin get assigned to bucket 1
    buckets = np.where(last_timestep < bins[..., 0], 1, buckets)
    # Values >= max_temp get assigned to the highest valid bucket
    buckets = np.where(last_timestep >= max_temp, bins.shape[-1], buckets)
    
    return buckets

# Ensure your dataset is chunked appropriately (adjust chunk sizes for your data)
# We keep time as a single chunk here to optimize rolling window operations
ds = ds.chunk({'time': -1, 'lat': 10, 'lon': 10})

# Create the 7-day rolling window
rolling_window = ds.t2m.rolling(time=7, center=False)

# Apply our binning function across all windows using apply_ufunc
ds_buckets = xr.apply_ufunc(
    compute_temp_buckets,
    rolling_window,
    input_core_dims=[['time']],  # We're operating over the time window dimension
    output_core_dims=[[]],  # Output has no time dimension (one value per window)
    vectorize=True,  # Vectorize over non-core dimensions (lat, lon)
    dask='parallelized',  # Use Dask to parallelize computation
    output_dtypes=[np.int32],  # Specify output data type
    dask_gufunc_kwargs={'allow_rechunk': True}  # Allow Dask to rechunk as needed
)

Why This Works Better Than Your Original Code #

1. No Manual Loops: apply_ufunc handles iterating over all rolling windows and spatial points automatically, eliminating the slow Python loop.
2. Dask Parallelism: With dask='parallelized', the computation is split across your Dask cluster (or local workers), drastically speeding up processing for large datasets.
3. Vectorized Operations: The pure NumPy function uses fast array operations instead of per-element processing, which is orders of magnitude faster than Python loops.
4. Cleaner, Maintainable Code: The logic is encapsulated in a single function, making it easier to debug and modify later.

Key Details to Note #

• Chunk Sizing: We chunked time as a single chunk to optimize rolling window operations (Dask handles rolling more efficiently on contiguous time chunks). Adjust lat/lon chunk sizes based on your available memory and cluster resources.
• Edge Case Handling: We explicitly handle cases where all values in a window are the same (to avoid empty bins) and fix boundary values that fall outside the generated bins.
• Result Consistency: You can verify the output matches your original code with:

  np.allclose(ds_buckets.compute(), ds_concat.values, equal_nan=True)
  

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