treecf 0.2.0__tar.gz → 0.2.2__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (64) hide show
  1. {treecf-0.2.0 → treecf-0.2.2}/PKG-INFO +12 -3
  2. {treecf-0.2.0 → treecf-0.2.2}/README.md +10 -1
  3. {treecf-0.2.0 → treecf-0.2.2}/pyproject.toml +2 -2
  4. {treecf-0.2.0 → treecf-0.2.2}/rust/Cargo.lock +1 -1
  5. {treecf-0.2.0 → treecf-0.2.2}/rust/Cargo.toml +1 -1
  6. {treecf-0.2.0 → treecf-0.2.2}/rust/src/exact/mod.rs +7 -0
  7. {treecf-0.2.0 → treecf-0.2.2}/rust/src/exact/search.rs +121 -9
  8. {treecf-0.2.0 → treecf-0.2.2}/rust/src/exact/test_support.rs +32 -0
  9. {treecf-0.2.0 → treecf-0.2.2}/rust/src/ga.rs +139 -21
  10. treecf-0.2.2/rust/src/interrupt.rs +27 -0
  11. {treecf-0.2.0 → treecf-0.2.2}/rust/src/lib.rs +1 -0
  12. {treecf-0.2.0 → treecf-0.2.2}/rust/src/py.rs +84 -5
  13. {treecf-0.2.0 → treecf-0.2.2}/rust/src/regions.rs +107 -11
  14. {treecf-0.2.0 → treecf-0.2.2}/src/treecf/__init__.py +4 -1
  15. {treecf-0.2.0 → treecf-0.2.2}/src/treecf/api.py +505 -32
  16. treecf-0.2.2/src/treecf/audit.py +525 -0
  17. {treecf-0.2.0 → treecf-0.2.2}/src/treecf/backends/exact.py +5 -0
  18. {treecf-0.2.0 → treecf-0.2.2}/src/treecf/backends/exact_rust.py +8 -1
  19. {treecf-0.2.0 → treecf-0.2.2}/src/treecf/batch.py +370 -17
  20. treecf-0.2.2/src/treecf/constraints/objects.py +142 -0
  21. {treecf-0.2.0 → treecf-0.2.2}/src/treecf/constraints/parser.py +26 -4
  22. {treecf-0.2.0 → treecf-0.2.2}/src/treecf/mining.py +117 -0
  23. treecf-0.2.2/src/treecf/plausibility.py +118 -0
  24. {treecf-0.2.0 → treecf-0.2.2}/src/treecf/regions.py +31 -3
  25. {treecf-0.2.0 → treecf-0.2.2}/src/treecf/targets.py +154 -5
  26. {treecf-0.2.0 → treecf-0.2.2}/src/treecf/viz.py +115 -4
  27. {treecf-0.2.0 → treecf-0.2.2}/src/treecf/viz_batch.py +82 -1
  28. treecf-0.2.0/src/treecf/constraints/objects.py +0 -86
  29. treecf-0.2.0/src/treecf/plausibility.py +0 -61
  30. {treecf-0.2.0 → treecf-0.2.2}/LICENSE +0 -0
  31. {treecf-0.2.0 → treecf-0.2.2}/rust/src/cells.rs +0 -0
  32. {treecf-0.2.0 → treecf-0.2.2}/rust/src/constraints.rs +0 -0
  33. {treecf-0.2.0 → treecf-0.2.2}/rust/src/exact/domains.rs +0 -0
  34. {treecf-0.2.0 → treecf-0.2.2}/rust/src/exact/orderpairs.rs +0 -0
  35. {treecf-0.2.0 → treecf-0.2.2}/rust/src/exact/propagation.rs +0 -0
  36. {treecf-0.2.0 → treecf-0.2.2}/rust/src/ir.rs +0 -0
  37. {treecf-0.2.0 → treecf-0.2.2}/src/treecf/_errors.py +0 -0
  38. {treecf-0.2.0 → treecf-0.2.2}/src/treecf/_json.py +0 -0
  39. {treecf-0.2.0 → treecf-0.2.2}/src/treecf/aim/__init__.py +0 -0
  40. {treecf-0.2.0 → treecf-0.2.2}/src/treecf/aim/cells.py +0 -0
  41. {treecf-0.2.0 → treecf-0.2.2}/src/treecf/backends/__init__.py +0 -0
  42. {treecf-0.2.0 → treecf-0.2.2}/src/treecf/backends/_exact_bounds.py +0 -0
  43. {treecf-0.2.0 → treecf-0.2.2}/src/treecf/backends/_exact_domains.py +0 -0
  44. {treecf-0.2.0 → treecf-0.2.2}/src/treecf/backends/_exact_orderpairs.py +0 -0
  45. {treecf-0.2.0 → treecf-0.2.2}/src/treecf/backends/_exact_propagation.py +0 -0
  46. {treecf-0.2.0 → treecf-0.2.2}/src/treecf/backends/genetic.py +0 -0
  47. {treecf-0.2.0 → treecf-0.2.2}/src/treecf/backends/genetic_rust.py +0 -0
  48. {treecf-0.2.0 → treecf-0.2.2}/src/treecf/backends/regions_rust.py +0 -0
  49. {treecf-0.2.0 → treecf-0.2.2}/src/treecf/constraints/__init__.py +0 -0
  50. {treecf-0.2.0 → treecf-0.2.2}/src/treecf/constraints/compile.py +0 -0
  51. {treecf-0.2.0 → treecf-0.2.2}/src/treecf/constraints/flatten.py +0 -0
  52. {treecf-0.2.0 → treecf-0.2.2}/src/treecf/ir/__init__.py +0 -0
  53. {treecf-0.2.0 → treecf-0.2.2}/src/treecf/ir/conformance.py +0 -0
  54. {treecf-0.2.0 → treecf-0.2.2}/src/treecf/ir/evaluate.py +0 -0
  55. {treecf-0.2.0 → treecf-0.2.2}/src/treecf/ir/flatten.py +0 -0
  56. {treecf-0.2.0 → treecf-0.2.2}/src/treecf/ir/model.py +0 -0
  57. {treecf-0.2.0 → treecf-0.2.2}/src/treecf/ir/parsers/__init__.py +0 -0
  58. {treecf-0.2.0 → treecf-0.2.2}/src/treecf/ir/parsers/catboost.py +0 -0
  59. {treecf-0.2.0 → treecf-0.2.2}/src/treecf/ir/parsers/json_dump.py +0 -0
  60. {treecf-0.2.0 → treecf-0.2.2}/src/treecf/ir/parsers/lightgbm.py +0 -0
  61. {treecf-0.2.0 → treecf-0.2.2}/src/treecf/ir/parsers/sklearn.py +0 -0
  62. {treecf-0.2.0 → treecf-0.2.2}/src/treecf/ir/parsers/xgboost.py +0 -0
  63. {treecf-0.2.0 → treecf-0.2.2}/src/treecf/objective.py +0 -0
  64. {treecf-0.2.0 → treecf-0.2.2}/src/treecf/py.typed +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: treecf
