allow --targetMate option for advancepvs.py script (vondele#47)

robertnurnberg · web-flow · commit bd660e04d4a9 · 2024-03-03T12:12:36.000+01:00
diff --git a/README.md b/README.md
@@ -29,3 +29,7 @@ puzzles from the information stored in `matetrackpv.epd`. For example, the file
 ```shell
 python advancepvs.py --plies 1 --mateType won && sed 's/; PV.*/;/' matedtrackpv.epd > matedtrack.epd
 ```
+Similarly, a file with only `#-10` positions can be created with the command
+```shell
+python advancepvs.py --targetMate -10 && grep 'bm #-10;' matedtrackpv.epd > mate-10.epd
+```
diff --git a/advancepvs.py b/advancepvs.py
@@ -22,10 +22,15 @@
         default=1,
         help="number of plies to advance",
     )
+    parser.add_argument(
+        "--targetMate",
+        type=int,
+        help="in each position advance enough plies to leave a mate-in-TARGETMATE (overrides --plies)",
+    )
     parser.add_argument(
         "--mateType",
         choices=["all", "won", "lost"],
-        default="won",
+        default="all",
         help="type of positions to advance from",
     )
     args = parser.parse_args()
@@ -45,13 +50,19 @@
 
     print(f"{len(fens)} FENs loaded...")
 
-    count = 0
+    count, plies = 0, args.plies
     with open(args.outFile, "w") as f:
         for fen, bm, pv, line in fens:
             plies_to_checkmate = 2 * bm - 1 if bm > 0 else -2 * bm
+            if args.targetMate:
+                m = args.targetMate
+                plies4m = 2 * m - 1 if m > 0 else -2 * m
+                plies = plies_to_checkmate - plies4m
+                if plies < 0:
+                    plies = plies_to_checkmate + 1
             if (
-                args.plies <= len(pv)
-                and args.plies < plies_to_checkmate
+                plies <= len(pv)
+                and plies < plies_to_checkmate
                 and (
                     args.mateType == "all"
                     or args.mateType == "won"
@@ -61,14 +72,17 @@
                 )
             ):
                 board = chess.Board(fen)
-                for move in pv[: args.plies]:
+                for move in pv[:plies]:
                     board.push(chess.Move.from_uci(move))
                     bm = -bm + (1 if bm > 0 else 0)
                 fen = board.epd()
-                pv = pv[args.plies :]
+                pv = pv[plies:]
                 f.write(f"{fen} bm #{bm}; PV: {' '.join(pv)};\n")
                 count += 1
             else:
                 f.write(line)
 
-    print(f"Positions in which we advanced {args.plies} plies: ", count)
+    if args.targetMate:
+        print(f"Number of #{args.targetMate} positions created: ", count)
+    else:
+        print(f"Positions in which we advanced {plies} plies: ", count)
