netcdf文件复制并修改

import numpy as np
from netCDF4 import Dataset
from datetime import datetime

# 获取当前时间
current_time = datetime.now()

# 格式化为只显示时间
formatted_time = current_time.strftime("%m-%d-%H-%M-%S")


def extend_list(lst, extension_count):
    """
    扩展列表两端，增加指定数量的元素
    :param lst: 原始列表
    :param extension_count: 扩展的元素数量
    :return: 扩展后的列表
    """
    original_type = type(lst[0])
    delta = lst[1] - lst[0]

    # 向左扩展
    left_extension = [original_type(lst[0] - delta * i) for i in range(1, extension_count + 1)]
    left_extension.reverse()

    # 向右扩展
    right_extension = [original_type(lst[-1] + delta * i) for i in range(1, extension_count + 1)]

    # 合并扩展后的结果
    extended_lst = np.concatenate([left_extension, lst, right_extension])

    return extended_lst


def copy_variable_attributes(src_var, dst_var):
    """
    复制变量的所有属性并设置分块大小
    :param src_var: 源变量
    :param dst_var: 目标变量
    """
    # 复制变量的属性
    for attr_name in src_var.ncattrs():
        if attr_name != '_FillValue':  # 跳过 _FillValue
            dst_var.setncattr(attr_name, src_var.getncattr(attr_name))


def copy_polar_stereographic(src_nc, dst_nc):
    """
    复制 polar_stereographic 变量及其属性
    :param src_nc: 源 NetCDF 文件
    :param dst_nc: 目标 NetCDF 文件
    """
    polar_stereographic_src = src_nc.variables['polar_stereographic']
    polar_stereographic_target = dst_nc.createVariable(
        'polar_stereographic', polar_stereographic_src.datatype, polar_stereographic_src.dimensions
    )

    # 复制所有属性
    for attr in polar_stereographic_src.ncattrs():
        polar_stereographic_target.setncattr(attr, polar_stereographic_src.getncattr(attr))

    # 复制数据（如果是一个简单的标量数据，可以直接复制）
    polar_stereographic_target[:] = polar_stereographic_src[:]

    x_max = int(dst_nc.variables['x'][:][-1] + 3125)
    y_yax = int(dst_nc.variables['y'][:][-1] + 3125)
    # 设置 GeoTransform（示例值，可以根据需要调整）
    polar_stereographic_target.GeoTransform = f'-{x_max} 6250 0 {y_yax} 0 -6250'


def apply_circle_mask_to_z(z):
    """
    在二维数组 z 上应用圆形掩膜
    :param z: 原始二维数据数组
    :return: 更新后的 z 数组
    """
    # 创建一个掩膜，初始化为 NaN
    # mask = np.full_like(z, np.nan, dtype=np.float32)
    mask = z
    # 获取圆心坐标
    center_x, center_y = z.shape[0] / 2, z.shape[1] / 2
    if center_x > center_y:
        radius = center_y
    else:
        radius = center_x
    # 遍历 z 数组，计算每个点到圆心的距离
    for i in range(z.shape[0]):  # 遍历二维数组的每一行
        for j in range(z.shape[1]):  # 遍历二维数组的每一列
            # 计算当前点 (i, j) 到圆心 (center_x, center_y) 的距离
            distance = np.sqrt((i - center_x) ** 2 + (j - center_y) ** 2)
            if distance > radius:
                mask[i, j] = np.nan
            # 如果点在圆内，设置值为 50，否则保持为掩膜空值
            # if distance <= radius:
            #     mask[i, j] = 50
    return mask


def create_and_save_nc(input_file_path, output_file_path, extension_count=100):
    """
    扩展原始 NetCDF 数据并保存到新文件
    :param input_file_path: 输入的 NetCDF 文件路径
    :param output_file_path: 输出的 NetCDF 文件路径
    :param extension_count: 要扩展的坐标数（默认扩展 100）
    """
    # 打开原始 NetCDF 文件并读取数据
    with Dataset(input_file_path, mode='r') as nc_file:
        # 读取现有的 x、y 和 z
        x = nc_file.variables['x'][:]
        y = nc_file.variables['y'][:]
        z = nc_file.variables['z'][:]
        # 获取现有的形状
        x_size, y_size = x.shape[0], y.shape[0]

        # 计算新坐标，增加指定数量的坐标点
        new_x = extend_list(x, extension_count)  # 在 x 轴增加指定数量的点
        new_y = extend_list(y, extension_count)  # 在 y 轴增加指定数量的点

        # 创建新的 z 数组，形状为新的 x 和 y 的交叉形状
        new_z = np.full((new_y.shape[0], new_x.shape[0]), np.ma.masked, dtype=np.float32)  # 新数组赋值为 NaN

        # 将现有的 z 数组放入新的 z 数组中，保持原有数据
        new_z[extension_count:y_size + extension_count, extension_count:x_size + extension_count] = z  # 现有数据放入新数组的中间

        # 应用圆形掩膜，更新 z 数据
        # new_z = apply_circle_mask_to_z(new_z)

        # 创建新的 NetCDF 文件来保存数据
        with Dataset(output_file_path, mode='w', format='NETCDF4') as new_nc_file:
            # 创建新的维度
            new_nc_file.createDimension('x', new_x.shape[0])
            new_nc_file.createDimension('y', new_y.shape[0])

            # 创建新的变量并保存数据
            x_var = new_nc_file.createVariable('x', 'f8', ('x',), zlib=False)
            y_var = new_nc_file.createVariable('y', 'f8', ('y',), zlib=False)
            # 设置分块大小
            chunk_sizes = (1, x_var.size)  # 根据你的需求设置块大小
            z_var = new_nc_file.createVariable('z', np.float32, ('y', 'x'),
                                               fill_value=np.float32(np.nan), chunksizes=chunk_sizes,
                                               zlib=True)

            # 将数据写入变量
            x_var[:] = new_x
            y_var[:] = new_y
            z_var[:] = new_z

            # 复制原始变量的属性并设置分块
            copy_variable_attributes(nc_file.variables['x'], x_var)
            copy_variable_attributes(nc_file.variables['y'], y_var)
            copy_variable_attributes(nc_file.variables['z'], z_var)

            # 复制空间参考信息
            copy_polar_stereographic(nc_file, new_nc_file)

    print(f"数据已保存到 {output_file_path}")


if __name__ == '__main__':
    input_file_path = r'F:\Polar\data\wind_speed_241130.nc'
    output_file_path = fr'F:\Polar\data\{formatted_time}-wind_speed_241130.nc'

    create_and_save_nc(input_file_path, output_file_path, extension_count=250)