treecf 0.2.4__tar.gz → 0.3.1__tar.gz

{treecf-0.2.4 → treecf-0.3.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: treecf
-Version: 0.2.4
+Version: 0.3.1
 Classifier: Development Status :: 4 - Beta
 Classifier: Intended Audience :: Science/Research
 Classifier: License :: OSI Approved :: MIT License
@@ -131,6 +131,12 @@ proved = exp.explain(x, target=t, backend="exact")      # proof="optimal", a cer
 boxed = exp.explain(x, target=t, region=True)            # res.region.describe() -> "utilization <= 0.4"
 ```
 
+## Contributing
+
+See [CONTRIBUTING.md](CONTRIBUTING.md) for dev setup, the test layers, and the
+project's hard invariants; report security issues privately per
+[SECURITY.md](SECURITY.md).
+
 ## License
 
 MIT

{treecf-0.2.4 → treecf-0.3.1}/README.md

@@ -65,6 +65,12 @@ proved = exp.explain(x, target=t, backend="exact")      # proof="optimal", a cer
 boxed = exp.explain(x, target=t, region=True)            # res.region.describe() -> "utilization <= 0.4"
 ```
 
+## Contributing
+
+See [CONTRIBUTING.md](CONTRIBUTING.md) for dev setup, the test layers, and the
+project's hard invariants; report security issues privately per
+[SECURITY.md](SECURITY.md).
+
 ## License
 
 MIT

{treecf-0.2.4 → treecf-0.3.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "treecf"
-version = "0.2.4"
+version = "0.3.1"
 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."
 readme = "README.md"
 license = { text = "MIT" }
@@ -121,10 +121,11 @@ ignore_missing_imports = true
 
 [tool.pytest.ini_options]
 testpaths = ["tests"]
-addopts = "-q --strict-markers -m 'not bench'"
+addopts = "-q --strict-markers -m 'not bench and not slow'"
 markers = [
     "bench: non-gating performance smoke benchmarks (run with -m bench)",
     "rust: cross-language tests requiring the dev Rust extension (maturin develop)",
+    "slow: docs snippet/structure harnesses (run with -m slow)",
 ]
 filterwarnings = [
     "error",

{treecf-0.2.4 → treecf-0.3.1}/rust/Cargo.lock

@@ -419,7 +419,7 @@ checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca"
 
 [[package]]
 name = "treecf-core"
-version = "0.2.4"
+version = "0.3.1"
 dependencies = [
  "numpy",
  "pyo3",

{treecf-0.2.4 → treecf-0.3.1}/rust/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "treecf-core"
-version = "0.2.4"
+version = "0.3.1"
 edition = "2021"
 # f64::next_down (cells.rs) stabilized in 1.86; pyo3 0.29 needs 1.83
 rust-version = "1.86"

{treecf-0.2.4 → treecf-0.3.1}/rust/src/cells.rs

@@ -152,7 +152,9 @@ pub fn feature_cells_joint(ensembles: &[&Ensemble]) -> Vec<Vec<Cell>> {
     let mut pairs: Vec<Vec<(f64, bool)>> = vec![Vec::new(); n_features];
     for ens in ensembles {
         for i in 0..ens.feature.len() {
-            if ens.feature[i] >= 0 {
+            // set-membership splits carry no threshold; their features partition
+            // into category blocks instead of interval cells
+            if ens.feature[i] >= 0 && ens.node_set[i] < 0 {
                 pairs[ens.feature[i] as usize].push((ens.threshold[i], ens.is_lt[i]));
             }
         }
@@ -160,6 +162,78 @@ pub fn feature_cells_joint(ensembles: &[&Ensemble]) -> Vec<Vec<Cell>> {
     pairs.iter().map(|p| build_cells(p)).collect()
 }
 
+/// Category blocks per feature across ensembles — port of
+/// `treecf.aim.cells.category_blocks`. Indexed by feature; numeric features get
+/// an empty vec. Two codes share a block iff no split in any ensemble separates
+/// them: set-membership splits partition by membership, numeric splits on a
+/// categorical feature (an isolation forest trained on raw codes) partition by
+/// threshold side. Blocks are ordered by smallest member (codes scan
+/// ascending), members ascend, and a block's representative is its first entry.
+///
+/// Panics where Python raises `ValueError` (cardinality disagreement), matching
+/// the `feature_cells_joint` convention.
+pub fn category_blocks_joint(ensembles: &[&Ensemble]) -> Vec<Vec<Vec<u32>>> {
+    let n_features = ensembles[0].n_features;
+    let mut cardinality = vec![0u32; n_features];
+    for ens in ensembles {
+        for (j, slot) in cardinality.iter_mut().enumerate() {
+            let k = ens.cardinality[j];
+            if k > 0 {
+                if *slot == 0 {
+                    *slot = k;
+                } else {
+                    assert_eq!(
+                        *slot, k,
+                        "ensembles disagree on the cardinality of feature {j}"
+                    );
+                }
+            }
+        }
+    }
+    let mut out = Vec::with_capacity(n_features);
+    for (j, &card) in cardinality.iter().enumerate() {
+        let k = card as usize;
+        if k == 0 {
+            out.push(Vec::new());
+            continue;
+        }
+        let mut signatures: Vec<Vec<bool>> = vec![Vec::new(); k];
+        for ens in ensembles {
+            for i in 0..ens.feature.len() {
+                if ens.feature[i] != j as i32 {
+                    continue;
+                }
+                if ens.node_set[i] >= 0 {
+                    for (code, sig) in signatures.iter_mut().enumerate() {
+                        sig.push(ens.set_contains(ens.node_set[i], code as f64));
+                    }
+                } else if ens.is_lt[i] {
+                    for (code, sig) in signatures.iter_mut().enumerate() {
+                        sig.push((code as f64) < ens.threshold[i]);
+                    }
+                } else {
+                    for (code, sig) in signatures.iter_mut().enumerate() {
+                        sig.push((code as f64) <= ens.threshold[i]);
+                    }
+                }
+            }
+        }
+        let mut blocks: Vec<Vec<u32>> = Vec::new();
+        let mut seen: Vec<usize> = Vec::new(); // block index -> a code carrying its signature
+        for code in 0..k {
+            match seen.iter().position(|&c| signatures[c] == signatures[code]) {
+                Some(b) => blocks[b].push(code as u32),
+                None => {
+                    seen.push(code);
+                    blocks.push(vec![code as u32]);
+                }
+            }
+        }
+        out.push(blocks);
+    }
+    out
+}
+
 /// Index of the unique cell containing `x` — port of `treecf.aim.cells.cell_index`.
 ///
 /// Panics where Python raises `ValueError`. The cells partition the whole line,
@@ -175,6 +249,7 @@ pub fn cell_index(cells: &[Cell], x: f64) -> usize {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::ir::Link;
 
     #[test]
     fn lt_le_collision_yields_singleton() {
@@ -252,4 +327,107 @@ mod tests {
         assert_eq!(cell_index(&cells, 1.0), 1);
         assert_eq!(cell_index(&cells, 1.5), 2);
     }
+
+    /// One set stump per word list, all on feature 0, cardinality `k`.
+    fn set_ens(sets: &[u64], k: u32) -> Ensemble {
+        let n = sets.len().max(1);
+        let mut feature = Vec::new();
+        let mut node_set = Vec::new();
+        let mut left = Vec::new();
+        let mut right = Vec::new();
+        let mut value = Vec::new();
+        let mut roots = Vec::new();
+        for (t, _) in sets.iter().enumerate() {
+            let base = (t * 3) as u32;
+            roots.push(base);
+            feature.extend_from_slice(&[0, -1, -1]);
+            node_set.extend_from_slice(&[t as i32, -1, -1]);
+            left.extend_from_slice(&[base + 1, 0, 0]);
+            right.extend_from_slice(&[base + 2, 0, 0]);
+            value.extend_from_slice(&[0.0, -1.0, 1.0]);
+        }
+        if sets.is_empty() {
+            roots.push(0);
+            feature.push(-1);
+            node_set.push(-1);
+            left.push(0);
+            right.push(0);
+            value.push(0.0);
+        }
+        let n_nodes = feature.len();
+        Ensemble::new(
+            feature,
+            vec![0.0; n_nodes],
+            vec![false; n_nodes],
+            vec![false; n_nodes],
+            left,
+            right,
+            value,
+            roots,
+            0.0,
+            Link::Identity,
+            n,
+        )
+        .unwrap()
+        .with_categories(
+            node_set,
+            (0..=sets.len() as u32).collect(),
+            sets.to_vec(),
+            {
+                let mut c = vec![0; n];
+                c[0] = k;
+                c
+            },
+        )
+        .unwrap()
+    }
+
+    #[test]
+    fn no_splits_one_block() {
+        let e = set_ens(&[], 4);
+        assert_eq!(e.category_blocks()[0], vec![vec![0, 1, 2, 3]]);
+    }
+
+    #[test]
+    fn one_set_two_blocks_numbered_by_smallest_member() {
+        let e = set_ens(&[0b101], 4); // {0, 2}
+        assert_eq!(e.category_blocks()[0], vec![vec![0, 2], vec![1, 3]]);
+    }
+
+    #[test]
+    fn two_sets_refine_to_singletons() {
+        let e = set_ens(&[0b101, 0b1100], 4); // {0,2} then {2,3}
+        assert_eq!(
+            e.category_blocks()[0],
+            vec![vec![0], vec![1], vec![2], vec![3]]
+        );
+    }
+
+    #[test]
+    fn joint_blocks_refine_and_numeric_splits_partition_codes() {
+        let a = set_ens(&[0b101], 4);
+        // a numeric stump on the same feature at 1.5 (an isolation forest would)
+        let b = {
+            let mut e = Ensemble::new(
+                vec![0, -1, -1],
+                vec![1.5, 0.0, 0.0],
+                vec![false, false, false],
+                vec![false, false, false],
+                vec![1, 0, 0],
+                vec![2, 0, 0],
+                vec![0.0, -1.0, 1.0],
+                vec![0],
+                0.0,
+                Link::Identity,
+                1,
+            )
+            .unwrap();
+            e = e
+                .with_categories(vec![-1, -1, -1], vec![0], vec![], vec![4])
+                .unwrap();
+            e
+        };
+        let joint = category_blocks_joint(&[&a, &b]);
+        assert_eq!(joint[0], vec![vec![0], vec![1], vec![2], vec![3]]);
+    }
 }

{treecf-0.2.4 → treecf-0.3.1}/rust/src/constraints.rs

@@ -35,6 +35,27 @@ pub struct Constraints {
     pub implications: Vec<(u32, f64, u32, f64)>, // cond_idx, cond_val, cons_idx, cons_val
     pub onehot: Vec<Vec<u32>>,
     pub allow_missing: Vec<(u32, f64, f64)>, // (feature, delta_to, delta_from), index-sorted
+    // (feature, allowed-code bitset words), index-sorted; an empty word list
+    // means the declared set is empty (nothing is allowed on that feature)
+    pub allowed_categories: Vec<(u32, Vec<u64>)>,
+}
+
+/// Membership of `v` in an allowed-code bitset: an integral code whose bit is
+/// set. Non-integral values are never members; codes beyond the words are not.
+#[inline]
+pub(crate) fn code_allowed(words: &[u64], v: f64) -> bool {
+    if !v.is_finite() {
+        return false;
+    }
+    let code = v as i64;
+    if code as f64 != v || code < 0 {
+        return false;
+    }
+    let word = (code >> 6) as usize;
+    if word >= words.len() {
+        return false;
+    }
+    (words[word] >> (code & 63)) & 1 == 1
 }
 
 /// Python `max(a, b)`: returns b only if b > a. Two consequences both mirrors
@@ -184,6 +205,12 @@ impl Constraints {
                 return false; // NaN sum compares false -> infeasible, like numpy
             }
         }
+        for (j, words) in &self.allowed_categories {
+            let v = row[*j as usize];
+            if !v.is_nan() && !code_allowed(words, v) {
+                return false;
+            }
+        }
         true
     }
 
@@ -239,6 +266,19 @@ impl Constraints {
                 row[j] = v;
             }
         }
+        for (j, words) in &self.allowed_categories {
+            let j = *j as usize;
+            // keep an allowed code; else the smallest allowed code (deterministic)
+            if words.is_empty() {
+                continue; // nothing legal to write; check rejects
+            }
+            if !row[j].is_nan() && !code_allowed(words, row[j]) {
+                if let Some(word) = words.iter().position(|&w| w != 0) {
+                    let bit = words[word].trailing_zeros() as usize;
+                    row[j] = (word * 64 + bit) as f64;
+                }
+            }
+        }
         // cyclic projection: boxes intersected with halfspaces, a few sweeps
         // (must mirror repair_matrix float-for-float: same term order, one
         // residual/denom division, lower-bound-first clipping)
@@ -345,6 +385,7 @@ mod tests {
             implications: vec![],
             onehot: vec![],
             allow_missing: vec![],
+            allowed_categories: vec![],
         }
     }
 

{treecf-0.2.4 → treecf-0.3.1}/rust/src/exact/domains.rs

@@ -12,7 +12,7 @@ use crate::ir::Ensemble;
 /// Python's `<`-based ordering: `-0.0` and `0.0` compare equal, so a stable sort
 /// leaves them in insertion order. Costs and sort values are never NaN here.
 #[inline]
-fn py_cmp(a: f64, b: f64) -> std::cmp::Ordering {
+pub(crate) fn py_cmp(a: f64, b: f64) -> std::cmp::Ordering {
     use std::cmp::Ordering;
     if a < b {
         Ordering::Less
@@ -27,6 +27,9 @@ fn py_cmp(a: f64, b: f64) -> std::cmp::Ordering {
 
 /// One feature's contribution to the objective — the per-feature term of
 /// `genetic.objective()`, same four cases, same multiply-then-divide order.
+/// A categorical change (`is_cat`) costs one flat unit in place of the
+/// absolute code distance; NaN transitions keep their declared deltas.
+#[allow(clippy::too_many_arguments)]
 fn term_cost(
     x_j: f64,
     r: f64,
@@ -35,6 +38,7 @@ fn term_cost(
     lam: f64,
     to_miss: f64,
     from_miss: f64,
+    is_cat: bool,
 ) -> f64 {
     let x_nan = x_j.is_nan();
     let r_nan = r.is_nan();
@@ -50,11 +54,12 @@ fn term_cost(
     if r == x_j {
         return 0.0;
     }
-    let delta = (r - x_j).abs();
+    let delta = if is_cat { 1.0 } else { (r - x_j).abs() };
     lam + (weight_j * delta) / sigma_j
 }
 
 /// Full-row objective, accumulated in ascending feature index.
+/// `cardinality[j] > 0` marks feature j categorical (flat change cost).
 pub(crate) fn cost_of_row(
     x: &[f64],
     row: &[f64],
@@ -62,11 +67,21 @@ pub(crate) fn cost_of_row(
     weights: &[f64],
     lam: f64,
     deltas: &[(f64, f64)],
+    cardinality: &[u32],
 ) -> f64 {
     let mut total = 0.0;
     for j in 0..x.len() {
         let (to_miss, from_miss) = deltas[j];
-        total += term_cost(x[j], row[j], weights[j], sigma[j], lam, to_miss, from_miss);
+        total += term_cost(
+            x[j],
+            row[j],
+            weights[j],
+            sigma[j],
+            lam,
+            to_miss,
+            from_miss,
+            cardinality[j] > 0,
+        );
     }
     total
 }
@@ -376,6 +391,7 @@ pub(crate) fn build_domains(
     weights: &[f64],
     lam: f64,
     policies: &[Option<ValuePolicy>],
+    blocks: &[Vec<Vec<u32>>],
 ) -> Vec<Vec<State>> {
     let (lo, hi, frozen) = cons.instance_bounds(x);
     let lo: Vec<f64> = lo
@@ -414,6 +430,28 @@ pub(crate) fn build_domains(
         let (to_miss, from_miss) = deltas[j];
         let (weight_j, sigma_j) = (weights[j], sigma[j]);
 
+        if !blocks[j].is_empty() {
+            let allowed = cons
+                .allowed_categories
+                .iter()
+                .find(|(idx, _)| *idx as usize == j)
+                .map(|(_, words)| words.as_slice());
+            domains.push(categorical_states(
+                x_j,
+                &blocks[j],
+                allowed,
+                frozen[j],
+                allow_j,
+                suppress_nan[j],
+                weight_j,
+                sigma_j,
+                lam,
+                to_miss,
+                from_miss,
+            ));
+            continue;
+        }
+
         if frozen[j] {
             let idx = if x_nan {
                 cells.len()
@@ -429,11 +467,19 @@ pub(crate) fn build_domains(
             if !x_nan {
                 // the pin fixes the only value the feature may take; going
                 // missing is a separate question AllowMissing still answers
-                let cost = term_cost(x_j, v, weight_j, sigma_j, lam, to_miss, from_miss);
+                let cost = term_cost(x_j, v, weight_j, sigma_j, lam, to_miss, from_miss, false);
                 let mut states = vec![State::new(v, cost, cell_index(cells, v), false)];
                 if allow_j && !suppress_nan[j] {
-                    let nan_cost =
-                        term_cost(x_j, f64::NAN, weight_j, sigma_j, lam, to_miss, from_miss);
+                    let nan_cost = term_cost(
+                        x_j,
+                        f64::NAN,
+                        weight_j,
+                        sigma_j,
+                        lam,
+                        to_miss,
+                        from_miss,
+                        false,
+                    );
                     states.push(State::new(f64::NAN, nan_cost, cells.len(), true));
                     sort_states(&mut states);
                 }
@@ -449,7 +495,7 @@ pub(crate) fn build_domains(
                 states.push(State::new(f64::NAN, 0.0, cells.len(), true));
             }
             if allow_j {
-                let cost = term_cost(x_j, v, weight_j, sigma_j, lam, to_miss, from_miss);
+                let cost = term_cost(x_j, v, weight_j, sigma_j, lam, to_miss, from_miss, false);
                 states.push(State::new(v, cost, cell_index(cells, v), false));
             }
             sort_states(&mut states);
@@ -492,7 +538,8 @@ pub(crate) fn build_domains(
                     if keep_added && val == x_j {
                         continue;
                     }
-                    let cost = term_cost(x_j, val, weight_j, sigma_j, lam, to_miss, from_miss);
+                    let cost =
+                        term_cost(x_j, val, weight_j, sigma_j, lam, to_miss, from_miss, false);
                     states.push(State::new(val, cost, local_idx, false));
                 }
                 continue;
@@ -504,7 +551,7 @@ pub(crate) fn build_domains(
                 if !iv.contains(val) || (keep_added && val == x_j) {
                     continue;
                 }
-                let cost = term_cost(x_j, val, weight_j, sigma_j, lam, to_miss, from_miss);
+                let cost = term_cost(x_j, val, weight_j, sigma_j, lam, to_miss, from_miss, false);
                 states.push(State::new(val, cost, local_idx, false));
                 added_here.push(val);
             }
@@ -523,14 +570,23 @@ pub(crate) fn build_domains(
             if added_here.contains(&r) {
                 continue; // the demanded value was this cell's nearest point too
             }
-            let cost = term_cost(x_j, r, weight_j, sigma_j, lam, to_miss, from_miss);
+            let cost = term_cost(x_j, r, weight_j, sigma_j, lam, to_miss, from_miss, false);
             let mut state = State::new(r, cost, local_idx, false);
             state.snapped = snapped;
             states.push(state);
         }
 
         if allow_j && !suppress_nan[j] {
-            let nan_cost = term_cost(x_j, f64::NAN, weight_j, sigma_j, lam, to_miss, from_miss);
+            let nan_cost = term_cost(
+                x_j,
+                f64::NAN,
+                weight_j,
+                sigma_j,
+                lam,
+                to_miss,
+                from_miss,
+                false,
+            );
             states.push(State::new(f64::NAN, nan_cost, cells.len(), true));
         }
 
@@ -575,15 +631,121 @@ fn referenced_features(cons: &Constraints) -> Vec<bool> {
     for &(j, _, _) in &cons.allow_missing {
         refs[j as usize] = true;
     }
+    for (j, _) in &cons.allowed_categories {
+        refs[*j as usize] = true;
+    }
     refs
 }
 
+/// Candidate states of a categorical feature — the mirror of
+/// `_categorical_states`. A block's representative is its smallest member the
+/// declared allowed set admits (cost is flat within a block, so the choice is
+/// free; smallest is the determinism rule); a block with no admissible member
+/// contributes no state, and an allowed set admitting nothing empties the
+/// whole domain — the same certified-infeasible signal a contradictory
+/// numeric pin produces. `cell_idx` carries the block index; the NaN state's
+/// sentinel index is the block count.
+#[allow(clippy::too_many_arguments)]
+fn categorical_states(
+    x_j: f64,
+    blocks_j: &[Vec<u32>],
+    allowed: Option<&[u64]>,
+    frozen_j: bool,
+    allow_j: bool,
+    suppressed: bool,
+    weight_j: f64,
+    sigma_j: f64,
+    lam: f64,
+    to_miss: f64,
+    from_miss: f64,
+) -> Vec<State> {
+    use crate::constraints::code_allowed;
+
+    let x_nan = x_j.is_nan();
+    let n_blocks = blocks_j.len();
+    let block_of = |code: f64| -> usize {
+        blocks_j
+            .iter()
+            .position(|block| block.contains(&(code as u32)))
+            .expect("code not covered by blocks (should be impossible)")
+    };
+
+    if frozen_j {
+        let idx = if x_nan { n_blocks } else { block_of(x_j) };
+        return vec![State::new(x_j, 0.0, idx, x_nan)];
+    }
+
+    if x_nan && !allow_j {
+        if suppressed {
+            return Vec::new();
+        }
+        return vec![State::new(x_j, 0.0, n_blocks, true)];
+    }
+
+    let is_allowed = |code: f64| -> bool {
+        match allowed {
+            None => true,
+            Some(words) => code_allowed(words, code),
+        }
+    };
+
+    let mut states: Vec<State> = Vec::new();
+    let mut keep_block = usize::MAX;
+    if !x_nan && is_allowed(x_j) {
+        keep_block = block_of(x_j);
+        states.push(State::new(x_j, 0.0, keep_block, false));
+    }
+
+    for (block_idx, block) in blocks_j.iter().enumerate() {
+        if block_idx == keep_block {
+            continue; // the unchanged factual already represents its block
+        }
+        let member = block.iter().copied().find(|&c| is_allowed(c as f64));
+        if let Some(code) = member {
+            let rep = code as f64;
+            let cost = term_cost(x_j, rep, weight_j, sigma_j, lam, to_miss, from_miss, true);
+            states.push(State::new(rep, cost, block_idx, false));
+        }
+    }
+
+    if allow_j && !suppressed {
+        let nan_cost = term_cost(
+            x_j,
+            f64::NAN,
+            weight_j,
+            sigma_j,
+            lam,
+            to_miss,
+            from_miss,
+            true,
+        );
+        states.push(State::new(f64::NAN, nan_cost, n_blocks, true));
+    }
+
+    sort_states(&mut states);
+    states
+}
+
 /// Search order: descending split count in the joint grid, ties ascending index.
 /// A feature with no split anywhere and no referencing constraint is left out —
 /// its domain is a single keep state, so it never needs to branch.
-pub(crate) fn feature_order(grids: &[Vec<Cell>], cons: &Constraints) -> Vec<usize> {
+pub(crate) fn feature_order(
+    grids: &[Vec<Cell>],
+    cons: &Constraints,
+    blocks: &[Vec<Vec<u32>>],
+) -> Vec<usize> {
     let referenced = referenced_features(cons);
-    let split_counts: Vec<usize> = grids.iter().map(|cells| cells.len() - 1).collect();
+    let split_counts: Vec<usize> = grids
+        .iter()
+        .enumerate()
+        .map(|(j, cells)| {
+            if blocks[j].is_empty() {
+                cells.len() - 1
+            } else {
+                blocks[j].len() - 1
+            }
+        })
+        .collect();
     let mut included: Vec<usize> = (0..grids.len())
         .filter(|&j| split_counts[j] > 0 || referenced[j])
         .collect();
@@ -1107,7 +1269,8 @@ mod tests {
         let mut cons = cons_base(4);
         cons.freeze = vec![3]; // referenced but split-free: kept, and last
         let grids = constraint_cells(&cons, &[&ens]);
-        assert_eq!(feature_order(&grids, &cons), vec![1, 2, 3]);
+        let blocks = crate::cells::category_blocks_joint(&[&ens]);
+        assert_eq!(feature_order(&grids, &cons, &blocks), vec![1, 2, 3]);
     }
 
     #[test]
@@ -1115,6 +1278,7 @@ mod tests {
         let ens = golden_ens();
         let cons = cons_base(2);
         let grids = constraint_cells(&cons, &[&ens]);
+        let blocks = crate::cells::category_blocks_joint(&[&ens]);
         let domains = build_domains(
             &grids,
             &[2.0, 0.0],
@@ -1123,8 +1287,9 @@ mod tests {
             &[1.0, 1.0],
             0.0,
             &no_policies(2),
+            &blocks,
         );
-        let order = feature_order(&grids, &cons);
+        let order = feature_order(&grids, &cons, &blocks);
         assert_eq!(h_suffix(&order, &domains), vec![0.0, 0.0]);
     }