【数值预报】按时间维度合并/重新生成nc、grib网格数据（按天、小时组织的文件合并成按月组织文件）

全球网格数据如grib、nc有些是按年组织的、有些是按月组织的、有些是按日组织的、有些是按小时组织的，然而这些在时间上都是一维形式，即普通的时间序列。

对于数值模式预报数据，如ECMWF、GRAPES、JMA、NCEP 等等，其数据在时间上是二维形式，即（起报时刻，超前序列）。举个简单例子，假设气象台对气温预报，星期一会预报星期二到星期五的气温，星期二会预测星期三到星期六的气温，星期三会预测星期四到星期日的气温，这里就可以看到，星期一到星期天产生的序列长度不是7，而是3 x 7 = 21，所以说预报数据是二维时间形式，这里的3就是起报时刻，7就是超前序列。

 一般二维预报数据是按月组织的，比如在ECMWF官网下载的数据集，下载后解析可以看到时间是二维的，如下图：

 一般读取这些网格数据用的是xarray，相比于netcdf4，xarray兼容grib和nc。使用xarray读取后，二维预报数据的维度（起报时刻，超前序列）会被解析为（time，step），比如下图中的time和step维度：

现在问题来了，有些全球网格的预报数据不是按月组织的，我想把它合成按月组织的怎么办？

例如中科院大气所的FGOALS模式，它是按天组织的，每天一个文件（一次预报），每次预报预测未来7天数据：

打开第一个文件，解析结果如上图右所示，发现是一维的预报数据，只有time维度，这与标准的二维预报数据不符。同时还发现，本来上图右的time维度应该是step，因为这是超前时间，超前7天，但由于是按天组织，退化成一维预报数据，导致用time维度替换了step维度。

所以一个大概合并流程要解决的问题包括：增加step维度，把原来的time维度替换成step维度，并将真正的起报时刻作为time维度。像什么xarray里的维度交换、更名等等操作全试了，都不行，因为这网格数据比想象的要复杂，它包含了坐标、索引等东西，这些东西是耦合的，所以不是交换维度、更名等操作能解决的。

下面给出一个个人的解决方案：

import numpy as np
import xarray as xr
import pandas as pd


def extract_all_element_data(total_ds):
    """获取dataset中所有要素对应的numpy数组 + 原有的属性信息，即元组 (nparr, attrs)
    """
    ele_data_dict = {}
    elements = list(total_ds.var())

    for ele in elements:
        ele_data_dict[ele] = (total_ds[ele].data, total_ds[ele].attrs)
    return ele_data_dict


def clear_dataset_dims(dataset:xr.Dataset):
    # 暂时只用去掉time维度
    dataset = dataset.drop_dims("time")
    return dataset


def get_idxs(raw_list, target_list):
    """得到target_list里的元素在raw_list里对应的下标
       Return: [idx1, idx2, ...], 其中len(target_list) == len(Return)
    """
    _dict, c = {}, 0
    for item in raw_list:
        _dict[item] = c
        c += 1
    ans = []
    for item in target_list:
        idx = _dict.pop(item)
        ans.append(idx)
    return ans


def get_step_coords(leading_times):
    """将[24, 48, 60, 72, 78...]等超前时间转换成xarray的坐标内容形式
       不用pandas直接生成序列的原因是这样写能处理非固定时间间隔的情况
    """
    pd_step = list(map(lambda x:"{} hour".format(int(x)), leading_times))
    step_coords = pd.TimedeltaIndex(pd_step)
    return step_coords


def get_time_coords(str_start_date:str, reftimes:list, days:int):
    """将[00:00:00, 16:00:00, 18:00:00, ...]等起报时间转换成xarray的坐标内容形式
       不用pandas直接生成序列的原因是这样写能处理非固定时间间隔的情况
    """
    import datetime
    start_date = datetime.datetime.strptime(str_start_date, "%Y-%m-%d")
    
    time_deltas = [datetime.timedelta(hours=hour) for hour in list(
        map(lambda x:float(x.split(':')[0]), reftimes))
    ]
    time_coords = []

    for i in range(days):
        day_i_reftimes = [start_date + i * datetime.timedelta(hours=24) + time_delta \
            for time_delta in time_deltas]
        time_coords.extend(day_i_reftimes)

    pd_time_coords = pd.DatetimeIndex(time_coords)
    return pd_time_coords


def get_missing_idx(idxs):
    max_id = max(idxs)
    for _idx in range(max_id + 1):
        if _idx not in idxs:
            return _idx
    return max_id + 1


def get_element_dims(total_ds:xr.Dataset):
    elements = list(total_ds.var())
    final_dim_name = None
    final_dim_shape = None
    for ele in elements:
        cur_dim_names = total_ds[ele].dims
        cur_dim_shape = total_ds[ele].shape
        if final_dim_name and final_dim_shape:
            assert final_dim_name == cur_dim_names
            assert final_dim_shape == cur_dim_shape
        final_dim_name = cur_dim_names
        final_dim_shape = cur_dim_shape
    assert final_dim_name is not None and final_dim_shape is not None

    return (final_dim_name, final_dim_shape)


def exchange_dims(raw_dim_list, idx_a, idx_b):
    tmp = raw_dim_list[idx_a]
    raw_dim_list[idx_a] = raw_dim_list[idx_b]
    raw_dim_list[idx_b] = tmp
    return raw_dim_list


