xarray+eofs：对EOF的简单实现

from eofs.standard import Eof
import xarray as xr
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt


def get_time_space(df, time_dim, lumped_space_dims):
    """
    Get a DataFrame with the dim: [row (e.g.,time), column (e.g, space)]

    Parameters
    ----------

    df : pandas.core.frame.DataFrame 
        | time | lat | lon | PM2.5 |
        | 2011 | 66  | 88  | 22    |
        | 2011 | 66  | 99  | 23    |
        | 2012 | 66  | 88  | 24    |
        | 2012 | 66  | 99  | 25    |
        
 
    time_dim:
        time dim name (str), e.g., "time"
    
    lumped_space_dims: 
        a list of the lumped space, e.g., ["lat","lon"]

    Returns
    -------
    dataframe:  
        with [row (e.g.,time), column (e.g, space)], e.g.,
        | time | loc1    | loc2    |
        |      | (66,88) | (66,99) |
        | 2011 | 22      | 24      |
        | 2012 | 23      | 25      |
    
    """
    return df.set_index([time_dim]+lumped_space_dims).unstack(lumped_space_dims)

import xarray as xr

# 数据
f = '/home/mw/input/cru3546/cru_ts4.07.2021.2022.pre.dat.nc'

# 打开数据集
ds = xr.open_dataset(f)
ds = ds.sel(lat=slice(20, 60), lon=slice(80, 130))
ds

df = ds.to_dataframe().reset_index() # get df from da
df=df.drop(columns=['stn'])
df

df_data = get_time_space(df, time_dim = "time", lumped_space_dims = ["lat","lon"])
display(df_data.head(5))
print("DataFrame Shape:",df_data.shape)

DataFrame Shape: (24, 8000)

import matplotlib.pyplot as plt

def visualization(da, pcs, eofs_da, evf):
    n = len(pcs[0])  # 获取模态数
    fig, axs = plt.subplots(n+1, 2, figsize=(20, 6*(n+1)))

    da.mean(dim=["lat","lon"]).plot(ax=axs[0, 0])
    axs[0, 0].set_title("average pre")
    da.mean(dim="time").plot(ax=axs[0, 1])
    axs[0, 1].set_title("average pre")

    for i in range(1, n+1):
        pc_i = xr.DataArray(pcs[:,i-1],dims=["time"])
        pc_i = pc_i.assign_coords(time=ds.time)
        eof_i = eofs_da.sel(mode=i-1)
        frac = str((evf[i-1] * 100).round(2))

        axs[i, 0].plot(pc_i)
        axs[i, 0].set_title("PC"+str(i)+" ("+frac+"%)")

        img = eof_i.plot(x='lon',y='lat',ax=axs[i, 1], vmin=-0.75, vmax=0.75, cmap="RdBu_r", add_colorbar=False)
        axs[i, 1].set_title("EOF"+str(i)+" ("+frac+"%)")
        fig.colorbar(img, ax=axs[i, 1])

    plt.tight_layout()
    plt.show()

from eofs.standard import Eof
from sklearn.preprocessing import StandardScaler
n=4
# 使用StandardScaler对数据进行标准化处理
scaler = StandardScaler()
df_data_scaled = scaler.fit_transform(df_data.values)

# 创建EOFs对象
solver = Eof(df_data_scaled)

# 计算主分量分数（PCs）
s_pcs = solver.pcs(npcs=4, pcscaling=2)

# 计算模态函数（EOFs）
s_eofs =solver.eofs(neofs=4, eofscaling=2)

s_eofs = s_eofs.reshape(4,100,80)

# 创建xarray数组
lon = ds.lon
lat = ds.lat
mode = np.arange(0,n)
s_eofs_ds = xr.DataArray(s_eofs,coords=[mode,lon, lat], dims=["mode", "lon", "lat"])
# 计算解释方差比例
s_evfs = solver.varianceFraction(neigs=4)


visualization(ds['pre'], s_pcs, s_eofs_ds, s_evfs)

/opt/conda/lib/python3.7/site-packages/sklearn/utils/extmath.py:770: RuntimeWarning: invalid value encountered in true_divide
  updated_mean = (last_sum + new_sum) / updated_sample_count
/opt/conda/lib/python3.7/site-packages/sklearn/utils/extmath.py:709: RuntimeWarning: Degrees of freedom <= 0 for slice.
  result = op(x, *args, **kwargs, dtype=np.float64)

coslat = np.array(np.sqrt(np.cos(np.deg2rad(ds.lat)))).reshape((1,80))


# 创建EOFs对象
solver1 = Eof(df_data_scaled.reshape(24,100,80),weights = coslat)

# 计算主分量分数（PCs）
s_pcs1 = solver1.pcs(npcs=4, pcscaling=2)

# 计算模态函数（EOFs）
s_eofs1 =solver.eofs(neofs=4, eofscaling=2)

s_eofs1 = s_eofs.reshape(4,100,80)

# 创建xarray数组
lon = ds.lon
lat = ds.lat
mode = np.arange(0,n)
s_eofs_ds1 = xr.DataArray(s_eofs1,coords=[mode,lon, lat], dims=["mode", "lon", "lat"])
# 计算解释方差比例
s_evfs1 = solver.varianceFraction(neigs=4)


visualization(ds['pre'], s_pcs1, s_eofs_ds1, s_evfs1)