Support backfilling cycle indexes for multiple measurement types

backend/LabelingPreTraining/detect_cycle_index.py

@@ -2,8 +2,9 @@
 
 The original full re-run encoded the first complete cycle of every file but
 never wrote back where that cycle starts and ends. This script re-detects the
-first complete cycle for every pressure file (rpm > 0) and records the real
-``sample_index`` bounds into two columns:
+first complete cycle for every file of a given measurement type (rpm > 0,
+selected with ``--measurement-type``) and records the real ``sample_index``
+bounds into two columns:
 
 * ``cycle_start`` / ``cycle_end`` >= 0 : first complete cycle, half-open
   ``[cycle_start, cycle_end)`` in actual sample_index values;
@@ -24,10 +25,14 @@ written (0/-1) is never re-selected on resume. Writes are flushed in one
 
 Usage:
     conda activate tspulse
-    python detect_cycle_index.py [--start-time ""] [--start-id 0]
-                                 [--batch 100] [--read-workers 4] [--limit 0]
-                                 [--report-every 100] [--write-batch 100]
-                                 [--dry-run]
+    python detect_cycle_index.py --measurement-type 位移
+    python detect_cycle_index.py --measurement-type 加速度
+    python detect_cycle_index.py --measurement-type 位移 --measurement-type 加速度
+    python detect_cycle_index.py --measurement-type 位移,加速度
+    python detect_cycle_index.py [--measurement-type 压力] [--start-time ""]
+                                 [--start-id 0] [--batch 100] [--read-workers 4]
+                                 [--limit 0] [--report-every 100]
+                                 [--write-batch 100] [--dry-run]
 """
 
 import argparse
@@ -42,7 +47,7 @@ sys.path.insert(0, str(BACKEND))
 
 from app.algorithms.cycles import detect_cycles  # noqa: E402
 import predict  # noqa: E402
-from predict import MEASUREMENT_TYPE, STAGE_LIMITS  # noqa: E402
+from predict import STAGE_LIMITS  # noqa: E402
 
 
 def _fetch_chunk(connection, file_ids: list[int], limit: int) -> dict[int, np.ndarray]:
@@ -185,6 +190,7 @@ def flush_bounds(connection, buffer: list[tuple[int, int, int]]) -> None:
 
 def fetch_batch(
     connection,
+    measurement_type: str,
     watermark_time: object | None,
     watermark_id: int,
     batch: int,
@@ -199,7 +205,7 @@ def fetch_batch(
                 ORDER BY sample_time ASC, id ASC
                 LIMIT %s
                 """,