def concat_to_forecast_data(start_date, days, reftimes, leading_times, filepaths, target_path):
    """将一维模式（超前时间）预报数据 转换成二维模式（起报, 超前）预报数据

    Args:
        start_date (str): '2020-01-01'，数据开始日期
        days (int): 61，所有数据覆盖的天数
        reftimes (list): ['00:00:00', '06:00:00', '18:00:00']，每天的起报时间
        leading_times (list): ['06', '12', '24',...] 或 ['6', '12', '24', ...] 或 [6, 12, 24, ...]，每报的超前时间
        filepaths (list): 待合并的nc或grib文件路径
        target_path (str): 合并后生成的nc文件路径
    """
    multi_ds = [xr.open_dataset(filepath) for filepath in filepaths]
    # 将所有文件的time全部积累到一起，本来ref, leading = (2, 7), 现在time维度为14
    total_ds = xr.concat(multi_ds, dim="time")
    
    nums_ref_perday, nums_leading = len(reftimes), len(leading_times)
    nums_ref = days * nums_ref_perday
    # 所有文件时间累积后应该与期望的时间数相等
    assert nums_ref * nums_leading == total_ds.dims["time"]
    
    # {"ele":nparr, ...}
    ele_data_dict = extract_all_element_data(total_ds)

    #raw_dims = dict(total_ds.dims) // dataset和要素对应的dataarray维度顺序还不一样，所以不能这么写
    raw_dim_name, raw_dim_shape = get_element_dims(total_ds)

    # 清除time维度等工作，不然直接给要素赋新值（reshape后的）会维度不匹配报错
    total_ds = clear_dataset_dims(total_ds)

    # 去除掉time等维度后剩余的维度
    rest_dims = dict(total_ds.dims)
    # raw_dims:(height, time, lon, lat), rest_dims:(height, lon, lat)
    # 则idxs:(0, 2, 3)
    idxs = get_idxs(raw_dim_name, list(rest_dims.keys()))

    total_ds.coords["step"] = get_step_coords(leading_times)
    total_ds.coords["time"] = get_time_coords(start_date, reftimes, days)

    for ele, ele_info in ele_data_dict.items():
        # (height, time, lon, lat) --> (time, step, height, lon, lat)，这里要利用好(0, 2, 3)
        # (0, 2, 3)消失的是1，即time维度，所以time变time，step，即（1）变成 (0, 1) ，剩下维度(0, 2, 3)放后排，变成(2, 3, 4) 
        nparr, attrs = ele_info
        raw_time_idx = get_missing_idx(idxs)

        final_dim_name = exchange_dims(list(raw_dim_name), 0, raw_time_idx)
        final_dim_shape = exchange_dims(list(raw_dim_shape), 0, raw_time_idx)
        rest_dim_names = list(final_dim_name)[slice(1, len(final_dim_name))]
        rest_dim_shape = list(final_dim_shape)[slice(1, len(final_dim_shape))]

        nparr = nparr.swapaxes(0, raw_time_idx)
        nparr = nparr.reshape(nums_ref, nums_leading, *rest_dim_shape)
        total_ds[ele] = (("time", "step", *rest_dim_names), nparr, attrs)
    
    total_ds.to_netcdf(target_path)


import os
import natsort

# 按天或按小时组织的预报数据目录
folder = "/Users/xxxx/data/FGOALS/u10"
filenames = os.listdir(folder)
filenames = list(filter(lambda x: x.endswith(".nc"), filenames))
filepaths = ["{}/{}".format(folder, filename) for filename in filenames]

# 保存os.listdir的文件顺序性
filepaths = natsort.natsorted(filepaths)

# 参数分别为预报文件起始时间、所有文件累计的天数、起报时刻、超前序列（小时）、待合成的filepaths、输出的filepath
concat_to_forecast_data('2020-01-01', days=59, reftimes=['00:00:00'], \
    leading_times=['24', '48', '72', '96', '120', '144', '168'], filepaths=filepaths, target_path="z.nc")

natsort是一个python包，直接pip安装即可，当然也可以不下载，只要你能保证filepaths的顺序即可

就上面例子而言，会输出一个“z.nc”文件，该文件包含了59个按天组织的文件的信息。

 同时打开z.nc和原始的预报数据，例如s2s_cas_2019_20200101_10u.nc，解析结果如下：

其中左图是原始数据，右图是合成数据，发现一摸一样，且左边只有一个time维度，但合成后的数据有time和step两个维度，说明合成成功。

进一步验证，查看其他时间的数据情况，如下视频 

 可观察到原始预报数据和合成后的数据始终是一致的，说明合成成功。

—————以下仅个人吐槽，与文无关，可以不看———————

这里默默吐槽下，也是对我导接的项目的小吐槽。一个项目说是给我们提供符合质量的数据，结果到项目结题了数据还没给全，真让我们乙方（我）心累。亲爱的读者，如果你恰好是我的数据提供方，还请用这个代码帮忙把FGOALS的数据合成为月数据的形式，缺失的直接复制近邻时间的文件当作本时刻的即可，比如1、2、3号缺了2号，直接拿1号数据复制一份文件，文件名改成2号的即可，这代码都可以合成，不甚感谢🙏。