Python 使用xarray根据经纬度获取nc文件子集

Python 使用xarray根据经纬度获取nc文件子集

使用xarray实现生成nc文件的子集，输入文件名和经纬度范围，保存成新的nc文件。

生成子集函数

import numpy as np
import xarray as xr

def subset_from_nc(filename: str, extent: list, save_filename: str = None) -> None:
    """
    根据经纬度范围获取nc文件子集，生成新nc文件
    :param filename: 输入文件名
    :param extent: 经纬度范围  [down up left right]
    :param save_filename: 输出文件名
    :return: None
    """
    dataset = xr.open_dataset(r'H:\test.nc')
    # 源文件需要有lat，lon两个变量名并且是升序排列（否则不适用searchsorted，需要其他方法获取索引），有些文件使用的是lat、lon全写。
    extent_index = [np.searchsorted(dataset['lat'].data, v) for v in extent[:2]]
    extent_index.extend(np.searchsorted(dataset['lon'].data, v) for v in extent[2:])
    dim_dict = {}
    for dim in dataset.dims:
        if dim == 'lat':
            dim_dict[f'{dim}'] = dataset[f'{dim}'][extent_index[0]:extent_index[1]+1]
        elif dim == 'lon':
            dim_dict[f'{dim}'] = dataset[f'{dim}'][extent_index[2]:extent_index[3]+1]
        else:
            dim_dict[f'{dim}'] = dataset[f'{dim}']

    var_dict = {}
    for var in dataset.variables:
        # 一维数据一般作为维度，所以这里不考虑，有特殊情况增加相应分支
        # 这里面的纬度名需要根据自己的数据相应调整，我测试的数据实际上是四维数据（time, depth, lat, lon）
        if len(dataset[f'{var}'].data.shape) == 2:
                var_dict[f'{var}'] = (('lat','lon'), dataset[f'{var}'].data[extent_index[0]:extent_index[1]+1, extent_index[2]:extent_index[3]+1])
        elif len(dataset[f'{var}'].data.shape) == 3:
                var_dict[f'{var}'] = (('depth', 'lat','lon'), dataset[f'{var}'].data[:, extent_index[0]:extent_index[1]+1, extent_index[2]:extent_index[3]+1])
        elif len(dataset[f'{var}'].data.shape) == 4:
                var_dict[f'{var}'] = (('time', 'depth', 'lat','lon'), dataset[f'{var}'].data[:, :, extent_index[0]:extent_index[1]+1, extent_index[2]:extent_index[3]+1])
        # 更高维的同理四维
        # ...
        else:
            pass
    # 新建数据集
    new_dataset = xr.Dataset(data_vars=var_dict, coords=dim_dict, attrs=dataset.attrs)
    # 写入变量属性
    for name in new_dataset.variables:
        new_dataset[f'{name}'].attrs = dataset[f'{name}'].attrs
    # 保存文件
    if save_filename is None:
        new_dataset.to_netcdf(filename.replace('.nc', '_subset[{}_{}_{}_{}].nc'.format(*extent)), format='NETCDF4', engine='netcdf4')
    else:
        new_dataset.to_netcdf(save_filename, format='NETCDF4', engine='netcdf4')
    dataset.close()

调用即可

subset_from_nc(r'H:test.nc', [10,50,100,150])