-                (MEASUREMENT_TYPE, batch),
+                (measurement_type, batch),
             )
         else:
             cursor.execute(
@@ -211,13 +217,15 @@ def fetch_batch(
                 ORDER BY sample_time ASC, id ASC
                 LIMIT %s
                 """,
-                (MEASUREMENT_TYPE, watermark_time, watermark_time, watermark_id, batch),
+                (measurement_type, watermark_time, watermark_time, watermark_id, batch),
             )
         return cursor.fetchall()
 
 
 def main() -> int:
     parser = argparse.ArgumentParser(description="回写首个完整周期的 sample_index 边界")
+    parser.add_argument("--measurement-type", action="append", default=[],
+                        help="要回写的测量类型，可多次指定或用逗号分隔（默认压力）")
     parser.add_argument("--start-time", type=str, default="",
                         help="起始水位时间(含)，留空从头开始")
     parser.add_argument("--start-id", type=int, default=0,
@@ -231,7 +239,18 @@ def main() -> int:
     parser.add_argument("--dry-run", action="store_true", help="只检测不回写数据库")
     args = parser.parse_args()
 
+    measurement_types: list[str] = []
+    for raw in (args.measurement_type or ["压力"]):
+        for part in str(raw).split(","):
+            part = part.strip()
+            if part and part not in measurement_types:
+                measurement_types.append(part)
+
     start_time = args.start_time.strip()
+    read_workers = max(1, min(args.read_workers, args.batch))
+    reader_connections = [predict.get_connection() for _ in range(read_workers)]
+
+    current_type = measurement_types[0]
     watermark_time: object | None = start_time or None
     watermark_id = args.start_id
     total = 0
@@ -240,75 +259,85 @@ def main() -> int:
     last_reported = 0
     write_buffer: list[tuple[int, int, int]] = []
 
-    read_workers = max(1, min(args.read_workers, args.batch))
-    reader_connections = [predict.get_connection() for _ in range(read_workers)]
-
-    def flush(connection) -> None:
-        if args.dry_run:
-            write_buffer.clear()
-        else:
-            flush_bounds(connection, write_buffer)
-
-    def summary_text() -> str:
-        return (
-            f"已处理 {total}, 有周期 {stats['有周期']}, 无周期 {stats['无周期']}, "
-            f"错误 {errors}, 水位 sample_time={watermark_time} id={watermark_id}"
-        )
-
-    print(f"水位起点: sample_time={start_time or '从头'}, id={args.start_id}, "
-          f"批大小: {args.batch}, 读并行: {read_workers} 连接, 窗口: {STAGE_LIMITS}"
+    print(f"测量类型: {', '.join(measurement_types)}, 水位起点: "
+          f"sample_time={start_time or '从头'}, id={args.start_id}, 批大小: {args.batch}, "
+          f"读并行: {read_workers} 连接, 窗口: {STAGE_LIMITS}"
           + (", 干跑(不回写)" if args.dry_run else ""))
 
     try:
-        while True:
-            connection = predict.get_connection()
-            try:
-                rows = fetch_batch(connection, watermark_time, watermark_id, args.batch)
-                if not rows:
-                    flush(connection)
-                    connection.close()
-                    print(f"无待处理文件，完成 [{summary_text()}]")
-                    return 0
+        for type_index, measurement_type in enumerate(measurement_types, start=1):
+            current_type = measurement_type
+            watermark_time = start_time or None
+            watermark_id = args.start_id
+            total = 0
+            stats = {"有周期": 0, "无周期": 0}
+            errors = 0
+            last_reported = 0
+            write_buffer = []
+
+            def flush(connection) -> None:
+                if args.dry_run:
+                    write_buffer.clear()
+                else:
+                    flush_bounds(connection, write_buffer)
 
-                resolved, batch_errors = resolve_batch_bounds(
-                    rows, reader_connections, read_workers,
+            def summary_text() -> str:
+                return (
+                    f"{measurement_type}: 已处理 {total}, 有周期 {stats['有周期']}, "
+                    f"无周期 {stats['无周期']}, 错误 {errors}, "
+                    f"水位 sample_time={watermark_time} id={watermark_id}"
                 )
-                errors += batch_errors
-
-                for row in rows:
-                    file_id = int(row["id"])
-                    total += 1
-                    bounds = resolved[file_id]
-                    if bounds is None:
-                        stats["无周期"] += 1
-                        write_buffer.append((file_id, -1, -1))
-                    else:
-                        stats["有周期"] += 1
-                        write_buffer.append((file_id, bounds[0], bounds[1]))
-                    if len(write_buffer) >= args.write_batch:
-                        flush(connection)
 
-                if total // args.report_every > last_reported:
-                    last_reported = total // args.report_every
-                    print(f"  汇总: {summary_text()}", flush=True)
+            print(f"\n开始处理类型 {type_index}/{len(measurement_types)}: {measurement_type}")
+            while True:
+                connection = predict.get_connection()
+                try:
+                    rows = fetch_batch(connection, measurement_type, watermark_time, watermark_id, args.batch)
+                    if not rows:
+                        flush(connection)
+                        connection.close()
+                        print(f"无待处理文件，完成 [{summary_text()}]")
+                        break
+
+                    resolved, batch_errors = resolve_batch_bounds(
+                        rows, reader_connections, read_workers,
+                    )
+                    errors += batch_errors
+
+                    for row in rows:
+                        file_id = int(row["id"])
+                        total += 1
+                        bounds = resolved[file_id]
+                        if bounds is None:
+                            stats["无周期"] += 1
+                            write_buffer.append((file_id, -1, -1))
+                        else:
+                            stats["有周期"] += 1
+                            write_buffer.append((file_id, bounds[0], bounds[1]))
+                        if len(write_buffer) >= args.write_batch:
+                            flush(connection)
+
+                    if total // args.report_every > last_reported:
+                        last_reported = total // args.report_every
+                        print(f"  汇总: {summary_text()}", flush=True)
+
+                    if args.limit and total >= args.limit:
+                        flush(connection)
+                        watermark_time = rows[-1]["sample_time"]
+                        watermark_id = int(rows[-1]["id"])
+                        connection.close()
+                        print(f"已达 --limit={args.limit}，退出")
+                        print(f"  汇总: {summary_text()}")
+                        break
 
-                if args.limit and total >= args.limit:
                     flush(connection)
                     watermark_time = rows[-1]["sample_time"]
                     watermark_id = int(rows[-1]["id"])
-                    connection.close()
-                    print(f"已达 --limit={args.limit}，退出")
-                    print(f"  汇总: {summary_text()}")
-                    return 0
-
-                flush(connection)
-                watermark_time = rows[-1]["sample_time"]
-                watermark_id = int(rows[-1]["id"])
-            finally:
-                try:
-                    connection.close()
-                except Exception:
-                    pass
+                finally:
+                    try:
+                        connection.close()
+                    except Exception:
+                        pass
 
     except KeyboardInterrupt:
         if write_buffer and not args.dry_run:
@@ -324,8 +353,8 @@ def main() -> int:
                 pass
         print(f"\n已中断。{summary_text()}")
         if watermark_time is not None:
-            print(f"续跑命令: python detect_cycle_index.py --start-time \"{watermark_time}\" "
-                  f"--start-id {watermark_id}")
+            print(f"续跑命令: python detect_cycle_index.py --measurement-type {current_type} "
+                  f"--start-time \"{watermark_time}\" --start-id {watermark_id}")
         return 0
 
     for connection in reader_connections: