feat: Expose advanced timeout options to tket2-py

tket2-py/src/passes.rs

@@ -102,18 +102,27 @@ fn lower_to_pytket<'py>(circ: &Bound<'py, PyAny>) -> PyResult<Bound<'py, PyAny>>
 /// optimising. This can be deactivated by setting `rebase` to `false`, in which
 /// case the circuit is expected to be in the Nam gate set.
 ///
-/// Will use at most `max_threads` threads (plus a constant) and take at most
-/// `timeout` seconds (plus a constant). Default to the number of cpus and
-/// 15min respectively.
+/// Will use at most `max_threads` threads (plus a constant). Defaults to the
+/// number of CPUs available.
+///
+/// The optimisation will terminate at the first of the following timeout
+/// criteria, if set:
+/// - `timeout` seconds (default: 15min) have elapsed since the start of the
+///    optimisation
+/// - `progress_timeout` (default: None) seconds have elapsed since progress
+///    in the cost function was last made
+/// - `max_circuit_cnt` (default: None) circuits have been explored.
 ///
 /// Log files will be written to the directory `log_dir` if specified.
 #[pyfunction]
+#[allow(clippy::too_many_arguments)]
 fn badger_optimise<'py>(
     circ: &Bound<'py, PyAny>,
     optimiser: &PyBadgerOptimiser,
     max_threads: Option<NonZeroUsize>,
     timeout: Option<u64>,
     progress_timeout: Option<u64>,
+    max_circuit_cnt: Option<usize>,
     log_dir: Option<PathBuf>,
     rebase: Option<bool>,
 ) -> PyResult<Bound<'py, PyAny>> {
@@ -165,6 +174,7 @@ fn badger_optimise<'py>(
                 progress_timeout,
                 n_threads: n_threads.try_into().unwrap(),
                 split_circuit: true,
+                max_circuit_cnt,
                 ..Default::default()
             };
             circ = optimiser.optimise(circ, log_file, options);

tket2-py/tket2/_tket2/passes.pyi

@@ -35,6 +35,7 @@ def badger_optimise(
     max_threads: int | None = None,
     timeout: int | None = None,
     progress_timeout: int | None = None,
+    max_circuit_cnt: int | None = None,
     log_dir: Path | None = None,
     rebase: bool | None = False,
 ) -> CircuitClass:
@@ -47,9 +48,16 @@ def badger_optimise(
     optimising. This can be deactivated by setting `rebase` to `false`, in which
     case the circuit is expected to be in the Nam gate set.
 
-    Will use at most `max_threads` threads (plus a constant) and take at most
-    `timeout` seconds (plus a constant). Default to the number of cpus and
-    15min respectively.
+    Will use at most `max_threads` threads (plus a constant). Defaults to the
+    number of CPUs available.
+
+    The optimisation will terminate at the first of the following timeout
+    criteria, if set:
+    - `timeout` seconds (default: 15min) have elapsed since the start of the
+      optimisation
+    - `progress_timeout` (default: None) seconds have elapsed since progress
+      in the cost function was last made
+    - `max_circuit_cnt` (default: None) circuits have been explored.
 
     Log files will be written to the directory `log_dir` if specified.
     """

tket2-py/tket2/passes.py

@@ -35,6 +35,7 @@ def badger_pass(
     max_threads: Optional[int] = None,
     timeout: Optional[int] = None,
     progress_timeout: Optional[int] = None,
+    max_circuit_cnt: Optional[int] = None,
     log_dir: Optional[Path] = None,
     rebase: bool = False,
 ) -> BasePass:
@@ -44,8 +45,9 @@ def badger_pass(
     `compile-rewriter <https://github.com/CQCL/tket2/tree/main/badger-optimiser>`_
     utility. If `rewriter` is not specified, a default one will be used.
 
-    The arguments `max_threads`, `timeout`, `log_dir` and `rebase` are optional
-    and will be passed on to the Badger optimiser if provided."""
+    The arguments `max_threads`, `timeout`, `progress_timeout`, `max_circuit_cnt`,
+    `log_dir` and `rebase` are optional and will be passed on to the Badger
+    optimiser if provided."""
     if rewriter is None:
         with resources.as_file(
             resources.files("tket2").joinpath("data/nam_6_3.rwr")
@@ -61,6 +63,7 @@ def apply(circuit: Circuit) -> Circuit:
             max_threads=max_threads,
             timeout=timeout,
             progress_timeout=progress_timeout,
+            max_circuit_cnt=max_circuit_cnt,
             log_dir=log_dir,
             rebase=rebase,
         )