Update stale docstrings, docs and tutorial outputs (#77)

README.md

@@ -92,7 +92,7 @@ installs the `sgtlearn` package plus native modules into `.venv`.
 
 ```bash
 pip install .
-pip install ".[dev]"   # dev extras (pytest, scikit-learn) only if needed
+pip install ".[dev]"   # dev extras (pytest, ruff, black, mypy, ...) only if needed
 ```
 
 > **Anaconda users:** Do not bootstrap the venv from an Anaconda Python.
@@ -102,27 +102,26 @@ pip install ".[dev]"   # dev extras (pytest, scikit-learn) only if needed
 > non-Anaconda Python — e.g. `uv venv --python 3.12 .venv` (downloads a
 > hermetic CPython), `pyenv`, or your distro's `python3`.
 
-## Build Workflow (scikit-build + CMake)
+## Build Workflow (scikit-build-core + CMake)
 
 `pip install .` drives this build path:
 
 1. `pyproject.toml` selects `scikit_build_core.build` as the backend.
 2. CMake is configured from `cpp/CMakeLists.txt`.
 3. Each file in `cpp/bindings/*.cpp` becomes one pybind11 module target.
-4. After each module is built, `pybind11-stubgen` generates a matching `.pyi`.
-5. The `.pyi` is generated and installed in the same location as the module `.so`.
+4. After each module is built, `pybind11-stubgen` generates a matching `.pyi` (best effort; a missing or failing `pybind11-stubgen` does not break the build).
+5. Any generated `.pyi` is installed in the same location as the module `.so`.
 
 ## C++ Folder Conventions
 
-- `cpp/include/sgtlearn/`: public headers for the core C++ API.
-- `cpp/src/`: internal C++ implementation for the core library.
+- `cpp/src/`: C++ headers and implementation for the core library (the public include directory of `sgtlearn_core`).
 - `cpp/bindings/`: pybind11 binding entrypoints; one `.cpp` file maps to one Python extension module.
 - `cpp/tests/`: C++ unit tests consumed by the `cpp_tests` executable target.
 
 ## CMake Targets
 
 - `sgtlearn_core` (static library): shared C++ logic used by Python modules and tests.
-- `<module_name>` (pybind11 module, one per file in `cpp/bindings/`): compiled extension modules installed into the package.
+- `<module_name>` (pybind11 module, one per file in `cpp/bindings/`): compiled extension modules installed as top-level modules next to the `sgtlearn` package.
 - `cpp_tests` (Catch2 executable): optional C++ test target, controlled by:
   - `-DSGTLEARN_BUILD_TESTS=ON` (build C++ tests)
   - `-DSGTLEARN_BUILD_TESTS=OFF` (default for `pip install`; the CMake option itself defaults to `ON`, but `pyproject.toml` overrides this so wheels don't ship test binaries)

ROADMAP.md

@@ -14,7 +14,7 @@
 - [ ] Exporting to Graphviz, planned for v0.4.0 ([#69](https://github.com/optimal-uoft/sgtlearn/issues/69))
 - [x] Feature importances matching scikit-learn's API for SGTs and SGT$_K$
 - [x] Adding TAO Refinement
-- [x] **Sklearn-style NaN support** (replaces current tail-bin placeholder): split search uses finite values only; each candidate is scored with missing sent left vs right—including an explicit missing-vs-non-missing split—with the winning direction stored per node (ties → right).
+- [x] **Sklearn-style NaN support**: split search uses finite values only, and each univariate node learns which child receives missing values.
 - [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
@@ -29,7 +29,13 @@ See the implementation specification in [#42](https://github.com/optimal-uoft/sg
 and the umbrella issue [#27](https://github.com/optimal-uoft/sgtlearn/issues/27).
 
 
+## v0.3.1
+- [x] Regularized best-first outer growth with `branching_penalty` and strict `max_leaf_nodes` budgets for multiway splits
+
+## Unreleased
+- [x] Text export of fitted trees with `export_text`
+
 ## v1.0.0
-- [ ] Cross-Feature Tree Binning (Similar to [DPDT](https://github.com/KohlerHECTOR/DPDTreeEstimator) 
+- [ ] Cross-Feature Tree Binning (Similar to [DPDT](https://github.com/KohlerHECTOR/DPDTreeEstimator)) 
 - [ ] Boosting
 - [ ] Default Optuna Support

cpp/README.md

@@ -4,17 +4,15 @@ This directory contains the C++ core library and pybind11 module definitions use
 
 ## Layout Conventions
 
-- `include/sgtlearn/`: public headers for `sgtlearn_core`.
-- `src/`: internal implementation for `sgtlearn_core`.
+- `src/`: headers (`.h`, `.tpp`) and sources (`.cpp`) for `sgtlearn_core`; `src/` is its public include directory.
 - `bindings/`: pybind11 module sources; each `*.cpp` file becomes one extension module target.
-- `../tests/cpp/`: C++ test files used by `cpp_tests`.
+- `tests/`: C++ test files used by `cpp_tests`.
 
 | Directory | Content Type | Visibility | Target |
 |:--|:--|:--|:--|
-| `include/sgtlearn/` | Public headers (`.hpp`) | Public | `sgtlearn_core` |
-| `src/` | Internal logic (`.cpp`) | Private | `sgtlearn_core` |
+| `src/` | Headers (`.h`, `.tpp`) and logic (`.cpp`) | Public | `sgtlearn_core` |
 | `bindings/` | Python wrappers (`.cpp`) | Private | `<module_name>` |
-| `../tests/cpp/` | Unit tests (`.cpp`) | Private | `cpp_tests` |
+| `tests/` | Unit tests (`.cpp`) | Private | `cpp_tests` |
 
 ## Targets Defined in `CMakeLists.txt`
 
@@ -33,10 +31,10 @@ This directory contains the C++ core library and pybind11 module definitions use
 For each pybind11 module target:
 
 1. the extension module (`.so`) is built;
-2. `pybind11_stubgen` runs as a post-build step;
+2. if `pybind11_stubgen` was importable at configure time, it runs as a non-fatal post-build step;
 3. the generated `.pyi` is written to the same output directory as that module.
 
-The install rules place the `.so` and matching `.pyi` together in the same install destination.
+The install rules place the `.so` and, when it was generated, the matching `.pyi` together in the same install destination.
 
 ## Common Commands
 
@@ -64,7 +62,7 @@ Skip C++ test target:
 cmake -S cpp -B build -DSGTLEARN_BUILD_TESTS=OFF
 ```
 
-Most developers should use `pip install .` from the repository root, which drives this CMake configuration via scikit-build.
+Most developers should use `pip install .` from the repository root, which drives this CMake configuration via scikit-build-core (with `-DSGTLEARN_BUILD_TESTS=OFF`).
 
 ## CLion Notes
 

cpp/bindings/Discretizers.cpp

@@ -1,6 +1,6 @@
 /**
  * @file Discretizers.cpp
- * @brief pybind11 module ``Discretizers``: univariate classification/regression discretizers.
+ * @brief pybind11 module ``Discretizers``: univariate and categorical classification/regression discretizers.
  */
 
 #include <carma>

cpp/bindings/ShapeGeneralizedTrees.cpp

@@ -86,7 +86,9 @@ PYBIND11_MODULE(ShapeGeneralizedTrees, m) {
            py::arg("features"),
            "Fit the routing tree. X is (n_samples, n_features) float32; y is "
            "uint class labels, 1-D (n_samples,) or 2-D (n_samples, n_outputs). "
-           "Optional sample_weight is 1-D float32.")
+           "Optional sample_weight is 1-D float32. features is a list of dicts "
+           "with 'type' ('continuous' or 'categorical') and 'indices' (the X "
+           "columns of each logical feature).")
       .def("predict", &ClassificationShapeGeneralizedTreePy::predict,
            py::arg("X"),
            "Predict class labels for X (shape (n_samples, n_features)). Returns "
@@ -114,7 +116,9 @@ PYBIND11_MODULE(ShapeGeneralizedTrees, m) {
           "feature_importance",
           &ClassificationShapeGeneralizedTreePy::featureImportance,
           "Normalized feature importances aligned with the features sequence "
-          "passed to fit. Available only after training.")
+          "passed to fit (all zeros if the tree has no split). Computed at fit "
+          "and not updated by TreeAlternatingOptimization. Available only "
+          "after training.")
       .def("tree_export", &ClassificationShapeGeneralizedTreePy::tree_export,
            "Return a flat snapshot of the fitted tree as a Python dict.");
 
@@ -152,16 +156,19 @@ PYBIND11_MODULE(ShapeGeneralizedTrees, m) {
            py::arg("pairwise_candidates") = 0,
            py::arg("pairwise_penalty") = 0.0,
            py::arg("branching_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.
-``absolute_error`` / ``mae`` skip coordinate descent by default.)")
+           R"(Regression tree. For each branch count, ``squared_error`` seeds bins with
+weighted k-means over bin means, starts from the inner tree's root split instead
+when that is no worse, and refines with coordinate descent, keeping the best
+feasible assignment per arity. ``absolute_error`` / ``mae`` score the root split
+and round-robin seeds and skip coordinate descent unless ``SGTLEARN_MAE_CD=1``.)")
       .def("fit", &RegressionShapeGeneralizedTreePy::fit, py::arg("X"),
            py::arg("y"), py::arg("sample_weight") = py::none(),
            py::arg("features"),
            "Fit the routing tree. X is (n_samples, n_features) float32; y is "
            "float32 targets, 1-D (n_samples,) or 2-D (n_samples, n_outputs). "
-           "Optional sample_weight is 1-D float32.")
+           "Optional sample_weight is 1-D float32. features is a list of dicts "
+           "with 'type' ('continuous' or 'categorical') and 'indices' (the X "
+           "columns of each logical feature).")
       .def("predict", &RegressionShapeGeneralizedTreePy::predict, py::arg("X"),
            "Predict targets for X. Returns (n_samples,) for a single output or "
            "(n_samples, n_outputs) for multi-output.")
@@ -179,7 +186,9 @@ otherwise the snapshot is restored and the branch objective is rebuilt.
           "feature_importance",
           &RegressionShapeGeneralizedTreePy::featureImportance,
           "Normalized feature importances aligned with the features sequence "
-          "passed to fit. Available only after training.")
+          "passed to fit (all zeros if the tree has no split). Computed at fit "
+          "and not updated by TreeAlternatingOptimization. Available only "
+          "after training.")
       .def_property_readonly(
           "leaf_regression_stats",
           &RegressionShapeGeneralizedTreePy::leafRegressionStats)

cpp/bindings/TreeAlternatingOptimization.cpp

@@ -167,8 +167,10 @@ PYBIND11_MODULE(TreeAlternatingOptimization, m) {
         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 "
+        "(n_samples, n_features) float32; y is class labels (uint) or float "
+        "targets matching the tree type, 1-D (n_samples,) or 2-D "
+        "(n_samples, n_outputs) with the fitted tree's number of outputs. "
+        "Optional sample_weight is 1-D float32. Runs up to n_runs bottom-up "
         "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 "

cpp/src/Estimators/RegressionShapeGeneralizedTree.cpp

@@ -2,8 +2,10 @@
  * @file Estimators/RegressionShapeGeneralizedTree.cpp
  * @brief Training, child partitioning, and prediction for the regression
  *        shape-generalized tree. Per-node inner fit: discretize -> search
- *        partition counts k in [2, numPartitions] -> coordinate descent (MSE)
- *        or round-robin (MAE) -> record best branch.
+ *        partition counts k in [2, numPartitions] -> k-means / root-split seed
+ *        plus coordinate descent (MSE), or root / round-robin seeds without
+ *        coordinate descent unless SGTLEARN_MAE_CD=1 (MAE) -> record best
+ *        branch per arity.
  */
 
 #include <memory>

cpp/src/Estimators/RegressionShapeGeneralizedTree.h

@@ -29,9 +29,9 @@
  * - **Outer growth** (`OuterTreeBuilder`): best-first regularized expansion;
  *   per-node split search and child creation via lambdas in ``fit``; commit
  *   step remains a local lambda in ``fit``.
- * - **Per-node split search** (``fit`` lambdas): for each
- *   discretize -> round-robin bin-to-partition seed -> ``coordinateDescent`` on
- *   ``SquaredError`` only; ``AbsoluteError`` keeps the round-robin map (no CD).
+ * - **Per-node split search** (``fit`` lambdas): for each candidate feature,
+ *   discretize -> seed bin-to-partition maps -> ``coordinateDescent``
+ *   (``SquaredError``; ``AbsoluteError`` only when ``SGTLEARN_MAE_CD=1``).
  * - **Leaf state**: per-leaf mean (squared error) or median (absolute error)
  *   plus optional ``[sum y, sum y^2]`` stats for squared error.
  * - **Inference**: `predict` walks `childIndices_` using
@@ -40,17 +40,21 @@
  * At every outer-tree node, for each candidate feature:
  *   1. **Discretize**: train univariate regression discretizer on the node's
  *      samples (per-bin stats and training column indices).
- *   2. **Initial assignment**: round-robin by discretizer bin index
- *      (``b % numPartitions``); regression does not use k-means seeding.
+ *   2. **Initial assignment** (for each k in [2, numPartitions]):
+ *      ``SquaredError`` seeds with weighted k-means over bin means and starts
+ *      from the inner tree's root split instead when that is no worse;
+ *      ``AbsoluteError`` scores a round-robin seed (``b % k``) and starts from
+ *      the root split.
  *   3. **Refinement**: ``SquaredError`` runs coordinate descent on the
- *      bin-to-partition map; if the post-CD objective clearly worsens vs the
- *      seed, restore the assignment snapshot and rebuild. ``AbsoluteError``
- *      keeps the round-robin seed (no CD).
+ *      bin-to-partition map; ``AbsoluteError`` skips it unless
+ *      ``SGTLEARN_MAE_CD=1``. Every scored assignment is observed and the
+ *      lowest-impurity feasible one is kept per occupied arity.
  *
  * The best-scoring feature wins; its inner discretizer + bin->partition
- * mapping become the routing rule for that node, producing `numPartitions`
- * children. Inner-node fitting matches the Python `BranchingTree` pattern;
- * the outer loop uses `OuterTreeBuilder` like Python's heap over
+ * mapping become the routing rule for that node, producing at most
+ * `numPartitions` children (empty partitions are compacted away).
+ * Inner-node fitting matches the Python `BranchingTree` pattern; the outer
+ * loop uses `OuterTreeBuilder` like Python's heap over
  * `best_impurity_decrease`.
  *
  * Inputs use Armadillo's column-major convention: X has shape

docs/api/ensemble.rst

@@ -18,19 +18,20 @@ Multi-output ``y`` is supported the same way as for single trees
 Forests accept ``branching_penalty`` (default ``0.0``) and the Shape²CART
 options ``pairwise_candidates`` and ``pairwise_penalty`` and forward them to
 every base estimator. Outer growth uses the same total sample-mass-weighted,
-output-averaged score, constant ``alpha``/``lambda``/``gamma`` costs, strict
+output-averaged score, constant ``min_impurity_decrease``,
+``branching_penalty`` and ``pairwise_penalty`` costs, strict
 actual-arity leaf budget, fixed-epsilon acceptance, and always-best-first order
 as single trees. Pair candidates are restricted to each node's
 ``max_features`` logical-feature subset. See :doc:`estimators` for the full
 scoring, candidate-count, pair-screening, categorical and joint missing-routing
 semantics, and the feature-importance warning for pair nodes.
 See :doc:`../tutorials/bivariate-branching` for a worked classifier example.
 
-:attr:`~sgtlearn.ensemble.RandomSGForestClassifier.mean_feature_importances_`
-and :attr:`~sgtlearn.ensemble.RandomSGForestClassifier.std_feature_importance_`
+:attr:`~sgtlearn.RandomSGForestClassifier.mean_feature_importances_`
+and :attr:`~sgtlearn.RandomSGForestClassifier.std_feature_importance_`
 (and the regressor counterparts) summarize per-tree
 :attr:`~sgtlearn.SGTClassifier.feature_importances_` across the forest,
-aligned with the shared :attr:`processed_features_`.
+aligned with the shared :attr:`~sgtlearn.RandomSGForestClassifier.processed_features_`.
 They are unavailable if any base tree has undergone positive-run TAO refinement.
 
 RandomSGForestClassifier
@@ -40,6 +41,7 @@ RandomSGForestClassifier
    :members:
    :inherited-members:
    :show-inheritance:
+   :exclude-members: mean_feature_importances_, std_feature_importance_
 
 RandomSGForestRegressor
 -----------------------
@@ -48,3 +50,4 @@ RandomSGForestRegressor
    :members:
    :inherited-members:
    :show-inheritance:
+   :exclude-members: mean_feature_importances_, std_feature_importance_