Python中多维数组Mann-Whitney U检验及xarray数组值修改问题

Hey there! Let's work through your xarray and Mann-Whitney U test problem together. I'll cover both an efficient way to run the test across all grid points and months, plus fix those assignment errors you hit.

Efficient Batch Processing: Avoid Manual Loops with xr.apply_ufunc #

Manual loops over grid points and months can be slow and error-prone. Instead, we can use xarray's apply_ufunc to vectorize the Mann-Whitney U test across all dimensions we care about. Here's how:

Step 1: Prep Your Data #

First, add a month coordinate to both your past and future temperature arrays so we can group data by month:

import numpy as np
import xarray as xr
import scipy.stats as sts

# Load your datasets (assuming fileHis and fileFut are already opened)
tp = fileHis['tas'].assign_coords(month=lambda x: x.time.dt.month)
tf = fileFut['tas'].assign_coords(month=lambda x: x.time.dt.month)

Step 2: Define a Helper Function for the Test #

Create a small function that takes two 1D time series (past and future for a single grid point) and returns the p-value from the Mann-Whitney U test. We'll add a check to handle all-NaN grid points gracefully:

def calculate_pval(past_ts, future_ts):
    # Skip grid points with no valid data
    if np.isnan(past_ts).all() or np.isnan(future_ts).all():
        return np.nan
    # Run the test and return the p-value
    stat, pval = sts.mannwhitneyu(past_ts, future_ts, alternative='two-sided')
    return pval

Step 3: Vectorize the Test Across All Dimensions #

Use xr.apply_ufunc to apply our helper function to every grid point and month. This handles all the broadcasting and dimension alignment automatically:

pvalue_mon = xr.apply_ufunc(
    calculate_pval,
    tp, tf,
    # Specify that each input uses 'time' as the core dimension to process
    input_core_dims=[['time'], ['time']],
    # The output is a scalar per grid point/month, so no core dimensions
    output_core_dims=[[]],
    # Auto-vectorize the function to handle rlon/rlat/month dimensions
    vectorize=True,
    # Allow dask for large datasets (optional but useful for big climate data)
    dask='allowed',
    # Set the output data type
    output_dtypes=[float]
).rename('pvalue')

# Ensure the month coordinate is labeled with 1-12
pvalue_mon = pvalue_mon.assign_coords(month=np.arange(1,13))

The resulting pvalue_mon will have dimensions matching your original data minus time, plus a month dimension (so (month, rlon, rlat) or similar, depending on your original dimension order). You can reorder dimensions with .transpose('rlon', 'rlat', 'month') if needed.

Fixing Assignment Errors in xarray #

Your earlier attempts to update single values failed because xarray doesn't allow direct assignment to sel() results, and update() expects an xarray object—not a single float. Here are two reliable ways to modify individual values:

Method 1: Label-Based Assignment with .loc #

Use .loc to target specific coordinates directly:

# Assign a p-value to the grid point (-6.375, 1.375) for month 1
pvalue_mon.loc[dict(rlon=-6.375, rlat=1.375, month=1)] = 0.05922372345359562

Method 2: Position-Based Assignment #

If you know the integer indices of the coordinates, you can index directly:

# Get boolean masks for the target coordinates
rlon_mask = pvalue_mon.rlon == -6.375
rlat_mask = pvalue_mon.rlat == 1.375
month_mask = pvalue_mon.month == 1

# Assign the value
pvalue_mon[month_mask, rlat_mask, rlon_mask] = 0.05922372345359562

Why Your Previous Code Failed #

• pvalue_mon.sel(...) = ...: Python doesn't let you assign values to the result of a function call (like sel()). .loc is the correct way to do label-based assignment in xarray.
• .update(result[1]): The update() method expects an xarray DataArray/Dataset, not a single numpy float. It's meant to merge entire datasets, not set individual values.

Full Example Recap #

Here's the complete workflow from start to finish:

import numpy as np
import xarray as xr
import scipy.stats as sts

# Load data (replace with your actual dataset paths)
fileHis = xr.open_dataset('historical_data.nc')
fileFut = xr.open_dataset('future_data.nc')

# Prep data with month coordinate
tp = fileHis['tas'].assign_coords(month=lambda x: x.time.dt.month)
tf = fileFut['tas'].assign_coords(month=lambda x: x.time.dt.month)

# Define test function
def calculate_pval(past_ts, future_ts):
    if np.isnan(past_ts).all() or np.isnan(future_ts).all():
        return np.nan
    stat, pval = sts.mannwhitneyu(past_ts, future_ts, alternative='two-sided')
    return pval

# Compute p-values across all grid points and months
pvalue_mon = xr.apply_ufunc(
    calculate_pval,
    tp, tf,
    input_core_dims=[['time'], ['time']],
    output_core_dims=[[]],
    vectorize=True,
    dask='allowed',
    output_dtypes=[float]
).rename('pvalue').assign_coords(month=np.arange(1,13))

# Example: Assign a specific p-value to a grid point and month
pvalue_mon.loc[dict(rlon=-6.375, rlat=1.375, month=1)] = 0.05922372345359562

# Save the result to a netCDF file (optional)
pvalue_mon.to_netcdf('monthly_pvalues.nc')

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