3
- Version: 0.2.0
3
+ Version: 0.2.2
4
4
  Classifier: Development Status :: 4 - Beta
5
5
  Classifier: Intended Audience :: Science/Research
6
6
  Classifier: License :: OSI Approved :: MIT License
@@ -48,7 +48,7 @@ Provides-Extra: sklearn
48
48
  Provides-Extra: viz
49
49
  Provides-Extra: xgboost
50
50
  License-File: LICENSE
51
- Summary: Constrained, threshold-aware counterfactual explanations for tree ensembles (XGBoost, LightGBM, CatBoost, sklearn) on a bundled Rust genetic engine.
51
+ Summary: Constrained, threshold-aware counterfactual explanations for tree ensembles (XGBoost, LightGBM, CatBoost, sklearn) — fast Rust genetic search, exact optimality proofs, certified infeasibility, and recourse regions.
52
52
  Keywords: counterfactual,xai,interpretability,recourse,gbdt,xgboost,lightgbm,catboost,credit-risk
53
53
  Author-email: Daniel Wlazlo <wlazlo.daniel@gmail.com>
54
54
  License: MIT
@@ -67,7 +67,7 @@ Project-URL: Issues, https://github.com/wlazlod/treecf/issues
67
67
  that the model's raw output lands in a target interval?"* — for XGBoost, LightGBM, CatBoost
68
68
  and scikit-learn tree ensembles.
69
69
 
