新增管道健康风险预测 api

timescaledb/composite_queries.py

@@ -1,10 +1,14 @@
+import time
 from typing import List, Optional, Any, Dict, Tuple
-from datetime import datetime
+from datetime import datetime, timedelta
 from psycopg import AsyncConnection
 import pandas as pd
+import numpy as np
 from api_ex.Fdataclean import clean_flow_data_df_kf
 from api_ex.Pdataclean import clean_pressure_data_df_km
+from api_ex.pipeline_health_analyzer import PipelineHealthAnalyzer
 
+from postgresql.internal_queries import InternalQueries
 from postgresql.scada_info import ScadaRepository as PostgreScadaRepository
 from timescaledb.schemas.realtime import RealtimeRepository
 from timescaledb.schemas.scheme import SchemeRepository
@@ -450,3 +454,158 @@ class CompositeQueries:
             return "success"
         except Exception as e:
             return f"error: {str(e)}"
+
+    @staticmethod
+    async def predict_pipeline_health(
+        timescale_conn: AsyncConnection,
+        db_name: str,
+        query_time: datetime,
+    ) -> List[Dict[str, Any]]:
+        """
+        预测管道健康状况
+
+        根据管网名称和当前时间，查询管道信息和实时数据，
+        使用随机生存森林模型预测管道的生存概率
+
+        Args:
+            timescale_conn: TimescaleDB 异步连接
+            db_name: 管网数据库名称
+            query_time: 查询时间
+            property_conditions: 可选的管道筛选条件，如 {"diameter": 300}
+
+        Returns:
+            预测结果列表，每个元素包含 link_id 和对应的生存函数
+
+        Raises:
+            ValueError: 当参数无效或数据不足时
+            FileNotFoundError: 当模型文件未找到时
+        """
+        try:
+            perf_start_time = time.time()
+            # 1. 准备时间范围（查询时间前后1秒）
+            start_time = query_time - timedelta(seconds=1)
+            end_time = query_time + timedelta(seconds=1)
+
+            # 2. 先查询流速数据（velocity），获取有数据的管道ID列表
+            velocity_data = await RealtimeRepository.get_links_field_by_time_range(
+                timescale_conn, start_time, end_time, "velocity"
+            )
+
+            if not velocity_data:
+                raise ValueError("未找到流速数据")
+
+            # 3. 只查询有流速数据的管道的基本信息
+            valid_link_ids = list(velocity_data.keys())
+
+            # 批量查询这些管道的详细信息
+            fields = ["id", "diameter", "node1", "node2"]
+            all_links = InternalQueries.get_links_by_property(
+                fields=fields,
+                db_name=db_name,
+            )
+
+            # 转换为字典以快速查找
+            links_dict = {link["id"]: link for link in all_links}
+
+            # 获取所有需要查询的节点ID
+            node_ids = set()
+            for link_id in valid_link_ids:
+                if link_id in links_dict:
+                    link = links_dict[link_id]
+                    node_ids.add(link["node1"])
+                    node_ids.add(link["node2"])
+
+            # 4. 批量查询压力数据（pressure）
+            pressure_data = await RealtimeRepository.get_nodes_field_by_time_range(
+                timescale_conn, start_time, end_time, "pressure"
+            )
+
+            # 5. 组合数据结构
+            materials = []
+            diameters = []
+            velocities = []
+            pressures = []
+            link_ids = []
+
+            for link_id in valid_link_ids:
+                # 跳过不在管道字典中的ID（如泵等其他元素）
+                if link_id not in links_dict:
+                    continue
+
+                link = links_dict[link_id]
+                diameter = link["diameter"]
+                node1 = link["node1"]
+                node2 = link["node2"]
+
+                # 获取流速数据
+                velocity_values = velocity_data[link_id]
+                velocity = velocity_values[-1]["value"] if velocity_values else 0
+
+                # 获取node1和node2的压力数据，计算平均值
+                node1_pressure = 0
+                node2_pressure = 0
+
+                if node1 in pressure_data and pressure_data[node1]:
+                    pressure_values = pressure_data[node1]
+                    node1_pressure = (
+                        pressure_values[-1]["value"] if pressure_values else 0
+                    )
+
+                if node2 in pressure_data and pressure_data[node2]:
+                    pressure_values = pressure_data[node2]
+                    node2_pressure = (
+                        pressure_values[-1]["value"] if pressure_values else 0
+                    )
+
+                # 计算平均压力
+                avg_pressure = (node1_pressure + node2_pressure) / 2
+
+                # 添加到列表
+                link_ids.append(link_id)
+                materials.append(7)  # 默认材料类型为7，可根据实际情况调整
+                diameters.append(diameter)
+                velocities.append(velocity)
+                pressures.append(avg_pressure)
+
+            if not link_ids:
+                raise ValueError("没有找到有效的管道数据用于预测")
+
+            # 6. 创建DataFrame
+            data = pd.DataFrame(
+                {
+                    "Material": materials,
+                    "Diameter": diameters,
+                    "Flow Velocity": velocities,
+                    "Pressure": pressures,
+                }
+            )
+
+            # 7. 使用PipelineHealthAnalyzer进行预测
+            analyzer = PipelineHealthAnalyzer(
+                model_path="api_ex/model/my_survival_forest_model_quxi.joblib"
+            )
+            survival_functions = analyzer.predict_survival(data)
+            print("预测管道健康耗时: {:.2f} 秒".format(time.time() - perf_start_time))
+            # 8. 组合结果
+            results = []
+            for i, link_id in enumerate(link_ids):
+                sf = survival_functions[i]
+                results.append(
+                    {
+                        "link_id": link_id,
+                        "diameter": diameters[i],
+                        "velocity": velocities[i],
+                        "pressure": pressures[i],
+                        "survival_function": {
+                            "x": sf.x.tolist(),  # 时间点（年）
+                            "y": sf.y.tolist(),  # 生存概率
+                            "a": float(sf.a),
+                            "b": float(sf.b),
+                        },
+                    }
+                )
+
+            return results
+
+        except Exception as e:
+            raise ValueError(f"管道健康预测失败: {str(e)}")

