气象大模型论文中评估指标的计算详解

import numpy as np
import xarray as xr

def compute_weighted_rmse(da_fc, da_true, mean_dims=xr.ALL_DIMS):
    """
    计算具有纬度加权的两个 xr.DataArrays 之间的 RMSE。

    参数:
        da_fc (xr.DataArray): 预测值。
        da_true (xr.DataArray): 真值。
        mean_dims: 要平均的维度
    返回:
        rmse: 纬度加权的均方根误差
    """
    error = da_fc - da_true
    weights_lat = np.cos(np.deg2rad(error.lat))
    weights_lat /= weights_lat.mean()
    rmse = np.sqrt(((error)**2 * weights_lat).mean(mean_dims))
    return rmse

def compute_weighted_acc(da_fc, da_true, mean_dims=xr.ALL_DIMS):
    """
    计算具有纬度加权的两个 xr.DataArrays 之间的 ACC。

    参数:
        da_fc (xr.DataArray): 预测值。
        da_true (xr.DataArray): 真值。
        mean_dims: 要平均的维度
    返回:
        acc: 纬度加权的异常相关系数
    """

    clim = da_true.mean('time')
    try:
        t = np.intersect1d(da_fc.time, da_true.time)
        fa = da_fc.sel(time=t) - clim
    except AttributeError:
        t = da_true.time.values
        fa = da_fc - clim
    a = da_true.sel(time=t) - clim

    weights_lat = np.cos(np.deg2rad(da_fc.lat))
    weights_lat /= weights_lat.mean()
    w = weights_lat

    fa_prime = fa - fa.mean()
    a_prime = a - a.mean()

    acc = (
            np.sum(w * fa_prime * a_prime) /
            np.sqrt(
                np.sum(w * fa_prime ** 2) * np.sum(w * a_prime ** 2)
            )
    )
    return acc

def compute_weighted_mae(da_fc, da_true, mean_dims=xr.ALL_DIMS):
    """
    计算具有纬度加权的两个 xr.DataArrays 之间的 MAE。
    参数:
        da_fc (xr.DataArray): 预测值。
        da_true (xr.DataArray): 真值。
        mean_dims: 要平均的维度
    返回:
        mae: 纬度加权的平均绝对误差
    """
    error = da_fc - da_true
    weights_lat = np.cos(np.deg2rad(error.lat))
    weights_lat /= weights_lat.mean()
    mae = (np.abs(error) * weights_lat).mean(mean_dims)
    return mae

import xarray as xr
import numpy as np

# 定义时间、纬度和经度的维度大小
time_size = 10
lat_size = 5
lon_size = 5

# 生成随机数据数组1
data_array1 = xr.DataArray(np.random.rand(time_size, lat_size, lon_size),
                           dims=('time', 'lat', 'lon'),
                           coords={'time': range(time_size), 'lat': range(lat_size), 'lon': range(lon_size)})

# 生成随机数据数组2
data_array2 = xr.DataArray(np.random.rand(time_size, lat_size, lon_size),
                           dims=('time', 'lat', 'lon'),
                           coords={'time': range(time_size), 'lat': range(lat_size), 'lon': range(lon_size)})

# 打印数据数组1
print("Data Array 1:")
print(data_array1)

# 打印数据数组2
print("Data Array 2:")
print(data_array2)

#rmse
print(compute_weighted_rmse(data_array1,data_array2))

#acc
print(compute_weighted_acc(data_array1,data_array2))

#mae
print(compute_weighted_mae(data_array1,data_array2))