Datawhale 春训营 AI + 航空安全 task2

气象数据特征工程与预测代码笔记

一、代码核心功能概述

本代码主要针对气象数据进行处理，核心功能包括数据特征工程、预测结果生成以及数据文件的批量处理。通过对原始气象数据张量进行特征提取与变换，生成新的特征工程数据，并基于此数据进行预测，最终将预测结果保存到指定目录。

二、关键变量定义

1. 数据切片范围
   • LAT_RANGE：定义纬度切片范围 slice(35, 81)，用于裁剪目标区域的纬度数据。
   • LON_RANGE：定义经度切片范围 slice(70, 141)，用于裁剪目标区域的经度数据。
2. 变量索引映射
   • VAR_INDEX_MAPPING：字典类型，存储气象变量（如 z、t、u 等）在数据张量通道维度中的起始索引，方便后续对不同变量数据的提取与处理。
3. 输出配置
   • OUTPUT_SETTINGS：列表类型，包含一系列元组，每个元组指定要输出的气象变量和对应的气压层（如 ('t', 200) 表示温度变量在 200hPa 气压层的数据），用于确定最终预测结果的内容。
4. 气压层列表
   • PRESSURE_VALUES：包含不同气压层数值的列表，如 [50, 100, 150, ..., 1000]，用于与 OUTPUT_SETTINGS 配合确定输出数据的气压层信息。
5. 噪声添加标志
   • IS_ADD_NOISE：布尔类型变量，决定是否在预测结果中添加随机乘性噪声，为模型结果增加一定随机性。

三、核心函数解析

1. get_output_channel_indices 函数

• 功能：根据 OUTPUT_SETTINGS 和 PRESSURE_VALUES 计算每个输出变量和气压层组合对应的全局索引，返回一个包含所有输出通道索引的列表。
• 实现步骤：遍历 OUTPUT_SETTINGS 中的每个元组，先在 PRESSURE_VALUES 中查找气压层对应的索引，再结合 VAR_INDEX_MAPPING 中变量的起始索引，计算出全局索引并添加到结果列表中。

2. create_advanced_features 函数

• 功能：基于原始输入张量创建高级特征，包括温度垂直梯度、风速、特定层级的温度和湿度等，并将这些特征与原始数据进行拼接。
• 实现步骤
  1. 计算输入张量在第 0 维上的均值，得到各变量的平均数据。
  2. 分别提取温度、U 风、V 风各层数据。
  3. 计算温度垂直梯度：通过对温度各层数据进行前向差分得到。
  4. 计算风速：选取 U 风和 V 风中间层数据，通过公式 $\sqrt{u^2+v^2}$ 计算。
  5. 提取特定层级的温度和湿度：如近地面温度和高空气压层的湿度数据。
  6. 将原始数据、计算得到的梯度、风速以及特定层级变量通过 torch.cat 进行拼接，得到最终的特征工程数据。

3. prediction_with_advanced_features 函数

• 功能：结合特征工程数据进行预测，生成符合输出要求的预测结果张量。
• 实现步骤
  1. 获取输出通道索引列表。
  2. 调用 create_advanced_features 函数生成特征工程数据。
  3. 使用 LAT_RANGE 和 LON_RANGE 对特征工程数据进行目标区域裁剪。
  4. 根据输出通道索引从裁剪后的数据中选取对应通道数据。
  5. 根据 IS_ADD_NOISE 标志判断是否添加随机乘性噪声。
  6. 对选取的数据进行维度扩展和复制，生成最终形状为 (1, 12, 30, H_cropped, W_cropped) 的预测结果张量。

4. process_input_files 函数

• 功能：批量处理输入文件夹中的 .pt 文件，对每个文件进行预测并将结果保存到输出文件夹。
• 实现步骤：首先创建输出文件夹（如果不存在），然后遍历输入文件夹中的所有 .pt 文件，加载文件数据，调用 prediction_with_advanced_features 函数进行预测，最后将预测结果保存到输出文件夹，并在处理过程中输出处理状态信息。

