feat: update prediction workbook output

app/prediction.py

@@ -58,6 +58,39 @@ FEATURES = [
 
 ID_COLUMN = "管道编号"
 PIPE_AGE_COLUMN = "管龄"
+MATERIAL_COLUMN = "管材"
+MATERIAL_CODE_OPTIONS = [
+    (1, "镀锌"),
+    (2, "钢塑"),
+    (3, "铝塑"),
+    (4, "PPR"),
+    (5, "PE"),
+    (6, "UPVC"),
+    (7, "铸铁"),
+    (8, "预应力"),
+    (9, "自应力"),
+    (10, "玻璃钢夹砂"),
+    (11, "钢管"),
+    (12, "钢套混凝土管"),
+    (13, "球墨铸铁"),
+    (14, "其他"),
+]
+MATERIAL_ALIAS_TO_CODE = {
+    str(code): code
+    for code, _ in MATERIAL_CODE_OPTIONS
+}
+MATERIAL_ALIAS_TO_CODE.update(
+    {
+        name.casefold(): code
+        for code, name in MATERIAL_CODE_OPTIONS
+    }
+)
+MATERIAL_ALIAS_TO_CODE.update(
+    {
+        f"{code}-{name}".casefold(): code
+        for code, name in MATERIAL_CODE_OPTIONS
+    }
+)
 SUPPORTED_EXTENSIONS = {".csv", ".xls", ".xlsx"}
 CHINESE_FONT_PROP = None
 
@@ -150,13 +183,21 @@ def secure_upload_name(original_filename: str, run_id: str) -> tuple[str, str]:
 def read_input_file(path: Path, suffix: str) -> pd.DataFrame:
     try:
         if suffix == ".csv":
-            return pd.read_csv(path)
-        return pd.read_excel(path)
+            return pd.read_csv(path, dtype={ID_COLUMN: "string"})
+        return pd.read_excel(path, dtype={ID_COLUMN: "string"})
     except Exception as exc:
         logging.exception("文件解析失败: %s", exc)
         raise PredictionError("文件解析失败，请检查编码或表格格式。")
 
 
+def normalize_id_value(value: Any, fallback: str) -> str:
+    if pd.isna(value):
+        return fallback
+    if isinstance(value, str):
+        return value.strip()
+    return str(value)
+
+
 def validate_input_frame(df: pd.DataFrame) -> None:
     if df.empty:
         raise PredictionError("上传文件没有可预测的数据。")
@@ -166,6 +207,28 @@ def validate_input_frame(df: pd.DataFrame) -> None:
         raise PredictionError(f"缺少必要字段: {', '.join(missing)}")
 
 
+def normalize_material_code(value: Any) -> int:
+    if pd.isna(value):
+        raise PredictionError("管材不能为空。")
+    if isinstance(value, str):
+        key = value.strip().casefold()
+    else:
+        numeric_value = float(value)
+        if not numeric_value.is_integer():
+            raise PredictionError(f"管材编码无效: {value}")
+        key = str(int(numeric_value))
+    code = MATERIAL_ALIAS_TO_CODE.get(key)
+    if code is None:
+        raise PredictionError(f"管材编码无效: {value}")
+    return code
+
+
+def prepare_model_features(df: pd.DataFrame) -> pd.DataFrame:
+    x_test = df[FEATURES].copy()
+    x_test[MATERIAL_COLUMN] = x_test[MATERIAL_COLUMN].map(normalize_material_code)
+    return x_test
+
+
 def grade_info(probability: float) -> tuple[str, str, str]:
     if probability <= 0.2:
         return ("I级", "管道安全风险十分严重，需立刻进行抢修或更新改造", "bg-dangerSoft text-dangerText")
@@ -318,7 +381,7 @@ def run_prediction(uploaded: FileStorage, user_id: int, model) -> PredictionArti
     try:
         df = read_input_file(original_path, suffix)
         validate_input_frame(df)
-        x_test = df[FEATURES].copy()
+        x_test = prepare_model_features(df)
         try:
             curves = model.predict_survival_function(x_test)
         except Exception as exc:
@@ -365,7 +428,7 @@ def render_survival_chart(df: pd.DataFrame, curves, image_path: Path) -> tuple[l
     for i, curve in enumerate(curves):
         times = [float(x) for x in list(curve.x)]
         probs = [float(y) for y in list(curve.y)]
-        pipe_id = str(df.iloc[i][ID_COLUMN]) if pd.notna(df.iloc[i][ID_COLUMN]) else f"Pipe_{i+1:03d}"
+        pipe_id = normalize_id_value(df.iloc[i][ID_COLUMN], f"Pipe_{i+1:03d}")
         pipe_age = f"{df.iloc[i][PIPE_AGE_COLUMN]} 年" if PIPE_AGE_COLUMN in df.columns and pd.notna(df.iloc[i][PIPE_AGE_COLUMN]) else "-"
         health_probability = interpolate_probability(times, probs, 10)
         remaining_life = estimate_remaining_life(times, probs)
@@ -386,15 +449,13 @@ def render_survival_chart(df: pd.DataFrame, curves, image_path: Path) -> tuple[l
             {
                 "管道编号": pipe_id,
                 "管龄": pipe_age,
-                "健康概率": health_probability,
-                "预计剩余寿命": remaining_life,
                 "健康等级": grade_label,
             }
         )
         plt.step(times, probs, where="post", linewidth=2, label=pipe_id)
 
     font_kwargs = chinese_font_kwargs()
-    plt.xlabel("预测时间轴（年）", **font_kwargs)
+    plt.xlabel("管龄（年）", **font_kwargs)
     plt.ylabel("生存概率", **font_kwargs)
     plt.title("预测分析图", **font_kwargs)
     plt.grid(alpha=0.18)
@@ -416,22 +477,58 @@ def write_prediction_workbook(
     summary_rows: list[dict[str, Any]],
     summary_sheet_rows: list[dict[str, Any]],
 ) -> None:
-    sample_data_rows: list[dict[str, Any]] = []
+    chart_times: set[float] = set()
+    chart_series: list[tuple[str, dict[float, float]]] = []
     for i, curve in enumerate(curves):
         times = [float(x) for x in list(curve.x)]
         probs = [float(y) for y in list(curve.y)]
         pipe_id = summary_rows[i]["pipe_id"]
-        for time, probability in zip(times, probs):
-            sample_data_rows.append(
-                {
-                    "管道编号": pipe_id,
-                    "样本序号": i + 1,
-                    "时间（年）": time,
-                    "生存概率": probability,
-                    "风险概率": 1 - probability,
-                }
-            )
+        chart_times.update(times)
+        chart_series.append((pipe_id, dict(zip(times, probs))))
+
+    sorted_times = sorted(chart_times)
+    sample_columns = ["管龄（年）", *[pipe_id for pipe_id, _ in chart_series]]
+    sample_data_rows = [
+        [time, *[series.get(time) for _, series in chart_series]]
+        for time in sorted_times
+    ]
 
     with pd.ExcelWriter(excel_path, engine="openpyxl") as writer:
-        pd.DataFrame(summary_sheet_rows).to_excel(writer, sheet_name="结果摘要", index=False)
-        pd.DataFrame(sample_data_rows).to_excel(writer, sheet_name="样本数据", index=False)
+        summary_df = pd.DataFrame(summary_sheet_rows, columns=[ID_COLUMN, PIPE_AGE_COLUMN, "健康等级"])
+        sample_df = pd.DataFrame(sample_data_rows, columns=sample_columns)
+        if ID_COLUMN in summary_df.columns:
+            summary_df[ID_COLUMN] = summary_df[ID_COLUMN].astype("string")
+
+        summary_df.to_excel(writer, sheet_name="结果摘要", index=False)
+        sample_df.to_excel(writer, sheet_name="样本数据", index=False)
+
+        sample_worksheet = writer.book["样本数据"]
+        sample_worksheet.freeze_panes = "A2"
+        sample_worksheet.auto_filter.ref = sample_worksheet.dimensions
+        sample_worksheet.column_dimensions["A"].width = 12
+        for column_cells in sample_worksheet.iter_cols(min_col=2, max_col=sample_worksheet.max_column):
+            header_cell = column_cells[0]
+            sample_worksheet.column_dimensions[header_cell.column_letter].width = max(12, len(str(header_cell.value)) + 2)
+            for cell in column_cells[1:]:
+                cell.number_format = "0.0%"
+
+        for sheet_name in ("结果摘要", "样本数据"):
+            worksheet = writer.book[sheet_name]
+            header_cells = next(worksheet.iter_rows(min_row=1, max_row=1), [])
+            id_column_index = None
+            for cell in header_cells:
+                if cell.value == ID_COLUMN:
+                    id_column_index = cell.column
+                    break
+            if id_column_index is None:
+                continue
+            for cell in worksheet.iter_cols(
+                min_col=id_column_index,
+                max_col=id_column_index,
+                min_row=2,
+                max_row=worksheet.max_row,
+            ):
+                for item in cell:
+                    if item.value is not None:
+                        item.value = str(item.value)
+                    item.number_format = "@"