70
- > Status: v0.1.0 on [PyPI](https://pypi.org/project/treecf/). See the [documentation](https://wlazlod.github.io/treecf/) for concepts and tutorials.
70
+ > On [PyPI](https://pypi.org/project/treecf/). See the [documentation](https://wlazlod.github.io/treecf/) for concepts and tutorials.
71
71
 
72
72
  ## Why another counterfactual package?
73
73
 
@@ -81,6 +81,12 @@ and scikit-learn tree ensembles.
81
81
  counterfactual in a median 0.24s versus 0.005s for the genetic heuristic, closing a
82
82
  median 14.33% cost gap the heuristic leaves on the table — measured on a 4-core dev
83
83
  machine (`scripts/bench_exact.py`).
84
+ - **Certified "no".** A completed exact search returns `Infeasible(proof="certified")` —
85
+ "no recourse exists within these constraints" becomes a provable statement, not a shrug
86
+ after a timeout.
87
+ - **Recourse regions.** Any verified counterfactual widens into a certified box — "reduce
88
+ utilization to ≤ 0.40", not "to 0.3972" — with every point in the box provably in-target
89
+ and constraint-feasible; works with every backend.
84
90
  - **Decision thresholds are first-class.** Targets are intervals on the raw model output —
85
91
  custom probability cutoffs, regression targets, and whole rating-grade ladders in one call.
86
92
  - **Real-world constraints.** Declarative layer for immutability, directionality, ranges,
@@ -113,6 +119,9 @@ exp = Explainer(
113
119
  ],
114
120
  )
115
121
  res = exp.explain(x, target=Target.probability(range=(0.0, 0.04)), seed=0)
122
+
123
+ proved = exp.explain(x, target=t, backend="exact") # proof="optimal", a certified "no", or a warned degrade
124
+ boxed = exp.explain(x, target=t, region=True) # res.region.describe() -> "utilization <= 0.4"
116
125
  ```
117
126
 
118
127
  ## License
@@ -6,7 +6,7 @@
6
6
  that the model's raw output lands in a target interval?"* — for XGBoost, LightGBM, CatBoost
7
7
  and scikit-learn tree ensembles.
8
8
 
9
- > Status: v0.1.0 on [PyPI](https://pypi.org/project/treecf/). See the [documentation](https://wlazlod.github.io/treecf/) for concepts and tutorials.
9
+ > On [PyPI](https://pypi.org/project/treecf/). See the [documentation](https://wlazlod.github.io/treecf/) for concepts and tutorials.
10
10
 
11
11
  ## Why another counterfactual package?
12
12
 
@@ -20,6 +20,12 @@ and scikit-learn tree ensembles.
20
20
  counterfactual in a median 0.24s versus 0.005s for the genetic heuristic, closing a
21
21
  median 14.33% cost gap the heuristic leaves on the table — measured on a 4-core dev
22
22
  machine (`scripts/bench_exact.py`).
23
+ - **Certified "no".** A completed exact search returns `Infeasible(proof="certified")` —
24
+ "no recourse exists within these constraints" becomes a provable statement, not a shrug
25
+ after a timeout.
26
+ - **Recourse regions.** Any verified counterfactual widens into a certified box — "reduce
27
+ utilization to ≤ 0.40", not "to 0.3972" — with every point in the box provably in-target
28
+ and constraint-feasible; works with every backend.
23
29
  - **Decision thresholds are first-class.** Targets are intervals on the raw model output —
24
30
  custom probability cutoffs, regression targets, and whole rating-grade ladders in one call.
25
31
  - **Real-world constraints.** Declarative layer for immutability, directionality, ranges,
@@ -52,6 +58,9 @@ exp = Explainer(
52
58
  ],
53
59
  )
54
60
  res = exp.explain(x, target=Target.probability(range=(0.0, 0.04)), seed=0)
61
+
62
+ proved = exp.explain(x, target=t, backend="exact") # proof="optimal", a certified "no", or a warned degrade
63
+ boxed = exp.explain(x, target=t, region=True) # res.region.describe() -> "utilization <= 0.4"
55
64
  ```
56
65
 
57
66
  ## License
@@ -1,7 +1,7 @@
1
1
  [project]
2
2
  name = "treecf"
3
- version = "0.2.0"
4
- description = "Constrained, threshold-aware counterfactual explanations for tree ensembles (XGBoost, LightGBM, CatBoost, sklearn) on a bundled Rust genetic engine."
3
+ version = "0.2.2"
4
+ description = "Constrained, threshold-aware counterfactual explanations for tree ensembles (XGBoost, LightGBM, CatBoost, sklearn) — fast Rust genetic search, exact optimality proofs, certified infeasibility, and recourse regions."
5
5
  readme = "README.md"
6
6
  license = { text = "MIT" }
7
7
  authors = [{ name = "Daniel Wlazlo", email = "wlazlo.daniel@gmail.com" }]
@@ -419,7 +419,7 @@ checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca"
419
419
 
420
420
  [[package]]
421
421
  name = "treecf-core"
422
- version = "0.2.0"
422
+ version = "0.2.2"
423
423
  dependencies = [
424
424
  "numpy",
425
425
  "pyo3",
@@ -1,6 +1,6 @@
1
1
  [package]
2
2
  name = "treecf-core"
3
- version = "0.2.0"
3
+ version = "0.2.2"
4
4
  edition = "2021"
5
5
  # f64::next_down (cells.rs) stabilized in 1.86; pyo3 0.29 needs 1.83
6
6
  rust-version = "1.86"
@@ -24,6 +24,12 @@
24
24
  //! all depend on the order nodes are visited in, so the search is
25
25
  //! single-threaded by construction and no rayon call may enter it. The
26
26
  //! RNG-free stages that `ga.rs` fans out have no analogue here.
27
+ //! The interrupt probe the search polls does not weaken this: it is asked
28
+ //! a question that depends on nothing in the search, it is asked on a
29
+ //! schedule fixed by the node counter alone, and the only thing it can do
30
+ //! is replace the whole answer with `Interrupted`. It never changes a
31
+ //! value that is returned — the same reasoning that already makes the
32
+ //! `time_budget_s` early break parity-neutral.
27
33
  //! 2. **Full re-summation.** Assigning a feature re-walks the trees that split
28
34
  //! on it from their roots, then the ensemble bracket is re-summed over every
29
35
  //! tree in ascending index (`base + tree_0 + tree_1 + ...`), never patched
@@ -70,6 +76,7 @@ pub(crate) mod search;
70
76
  #[cfg(test)]
71
77
  pub(crate) mod test_support;
72
78
 
79
+ pub use crate::interrupt::SearchOutcome;
73
80
  pub use domains::constraint_cells;
74
81
  pub use search::{solve_exact, ExactParams, ExactResult, ExactStats};
75
82
 
@@ -15,12 +15,18 @@ use crate::exact::domains::{
15
15
  use crate::exact::orderpairs::{achievable_bounds, boundary_candidates, intersect_cell};
16
16
  use crate::exact::propagation::{PropFrame, Propagation};
17
17
  use crate::exact::ValuePolicy;
18
+ use crate::interrupt::{InterruptProbe, SearchOutcome};
18
19
  use crate::ir::Ensemble;
19
20
 
20
21
  /// The tolerance `check_matrix` allows a linear constraint; an order pair counts
21
22
  /// as broken exactly when the arbiter would reject it.
22
23
  const LINEAR_SLACK: f64 = 1e-9;
23
24
 
25
+ /// How many expanded nodes between two interrupt polls. Asking is cheap but not
26
+ /// free, and a search that gets nowhere near this many nodes finishes fast
27
+ /// enough that nobody reaches for the keyboard.
28
+ const SIGNAL_CHECK_INTERVAL: u64 = 1 << 18;
29
+
24
30
  /// The exact set of counters `solve_exact` reports — the seven keys of Python's
25
31
  /// `_stats`. `nodes_pruned_score` counts branches the ensemble can no longer
26
32
  /// bring into the target (the plausibility bound counts here too); every other
@@ -661,7 +667,13 @@ fn stats(
661
667
  /// `completed` false only means the search never settled the whole space.
662
668
  ///
663
669
  /// `Err` mirrors Python's `ConstraintValidationError` for a multi-feature Linear
664
- /// outside the canonical order-pair shape.
670
+ /// outside the canonical order-pair shape — it stays validation-only, so a
671
+ /// caller can keep telling a bad constraint apart from a stopped search.
672
+ ///
673
+ /// `probe` is asked every `SIGNAL_CHECK_INTERVAL` expanded nodes whether to
674
+ /// stop. Answering yes throws the search away, incumbent included, and returns
675
+ /// `SearchOutcome::Interrupted`; answering no leaves every returned number
676
+ /// exactly as it would be without a probe at all.
665
677
  #[allow(clippy::too_many_arguments)]
666
678
  pub fn solve_exact(
667
679
  ens: &Ensemble,
@@ -675,7 +687,8 @@ pub fn solve_exact(
675
687
  plausibility: Option<(&Ensemble, f64)>,
676
688
  params: &ExactParams,
677
689
  incumbent: Option<(f64, &[f64])>,
678
- ) -> Result<ExactResult, String> {
690
+ probe: InterruptProbe<'_>,
691
+ ) -> Result<SearchOutcome<ExactResult>, String> {
679
692
  let start = Instant::now();
680
693
  let order_pairs = validate(cons)?;
681
694
  let (lo_t, hi_t) = interval;
@@ -685,13 +698,13 @@ pub fn solve_exact(
685
698
 
686
699
  // (a) The factual itself: nothing is ever cheaper than not moving at all.
687
700
  if accepts(ens, if_ens, min_total_path, cons, x, lo_t, hi_t, x) {
688
- return Ok(ExactResult {
701
+ return Ok(SearchOutcome::Done(ExactResult {
689
702
  x_cf: Some(x.to_vec()),
690
703
  proof: "optimal",
691
704
  stats: stats(0, 0, 0, 0.0, gap, true, false),
692
705
  snapped: Vec::new(),
693
706
  distance: Some(0.0),
694
- });
707
+ }));
695
708
  }
696
709
 
697
710
  let ensembles: Vec<&Ensemble> = match if_ens {
@@ -704,13 +717,13 @@ pub fn solve_exact(
704
717
  if order.iter().any(|&j| domains[j].is_empty()) {
705
718
  // Contradictory constraints left a feature with no legal value at all:
706
719
  // nothing to search, and nothing can be feasible.
707
- return Ok(ExactResult {
720
+ return Ok(SearchOutcome::Done(ExactResult {
708
721
  x_cf: None,
709
722
  proof: "optimal",
710
723
  stats: stats(0, 0, 0, f64::INFINITY, gap, true, false),
711
724
  snapped: Vec::new(),
712
725
  distance: None,
713
- });
726
+ }));
714
727
  }
715
728
  let h_suffix = h_suffix(&order, &domains);
716
729
 
@@ -892,6 +905,12 @@ pub fn solve_exact(
892
905
  completed = false;
893
906
  break;
894
907
  }
908
+ // Every node reaching this line is expanded below, so each count is
909
+ // asked about at most once and a search of fewer than
910
+ // SIGNAL_CHECK_INTERVAL nodes is never asked at all.
911
+ if nodes_expanded > 0 && nodes_expanded % SIGNAL_CHECK_INTERVAL == 0 && probe() {
912
+ return Ok(SearchOutcome::Interrupted);
913
+ }
895
914
 
896
915
  nodes_expanded += 1;
897
916
  let j = ctx.order[k];
@@ -1081,7 +1100,7 @@ pub fn solve_exact(
1081
1100
  }
1082
1101
  }
1083
1102
 
1084
- Ok(ExactResult {
1103
+ Ok(SearchOutcome::Done(ExactResult {
1085
1104
  distance: incumbent_row.as_ref().map(|_| incumbent_cost),
1086
1105
  x_cf: incumbent_row,
1087
1106
  proof,
@@ -1095,7 +1114,7 @@ pub fn solve_exact(
1095
1114
  warm_start_used,
1096
1115
  ),
1097
1116
  snapped,
1098
- })
1117
+ }))
1099
1118
  }
1100
1119
 
1101
1120
  #[cfg(test)]
@@ -1104,6 +1123,95 @@ mod tests {
1104
1123
  use crate::constraints::{LinearC, LIN_LE, POLICY_SATISFIED};
1105
1124
  use crate::exact::test_support::*;
1106
1125
 
1126
+ // -------------------------------------------------------- interrupt probe ---
1127
+
1128
+ /// A search that stays under the polling interval is never asked anything,
1129
+ /// so a probe that would have said stop changes nothing: the answer is the
1130
+ /// same bits the plain solve gives.
1131
+ #[test]
1132
+ fn a_short_search_is_never_asked_and_answers_the_same_bits() {
1133
+ let ens = stumps(&[(0, 1.0, true, 0.0, 1.0), (1, 0.5, true, 0.0, 0.5)], 2);
1134
+ let cons = cons_base(2);
1135
+ let x = [0.0, 0.0];
1136
+ let plain = solve(
1137
+ &ens,
1138
+ &x,
1139
+ (1.0, 2.0),
1140
+ &cons,
1141
+ 0.0,
1142
+ &no_policies(2),
1143
+ &ExactParams::default(),
1144
+ None,
1145
+ );
1146
+ let mut polls = 0usize;
1147
+ let outcome = solve_probed(
1148
+ &ens,
1149
+ &x,
1150
+ (1.0, 2.0),
1151
+ &cons,
1152
+ 0.0,
1153
+ &no_policies(2),
1154
+ &ExactParams::default(),
1155
+ None,
1156
+ &mut || {
1157
+ polls += 1;
1158
+ true // would stop the search — if it were ever asked
1159
+ },
1160
+ );
1161
+ let SearchOutcome::Done(probed) = outcome else {
1162
+ panic!("a search of a few nodes must not be interrupted")
1163
+ };
1164
+ assert_eq!(polls, 0);
1165
+ assert!(probed.stats.nodes_expanded > 0);
1166
+ assert_eq!(
1167
+ bits_of(probed.x_cf.as_ref().unwrap()),
1168
+ bits_of(plain.x_cf.as_ref().unwrap())
1169
+ );
1170
+ assert_eq!(
1171
+ probed.distance.unwrap().to_bits(),
1172
+ plain.distance.unwrap().to_bits()
1173
+ );
1174
+ assert_eq!(probed.proof, plain.proof);
1175
+ assert_eq!(probed.stats, plain.stats);
1176
+ assert_eq!(probed.snapped, plain.snapped);
1177
+ }
1178
+
1179
+ /// Twenty levers, each worth one point, and a target halfway between two
1180
+ /// whole numbers: no assignment can ever land in it, so no incumbent is
1181
+ /// found, nothing is cut on cost, and the search runs well past the
1182
+ /// polling interval. The warm start it was handed is a result it could
1183
+ /// have returned — a probe that says stop drops that too and reports only
1184
+ /// that it stopped.
1185
+ #[test]
1186
+ fn a_long_search_drops_even_its_warm_start_when_the_probe_says_stop() {
1187
+ let specs: Vec<(i32, f64, bool, f64, f64)> =
1188
+ (0..20).map(|j| (j, 1.0, true, 0.0, 1.0)).collect();
1189
+ let ens = stumps(&specs, 20);
1190
+ let cons = cons_base(20);
1191
+ let x = vec![0.0; 20];
1192
+ let params = ExactParams {
1193
+ time_budget_s: 1e9, // the probe, not the clock, must end this one
1194
+ ..ExactParams::default()
1195
+ };
1196
+ let mut polls = 0usize;
1197
+ let outcome = solve_probed(
1198
+ &ens,
1199
+ &x,
1200
+ (10.5, 10.5),
1201
+ &cons,
1202
+ 0.0,
1203
+ &no_policies(20),
1204
+ &params,
1205
+ Some((1e12, x.as_slice())), // too dear to prune anything with
1206
+ &mut || {
1207
+ polls += 1;
1208
+ true
1209
+ },
1210
+ );
1211
+ assert!(matches!(outcome, SearchOutcome::Interrupted));
1212
+ assert_eq!(polls, 1);
1213
+ }
1214
+
1107
1215
  // ---------------------------------------------------- ensemble bounds ---
1108
1216
 
1109
1217
  #[test]
@@ -1760,7 +1868,7 @@ mod tests {
1760
1868
  let ens = stumps(&[(0, 1.0, true, 0.0, 1.0)], 2);
1761
1869
  let if_ens = stumps(&[(1, 0.5, true, 2.0, 5.0)], 2);
1762
1870
  let cons = cons_base(2);
1763
- let result = solve_exact(
1871
+ let outcome = solve_exact(
1764
1872
  &ens,
1765
1873
  &[0.0, 0.0],
1766
1874
  (1.0, 2.0),
@@ -1772,8 +1880,12 @@ mod tests {
1772
1880
  Some((&if_ens, 4.0)),
1773
1881
  &ExactParams::default(),
1774
1882
  None,
1883
+ &mut || false,
1775
1884
  )
1776
1885
  .unwrap();
1886
+ let SearchOutcome::Done(result) = outcome else {
1887
+ unreachable!("no-op probe never interrupts")
1888
+ };
1777
1889
  assert_eq!(
1778
1890
  bits_of(result.x_cf.as_ref().unwrap()),
1779
1891
  vec![0x3ff0000000000000, 0x3fe0000000000000]
@@ -6,6 +6,7 @@ use crate::constraints::Constraints;
6
6
  use crate::exact::domains::{build_domains, constraint_cells, State};
7
7
  use crate::exact::search::solve_exact;
8
8
  use crate::exact::{ExactParams, ExactResult, ValuePolicy};
9
+ use crate::interrupt::{InterruptProbe, SearchOutcome};
9
10
  use crate::ir::{Ensemble, Link};
10
11
 
11
12
  // ------------------------------------------------------------ builders ---
@@ -139,6 +140,36 @@ pub(crate) fn solve(
139
140
  params: &ExactParams,
140
141
  incumbent: Option<(f64, &[f64])>,
141
142
  ) -> ExactResult {
143
+ match solve_probed(
144
+ ens,
145
+ x,
146
+ interval,
147
+ cons,
148
+ lam,
149
+ policies,
150
+ params,
151
+ incumbent,
152
+ &mut || false,
153
+ ) {
154
+ SearchOutcome::Done(result) => result,
155
+ SearchOutcome::Interrupted => unreachable!("no-op probe never interrupts"),
156
+ }
157
+ }
158
+
159
+ /// `solve` with a probe of the caller's own — for the tests that care what the
160
+ /// probe is asked and what happens when it says yes.
161
+ #[allow(clippy::too_many_arguments)]
162
+ pub(crate) fn solve_probed(
163
+ ens: &Ensemble,
164
+ x: &[f64],
165
+ interval: (f64, f64),
166
+ cons: &Constraints,
167
+ lam: f64,
168
+ policies: &[Option<ValuePolicy>],
169
+ params: &ExactParams,
170
+ incumbent: Option<(f64, &[f64])>,
171
+ probe: InterruptProbe<'_>,
172
+ ) -> SearchOutcome<ExactResult> {
142
173
  let p = ens.n_features;
143
174
  solve_exact(
144
175
  ens,
@@ -152,6 +183,7 @@ pub(crate) fn solve(
152
183
  None,
153
184
  params,
154
185
  incumbent,
186
+ probe,
155
187
  )
156
188
  .unwrap()
157
189
  }
@@ -5,6 +5,10 @@
5
5
  //! Rayon parallelizes only RNG-free stages (fitness/check/repair), so results
6
6
  //! are identical across thread counts by construction — the audit showed child
7
7
  //! creation is 1-4 % of wall time, so sequential variation costs little.
8
+ //! `solve_genetic_batch` fans its tasks out a chunk at a time rather than all
9
+ //! at once; that argument is untouched, since each task carries its own seed
10
+ //! and reads nothing from its neighbours, so which tasks travel together and
11
+ //! on which thread cannot change any of their answers.
8
12
 
9
13
  use std::time::Instant;
10
14
 
@@ -14,8 +18,14 @@ use rand_distr::{Distribution, Normal};
14
18
  use rand_pcg::Pcg64Mcg;
15
19
 
16
20
  use crate::constraints::Constraints;
21
+ use crate::interrupt::{InterruptProbe, SearchOutcome};
17
22
  use crate::ir::Ensemble;
18
23
 
24
+ /// How many batch tasks are fanned out between two interrupt polls. Big enough
25
+ /// that the rayon fan-out still has plenty to chew on, small enough that a
26
+ /// large batch comes back to the calling thread often.
27
+ const SIGNAL_CHECK_CHUNK: usize = 256;
28
+
19
29
  pub struct GaParams {
20
30
  pub population: usize,
21
31
  pub max_generations: usize,
@@ -241,6 +251,11 @@ pub fn solve_genetic(
241
251
  /// Independent GA searches for a batch of `(row_index, seed)` tasks, fanned
242
252
  /// out with rayon. Every task is independently seeded, so the output is
243
253
  /// thread-count-independent by construction; order follows `tasks`.
254
+ ///
255
+ /// The fan-out runs a chunk of `SIGNAL_CHECK_CHUNK` tasks at a time and asks
256
+ /// `probe` on the calling thread before each chunk. Answering yes drops the
257
+ /// chunks already finished and returns `SearchOutcome::Interrupted`; answering
258
+ /// no gives the same vector, task for task, as one undivided fan-out.
244
259
  #[allow(clippy::too_many_arguments)]
245
260
  pub fn solve_genetic_batch(
246
261
  ens: &Ensemble,
@@ -255,30 +270,39 @@ pub fn solve_genetic_batch(
255
270
  background: Option<(&[f64], usize)>,
256
271
  plausibility: Option<(&Ensemble, f64)>,
257
272
  params: &GaParams,
258
- ) -> Vec<GaResult> {
273
+ probe: InterruptProbe<'_>,
274
+ ) -> SearchOutcome<Vec<GaResult>> {
259
275
  use rayon::prelude::*;
260
276
  let p = ens.n_features;
261
277
  // Warm the cell caches up front so parallel tasks don't build them twice.
262
278
  let _ = ens.feature_cells();
263
- tasks
264
- .par_iter()
265
- .map(|&(row, seed)| {
266
- solve_genetic(
267
- ens,
268
- &xs[row * p..(row + 1) * p],
269
- lo_t,
270
- hi_t,
271
- cons,
272
- sigma,
273
- weights,
274
- lam,
275
- background,
276
- plausibility,
277
- Some(seed),
278
- params,
279
- )
280
- })
281
- .collect()
279
+ let mut results: Vec<GaResult> = Vec::with_capacity(tasks.len());
280
+ for chunk in tasks.chunks(SIGNAL_CHECK_CHUNK) {
281
+ if probe() {
282
+ return SearchOutcome::Interrupted;
283
+ }
284
+ let done: Vec<GaResult> = chunk
285
+ .par_iter()
286
+ .map(|&(row, seed)| {
287
+ solve_genetic(
288
+ ens,
289
+ &xs[row * p..(row + 1) * p],
290
+ lo_t,
291
+ hi_t,
292
+ cons,
293
+ sigma,
294
+ weights,
295
+ lam,
296
+ background,
297
+ plausibility,
298
+ Some(seed),
299
+ params,
300
+ )
301
+ })
302
+ .collect();
303
+ results.extend(done);
304
+ }
305
+ SearchOutcome::Done(results)
282
306
  }
283
307
 
284
308
  fn pin_fixed(pop: &mut [f64], n_rows: usize, p: usize, fixed: &[bool], x: &[f64]) {
@@ -516,7 +540,7 @@ mod tests {
516
540
  let cons = empty_constraints(2);
517
541
  let xs = [0.0, 0.0, -1.0, 2.0];
518
542
  let tasks = [(0usize, 1u64), (0, 2), (1, 3), (1, 1)];
519
- let batch = solve_genetic_batch(
543
+ let outcome = solve_genetic_batch(
520
544
  &ens,
521
545
  &xs,
522
546
  &tasks,
@@ -529,7 +553,11 @@ mod tests {
529
553
  None,
530
554
  None,
531
555
  &params(),
556
+ &mut || false,
532
557
  );
558
+ let SearchOutcome::Done(batch) = outcome else {
559
+ unreachable!("no-op probe never interrupts")
560
+ };
533
561
  for (result, &(row, seed)) in batch.iter().zip(&tasks) {
534
562
  let single = solve_genetic(
535
563
  &ens,
@@ -550,6 +578,96 @@ mod tests {
550
578
  }
551
579
  }
552
580
 
581
+ /// The batch asks before it starts, so a probe that says stop right away
582
+ /// costs one question and no search at all.
583
+ #[test]
584
+ fn batch_stops_before_the_first_chunk_when_the_probe_says_stop() {
585
+ let ens = stump();
586
+ let cons = empty_constraints(2);
587
+ let xs = [0.0, 0.0];
588
+ let tasks = [(0usize, 1u64), (0, 2)];
589
+ let mut polls = 0usize;
590
+ let outcome = solve_genetic_batch(
591
+ &ens,
592
+ &xs,
593
+ &tasks,
594
+ 0.5,
595
+ f64::INFINITY,
596
+ &cons,
597
+ &[1.0, 1.0],
598
+ &[1.0, 1.0],
599
+ 0.05,
600
+ None,
601
+ None,
602
+ &params(),
603
+ &mut || {
604
+ polls += 1;
605
+ true
606
+ },
607
+ );
608
+ assert!(matches!(outcome, SearchOutcome::Interrupted));
609
+ assert_eq!(polls, 1);
610
+ }
611
+
612
+ /// Chunking is a polling schedule and nothing else: 300 tasks are two
613
+ /// chunks and therefore two questions, and every task still comes back
614
+ /// with exactly what its own single solve produces.
615
+ #[test]
616
+ fn batch_polls_once_per_chunk_and_answers_task_for_task() {
617
+ let ens = stump();
618
+ let cons = empty_constraints(2);
619
+ let xs = [0.0, 0.0, -1.0, 2.0];
620
+ let tasks: Vec<(usize, u64)> = (0..300).map(|i| (i % 2, i as u64 + 1)).collect();
621
+ let small = GaParams {
622
+ population: 8,
623
+ max_generations: 3,
624
+ stall_generations: 2,
625
+ ..params()
626
+ };
627
+ let mut polls = 0usize;
628
+ let outcome = solve_genetic_batch(
629
+ &ens,
630
+ &xs,
631
+ &tasks,
632
+ 0.5,
633
+ f64::INFINITY,
634
+ &cons,
635
+ &[1.0, 1.0],
636
+ &[1.0, 1.0],
637
+ 0.05,
638
+ None,
639
+ None,
640
+ &small,
641
+ &mut || {
642
+ polls += 1;
643
+ false
644
+ },
645
+ );
646
+ let SearchOutcome::Done(batch) = outcome else {
647
+ unreachable!("this probe never says stop")
648
+ };
649
+ assert_eq!(polls, 2);
650
+ assert_eq!(batch.len(), tasks.len());
651
+ for (result, &(row, seed)) in batch.iter().zip(&tasks) {
652
+ let single = solve_genetic(
653
+ &ens,
654
+ &xs[row * 2..(row + 1) * 2],
655
+ 0.5,
656
+ f64::INFINITY,
657
+ &cons,
658
+ &[1.0, 1.0],
659
+ &[1.0, 1.0],
660
+ 0.05,
661
+ None,
662
+ None,
663
+ Some(seed),
664
+ &small,
665
+ );
666
+ assert_eq!(result.generations, single.generations);
667
+ assert_eq!(result.x_cf, single.x_cf);
668
+ }
669
+ }
670
+
553
671
  #[test]
554
672
  fn frozen_target_is_infeasible() {
555
673
  let ens = stump();
@@ -0,0 +1,27 @@
1
+ //! The stop question a long-running search polls, and the answer shape it
2
+ //! reports back with.
3
+ //!
4
+ //! A probe is a caller-supplied closure that takes no arguments, reads no
5
+ //! search state and returns no value beyond "stop now?". A search polls it on
6
+ //! a schedule fixed by its own work counters — never on anything derived from
7
+ //! the data it is searching — and the only thing a `true` answer can do is
8
+ //! abandon the search. So a probe that always answers `false` leaves every
9
+ //! number a search computes exactly as it was before the probe existed.
10
+
11
+ /// The stop question: `true` means abandon the search.
12
+ ///
13
+ /// It is a `FnMut` because the real one carries state of its own (the pending
14
+ /// error a caller wants to raise afterwards); nothing here depends on that.
15
+ pub type InterruptProbe<'a> = &'a mut dyn FnMut() -> bool;
16
+
17
+ /// What a probed search answers with.
18
+ ///
19
+ /// `Interrupted` deliberately carries nothing. Whatever the search had found
20
+ /// when the probe fired — an incumbent, a half-grown box, the results of the
21
+ /// chunks already done — is dropped here, so no caller can mistake an
22
+ /// abandoned search for a finished one.
23
+ #[derive(Debug)]
24
+ pub enum SearchOutcome<T> {
25
+ Done(T),
26
+ Interrupted,
27
+ }
@@ -4,6 +4,7 @@ pub mod cells;
4
4
  pub mod constraints;
5
5
  pub mod exact;
6
6
  pub mod ga;
7
+ pub mod interrupt;
7
8
  pub mod ir;
8
9
  pub mod regions;
9
10