OrgTJWater/pipeline-lifetime
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update prediction results workflow

Browse Source
This commit is contained in:
jiang
2026-07-20 13:47:39 +08:00
parent 8d347cfbc2
commit a20339d35e
10 changed files with 239 additions and 207 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_prediction.py
+72 -7
View File
@@ -7,8 +7,10 @@ from unittest.mock import patch
import pandas as pd
from openpyxl import load_workbook
from openpyxl import Workbook
from app.prediction import (
CHART_DISPLAY_LIMIT,
FEATURES,
ID_COLUMN,
INPUT_COLUMNS,
@@ -70,7 +72,8 @@ class PredictionHelpersTest(unittest.TestCase):
worksheet = workbook["Template"]
headers = [cell.value for cell in next(worksheet.iter_rows(min_row=1, max_row=1))]
self.assertEqual(headers, INPUT_COLUMNS)
self.assertEqual(len(headers), len(INPUT_COLUMNS))
self.assertEqual(set(headers), set(INPUT_COLUMNS))
def test_read_input_file_preserves_text_pipe_ids(self) -> None:
with TemporaryDirectory() as temp_dir:
@@ -81,6 +84,36 @@ class PredictionHelpersTest(unittest.TestCase):
self.assertEqual(df[ID_COLUMN].tolist(), ["00123"])
def test_read_input_file_calculates_defects_from_template_detail_sheet(self) -> None:
with TemporaryDirectory() as temp_dir:
output_path = Path(temp_dir) / "input.xlsx"
workbook = Workbook()
template = workbook.active
template.title = "Template"
template.append(INPUT_COLUMNS)
template.append(["001", 5, 0, 5, 1, 100, 1.2, 0.4, 20, 800, 1, "=缺陷计算!E2", "=缺陷计算!J2"])
detail = workbook.create_sheet("缺陷计算")
detail.append([
ID_COLUMN,
"泄漏\n权重0.5",
"腐蚀\n权重0.4",
"管瘤\n权重0.1",
"结构性缺陷值",
"气囊\n权重0.5",
"杂质\n权重0.2",
"异物\n权重0.2",
"不明连接\n权重0.1",
"功能性缺陷值",
])
detail.append(["001", "严重", "中度", "轻度", None, "轻度", "无", "中度", "严重", None])
workbook.save(output_path)
df = read_input_file(output_path, ".xlsx")
self.assertAlmostEqual(float(df.loc[0, "结构缺陷"]), 3.8)
self.assertAlmostEqual(float(df.loc[0, "功能缺陷"]), 1.6)
def test_prepare_model_features_maps_material_aliases_to_codes(self) -> None:
rows = []
for value in ["镀锌", "2-钢塑", 13]:
@@ -110,8 +143,8 @@ class PredictionHelpersTest(unittest.TestCase):
]
summary_rows = [{"pipe_id": "P001"}, {"pipe_id": "P002"}]
summary_sheet_rows = [
{ID_COLUMN: "P001", PIPE_AGE_COLUMN: "10 年", "健康概率": 0.7, "预计剩余寿命": 2.0, "健康等级": "IV级"},
{ID_COLUMN: "P002", PIPE_AGE_COLUMN: "12 年", "健康概率": 0.6, "预计剩余寿命": 1.5, "健康等级": "III级"},
{ID_COLUMN: "P001", PIPE_AGE_COLUMN: "10 年", "健康风险值": 0.3},
{ID_COLUMN: "P002", PIPE_AGE_COLUMN: "12 年", "健康风险值": 0.4},
]
with TemporaryDirectory() as temp_dir:
@@ -122,7 +155,7 @@ class PredictionHelpersTest(unittest.TestCase):
self.assertEqual(workbook.sheet_names, ["结果摘要", "样本数据"])
summary_data = pd.read_excel(output_path, sheet_name="结果摘要")
self.assertEqual(summary_data.columns.tolist(), [ID_COLUMN, PIPE_AGE_COLUMN, "健康等级"])
self.assertEqual(summary_data.columns.tolist(), [ID_COLUMN, PIPE_AGE_COLUMN, "健康风险值"])
sample_data = pd.read_excel(output_path, sheet_name="样本数据")
self.assertEqual(sample_data.columns.tolist(), ["管龄（年）", "P001", "P002"])
@@ -140,7 +173,7 @@ class PredictionHelpersTest(unittest.TestCase):
def test_prediction_workbook_writes_pipe_ids_as_excel_text(self) -> None:
curves = [DummyCurve([1], [0.9])]
summary_rows = [{"pipe_id": "00123"}]
summary_sheet_rows = [{ID_COLUMN: "00123", PIPE_AGE_COLUMN: "10 年", "健康等级": "V级"}]
summary_sheet_rows = [{ID_COLUMN: "00123", PIPE_AGE_COLUMN: "10 年", "健康风险值": 0.1}]
with TemporaryDirectory() as temp_dir:
output_path = Path(temp_dir) / "prediction.xlsx"
@@ -161,13 +194,45 @@ class PredictionHelpersTest(unittest.TestCase):
with TemporaryDirectory() as temp_dir:
output_path = Path(temp_dir) / "chart.png"
with patch("app.prediction.plt.xlabel") as xlabel, patch("app.prediction.plt.ylabel") as ylabel:
render_survival_chart(df, [DummyCurve([1, 2], [0.9, 0.7])], output_path)
with (
patch("app.prediction.plt.xlabel") as xlabel,
patch("app.prediction.plt.ylabel") as ylabel,
patch("app.prediction.plt.figtext") as figtext,
patch("app.prediction.plt.title") as title,
):
summary_rows, _ = render_survival_chart(df, [DummyCurve([1, 10, 12], [0.9, 0.7, 0.4])], output_path)
self.assertAlmostEqual(summary_rows[0]["health_probability"], 0.4)
self.assertAlmostEqual(summary_rows[0]["health_risk"], 0.6)
self.assertEqual(summary_rows[0]["grade_label"], "II级")
xlabel.assert_called_once()
self.assertEqual(xlabel.call_args.args[0], "管龄（年）")
ylabel.assert_called_once()
self.assertEqual(ylabel.call_args.args[0], "健康风险")
title.assert_called_once()
self.assertEqual(title.call_args.args[0], "管道的剩余寿命分析图")
figtext.assert_called_once()
self.assertIn(f"前{CHART_DISPLAY_LIMIT}条管道的示例数据", figtext.call_args.args[2])
def test_survival_chart_only_plots_first_ten_pipes(self) -> None:
sample_count = CHART_DISPLAY_LIMIT + 2
df = pd.DataFrame(
{
ID_COLUMN: [f"P{i:03d}" for i in range(sample_count)],
PIPE_AGE_COLUMN: [12] * sample_count,
}
)
curves = [DummyCurve([1, 10, 12], [0.9, 0.7, 0.4]) for _ in range(sample_count)]
with TemporaryDirectory() as temp_dir:
output_path = Path(temp_dir) / "chart.png"
with patch("app.prediction.plt.step") as step, patch("app.prediction.plt.legend"):
summary_rows, summary_sheet_rows = render_survival_chart(df, curves, output_path)
self.assertEqual(step.call_count, CHART_DISPLAY_LIMIT)
self.assertEqual(len(summary_rows), sample_count)
self.assertEqual(len(summary_sheet_rows), sample_count)
if __name__ == "__main__":