四、主程序逻辑

在 if __name__ == "__main__" 代码块中，首先设置随机种子（random.seed(42)、np.random.seed(42)、torch.manual_seed(42)）以确保结果的可复现性，然后调用 process_input_files 函数开始处理数据文件。处理完成后，检查输出文件夹中是否存在 001.pt 文件，如果存在则加载文件并验证输出张量的形状是否符合预期，否则提示未找到对应文件。

五、应用与拓展思考

1. 当前应用：本代码适用于基于气象数据的初步预测任务，通过简单的特征工程提升数据的信息量，为后续预测模型提供更丰富的输入特征。
2. 拓展方向：
   • 可以进一步优化特征工程方法，例如采用更复杂的梯度计算方式、添加更多类型的气象特征等。
   • 结合深度学习模型（如卷积神经网络、循环神经网络等）替代当前简单的预测逻辑，提高预测的准确性和泛化能力。
   • 增加数据预处理步骤，如数据归一化、异常值处理等，提升数据质量。

以下是代码：

import torch
import os
from pathlib import Path
import random
import numpy as np

# 定义纬度切片范围
LAT_RANGE = slice(35, 81)
# 定义经度切片范围
LON_RANGE = slice(70, 141)
# 变量组字典，存储各变量的起始索引
VAR_INDEX_MAPPING = {
    'z': 0, 't': 13, 'u': 26, 'v': 39, 'q': 52,
    'ciwc': 65, 'clwc': 78, 'crwc': 91, 'cswc': 104
}
# 输出配置，指定要输出的变量和气压层
OUTPUT_SETTINGS = [
    ('t', 200), ('t', 500), ('t', 700), ('t', 850), ('t', 1000),
    ('q', 200), ('q', 500), ('q', 700), ('q', 850), ('q', 1000),
    ('ciwc', 200), ('ciwc', 500), ('ciwc', 700), ('ciwc', 850), ('ciwc', 1000),
    ('clwc', 200), ('clwc', 500), ('clwc', 700), ('clwc', 850), ('clwc', 1000),
    ('crwc', 200), ('crwc', 500), ('crwc', 700), ('crwc', 850), ('crwc', 1000),
    ('cswc', 200), ('cswc', 500), ('cswc', 700), ('cswc', 850), ('cswc', 1000)
]
# 气压层列表
PRESSURE_VALUES = [50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000]
# 是否添加噪声的标志
IS_ADD_NOISE = True

def get_output_channel_indices():
    """
    获取输出通道的索引。
    根据输出配置和气压层列表，计算每个输出变量和气压层组合对应的全局索引。
    """
    indices = []
    for var, pressure in OUTPUT_SETTINGS:
        try:
            # 查找气压层在气压层列表中的索引
            pressure_index = PRESSURE_VALUES.index(pressure)
        except ValueError:
            raise ValueError(f"气压层 {pressure} 未在定义的 PRESSURE_VALUES 中。")
        # 计算全局索引
        global_index = VAR_INDEX_MAPPING[var] + pressure_index
        indices.append(global_index)
    return indices

