feat: add bivariate Shape²CART branching

README.md

@@ -4,18 +4,17 @@
 
 `sgtlearn` is a Python package for learning [Shape Generalized Trees (SGTs)](https://neurips.cc/virtual/2025/loc/san-diego/poster/115950).
 
-- 🌳 **Shape Generalized Trees (SGTs):** A class of decision trees where each node applies a learnable, axis-aligned shape function to a feature for non-linear and interpretable splits.
+- 🌳 **Shape Generalized Trees (SGTs):** A class of decision trees where each node applies a learnable, axis-aligned shape function to one or two logical features for non-linear and interpretable splits.
 - 👁 **Interpretability:** Each node's shape function can be visualized directly.
 - ⚡ **ShapeCART Algorithm:** An efficient induction method for learning SGTs from data.
 - 🔀 **Extensions:**
   - **Shape²GT (S²GT):** Bivariate shape functions for richer splits.
   - **SGT<sub>K</sub>:** Multi-way branching generalization.
   - **Shape²CART & ShapeCART<sub>K</sub>:** Algorithms for learning S²GTs and SGT<sub>K</sub>s.
 
+
 > [!NOTE]
-> This codebase is an efficient, but working implementation of the algorithms in the paper "Empowering Decision Trees via Shape Function Branching". Please refer to the [ROADMAP](ROADMAP.md) for a detailed list of features that are currently implemented and those that are planned for future releases. For the canonical code base for the paper, please refer to https://github.com/optimal-uoft/Empowering-DTs-via-Shape-Functions. Features in the paper that are not yet implemented in this codebase include: 
-> * Bivariate shape functions (Shape$^2$CART) + Higher branching factors for bivariate splits (Shape$^2$SGT$_K$)
-> * Visualization for bivariate splits (ex. contour plots)
+> This codebase is an efficient implementation of the algorithms in "Empowering Decision Trees via Shape Function Branching." See the [ROADMAP](ROADMAP.md) for implementation status and the [canonical research code](https://github.com/optimal-uoft/Empowering-DTs-via-Shape-Functions) for the paper's original implementation.
 
 ## Installation
 

ROADMAP.md

@@ -17,6 +17,12 @@
 - [x] **NaN routing at predict**: if training saw missing at that split, follow the stored direction; otherwise route to the majority child.
 
 ## v0.3.0
-- [ ] multioutput support
-- [ ] Shape$^2$CART
-- [ ] Shape$^2$CART Random Forest Ensembling  
+- [x] multioutput support
+- [x] Opt-in Shape$^2$CART for SGT estimators, including continuous/categorical
+  pairs, joint missing routing, and multiway branching ([tutorial](https://sgtlearn.readthedocs.io/en/latest/tutorials/bivariate-branching.html))
+- [x] Shape$^2$CART Random Forest Ensembling
+- [x] Pair-aware TAO refinement ([#48](https://github.com/optimal-uoft/sgtlearn/issues/48))
+- [x] Shape$^2$CART routing heatmap visualization ([#28](https://github.com/optimal-uoft/sgtlearn/issues/28))
+
+See the implementation specification in [#42](https://github.com/optimal-uoft/sgtlearn/issues/42)
+and the umbrella issue [#27](https://github.com/optimal-uoft/sgtlearn/issues/27).

cpp/CMakeLists.txt

@@ -77,6 +77,10 @@ add_library(sgtlearn_core STATIC
         src/Splitters/categorical/CategoricalRegressionSplitter.h
         src/Splitters/categorical/CategoricalRegressionSplitter.cpp
         src/Discretizers/univariate/UnivariateClassificationDiscretizer.cpp
+        src/Discretizers/pair/PairClassificationDiscretizer.h
+        src/Discretizers/pair/PairClassificationDiscretizer.cpp
+        src/Discretizers/pair/PairRegressionDiscretizer.h
+        src/Discretizers/pair/PairRegressionDiscretizer.cpp
         src/Splitters/univariate/SquaredErrorSplitter.h
         src/Splitters/univariate/SquaredErrorSplitter.cpp
         src/Discretizers/univariate/UnivariateRegressionDiscretizer.cpp
@@ -250,4 +254,4 @@ if (SGTLEARN_BUILD_TESTS)
 
 endif ()
 
-# endregion
+# endregion

cpp/bindings/ShapeGeneralizedTrees.cpp

@@ -52,15 +52,17 @@ PYBIND11_MODULE(ShapeGeneralizedTrees, m) {
                        size_t coordinate_descent_max_iters,
                        size_t coordinate_descent_patience,
                        bool coordinate_descent_smart_init, uint64_t random_state,
-                       py::object max_features) {
+                       py::object max_features, size_t pairwise_candidates,
+                       double pairwise_penalty) {
              return ClassificationShapeGeneralizedTreePy(
                  std::move(criterion), std::move(num_classes), num_partitions,
                  outer_min_leaf_size, outer_min_gain_split, outer_max_depth,
                  outer_max_leaf_nodes, inner_min_leaf_size, inner_min_gain_split,
                  inner_max_depth, inner_max_leaf_nodes,
                  coordinate_descent_max_iters, coordinate_descent_patience,
                  coordinate_descent_smart_init, random_state,
-                 std::move(max_features));
+                 std::move(max_features), pairwise_candidates,
+                 pairwise_penalty);
            }),
            py::arg("criterion") = "gini", py::arg("num_classes"),
            py::arg("num_partitions") = 2,
@@ -76,7 +78,9 @@ PYBIND11_MODULE(ShapeGeneralizedTrees, m) {
            py::arg("coordinate_descent_patience") = 5,
            py::arg("coordinate_descent_smart_init") = true,
            py::arg("random_state") = 42,
-           py::arg("max_features") = py::none())
+           py::arg("max_features") = py::none(),
+           py::arg("pairwise_candidates") = 0,
+           py::arg("pairwise_penalty") = 0.0)
       .def("fit", &ClassificationShapeGeneralizedTreePy::fit, py::arg("X"),
            py::arg("y"), py::arg("sample_weight") = py::none(),
            py::arg("features"),
@@ -104,6 +108,8 @@ PYBIND11_MODULE(ShapeGeneralizedTrees, m) {
           &ClassificationShapeGeneralizedTreePy::classesPerOutput)
       .def_property_readonly(
           "is_fitted", &ClassificationShapeGeneralizedTreePy::isFitted)
+      .def_property_readonly(
+          "has_pair_nodes", &ClassificationShapeGeneralizedTreePy::hasPairNodes)
       .def_property_readonly(
           "feature_importance",
           &ClassificationShapeGeneralizedTreePy::featureImportance,
@@ -122,14 +128,16 @@ PYBIND11_MODULE(ShapeGeneralizedTrees, m) {
                        size_t coordinate_descent_max_iters,
                        size_t coordinate_descent_patience,
                        bool coordinate_descent_smart_init, uint64_t random_state,
-                       py::object max_features) {
+                       py::object max_features, size_t pairwise_candidates,
+                       double pairwise_penalty) {
              return RegressionShapeGeneralizedTreePy(
                  std::move(criterion), num_partitions, outer_min_leaf_size,
                  outer_min_gain_split, outer_max_depth, outer_max_leaf_nodes,
                  inner_min_leaf_size, inner_min_gain_split, inner_max_depth,
                  inner_max_leaf_nodes, coordinate_descent_max_iters,
                  coordinate_descent_patience, coordinate_descent_smart_init,
-                 random_state, std::move(max_features));
+                 random_state, std::move(max_features), pairwise_candidates,
+                 pairwise_penalty);
            }),
            py::arg("criterion") = "squared_error",
            py::arg("num_partitions") = 2, py::arg("outer_min_leaf_size") = 1,
@@ -142,6 +150,8 @@ PYBIND11_MODULE(ShapeGeneralizedTrees, m) {
            py::arg("coordinate_descent_patience") = 5,
            py::arg("coordinate_descent_smart_init") = true,
            py::arg("random_state") = 42, py::arg("max_features") = py::none(),
+           py::arg("pairwise_candidates") = 0,
+           py::arg("pairwise_penalty") = 0.0,
            R"(Regression tree: inner bins are round-robin seeded. ``squared_error`` runs
 coordinate descent and keeps the map only if branch MSE improves clearly vs the seed;
 otherwise the snapshot is restored and the branch objective is rebuilt.
@@ -164,6 +174,8 @@ is accepted for API parity with ClassificationShapeGeneralizedTree but ignored.)
                              &RegressionShapeGeneralizedTreePy::nOutputs)
       .def_property_readonly("is_fitted",
                              &RegressionShapeGeneralizedTreePy::isFitted)
+      .def_property_readonly(
+          "has_pair_nodes", &RegressionShapeGeneralizedTreePy::hasPairNodes)
       .def_property_readonly(
           "feature_importance",
           &RegressionShapeGeneralizedTreePy::featureImportance,

cpp/bindings/TreeAlternatingOptimization.cpp

@@ -23,6 +23,7 @@
 #include "algorithms/TAO/TreeAlternatingOptimization.h"
 
 #include <armadillo>
+#include <cmath>
 #include <cstddef>
 #include <memory>
 #include <optional>
@@ -130,13 +131,17 @@ TaoRunContext TaoRunContext::make(py::object tree, const py::array &X,
 void TreeAlternatingOptimization(py::object tree, const py::array &X,
                                  const py::array &y,
                                  py::object sample_weight = py::none(),
-                                 size_t n_runs = 10, double lambda_ = 0.0) {
+                                 size_t n_runs = 10, double lambda_ = 0.0,
+                                 double tao_pair_scale = 1.1) {
+  if (!std::isfinite(tao_pair_scale) || tao_pair_scale < 0.0)
+    throw std::invalid_argument(
+        "tao_pair_scale must be finite and non-negative");
   TaoRunContext ctx = TaoRunContext::make(tree, X, y, sample_weight);
   if (!ctx.isFitted())
     throw std::logic_error("TreeAlternatingOptimization: model is not fitted");
 
   py::gil_scoped_release release;
-  tao::optimize(ctx.adapter(), n_runs, lambda_);
+  tao::optimize(ctx.adapter(), n_runs, lambda_, tao_pair_scale);
 }
 
 } // namespace
@@ -159,11 +164,13 @@ PYBIND11_MODULE(TreeAlternatingOptimization, m) {
   m.def("TreeAlternatingOptimization", &TreeAlternatingOptimization,
         py::arg("tree"), py::arg("X"), py::arg("y"),
         py::arg("sample_weight") = py::none(), py::arg("n_runs") = 10,
-        py::arg("lambda_") = 0.0,
+        py::arg("lambda_") = 0.0, py::arg("tao_pair_scale") = 1.1,
         "Refine a fitted ClassificationShapeGeneralizedTree or "
         "RegressionShapeGeneralizedTree in place. X is "
         "(n_samples, n_features) float32; y is 1-D class labels (uint) or "
         "float targets matching the tree type. Runs up to n_runs bottom-up "
-        "sweeps; lambda_ penalizes non-constant routing splits by "
-        "lambda_ * totalSampleWeight in weighted reward units.");
+        "sweeps. In weighted reward units, lambda_ penalizes single-feature "
+        "routers by lambda_ * nodeSampleCount and pair routers by "
+        "tao_pair_scale * lambda_ * nodeSampleCount; dummy routers are "
+        "unpenalized.");
 }