timescaledb/internal_queries.py

@@ -9,8 +9,6 @@ from timescaledb.schemas.scada import ScadaRepository
 import psycopg
 import time
 
-# 内部使用存储类
-
 
 class InternalStorage:
     @staticmethod
@@ -78,6 +76,8 @@ class InternalStorage:
                 else:
                     raise  # 达到最大重试次数后抛出异常
 
+
+class InternalQueries:
     @staticmethod
     def query_scada_by_ids_time(
         device_ids: List[str],

timescaledb/router.py

@@ -594,3 +594,39 @@ async def clean_scada_data(
         )
     except ValueError as e:
         raise HTTPException(status_code=400, detail=str(e))
+
+
+@router.get("/composite/pipeline-health-prediction")
+async def predict_pipeline_health(
+    query_time: datetime = Query(..., description="查询时间"),
+    db_name: str = Query(..., description="管网数据库名称"),
+    timescale_conn: AsyncConnection = Depends(get_database_connection),
+):
+    """
+    预测管道健康状况
+
+    根据管网名称和当前时间，查询管道信息和实时数据，
+    使用随机生存森林模型预测管道的生存概率
+
+    Args:
+        query_time: 查询时间
+        db_name: 管网数据库名称
+
+    Returns:
+        预测结果列表，每个元素包含 link_id 和对应的生存函数
+    """
+    try:
+        result = await CompositeQueries.predict_pipeline_health(
+            timescale_conn, db_name, query_time
+        )
+
+        return {
+            "success": True,
+            "result": result,
+        }
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail=str(e))
+    except FileNotFoundError as e:
+        raise HTTPException(status_code=404, detail=str(e))
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"内部服务器错误: {str(e)}")