def create_advanced_features(input_tensor):
    """
    基于原始输入张量创建一些高级特征。
    包含温度垂直梯度、风速、特定层级的温度和湿度等特征。
    """
    # 计算输入张量在第 0 维上的均值，得到 (C, H, W) 形状的张量
    mean_tensor = input_tensor.squeeze(0).mean(dim=0)
    # 提取温度各层数据
    temperature_layers = mean_tensor[VAR_INDEX_MAPPING['t']:VAR_INDEX_MAPPING['u']]
    # 提取 U 风各层数据
    u_wind_layers = mean_tensor[VAR_INDEX_MAPPING['u']:VAR_INDEX_MAPPING['v']]
    # 提取 V 风各层数据
    v_wind_layers = mean_tensor[VAR_INDEX_MAPPING['v']:VAR_INDEX_MAPPING['q']]

    # 1. 计算温度垂直梯度 (dT/dp) - 简单的前向差分
    temperature_gradient = temperature_layers[1:] - temperature_layers[:-1]

    # 2. 计算风速 (sqrt(u^2 + v^2)) - 仅计算中间层的风速
    mid_layer_wind_speed = torch.sqrt(u_wind_layers[6:7]**2 + v_wind_layers[6:7]**2)

    # 3. 提取特定层级的温度和湿度 (例如，近地面和高空)
    surface_temperature = temperature_layers[12:13]
    high_altitude_humidity = mean_tensor[VAR_INDEX_MAPPING['q']:VAR_INDEX_MAPPING['ciwc']][0:1]

    # 将原始数据、梯度、风速和特定层级变量拼接在一起
    combined_features = torch.cat([mean_tensor,
                                   torch.nn.functional.interpolate(temperature_gradient.unsqueeze(0), size=(mean_tensor.shape[1], mean_tensor.shape[2]), mode='bilinear', align_corners=False).squeeze(0),
                                   torch.nn.functional.interpolate(mid_layer_wind_speed.unsqueeze(0), size=(mean_tensor.shape[1], mean_tensor.shape[2]), mode='bilinear', align_corners=False).squeeze(0),
                                   torch.nn.functional.interpolate(surface_temperature.unsqueeze(0), size=(mean_tensor.shape[1], mean_tensor.shape[2]), mode='bilinear', align_corners=False).squeeze(0),
                                   torch.nn.functional.interpolate(high_altitude_humidity.unsqueeze(0), size=(mean_tensor.shape[1], mean_tensor.shape[2]), mode='bilinear', align_corners=False).squeeze(0)], dim=0)

    return combined_features

def prediction_with_advanced_features(input_tensor):
    """
    对输入数据进行预测，包含高级特征工程。
    """
    # 获取输出通道的索引
    output_channel_indices = get_output_channel_indices()

    # 创建高级特征数据
    advanced_features = create_advanced_features(input_tensor)

    # 进行目标区域裁剪
    cropped_tensor = advanced_features[:, LAT_RANGE, LON_RANGE]

    # 选择原始输出通道对应的特征进行后续处理
    output_core = cropped_tensor[output_channel_indices, :, :]

    # 添加随机乘性噪声（选择性激活）
    if IS_ADD_NOISE and random.randint(0, 2) == 1:
        noise = 1 + 0.0222 * torch.rand_like(output_core)
        noise = torch.round(noise * 10000) / 10000
        output_core = output_core * noise

    # 扩展并复制时间维度：得到形状 (1, 12, 30, H_cropped, W_cropped)
    output = output_core.unsqueeze(0).unsqueeze(0).expand(1, 12, -1, -1, -1)

    return output.half()

def process_input_files(input_folder="input", output_folder="output"):
    """
    处理输入文件夹中的 .pt 文件，并将预测结果保存到输出文件夹。
    """
    output_folder = Path(output_folder)
    output_folder.mkdir(parents=True, exist_ok=True)

    input_folder = Path(input_folder)
    for file in input_folder.glob("*.pt"):
        output_file = output_folder / file.name
        try:
            input_data = torch.load(file)
            prediction = prediction_with_advanced_features(input_data)
            torch.save(prediction, output_file)
            print(f"成功处理: {file.name}")
        except Exception as e:
            print(f"处理失败 {file.name}: {str(e)}")

if __name__ == "__main__":
    random.seed(42)
    np.random.seed(42)
    torch.manual_seed(42)
    process_input_files()
    sample_file = Path("output") / "001.pt"
    if sample_file.exists():
        sample = torch.load(sample_file)
        print(f"输出形状验证: {sample.shape} (应为 torch.Size([1, 12, 30, {sample.shape[3]}, {sample.shape[4]}]))")
    else:
        print("未找到 '001.pt' 文件，请检查 output 目录。")