From b7fce2e5aae961ea3182f706eb6a24fd355babd2 Mon Sep 17 00:00:00 2001
From: Dustella <fdnoaivj@outlook.com>
Date: Fri, 24 Jan 2025 10:58:32 +0800
Subject: [PATCH] fix: balloon year last 3 modes

---
 balloon/plot_year_backend.py | 75 ++++++++++++++++++++++++++++++++----
 1 file changed, 67 insertions(+), 8 deletions(-)

diff --git a/balloon/plot_year_backend.py b/balloon/plot_year_backend.py
index 39a058a..dd75a83 100644
--- a/balloon/plot_year_backend.py
+++ b/balloon/plot_year_backend.py
@@ -164,15 +164,61 @@ def plot_MFv(filtered_df1, ax):
     ax.set_ylabel("Occurrence(%)")
 
 
-def plot_horizontal_propagation(filtered_df, ax, season):
-    season_data = filtered_df[filtered_df["season"] == season]
-    windrose = WindroseAxes.from_ax(ax)
-    ax.set_title(season, fontsize=18)
-    windrose.bar(season_data["b"], np.ones_like(
-        season_data["b"]), normed=False)
+def plot_horizontal_propagation(filtered_df:pd.DataFrame,  season):
+    def get_season(date):
+        month = date.month
+        if month in [12, 1, 2]:
+            return "Winter"
+        elif month in [3, 4, 5]:
+            return "Spring"
+        elif month in [6, 7, 8]:
+            return "Summer"
+        else:
+            return "Fall"
+    fig = plt.figure(figsize=(36, 20))
+    gs = gridspec.GridSpec(2, 2)
+    ax10 = []
+    for i in range(4):
+        ax10.append(fig.add_subplot(gs[i//2, i % 2 ], projection="windrose"))
+    
+    data = filtered_df.copy()
+    
+    data.loc[:, "date"] = data["file_name"].str[-28:-16]
+    filtered_df = data[["date"] + [col for col in data.columns if col != "file_name" and col != "date"]]
+    filtered_df.reset_index(drop=True, inplace=True)
+
+    filtered_df = filtered_df.drop_duplicates(subset="date", keep="last")  # 使用 drop_duplicates 函数，保留每组重复中的最后一个记录
+
+    
+    filtered_df["date1"] = filtered_df["date"].str.split("T").str[0]  # 再加一列，只保留日期部分
+    filtered_df["date1"] = pd.to_datetime(filtered_df["date1"], format="%Y%m%d")  # 确保 'date1' 列是日期格式
+    filtered_df["season"] = filtered_df["date1"].apply(get_season)  # 添加季节列
+    seasons = ["Winter", "Spring", "Summer", "Fall"]
+    for ax, season in zip(ax10, seasons):
+        season_data = filtered_df[filtered_df["season"] == season]
+        windrose = WindroseAxes.from_ax(ax)
+        ax.set_title(season, fontsize=18)
+        windrose.bar(season_data["b"], np.ones_like(season_data["b"]), normed=False)  # normed=True表示占每个季节的占比
+    # season_data = filtered_df[filtered_df["season"] == season]
+    # windrose = WindroseAxes.from_ax(ax)
+    # ax.set_title(season, fontsize=18)
+    # windrose.bar(season_data["b"], np.ones_like(
+    #     season_data["b"]), normed=False)
 
 
 def plot_vertical_propagation(filtered_df, ax):
+    data = filtered_df.copy()
+    
+    data.loc[:, "date"] = data["file_name"].str[-28:-16]
+    filtered_df = data[["date"] + [col for col in data.columns if col != "file_name" and col != "date"]]
+    filtered_df.reset_index(drop=True, inplace=True)
+
+    filtered_df = filtered_df.drop_duplicates(subset="date", keep="last")  # 使用 drop_duplicates 函数，保留每组重复中的最后一个记录
+
+    
+    filtered_df["date1"] = filtered_df["date"].str.split("T").str[0]  # 再加一列，只保留日期部分
+    filtered_df["date1"] = pd.to_datetime(filtered_df["date1"], format="%Y%m%d")  # 确保 'date1' 列是日期格式
+ 
     filtered_df.set_index("date1", inplace=True)
     monthly_stats_df = (
         filtered_df.groupby([filtered_df.index.month, filtered_df.index.year])
@@ -195,6 +241,19 @@ def plot_vertical_propagation(filtered_df, ax):
 
 
 def plot_energy_distribution(filtered_df, ax):
+    data = filtered_df.copy()
+    
+    data.loc[:, "date"] = data["file_name"].str[-28:-16]
+    filtered_df = data[["date"] + [col for col in data.columns if col != "file_name" and col != "date"]]
+    filtered_df.reset_index(drop=True, inplace=True)
+
+    filtered_df = filtered_df.drop_duplicates(subset="date", keep="last")  # 使用 drop_duplicates 函数，保留每组重复中的最后一个记录
+
+    
+    filtered_df["date1"] = filtered_df["date"].str.split("T").str[0]  # 再加一列，只保留日期部分
+    filtered_df["date1"] = pd.to_datetime(filtered_df["date1"], format="%Y%m%d")  # 确保 'date1' 列是日期格式
+ 
+    
     filtered_df.reset_index(inplace=True)
     filtered_df["year_month"] = filtered_df["date1"].dt.to_period("M")
     monthly_avg = filtered_df.groupby("year_month")[["Ek", "E_p"]].mean()
@@ -251,7 +310,7 @@ def render_based_on_mode(df, mode, season=None):
     elif mode == "经向动量通量统计结果":
         plot_MFv(df, ax)
     elif mode == "horizontal propagation":
-        plot_horizontal_propagation(df, ax, season)
+        plot_horizontal_propagation(df, season="Fall")
     elif mode == "每月上传/下传重力波占比":
         plot_vertical_propagation(df, ax)
     elif mode == "动能和势能分布情况":
@@ -259,7 +318,7 @@ def render_based_on_mode(df, mode, season=None):
     else:
         raise ValueError("Invalid mode")
     plt.rcParams["font.sans-serif"] = ["SimHei"]  # 显示中文
-    fig.savefig(buf)
+    plt.savefig(buf)
     buf.seek(0)
     plt.close()
     return buf