214 lines
8.1 KiB
Python
214 lines
8.1 KiB
Python
from io import BytesIO
|
||
import os
|
||
import pandas as pd
|
||
import numpy as np
|
||
import matplotlib.pyplot as plt
|
||
from scipy.optimize import curve_fit
|
||
import seaborn as sns
|
||
|
||
# 解决绘图中中文不能显示的问题
|
||
import matplotlib
|
||
matplotlib.rcParams['font.sans-serif'] = ['SimHei']
|
||
matplotlib.rcParams['axes.unicode_minus'] = False
|
||
|
||
# ------------------------------ 模型定义 ------------------------------
|
||
|
||
|
||
def tidal_model(t, v0, a1, b1, a2, b2, a3, b3, a4, b4):
|
||
"""潮汐波模型"""
|
||
T1, T2, T3, T4 = 6, 8, 12, 24 # 周期以小时为单位
|
||
return (v0 + a1 * np.sin((2 * np.pi / T1) * t + b1)
|
||
+ a2 * np.sin((2 * np.pi / T2) * t + b2)
|
||
+ a3 * np.sin((2 * np.pi / T3) * t + b3)
|
||
+ a4 * np.sin((2 * np.pi / T4) * t + b4))
|
||
|
||
|
||
def planetary_model(t, v0, a, b, period):
|
||
"""行星波模型"""
|
||
return v0 + a * np.sin((2 * np.pi / period) * t + b)
|
||
|
||
# ------------------------------ 数据处理工具函数 ------------------------------
|
||
|
||
|
||
def preprocess_data(file_dir, year):
|
||
"""读取数据文件并合并成完整的 DataFrame"""
|
||
file_list = [f for f in os.listdir(file_dir) if f.endswith('.txt')]
|
||
final_df = pd.DataFrame()
|
||
|
||
for file_name in file_list:
|
||
file_path = os.path.join(file_dir, file_name)
|
||
df = pd.read_csv(file_path, sep='\s+')
|
||
|
||
# 替换异常值为 NaN
|
||
df.loc[df['uwind'] == 1000000, 'uwind'] = np.nan
|
||
df.loc[df['vwind'] == 1000000, 'vwind'] = np.nan
|
||
|
||
# 筛选所需列并重命名
|
||
date_part = os.path.splitext(file_name)[0][-8:]
|
||
df = df[['height', 'time', 'uwind', 'vwind']].rename(columns={
|
||
'uwind': f'{date_part}_uwind',
|
||
'vwind': f'{date_part}_vwind'
|
||
})
|
||
|
||
if final_df.empty:
|
||
final_df = df
|
||
else:
|
||
final_df = pd.merge(
|
||
final_df, df, on=['height', 'time'], how='outer')
|
||
|
||
# 补全缺失列并排序
|
||
all_dates = pd.date_range(
|
||
f'{year}-01-01', f'{year}-12-31').strftime('%Y%m%d')
|
||
expected_columns = [f'{date}_uwind' for date in all_dates] + \
|
||
[f'{date}_vwind' for date in all_dates]
|
||
for col in expected_columns:
|
||
if col not in final_df.columns:
|
||
final_df[col] = np.nan
|
||
final_df = final_df[['height', 'time'] + sorted(expected_columns)]
|
||
|
||
return final_df
|
||
|
||
# ------------------------------ 拟合工具函数 ------------------------------
|
||
|
||
|
||
def fit_wind_speed(data, model_func, initial_guess, bounds, window, num_params):
|
||
"""对风速数据进行滑动窗口拟合"""
|
||
# replace "_uwind" and "_vwind" in date column
|
||
data["date"] = data["date"].str.replace(
|
||
"_uwind", "").str.replace("_vwind", "")
|
||
data["date"] = pd.to_datetime(data["date"])
|
||
data.set_index("date", inplace=True)
|
||
|
||
unique_dates = np.unique(data.index.date)
|
||
params = []
|
||
|
||
for date in unique_dates:
|
||
start_date = pd.to_datetime(date) - pd.Timedelta(days=(window - 1) / 2)
|
||
end_date = pd.to_datetime(date) + pd.Timedelta(days=(window - 1) / 2)
|
||
window_data = data[start_date:end_date]
|
||
|
||
t = np.arange(len(window_data)) + 1
|
||
y = window_data['wind_speed'].values
|
||
|
||
valid_mask = ~np.isnan(y)
|
||
t = t[valid_mask]
|
||
y = y[valid_mask]
|
||
|
||
if len(y) < num_params * 6:
|
||
params.append([None] * len(initial_guess))
|
||
continue
|
||
|
||
try:
|
||
popt, _ = curve_fit(
|
||
model_func, t, y, p0=initial_guess, bounds=bounds)
|
||
params.append(popt)
|
||
except RuntimeError:
|
||
params.append([None] * len(initial_guess))
|
||
|
||
params_df = pd.DataFrame(params, columns=[
|
||
"v0"] + [f"param_{i}" for i in range(1, len(initial_guess))], index=unique_dates)
|
||
return params_df
|
||
|
||
# ------------------------------ 热力图绘制函数 ------------------------------
|
||
|
||
|
||
def plot_heatmaps(result_dict, param, title, vmax):
|
||
"""绘制热力图"""
|
||
fig, axes = plt.subplots(1, 2, figsize=(18, 8))
|
||
fig.suptitle(title, fontsize=20)
|
||
|
||
for i, (wind_type, data) in enumerate(result_dict.items()):
|
||
heatmap_data = data[param].unstack(level=0)
|
||
heatmap_data.index = pd.to_datetime(heatmap_data.index)
|
||
heatmap_data = heatmap_data.iloc[:, ::-1]
|
||
|
||
sns.heatmap(
|
||
heatmap_data.T,
|
||
cmap="viridis",
|
||
vmax=vmax,
|
||
cbar_kws={'label': '振幅 (m/s)'},
|
||
yticklabels=[
|
||
f'{height / 1000:.1f}' for height in heatmap_data.columns],
|
||
ax=axes[i]
|
||
)
|
||
axes[i].set_title(f'{wind_type}振幅变化', fontsize=14)
|
||
axes[i].set_xlabel('日期', fontsize=12)
|
||
axes[i].set_ylabel('高度 (km)', fontsize=12)
|
||
|
||
plt.tight_layout(rect=[0, 0, 1, 0.95])
|
||
# plt.show()
|
||
|
||
|
||
# ------------------------------ 主逻辑 ------------------------------
|
||
def final_plot_v2(year, station, model_name):
|
||
"""
|
||
主逻辑函数，用于处理数据并绘制选定模型的热力图。
|
||
|
||
参数:
|
||
- year: int，年份
|
||
- station: str，站点名称
|
||
- model_name: str，模型名称（如 '潮汐波', '2日行星波', '5日行星波', '10日行星波', '16日行星波'）
|
||
"""
|
||
# 配置文件路径
|
||
file_dir = rf'./radar/data\{station}\{year}'
|
||
output_dir = rf'./radar/tmp\{station}\{year}\时空'
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
|
||
# 数据预处理
|
||
final_df = preprocess_data(file_dir, year)
|
||
uwind_data = final_df[["height", "time"] +
|
||
[col for col in final_df.columns if "_uwind" in col]]
|
||
vwind_data = final_df[["height", "time"] +
|
||
[col for col in final_df.columns if "_vwind" in col]]
|
||
heights = uwind_data["height"].unique()
|
||
|
||
# 模型配置
|
||
models = {
|
||
"潮汐波": {"func": tidal_model, "params": 9, "initial": [0, 1, 0, 1, 0, 1, 0, 1, 0], "bounds": ([-np.inf, 0, -np.inf, 0, -np.inf, 0, -np.inf, 0, -np.inf], [np.inf] * 9), "window": 5},
|
||
"2日行星波": {"func": lambda t, v0, a, b: planetary_model(t, v0, a, b, 48), "params": 3, "initial": [0, 1, 0], "bounds": ([-np.inf, 0, -np.inf], [np.inf, np.inf, np.inf]), "window": 5},
|
||
"5日行星波": {"func": lambda t, v0, a, b: planetary_model(t, v0, a, b, 120), "params": 3, "initial": [0, 1, 0], "bounds": ([-np.inf, 0, -np.inf], [np.inf, np.inf, np.inf]), "window": 15},
|
||
"10日行星波": {"func": lambda t, v0, a, b: planetary_model(t, v0, a, b, 240), "params": 3, "initial": [0, 1, 0], "bounds": ([-np.inf, 0, -np.inf], [np.inf, np.inf, np.inf]), "window": 29},
|
||
"16日行星波": {"func": lambda t, v0, a, b: planetary_model(t, v0, a, b, 384), "params": 3, "initial": [0, 1, 0], "bounds": ([-np.inf, 0, -np.inf], [np.inf, np.inf, np.inf]), "window": 49}
|
||
}
|
||
|
||
# 检查模型名称是否有效
|
||
if model_name not in models:
|
||
raise ValueError(
|
||
f"无效的模型名称: '{model_name}'。可选模型为: {list(models.keys())}")
|
||
|
||
# 获取模型配置
|
||
config = models[model_name]
|
||
|
||
# 拟合纬向风和经向风
|
||
u_result_dict, v_result_dict = {}, {}
|
||
for height in heights:
|
||
# 纬向风数据
|
||
height_data_u = uwind_data[uwind_data["height"] == height].melt(
|
||
id_vars=["time", "height"], var_name="date", value_name="wind_speed")
|
||
u_result_dict[height] = fit_wind_speed(
|
||
height_data_u, config["func"], config["initial"], config["bounds"], config["window"], config["params"])
|
||
|
||
# 经向风数据
|
||
height_data_v = vwind_data[vwind_data["height"] == height].melt(
|
||
id_vars=["time", "height"], var_name="date", value_name="wind_speed")
|
||
v_result_dict[height] = fit_wind_speed(
|
||
height_data_v, config["func"], config["initial"], config["bounds"], config["window"], config["params"])
|
||
|
||
# 汇总结果
|
||
u_final_result = pd.concat(u_result_dict, names=["height", "date"])
|
||
v_final_result = pd.concat(v_result_dict, names=["height", "date"])
|
||
|
||
# 绘制热力图
|
||
plot_heatmaps(
|
||
{"纬向风": u_final_result, "经向风": v_final_result},
|
||
param="param_1", # 振幅参数
|
||
title=f'{year}年{model_name}振幅时空变化',
|
||
vmax=100,
|
||
)
|
||
|
||
buffer = BytesIO()
|
||
plt.savefig(buffer, format='png')
|
||
buffer.seek(0)
|
||
plt.close()
|
||
return buffer
|