perpetual 1.0.35__cp314-cp314-macosx_11_0_arm64.whl

perpetual/booster.py

@@ -0,0 +1,1915 @@
+import inspect
+import json
+import warnings
+from types import FunctionType
+from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Union, cast
+
+import numpy as np
+from typing_extensions import Self
+
+from perpetual.data import Node
+from perpetual.perpetual import (
+    MultiOutputBooster as CrateMultiOutputBooster,  # type: ignore
+)
+from perpetual.perpetual import (
+    PerpetualBooster as CratePerpetualBooster,  # type: ignore
+)
+from perpetual.serialize import BaseSerializer, ObjectSerializer
+from perpetual.types import BoosterType, MultiOutputBoosterType
+from perpetual.utils import (
+    CONTRIBUTION_METHODS,
+    convert_input_array,
+    convert_input_frame,
+    convert_input_frame_columnar,
+    transform_input_frame,
+    transform_input_frame_columnar,
+    type_df,
+)
+
+
+class PerpetualBooster:
+    # Define the metadata parameters
+    # that are present on all instances of this class
+    # this is useful for parameters that should be
+    # attempted to be loaded in and set
+    # as attributes on the booster after it is loaded.
+    metadata_attributes: Dict[str, BaseSerializer] = {
+        "feature_names_in_": ObjectSerializer(),
+        "n_features_": ObjectSerializer(),
+        "feature_importance_method": ObjectSerializer(),
+        "cat_mapping": ObjectSerializer(),
+        "classes_": ObjectSerializer(),
+    }
+
+    def __init__(
+        self,
+        *,
+        objective: Union[
+            str, Tuple[FunctionType, FunctionType, FunctionType]
+        ] = "LogLoss",
+        budget: float = 0.5,
+        num_threads: Optional[int] = None,
+        monotone_constraints: Union[Dict[Any, int], None] = None,
+        force_children_to_bound_parent: bool = False,
+        missing: float = np.nan,
+        allow_missing_splits: bool = True,
+        create_missing_branch: bool = False,
+        terminate_missing_features: Optional[Iterable[Any]] = None,
+        missing_node_treatment: str = "None",
+        log_iterations: int = 0,
+        feature_importance_method: str = "Gain",
+        quantile: Optional[float] = None,
+        reset: Optional[bool] = None,
+        categorical_features: Union[Iterable[int], Iterable[str], str, None] = "auto",
+        timeout: Optional[float] = None,
+        iteration_limit: Optional[int] = None,
+        memory_limit: Optional[float] = None,
+        stopping_rounds: Optional[int] = None,
+        max_bin: int = 256,
+        max_cat: int = 1000,
+    ):
+        """
+        Gradient Boosting Machine with Perpetual Learning.
+
+        A self-generalizing gradient boosting machine that doesn't need hyperparameter optimization.
+        It automatically finds the best configuration based on the provided budget.
+
+        Parameters
+        ----------
+        objective : str or tuple, default="LogLoss"
+            Learning objective function to be used for optimization. Valid options are:
+
+            - "LogLoss": logistic loss for binary classification.
+            - "SquaredLoss": squared error for regression.
+            - "QuantileLoss": quantile error for quantile regression.
+            - "HuberLoss": Huber loss for robust regression.
+            - "AdaptiveHuberLoss": adaptive Huber loss for robust regression.
+            - "ListNetLoss": ListNet loss for ranking.
+            - custom objective: a tuple of (grad, hess, init) functions.
+
+        budget : float, default=0.5
+            A positive number for fitting budget. Increasing this number will more likely result
+            in more boosting rounds and increased predictive power.
+        num_threads : int, optional
+            Number of threads to be used during training and prediction.
+        monotone_constraints : dict, optional
+            Constraints to enforce a specific relationship between features and target.
+            Keys are feature indices or names, values are -1, 1, or 0.
+        force_children_to_bound_parent : bool, default=False
+            Whether to restrict children nodes to be within the parent's range.
+        missing : float, default=np.nan
+            Value to consider as missing data.
+        allow_missing_splits : bool, default=True
+            Whether to allow splits that separate missing from non-missing values.
+        create_missing_branch : bool, default=False
+            Whether to create a separate branch for missing values (ternary trees).
+        terminate_missing_features : iterable, optional
+            Features for which missing branches will always be terminated if
+            ``create_missing_branch`` is True.
+        missing_node_treatment : str, default="None"
+            How to handle weights for missing nodes if ``create_missing_branch`` is True.
+            Options: "None", "AssignToParent", "AverageLeafWeight", "AverageNodeWeight".
+        log_iterations : int, default=0
+            Logging frequency (every N iterations). 0 disables logging.
+        feature_importance_method : str, default="Gain"
+            Method for calculating feature importance. Options: "Gain", "Weight", "Cover",
+            "TotalGain", "TotalCover".
+        quantile : float, optional
+            Target quantile for quantile regression (objective="QuantileLoss").
+        reset : bool, optional
+            Whether to reset the model or continue training on subsequent calls to fit.
+        categorical_features : str or iterable, default="auto"
+            Feature indices or names to treat as categorical.
+        timeout : float, optional
+            Time limit for fitting in seconds.
+        iteration_limit : int, optional
+            Maximum number of boosting iterations.
+        memory_limit : float, optional
+            Memory limit for training in GB.
+        stopping_rounds : int, optional
+            Early stopping rounds.
+        max_bin : int, default=256
+            Maximum number of bins for feature discretization.
+        max_cat : int, default=1000
+            Maximum unique categories before a feature is treated as numerical.
+
+        Attributes
+        ----------
+        feature_names_in_ : list of str
+            Names of features seen during :meth:`fit`.
+        n_features_ : int
+            Number of features seen during :meth:`fit`.
+        classes_ : list
+            Class labels for classification tasks.
+        feature_importances_ : ndarray of shape (n_features,)
+            Feature importances calculated via ``feature_importance_method``.
+
+        See Also
+        --------
+        perpetual.sklearn.PerpetualClassifier : Scikit-learn compatible classifier.
+        perpetual.sklearn.PerpetualRegressor : Scikit-learn compatible regressor.
+
+        Examples
+        --------
+        Basic usage for binary classification:
+
+        >>> from perpetual import PerpetualBooster
+        >>> from sklearn.datasets import make_classification
+        >>> X, y = make_classification(n_samples=1000, n_features=20)
+        >>> model = PerpetualBooster(objective="LogLoss")
+        >>> model.fit(X, y)
+        >>> preds = model.predict(X[:5])
+        """
+
+        terminate_missing_features_ = (
+            set() if terminate_missing_features is None else terminate_missing_features
+        )
+        monotone_constraints_ = (
+            {} if monotone_constraints is None else monotone_constraints
+        )
+
+        if isinstance(objective, str):
+            self.objective = objective
+            self.loss = None
+            self.grad = None
+            self.init = None
+        else:
+            self.objective = None
+            self.loss = objective[0]
+            self.grad = objective[1]
+            self.init = objective[2]
+        self.budget = budget
+        self.num_threads = num_threads
+        self.monotone_constraints = monotone_constraints_
+        self.force_children_to_bound_parent = force_children_to_bound_parent
+        self.allow_missing_splits = allow_missing_splits
+        self.missing = missing
+        self.create_missing_branch = create_missing_branch
+        self.terminate_missing_features = terminate_missing_features_
+        self.missing_node_treatment = missing_node_treatment
+        self.log_iterations = log_iterations
+        self.feature_importance_method = feature_importance_method
+        self.quantile = quantile
+        self.reset = reset
+        self.categorical_features = categorical_features
+        self.timeout = timeout
+        self.iteration_limit = iteration_limit
+        self.memory_limit = memory_limit
+        self.stopping_rounds = stopping_rounds
+        self.max_bin = max_bin
+        self.max_cat = max_cat
+
+        booster = CratePerpetualBooster(
+            objective=self.objective,
+            budget=self.budget,
+            max_bin=self.max_bin,
+            num_threads=self.num_threads,
+            monotone_constraints=dict(),
+            force_children_to_bound_parent=self.force_children_to_bound_parent,
+            missing=self.missing,
+            allow_missing_splits=self.allow_missing_splits,
+            create_missing_branch=self.create_missing_branch,
+            terminate_missing_features=set(),
+            missing_node_treatment=self.missing_node_treatment,
+            log_iterations=self.log_iterations,
+            quantile=self.quantile,
+            reset=self.reset,
+            categorical_features=set(),
+            timeout=self.timeout,
+            iteration_limit=self.iteration_limit,
+            memory_limit=self.memory_limit,
+            stopping_rounds=self.stopping_rounds,
+            loss=self.loss,
+            grad=self.grad,
+            init=self.init,
+        )
+        self.booster = cast(BoosterType, booster)
+
+    def fit(self, X, y, sample_weight=None, group=None) -> Self:
+        """
+        Fit the gradient booster on a provided dataset.
+
+        Parameters
+        ----------
+        X : array-like of shape (n_samples, n_features)
+            Training data. Can be a Polars or Pandas DataFrame, or a 2D Numpy array.
+            Polars DataFrames use a zero-copy columnar path for efficiency.
+        y : array-like of shape (n_samples,) or (n_samples, n_targets)
+            Target values.
+        sample_weight : array-like of shape (n_samples,), optional
+            Individual weights for each sample. If None, all samples are weighted equally.
+        group : array-like, optional
+            Group labels for ranking objectives.
+
+        Returns
+        -------
+        self : object
+            Returns self.
+        """
+
+        # Check if input is a Polars DataFrame for zero-copy columnar path
+        is_polars = type_df(X) == "polars_df"
+
+        if is_polars:
+            # Use columnar path for Polars DataFrames (true zero-copy)
+            (
+                features_,
+                columns,  # list of 1D arrays instead of flat_data
+                masks,
+                rows,
+                cols,
+                categorical_features_,
+                cat_mapping,
+            ) = convert_input_frame_columnar(X, self.categorical_features, self.max_cat)
+        else:
+            # Use existing flat path for pandas and numpy
+            (
+                features_,
+                flat_data,
+                rows,
+                cols,
+                categorical_features_,
+                cat_mapping,
+            ) = convert_input_frame(X, self.categorical_features, self.max_cat)
+
+        self.n_features_ = cols
+        self.cat_mapping = cat_mapping
+        self.feature_names_in_ = features_
+
+        y_, classes_ = convert_input_array(y, self.objective, is_target=True)
+        self.classes_ = np.array(classes_).tolist()
+
+        if sample_weight is None:
+            sample_weight_ = None
+        else:
+            sample_weight_, _ = convert_input_array(sample_weight, self.objective)
+
+        if group is None:
+            group_ = None
+        else:
+            group_, _ = convert_input_array(group, self.objective, is_int=True)
+
+        # Convert the monotone constraints into the form needed
+        # by the rust code.
+        crate_mc = self._standardize_monotonicity_map(X)
+        crate_tmf = self._standardize_terminate_missing_features(X)
+
+        if (len(classes_) <= 2) or (
+            len(classes_) > 1 and self.objective == "SquaredLoss"
+        ):
+            booster = CratePerpetualBooster(
+                objective=self.objective,
+                budget=self.budget,
+                max_bin=self.max_bin,
+                num_threads=self.num_threads,
+                monotone_constraints=crate_mc,
+                force_children_to_bound_parent=self.force_children_to_bound_parent,
+                missing=self.missing,
+                allow_missing_splits=self.allow_missing_splits,
+                create_missing_branch=self.create_missing_branch,
+                terminate_missing_features=crate_tmf,
+                missing_node_treatment=self.missing_node_treatment,
+                log_iterations=self.log_iterations,
+                quantile=self.quantile,
+                reset=self.reset,
+                categorical_features=categorical_features_,
+                timeout=self.timeout,
+                iteration_limit=self.iteration_limit,
+                memory_limit=self.memory_limit,
+                stopping_rounds=self.stopping_rounds,
+                loss=self.loss,
+                grad=self.grad,
+                init=self.init,
+            )
+            self.booster = cast(BoosterType, booster)
+        else:
+            booster = CrateMultiOutputBooster(
+                n_boosters=len(classes_),
+                objective=self.objective,
+                budget=self.budget,
+                max_bin=self.max_bin,
+                num_threads=self.num_threads,
+                monotone_constraints=crate_mc,
+                force_children_to_bound_parent=self.force_children_to_bound_parent,
+                missing=self.missing,
+                allow_missing_splits=self.allow_missing_splits,
+                create_missing_branch=self.create_missing_branch,
+                terminate_missing_features=crate_tmf,
+                missing_node_treatment=self.missing_node_treatment,
+                log_iterations=self.log_iterations,
+                quantile=self.quantile,
+                reset=self.reset,
+                categorical_features=categorical_features_,
+                timeout=self.timeout,
+                iteration_limit=self.iteration_limit,
+                memory_limit=self.memory_limit,
+                stopping_rounds=self.stopping_rounds,
+                loss=self.loss,
+                grad=self.grad,
+                init=self.init,
+            )
+            self.booster = cast(MultiOutputBoosterType, booster)
+
+        self._set_metadata_attributes("n_features_", self.n_features_)
+        self._set_metadata_attributes("cat_mapping", self.cat_mapping)
+        self._set_metadata_attributes("feature_names_in_", self.feature_names_in_)
+        self._set_metadata_attributes(
+            "feature_importance_method", self.feature_importance_method
+        )
+        self._set_metadata_attributes("classes_", self.classes_)
+
+        self.categorical_features = categorical_features_
+
+        if is_polars:
+            # Use columnar fit for Polars (zero-copy)
+            self.booster.fit_columnar(
+                columns=columns,
+                masks=masks,
+                rows=rows,
+                y=y_,
+                sample_weight=sample_weight_,  # type: ignore
+                group=group_,
+            )
+        else:
+            # Use standard fit for pandas/numpy
+            self.booster.fit(
+                flat_data=flat_data,
+                rows=rows,
+                cols=cols,
+                y=y_,
+                sample_weight=sample_weight_,  # type: ignore
+                group=group_,
+            )
+
+        return self
+
+    def prune(self, X, y, sample_weight=None, group=None) -> Self:
+        """
+        Prune the gradient booster on a provided dataset.
+
+        This removes nodes that do not contribute to a reduction in loss on the provided
+        validation set.
+
+        Parameters
+        ----------
+        X : array-like of shape (n_samples, n_features)
+            Validation data.
+        y : array-like of shape (n_samples,)
+            Validation targets.
+        sample_weight : array-like of shape (n_samples,), optional
+            Weights for validation samples.
+        group : array-like, optional
+            Group labels for ranking objectives.
+
+        Returns
+        -------
+        self : object
+            Returns self.
+        """
+
+        _, flat_data, rows, cols = transform_input_frame(X, self.cat_mapping)
+
+        y_, _ = convert_input_array(y, self.objective)
+
+        if sample_weight is None:
+            sample_weight_ = None
+        else:
+            sample_weight_, _ = convert_input_array(sample_weight, self.objective)
+
+        if group is None:
+            group_ = None
+        else:
+            group_, _ = convert_input_array(group, self.objective, is_int=True)
+
+        self.booster.prune(
+            flat_data=flat_data,
+            rows=rows,
+            cols=cols,
+            y=y_,
+            sample_weight=sample_weight_,  # type: ignore
+            group=group_,
+        )
+
+        return self
+
+    def calibrate(
+        self, X_train, y_train, X_cal, y_cal, alpha, sample_weight=None, group=None
+    ) -> Self:
+        """
+        Calibrate the gradient booster for prediction intervals.
+
+        Uses the provided training and calibration sets to compute scaling factors
+        for intervals.
+
+        Parameters
+        ----------
+        X_train : array-like
+            Data used to train the base model.
+        y_train : array-like
+            Targets for training data.
+        X_cal : array-like
+            Independent calibration dataset.
+        y_cal : array-like
+            Targets for calibration data.
+        alpha : float or array-like
+            Significance level(s) for the intervals (1 - coverage).
+        sample_weight : array-like, optional
+            Sample weights.
+        group : array-like, optional
+            Group labels.
+
+        Returns
+        -------
+        self : object
+            Returns self.
+        """
+
+        is_polars = type_df(X_train) == "polars_df"
+        if is_polars:
+            features_train, cols_train, masks_train, rows_train, _ = (
+                transform_input_frame_columnar(X_train, self.cat_mapping)
+            )
+            self._validate_features(features_train)
+            features_cal, cols_cal, masks_cal, rows_cal, _ = (
+                transform_input_frame_columnar(X_cal, self.cat_mapping)
+            )
+            # Use columnar calibration
+            y_train_, _ = convert_input_array(y_train, self.objective)
+            y_cal_, _ = convert_input_array(y_cal, self.objective)
+            if sample_weight is None:
+                sample_weight_ = None
+            else:
+                sample_weight_, _ = convert_input_array(sample_weight, self.objective)
+
+            self.booster.calibrate_columnar(
+                columns=cols_train,
+                masks=masks_train,
+                rows=rows_train,
+                y=y_train_,
+                columns_cal=cols_cal,
+                masks_cal=masks_cal,
+                rows_cal=rows_cal,
+                y_cal=y_cal_,
+                alpha=np.array(alpha),
+                sample_weight=sample_weight_,  # type: ignore
+                group=group,
+            )
+        else:
+            _, flat_data_train, rows_train, cols_train = transform_input_frame(
+                X_train, self.cat_mapping
+            )
+
+            y_train_, _ = convert_input_array(y_train, self.objective)
+
+            _, flat_data_cal, rows_cal, cols_cal = transform_input_frame(
+                X_cal, self.cat_mapping
+            )
+
+            y_cal_, _ = convert_input_array(y_cal, self.objective)
+
+            if sample_weight is None:
+                sample_weight_ = None
+            else:
+                sample_weight_, _ = convert_input_array(sample_weight, self.objective)
+
+            self.booster.calibrate(
+                flat_data=flat_data_train,
+                rows=rows_train,
+                cols=cols_train,
+                y=y_train_,
+                flat_data_cal=flat_data_cal,
+                rows_cal=rows_cal,
+                cols_cal=cols_cal,
+                y_cal=y_cal_,
+                alpha=np.array(alpha),
+                sample_weight=sample_weight_,  # type: ignore
+                group=group,
+            )
+
+        return self
+
+    def _validate_features(self, features: List[str]):
+        if len(features) > 0 and hasattr(self, "feature_names_in_"):
+            if features[0] != "0" and self.feature_names_in_[0] != "0":
+                if features != self.feature_names_in_:
+                    raise ValueError(
+                        f"Columns mismatch between data {features} passed, and data {self.feature_names_in_} used at fit."
+                    )
+
+    def predict_intervals(self, X, parallel: Union[bool, None] = None) -> dict:
+        """
+        Predict intervals with the fitted booster on new data.
+
+        Parameters
+        ----------
+        X : array-like of shape (n_samples, n_features)
+            New data for prediction.
+        parallel : bool, optional
+            Whether to run prediction in parallel. If None, uses class default.
+
+        Returns
+        -------
+        intervals : dict
+            A dictionary containing lower and upper bounds for the specified alpha levels.
+        """
+        is_polars = type_df(X) == "polars_df"
+        if is_polars:
+            features_, columns, masks, rows, cols = transform_input_frame_columnar(
+                X, self.cat_mapping
+            )
+            self._validate_features(features_)
+            return self.booster.predict_intervals_columnar(
+                columns=columns, masks=masks, rows=rows, parallel=parallel
+            )
+
+        features_, flat_data, rows, cols = transform_input_frame(X, self.cat_mapping)
+        self._validate_features(features_)
+
+        return self.booster.predict_intervals(
+            flat_data=flat_data,
+            rows=rows,
+            cols=cols,
+            parallel=parallel,
+        )
+
+    def predict(self, X, parallel: Union[bool, None] = None) -> np.ndarray:
+        """
+        Predict with the fitted booster on new data.
+
+        Parameters
+        ----------
+        X : array-like of shape (n_samples, n_features)
+            Input features.
+        parallel : bool, optional
+            Whether to run prediction in parallel.
+
+        Returns
+        -------
+        predictions : ndarray of shape (n_samples,)
+            The predicted values (log-odds for classification, raw values for regression).
+        """
+        is_polars = type_df(X) == "polars_df"
+        if is_polars:
+            features_, columns, masks, rows, cols = transform_input_frame_columnar(
+                X, self.cat_mapping
+            )
+        else:
+            features_, flat_data, rows, cols = transform_input_frame(
+                X, self.cat_mapping
+            )
+        self._validate_features(features_)
+
+        if len(self.classes_) == 0:
+            if is_polars:
+                return self.booster.predict_columnar(
+                    columns=columns, masks=masks, rows=rows, parallel=parallel
+                )
+            return self.booster.predict(
+                flat_data=flat_data, rows=rows, cols=cols, parallel=parallel
+            )
+        elif len(self.classes_) == 2:
+            if is_polars:
+                return np.rint(
+                    self.booster.predict_proba_columnar(
+                        columns=columns, masks=masks, rows=rows, parallel=parallel
+                    )
+                ).astype(int)
+            return np.rint(
+                self.booster.predict_proba(
+                    flat_data=flat_data, rows=rows, cols=cols, parallel=parallel
+                )
+            ).astype(int)
+        else:
+            if is_polars:
+                preds = self.booster.predict_columnar(
+                    columns=columns, masks=masks, rows=rows, parallel=parallel
+                )
+            else:
+                preds = self.booster.predict(
+                    flat_data=flat_data, rows=rows, cols=cols, parallel=parallel
+                )
+            preds_matrix = preds.reshape((-1, len(self.classes_)), order="F")
+            indices = np.argmax(preds_matrix, axis=1)
+            return np.array([self.classes_[i] for i in indices])
+
+    def predict_proba(self, X, parallel: Union[bool, None] = None) -> np.ndarray:
+        """
+        Predict class probabilities with the fitted booster on new data.
+
+        Only valid for classification tasks.
+
+        Parameters
+        ----------
+        X : array-like of shape (n_samples, n_features)
+            Input features.
+        parallel : bool, optional
+            Whether to run prediction in parallel.
+
+        Returns
+        -------
+        probabilities : ndarray of shape (n_samples, n_classes)
+            The class probabilities.
+        """
+        is_polars = type_df(X) == "polars_df"
+        if is_polars:
+            features_, columns, masks, rows, cols = transform_input_frame_columnar(
+                X, self.cat_mapping
+            )
+        else:
+            features_, flat_data, rows, cols = transform_input_frame(
+                X, self.cat_mapping
+            )
+        self._validate_features(features_)
+
+        if len(self.classes_) > 2:
+            if is_polars:
+                probabilities = self.booster.predict_proba_columnar(
+                    columns=columns, masks=masks, rows=rows, parallel=parallel
+                )
+            else:
+                probabilities = self.booster.predict_proba(
+                    flat_data=flat_data, rows=rows, cols=cols, parallel=parallel
+                )
+            return probabilities.reshape((-1, len(self.classes_)), order="C")
+        elif len(self.classes_) == 2:
+            if is_polars:
+                probabilities = self.booster.predict_proba_columnar(
+                    columns=columns, masks=masks, rows=rows, parallel=parallel
+                )
+            else:
+                probabilities = self.booster.predict_proba(
+                    flat_data=flat_data, rows=rows, cols=cols, parallel=parallel
+                )
+            return np.concatenate(
+                [(1.0 - probabilities).reshape(-1, 1), probabilities.reshape(-1, 1)],
+                axis=1,
+            )
+        else:
+            warnings.warn(
+                f"predict_proba not implemented for regression. n_classes = {len(self.classes_)}"
+            )
+            return np.ones((rows, 1))
+
+    def predict_log_proba(self, X, parallel: Union[bool, None] = None) -> np.ndarray:
+        """
+        Predict class log-probabilities with the fitted booster on new data.
+
+        Only valid for classification tasks.
+
+        Parameters
+        ----------
+        X : array-like of shape (n_samples, n_features)
+            Input features.
+        parallel : bool, optional
+            Whether to run prediction in parallel.
+
+        Returns
+        -------
+        log_probabilities : ndarray of shape (n_samples, n_classes)
+            The log-probabilities of each class.
+        """
+        is_polars = type_df(X) == "polars_df"
+        if is_polars:
+            features_, columns, masks, rows, cols = transform_input_frame_columnar(
+                X, self.cat_mapping
+            )
+        else:
+            features_, flat_data, rows, cols = transform_input_frame(
+                X, self.cat_mapping
+            )
+        self._validate_features(features_)
+
+        if len(self.classes_) > 2:
+            if is_polars:
+                preds = self.booster.predict_columnar(
+                    columns=columns, masks=masks, rows=rows, parallel=parallel
+                )
+            else:
+                preds = self.booster.predict(
+                    flat_data=flat_data,
+                    rows=rows,
+                    cols=cols,
+                    parallel=parallel,
+                )
+            return preds.reshape((-1, len(self.classes_)), order="F")
+        elif len(self.classes_) == 2:
+            if is_polars:
+                return self.booster.predict_columnar(
+                    columns=columns, masks=masks, rows=rows, parallel=parallel
+                )
+            return self.booster.predict(
+                flat_data=flat_data,
+                rows=rows,
+                cols=cols,
+                parallel=parallel,
+            )
+        else:
+            warnings.warn("predict_log_proba not implemented for regression.")
+            return np.ones((rows, 1))
+
+    def predict_nodes(self, X, parallel: Union[bool, None] = None) -> List:
+        """
+        Predict leaf node indices with the fitted booster on new data.
+
+        Parameters
+        ----------
+        X : array-like of shape (n_samples, n_features)
+            Input features.
+        parallel : bool, optional
+            Whether to run prediction in parallel.
+
+        Returns
+        -------
+        node_indices : list of ndarray
+            A list where each element corresponds to a tree and contains node indices
+            for each sample.
+        """
+        is_polars = type_df(X) == "polars_df"
+        if is_polars:
+            features_, columns, masks, rows, cols = transform_input_frame_columnar(
+                X, self.cat_mapping
+            )
+            self._validate_features(features_)
+            return self.booster.predict_nodes_columnar(
+                columns=columns, masks=masks, rows=rows, parallel=parallel
+            )
+
+        features_, flat_data, rows, cols = transform_input_frame(X, self.cat_mapping)
+        self._validate_features(features_)
+
+        return self.booster.predict_nodes(
+            flat_data=flat_data, rows=rows, cols=cols, parallel=parallel
+        )
+
+    @property
+    def feature_importances_(self) -> np.ndarray:
+        vals = self.calculate_feature_importance(
+            method=self.feature_importance_method, normalize=True
+        )
+        if hasattr(self, "feature_names_in_"):
+            vals = cast(Dict[str, float], vals)
+            return np.array([vals.get(ft, 0.0) for ft in self.feature_names_in_])
+        else:
+            vals = cast(Dict[int, float], vals)
+            return np.array([vals.get(ft, 0.0) for ft in range(self.n_features_)])
+
+    def predict_contributions(
+        self, X, method: str = "Average", parallel: Union[bool, None] = None
+    ) -> np.ndarray:
+        """
+        Predict feature contributions (SHAP-like values) for new data.
+
+        Parameters
+        ----------
+        X : array-like of shape (n_samples, n_features)
+            Input features.
+        method : str, default="Average"
+            Method to calculate contributions. Options:
+
+            - "Average": Internal node averages.
+            - "Shapley": Exact tree SHAP values.
+            - "Weight": Saabas-style leaf weights.
+            - "BranchDifference": Difference between chosen and other branch.
+            - "MidpointDifference": Weighted difference between branches.
+            - "ModeDifference": Difference from the most frequent node.
+            - "ProbabilityChange": Change in probability (LogLoss only).
+
+        parallel : bool, optional
+            Whether to run prediction in parallel.
+
+        Returns
+        -------
+        contributions : ndarray of shape (n_samples, n_features + 1)
+            The contribution of each feature to the prediction. The last column
+            is the bias term.
+        """
+        is_polars = type_df(X) == "polars_df"
+        if is_polars:
+            features_, columns, masks, rows, cols = transform_input_frame_columnar(
+                X, self.cat_mapping
+            )
+            self._validate_features(features_)
+            contributions = self.booster.predict_contributions_columnar(
+                columns=columns,
+                masks=masks,
+                rows=rows,
+                method=CONTRIBUTION_METHODS.get(method, method),
+                parallel=parallel,
+            )
+        else:
+            features_, flat_data, rows, cols = transform_input_frame(
+                X, self.cat_mapping
+            )
+            self._validate_features(features_)
+
+            contributions = self.booster.predict_contributions(
+                flat_data=flat_data,
+                rows=rows,
+                cols=cols,
+                method=CONTRIBUTION_METHODS.get(method, method),
+                parallel=parallel,
+            )
+
+        if len(self.classes_) > 2:
+            return (
+                np.reshape(contributions, (len(self.classes_), rows, cols + 1))
+                .transpose(1, 0, 2)
+                .reshape(rows, -1)
+            )
+        return np.reshape(contributions, (rows, cols + 1))
+
+    def partial_dependence(
+        self,
+        X,
+        feature: Union[str, int],
+        samples: Optional[int] = 100,
+        exclude_missing: bool = True,
+        percentile_bounds: Tuple[float, float] = (0.2, 0.98),
+    ) -> np.ndarray:
+        """
+        Calculate the partial dependence values of a feature.
+
+        For each unique value of the feature, this gives the estimate of the predicted
+        value for that feature, with the effects of all other features averaged out.
+
+        Parameters
+        ----------
+        X : array-like
+            Data used to calculate partial dependence. Should be the same format
+            as passed to :meth:`fit`.
+        feature : str or int
+            The feature for which to calculate partial dependence.
+        samples : int, optional, default=100
+            Number of evenly spaced samples to select. If None, all unique values are used.
+        exclude_missing : bool, optional, default=True
+            Whether to exclude missing values from the calculation.
+        percentile_bounds : tuple of float, optional, default=(0.2, 0.98)
+            Lower and upper percentiles for sample selection.
+
+        Returns
+        -------
+        pd_values : ndarray of shape (n_samples, 2)
+            The first column contains the feature values, and the second column
+            contains the partial dependence values.
+
+        Examples
+        --------
+        >>> import matplotlib.pyplot as plt
+        >>> pd_values = model.partial_dependence(X, feature="age")
+        >>> plt.plot(pd_values[:, 0], pd_values[:, 1])
+        """
+        if isinstance(feature, str):
+            is_polars = type_df(X) == "polars_df"
+            if not (type_df(X) == "pandas_df" or is_polars):
+                raise ValueError(
+                    "If `feature` is a string, then the object passed as `X` must be a pandas or polars DataFrame."
+                )
+            if is_polars:
+                values = X[feature].to_numpy()
+            else:
+                values = X.loc[:, feature].to_numpy()
+
+            if hasattr(self, "feature_names_in_") and self.feature_names_in_[0] != "0":
+                [feature_idx] = [
+                    i for i, v in enumerate(self.feature_names_in_) if v == feature
+                ]
+            else:
+                w_msg = (
+                    "No feature names were provided at fit, but feature was a string, attempting to "
+                    + "determine feature index from DataFrame Column, "
+                    + "ensure columns are the same order as data passed when fit."
+                )
+                warnings.warn(w_msg)
+                features = X.columns if is_polars else X.columns.to_list()
+                [feature_idx] = [i for i, v in enumerate(features) if v == feature]
+        elif isinstance(feature, int):
+            feature_idx = feature
+            if type_df(X) == "pandas_df":
+                values = X.to_numpy()[:, feature]
+            elif type_df(X) == "polars_df":
+                values = X.to_numpy(allow_copy=False)[:, feature]
+            else:
+                values = X[:, feature]
+        else:
+            raise ValueError(
+                f"The parameter `feature` must be a string, or an int, however an object of type {type(feature)} was passed."
+            )
+        min_p, max_p = percentile_bounds
+        values = values[~(np.isnan(values) | (values == self.missing))]
+        if samples is None:
+            search_values = np.sort(np.unique(values))
+        else:
+            # Exclude missing from this calculation.
+            search_values = np.quantile(values, np.linspace(min_p, max_p, num=samples))
+
+        # Add missing back, if they wanted it...
+        if not exclude_missing:
+            search_values = np.append([self.missing], search_values)
+
+        res = []
+        for v in search_values:
+            res.append(
+                (v, self.booster.value_partial_dependence(feature=feature_idx, value=v))
+            )
+        return np.array(res)
+
+    def calculate_feature_importance(
+        self, method: str = "Gain", normalize: bool = True
+    ) -> Union[Dict[int, float], Dict[str, float]]:
+        """
+        Calculate feature importance for the model.
+
+        Parameters
+        ----------
+        method : str, optional, default="Gain"
+            Importance method. Options:
+
+            - "Weight": Number of times a feature is used in splits.
+            - "Gain": Average improvement in loss brought by a feature.
+            - "Cover": Average number of samples affected by splits on a feature.
+            - "TotalGain": Total improvement in loss brought by a feature.
+            - "TotalCover": Total number of samples affected by splits on a feature.
+
+        normalize : bool, optional, default=True
+            Whether to normalize importance scores to sum to 1.
+
+        Returns
+        -------
+        importance : dict
+            A dictionary mapping feature names (or indices) to importance scores.
+        """
+        importance_: Dict[int, float] = self.booster.calculate_feature_importance(
+            method=method,
+            normalize=normalize,
+        )
+        if hasattr(self, "feature_names_in_"):
+            feature_map: Dict[int, str] = {
+                i: f for i, f in enumerate(self.feature_names_in_)
+            }
+            return {feature_map[i]: v for i, v in importance_.items()}
+        return importance_
+
+    def text_dump(self) -> List[str]:
+        """
+        Return the booster model in a human-readable text format.
+
+        Returns
+        -------
+        dump : list of str
+            A list where each string represents a tree in the ensemble.
+        """
+        return self.booster.text_dump()
+
+    def json_dump(self) -> str:
+        """
+        Return the booster model in JSON format.
+
+        Returns
+        -------
+        dump : str
+            The JSON representation of the model.
+        """
+        return self.booster.json_dump()
+
+    @classmethod
+    def load_booster(cls, path: str) -> Self:
+        """
+        Load a booster model from a file.
+
+        Parameters
+        ----------
+        path : str
+            Path to the saved booster (JSON format).
+
+        Returns
+        -------
+        model : PerpetualBooster
+            The loaded booster object.
+        """
+        try:
+            booster = CratePerpetualBooster.load_booster(str(path))
+        except ValueError:
+            booster = CrateMultiOutputBooster.load_booster(str(path))
+
+        params = booster.get_params()
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            c = cls(**params)
+        c.booster = booster
+        for m in c.metadata_attributes:
+            try:
+                m_ = c._get_metadata_attributes(m)
+                setattr(c, m, m_)
+                # If "feature_names_in_" is present, we know a
+                # pandas dataframe was used for fitting, in this case
+                # get back the original monotonicity map, with the
+                # feature names as keys.
+                if m == "feature_names_in_" and c.feature_names_in_[0] != "0":
+                    if c.monotone_constraints is not None:
+                        c.monotone_constraints = {
+                            ft: c.monotone_constraints[i]
+                            for i, ft in enumerate(c.feature_names_in_)
+                        }
+            except KeyError:
+                pass
+        return c
+
+    def save_booster(self, path: str):
+        """
+        Save the booster model to a file.
+
+        The model is saved in a JSON-based format.
+
+        Parameters
+        ----------
+        path : str
+            Path where the model will be saved.
+        """
+        self.booster.save_booster(str(path))
+
+    def _standardize_monotonicity_map(
+        self,
+        X,
+    ) -> Dict[int, Any]:
+        if isinstance(X, np.ndarray):
+            return self.monotone_constraints
+        else:
+            feature_map = {f: i for i, f in enumerate(X.columns)}
+            return {feature_map[f]: v for f, v in self.monotone_constraints.items()}
+
+    def _standardize_terminate_missing_features(
+        self,
+        X,
+    ) -> Set[int]:
+        if isinstance(X, np.ndarray):
+            return set(self.terminate_missing_features)
+        else:
+            feature_map = {f: i for i, f in enumerate(X.columns)}
+            return set(feature_map[f] for f in self.terminate_missing_features)
+
+    def insert_metadata(self, key: str, value: str):
+        """
+        Insert metadata into the model.
+
+        Metadata is saved alongside the model and can be retrieved later.
+
+        Parameters
+        ----------
+        key : str
+            The key for the metadata item.
+        value : str
+            The value for the metadata item.
+        """
+        self.booster.insert_metadata(key=key, value=value)
+
+    def get_metadata(self, key: str) -> str:
+        """
+        Get metadata associated with a given key.
+
+        Parameters
+        ----------
+        key : str
+            The key to look up in the metadata.
+
+        Returns
+        -------
+        value : str
+            The value associated with the key.
+        """
+        v = self.booster.get_metadata(key=key)
+        return v
+
+    def _set_metadata_attributes(self, key: str, value: Any) -> None:
+        value_ = self.metadata_attributes[key].serialize(value)
+        self.insert_metadata(key=key, value=value_)
+
+    def _get_metadata_attributes(self, key: str) -> Any:
+        value = self.get_metadata(key)
+        return self.metadata_attributes[key].deserialize(value)
+
+    @property
+    def base_score(self) -> Union[float, Iterable[float]]:
+        """
+        The base score(s) of the model.
+
+        Returns
+        -------
+        score : float or iterable of float
+            The initial prediction value(s) of the model.
+        """
+        return self.booster.base_score
+
+    @property
+    def number_of_trees(self) -> Union[int, Iterable[int]]:
+        """
+        The number of trees in the ensemble.
+
+        Returns
+        -------
+        n_trees : int or iterable of int
+            Total number of trees.
+        """
+        return self.booster.number_of_trees
+
+    # Make picklable with getstate and setstate
+    def __getstate__(self) -> Dict[Any, Any]:
+        booster_json = self.json_dump()
+        # Delete booster
+        # Doing it like this, so it doesn't delete it globally.
+        res = {k: v for k, v in self.__dict__.items() if k != "booster"}
+        res["__booster_json_file__"] = booster_json
+        return res
+
+    def __setstate__(self, d: Dict[Any, Any]) -> None:
+        # Load the booster object the pickled JSon string.
+        try:
+            booster_object = CratePerpetualBooster.from_json(d["__booster_json_file__"])
+        except ValueError:
+            booster_object = CrateMultiOutputBooster.from_json(
+                d["__booster_json_file__"]
+            )
+        d["booster"] = booster_object
+        # Are there any new parameters, that need to be added to the python object,
+        # that would have been loaded in as defaults on the json object?
+        # This makes sure that defaults set with a serde default function get
+        # carried through to the python object.
+        for p, v in booster_object.get_params().items():
+            if p not in d:
+                d[p] = v
+        del d["__booster_json_file__"]
+        self.__dict__ = d
+
+    # Functions for scikit-learn compatibility, will feel out adding these manually,
+    # and then if that feels too unwieldy will add scikit-learn as a dependency.
+    def get_params(self, deep=True) -> Dict[str, Any]:
+        """
+        Get parameters for this booster.
+
+        Parameters
+        ----------
+        deep : bool, default=True
+            Currently ignored, exists for scikit-learn compatibility.
+
+        Returns
+        -------
+        params : dict
+            Parameter names mapped to their values.
+        """
+        args = inspect.getfullargspec(PerpetualBooster).kwonlyargs
+        return {param: getattr(self, param) for param in args}
+
+    def set_params(self, **params: Any) -> Self:
+        """
+        Set parameters for this booster.
+
+        Parameters
+        ----------
+        **params : dict
+            Booster parameters.
+
+        Returns
+        -------
+        self : object
+            Returns self.
+        """
+        old_params = self.get_params()
+        old_params.update(params)
+        PerpetualBooster.__init__(self, **old_params)
+        return self
+
+    def get_node_lists(self, map_features_names: bool = True) -> List[List[Node]]:
+        """
+        Return tree structures as lists of node objects.
+
+        Parameters
+        ----------
+        map_features_names : bool, default=True
+            Whether to use feature names instead of indices.
+
+        Returns
+        -------
+        trees : list of list of Node
+            Each inner list represents a tree.
+        """
+        dump = json.loads(self.json_dump())
+        if "trees" in dump:
+            all_booster_trees = [dump["trees"]]
+        else:
+            # Multi-output
+            all_booster_trees = [b["trees"] for b in dump["boosters"]]
+
+        feature_map: Union[Dict[int, str], Dict[int, int]]
+        leaf_split_feature: Union[str, int]
+        if map_features_names and hasattr(self, "feature_names_in_"):
+            feature_map = {i: ft for i, ft in enumerate(self.feature_names_in_)}
+            leaf_split_feature = ""
+        else:
+            feature_map = {i: i for i in range(self.n_features_)}
+            leaf_split_feature = -1
+
+        trees = []
+        for booster_trees in all_booster_trees:
+            for t in booster_trees:
+                nodes = []
+                for node in t["nodes"].values():
+                    if not node["is_leaf"]:
+                        node["split_feature"] = feature_map[node["split_feature"]]
+                    else:
+                        node["split_feature"] = leaf_split_feature
+                    nodes.append(Node(**node))
+                trees.append(nodes)
+        return trees
+
+    def trees_to_dataframe(self) -> Any:
+        """
+        Return the tree structures as a DataFrame.
+
+        Returns
+        -------
+        df : DataFrame
+            A Polars or Pandas DataFrame containing tree information.
+        """
+
+        def node_to_row(
+            n: Node,
+            tree_n: int,
+        ) -> Dict[str, Any]:
+            def _id(i: int) -> str:
+                return f"{tree_n}-{i}"
+
+            return dict(
+                Tree=tree_n,
+                Node=n.num,
+                ID=_id(n.num),
+                Feature="Leaf" if n.is_leaf else str(n.split_feature),
+                Split=None if n.is_leaf else n.split_value,
+                Yes=None if n.is_leaf else _id(n.left_child),
+                No=None if n.is_leaf else _id(n.right_child),
+                Missing=None if n.is_leaf else _id(n.missing_node),
+                Gain=n.weight_value if n.is_leaf else n.split_gain,
+                Cover=n.hessian_sum,
+                Left_Cats=n.left_cats,
+                Right_Cats=n.right_cats,
+            )
+
+        # Flatten list of lists using list comprehension
+        vals = [
+            node_to_row(n, i)
+            for i, tree in enumerate(self.get_node_lists())
+            for n in tree
+        ]
+
+        try:
+            import polars as pl
+
+            return pl.from_records(vals).sort(
+                ["Tree", "Node"], descending=[False, False]
+            )
+        except ImportError:
+            import pandas as pd
+
+            return pd.DataFrame.from_records(vals).sort_values(
+                ["Tree", "Node"], ascending=[True, True]
+            )
+
+    def _to_xgboost_json(self) -> Dict[str, Any]:
+        """Convert the Perpetual model to an XGBoost JSON model structure."""
+
+        # Check if it's a multi-output model
+        is_multi = len(self.classes_) > 2
+
+        # Get raw dump
+        raw_dump = json.loads(self.json_dump())
+
+        # Initialize XGBoost structure
+        xgb_json = {
+            "learner": {
+                "attributes": {},
+                "feature_names": [],
+                "feature_types": [],
+                "gradient_booster": {
+                    "model": {
+                        "gbtree_model_param": {
+                            "num_parallel_tree": "1",
+                        },
+                        "trees": [],
+                        "tree_info": [],
+                        "iteration_indptr": [],
+                        "cats": {
+                            "enc": [],
+                            "feature_segments": [],
+                            "sorted_idx": [],
+                        },
+                    },
+                    "name": "gbtree",
+                },
+                "learner_model_param": {
+                    "boost_from_average": "1",
+                    "num_feature": str(self.n_features_),
+                },
+                "objective": {
+                    "name": "binary:logistic",
+                },
+            },
+            "version": [3, 1, 3],  # Use a reasonably recent version
+        }
+
+        # Fill feature names if available
+        if hasattr(self, "feature_names_in_"):
+            xgb_json["learner"]["feature_names"] = self.feature_names_in_
+            xgb_json["learner"]["feature_types"] = ["float"] * self.n_features_
+        else:
+            xgb_json["learner"]["feature_names"] = [
+                f"f{i}" for i in range(self.n_features_)
+            ]
+            xgb_json["learner"]["feature_types"] = ["float"] * self.n_features_
+
+        # Objective and Base Score Handling
+        if is_multi:
+            # Multi-class
+            n_classes = len(self.classes_)
+            xgb_json["learner"]["objective"]["name"] = "multi:softprob"
+            xgb_json["learner"]["objective"]["softmax_multiclass_param"] = {
+                "num_class": str(n_classes)
+            }
+            xgb_json["learner"]["learner_model_param"]["num_class"] = str(n_classes)
+            xgb_json["learner"]["learner_model_param"]["num_target"] = "1"
+
+            # Base score vector [0.5, 0.5, ...]
+            # 5.0E-1
+            base_score_str = ",".join(["5.0E-1"] * n_classes)
+            xgb_json["learner"]["learner_model_param"]["base_score"] = (
+                f"[{base_score_str}]"
+            )
+
+            boosters = raw_dump["boosters"]
+
+            trees = []
+            tree_info = []
+            # For multi-class, we need to interleave trees if we want to follow XGBoost structure perfectly?
+            # Or can we just dump them? iteration_indptr depends on this.
+            # XGBoost expects trees for iteration i to be contiguous.
+            # Perpetual stores boosters separately.
+            # Booster 0 has trees for class 0. Booster 1 for class 1.
+            # We need to rearrange them: Round 0 (C0, C1, C2), Round 1 (C0, C1, C2)...
+
+            # Assuming all boosters have same number of trees?
+            num_trees_per_booster = [len(b["trees"]) for b in boosters]
+            max_trees = max(num_trees_per_booster) if num_trees_per_booster else 0
+
+            # Iteration pointers: 0, 3, 6...
+            # But what if some booster has fewer trees? (Early stopping might cause this?)
+            # Perpetual implementation usually stops all or none?
+            # "MultiOutputBooster::fit" trains them sequentially but they might have different tree counts if EarlyStopping is per-booster.
+            # But XGBoost expects consistent num_class trees per round (or use "multi:softprob"?)
+            # If we just list them, XGBoost might get confused if we don't align them.
+
+            # Let's try to align them by round.
+
+            iteration_indptr = [0]
+            current_ptr = 0
+
+            for round_idx in range(max_trees):
+                # For each class
+                for group_id, booster_dump in enumerate(boosters):
+                    booster_trees = booster_dump["trees"]
+                    if round_idx < len(booster_trees):
+                        tree = booster_trees[round_idx]
+                        base_score = booster_dump["base_score"]
+
+                        xgb_tree = self._convert_tree(tree, current_ptr)
+
+                        if round_idx == 0:
+                            self._adjust_tree_leaves(xgb_tree, base_score)
+
+                        trees.append(xgb_tree)
+                        tree_info.append(group_id)
+                        current_ptr += 1
+                    else:
+                        # Missing tree for this class in this round?
+                        # Should we insert a dummy tree (0 prediction)?
+                        # For now, let's assume balanced trees or hope XGB handles it.
+                        # If we skip, tree_info tracks class.
+                        pass
+
+                iteration_indptr.append(current_ptr)
+
+            xgb_json["learner"]["gradient_booster"]["model"]["trees"] = trees
+            xgb_json["learner"]["gradient_booster"]["model"]["tree_info"] = tree_info
+            xgb_json["learner"]["gradient_booster"]["model"]["gbtree_model_param"][
+                "num_trees"
+            ] = str(len(trees))
+            xgb_json["learner"]["gradient_booster"]["model"]["iteration_indptr"] = (
+                iteration_indptr
+            )
+
+        else:
+            # Binary or Regression
+            if self.objective == "LogLoss":
+                xgb_json["learner"]["objective"]["name"] = "binary:logistic"
+                xgb_json["learner"]["objective"]["reg_loss_param"] = {
+                    "scale_pos_weight": "1"
+                }
+                xgb_json["learner"]["learner_model_param"]["num_class"] = "0"
+                xgb_json["learner"]["learner_model_param"]["num_target"] = "1"
+
+                # Base Score
+                base_score_val = 1.0 / (1.0 + np.exp(-raw_dump["base_score"]))
+                xgb_json["learner"]["learner_model_param"]["base_score"] = (
+                    f"[{base_score_val:.6E}]"
+                )
+
+            elif self.objective == "SquaredLoss":
+                xgb_json["learner"]["objective"]["name"] = "reg:squarederror"
+                xgb_json["learner"]["objective"]["reg_loss_param"] = {}
+                xgb_json["learner"]["learner_model_param"]["num_class"] = "0"
+                xgb_json["learner"]["learner_model_param"]["num_target"] = "1"
+                xgb_json["learner"]["learner_model_param"]["base_score"] = (
+                    f"[{raw_dump['base_score']:.6E}]"
+                )
+            else:
+                warnings.warn(
+                    f"Objective {self.objective} not explicitly supported for XGBoost export. Defaulting to reg:squarederror."
+                )
+                xgb_json["learner"]["objective"]["name"] = "reg:squarederror"
+                xgb_json["learner"]["objective"]["reg_loss_param"] = {}
+                xgb_json["learner"]["learner_model_param"]["num_class"] = "0"
+                xgb_json["learner"]["learner_model_param"]["num_target"] = "1"
+                xgb_json["learner"]["learner_model_param"]["base_score"] = (
+                    f"[{raw_dump['base_score']:.6E}]"
+                )
+
+            trees = []
+            tree_info = []
+            for tree_idx, tree in enumerate(raw_dump["trees"]):
+                xgb_tree = self._convert_tree(tree, tree_idx)
+                trees.append(xgb_tree)
+                tree_info.append(0)
+
+            xgb_json["learner"]["gradient_booster"]["model"]["trees"] = trees
+            xgb_json["learner"]["gradient_booster"]["model"]["tree_info"] = tree_info
+            xgb_json["learner"]["gradient_booster"]["model"]["gbtree_model_param"][
+                "num_trees"
+            ] = str(len(trees))
+            xgb_json["learner"]["gradient_booster"]["model"]["iteration_indptr"] = list(
+                range(len(trees) + 1)
+            )
+
+        return xgb_json
+
+    def _convert_tree(self, tree: Dict[str, Any], group_id: int) -> Dict[str, Any]:
+        """Convert a single Perpetual tree to XGBoost dictionary format."""
+
+        nodes_dict = tree["nodes"]
+        # Convert keys to int and sort
+        sorted_keys = sorted(nodes_dict.keys(), key=lambda x: int(x))
+
+        # Mapping from Perpetual ID (int) to XGBoost Array Index (int)
+        node_map = {int(k): i for i, k in enumerate(sorted_keys)}
+
+        num_nodes = len(sorted_keys)
+        # print(f"DEBUG: Converting tree group={group_id}. num_nodes={num_nodes}")
+
+        left_children = [-1] * num_nodes
+        right_children = [-1] * num_nodes
+        parents = [2147483647] * num_nodes
+        split_indices = [0] * num_nodes
+        split_conditions = [0.0] * num_nodes
+        split_type = [0] * num_nodes
+        sum_hessian = [0.0] * num_nodes
+        loss_changes = [0.0] * num_nodes
+        base_weights = [0.0] * num_nodes
+        default_left = [0] * num_nodes
+
+        categories = []
+        categories_nodes = []
+        categories_segments = []
+        categories_sizes = []
+
+        for i, k in enumerate(sorted_keys):
+            node = nodes_dict[k]
+            nid = int(node["num"])
+            idx = node_map[nid]
+
+            # print(f"  DEBUG: Node {i} nid={nid} idx={idx}")
+
+            sum_hessian[idx] = node["hessian_sum"]
+            base_weights[idx] = node["weight_value"]
+            loss_changes[idx] = node.get("split_gain", 0.0)
+
+            if node["is_leaf"]:
+                left_children[idx] = -1
+                right_children[idx] = -1
+                split_indices[idx] = 0
+                split_conditions[idx] = node["weight_value"]
+            else:
+                left_id = node["left_child"]
+                right_id = node["right_child"]
+
+                left_idx = node_map[left_id]
+                right_idx = node_map[right_id]
+
+                left_children[idx] = left_idx
+                right_children[idx] = right_idx
+                parents[left_idx] = idx
+                parents[right_idx] = idx
+
+                split_indices[idx] = node["split_feature"]
+                split_conditions[idx] = node["split_value"]
+
+                # Missing handling
+                # If missing_node goes left
+                if node["missing_node"] == left_id:
+                    default_left[idx] = 1
+                else:
+                    default_left[idx] = 0
+
+                if (
+                    "left_cats" in node
+                    and node["left_cats"] is not None
+                    and len(node["left_cats"]) > 0
+                ):
+                    # It's a categorical split
+                    cats = node["left_cats"]
+                    # XGBoost uses split_type=1 for categorical?
+                    # Or just presence in categories_nodes?
+                    # Docs say: split_type [default=0]: 0=numerical, 1=categorical
+                    split_type[idx] = 1
+
+                    # Update categorical arrays
+                    categories_nodes.append(idx)
+                    categories_sizes.append(len(cats))
+                    # Segment is start index.
+                    # If this is the first one, 0. Else prev_segment + prev_size?
+                    # Actually valid XGBoost format usually has segments as exclusive scan.
+                    # [0, len0, len0+len1, ...]
+                    # Wait, segments length should be same as nodes?
+                    # Let's check logic:
+                    # segments[i] points to start of cats for node i (in categories_nodes)
+
+                    next_segment = (
+                        (categories_segments[-1] + categories_sizes[-2])
+                        if categories_segments
+                        else 0
+                    )
+                    categories_segments.append(next_segment)
+
+                    categories.extend(sorted(cats))
+
+                    # split_condition for categorical is usually NaN or special?
+                    # XGBoost JSON parser might ignore it if type is categorical
+                    # But often it is set to something.
+
+        return {
+            "base_weights": base_weights,
+            "default_left": default_left,
+            "id": group_id,
+            "left_children": left_children,
+            "loss_changes": loss_changes,
+            "parents": parents,
+            "right_children": right_children,
+            "split_conditions": split_conditions,
+            "split_indices": split_indices,
+            "split_type": split_type,
+            "sum_hessian": sum_hessian,
+            "tree_param": {
+                "num_deleted": "0",
+                "num_feature": str(self.n_features_),
+                "num_nodes": str(num_nodes),
+                "size_leaf_vector": "1",
+            },
+            "categories": categories,
+            "categories_nodes": categories_nodes,
+            "categories_segments": categories_segments,
+            "categories_sizes": categories_sizes,
+        }
+
+    def _adjust_tree_leaves(self, xgb_tree: Dict[str, Any], adjustment: float):
+        """Add adjustment value to all leaves in an XGBoost tree dict."""
+        left_children = xgb_tree["left_children"]
+        split_conditions = xgb_tree["split_conditions"]
+        base_weights = xgb_tree["base_weights"]
+
+        for i, left in enumerate(left_children):
+            if left == -1:  # Leaf
+                split_conditions[i] += adjustment
+                base_weights[i] += adjustment
+
+    def save_as_xgboost(self, path: str):
+        """
+        Save the model in XGBoost JSON format.
+
+        Parameters
+        ----------
+        path : str
+            The path where the XGBoost-compatible model will be saved.
+        """
+        xgboost_json = self._to_xgboost_json()
+        with open(path, "w") as f:
+            json.dump(xgboost_json, f, indent=2)
+
+    def save_as_onnx(self, path: str, name: str = "perpetual_model"):
+        """
+        Save the model in ONNX format.
+
+        Parameters
+        ----------
+        path : str
+            The path where the ONNX model will be saved.
+        name : str, optional, default="perpetual_model"
+            The name of the graph in the exported model.
+        """
+        import json
+
+        import onnx
+        from onnx import TensorProto, helper
+
+        raw_dump = json.loads(self.json_dump())
+        is_classifier = len(self.classes_) >= 2
+        is_multi = is_classifier and len(self.classes_) > 2
+        n_classes = len(self.classes_) if is_classifier else 1
+
+        if "trees" in raw_dump:
+            booster_data = [{"trees": raw_dump["trees"]}]
+        else:
+            booster_data = raw_dump["boosters"]
+
+        feature_map_inverse = (
+            {v: k for k, v in enumerate(self.feature_names_in_)}
+            if hasattr(self, "feature_names_in_")
+            else None
+        )
+
+        nodes_treeids = []
+        nodes_nodeids = []
+        nodes_featureids = []
+        nodes_values = []
+        nodes_modes = []
+        nodes_truenodeids = []
+        nodes_falsenodeids = []
+        nodes_missing_value_tracks_true = []
+
+        target_treeids = []
+        target_nodeids = []
+        target_ids = []
+        target_weights = []
+
+        # Base score handling
+        base_score = self.base_score
+        if is_classifier:
+            if is_multi:
+                base_values = [float(b) for b in base_score]
+            else:
+                base_values = [float(base_score)]
+        else:
+            base_values = [float(base_score)]
+
+        global_tree_idx = 0
+        for b_idx, booster in enumerate(booster_data):
+            for tree_data in booster["trees"]:
+                nodes_dict = tree_data["nodes"]
+                node_keys = sorted(nodes_dict.keys(), key=lambda x: int(x))
+
+                node_id_to_idx = {}
+                for i, k in enumerate(node_keys):
+                    node_id_to_idx[int(k)] = i
+
+                for k in node_keys:
+                    node_dict = nodes_dict[k]
+                    nid = int(node_dict["num"])
+                    idx_for_onnx = node_id_to_idx[nid]
+
+                    nodes_treeids.append(global_tree_idx)
+                    nodes_nodeids.append(idx_for_onnx)
+
+                    if node_dict["is_leaf"]:
+                        nodes_modes.append("LEAF")
+                        nodes_featureids.append(0)
+                        nodes_values.append(0.0)
+                        nodes_truenodeids.append(0)
+                        nodes_falsenodeids.append(0)
+                        nodes_missing_value_tracks_true.append(0)
+
+                        target_treeids.append(global_tree_idx)
+                        target_nodeids.append(idx_for_onnx)
+                        target_ids.append(b_idx if is_multi else 0)
+                        target_weights.append(float(node_dict["weight_value"]))
+                    else:
+                        nodes_modes.append("BRANCH_LT")
+                        feat_val = node_dict["split_feature"]
+                        f_idx = 0
+                        if isinstance(feat_val, int):
+                            f_idx = feat_val
+                        elif feature_map_inverse and feat_val in feature_map_inverse:
+                            f_idx = feature_map_inverse[feat_val]
+                        elif isinstance(feat_val, str) and feat_val.isdigit():
+                            f_idx = int(feat_val)
+
+                        nodes_featureids.append(f_idx)
+                        nodes_values.append(float(node_dict["split_value"]))
+
+                        tracks_true = 0
+                        if node_dict["missing_node"] == node_dict["left_child"]:
+                            tracks_true = 1
+                        nodes_missing_value_tracks_true.append(tracks_true)
+
+                        nodes_truenodeids.append(
+                            node_id_to_idx[int(node_dict["left_child"])]
+                        )
+                        nodes_falsenodeids.append(
+                            node_id_to_idx[int(node_dict["right_child"])]
+                        )
+
+                global_tree_idx += 1
+
+        input_name = "input"
+        input_type = helper.make_tensor_value_info(
+            input_name, TensorProto.FLOAT, [None, self.n_features_]
+        )
+
+        raw_scores_name = "raw_scores"
+        reg_node = helper.make_node(
+            "TreeEnsembleRegressor",
+            inputs=[input_name],
+            outputs=[raw_scores_name],
+            domain="ai.onnx.ml",
+            nodes_treeids=nodes_treeids,
+            nodes_nodeids=nodes_nodeids,
+            nodes_featureids=nodes_featureids,
+            nodes_values=nodes_values,
+            nodes_modes=nodes_modes,
+            nodes_truenodeids=nodes_truenodeids,
+            nodes_falsenodeids=nodes_falsenodeids,
+            nodes_missing_value_tracks_true=nodes_missing_value_tracks_true,
+            target_treeids=target_treeids,
+            target_nodeids=target_nodeids,
+            target_ids=target_ids,
+            target_weights=target_weights,
+            base_values=base_values,
+            n_targets=n_classes if is_multi else 1,
+            name="PerpetualTreeEnsemble",
+        )
+
+        ops = [reg_node]
+        if is_classifier:
+            # Prepare class labels mapping
+            classes = self.classes_
+            if all(isinstance(c, (int, np.integer)) for c in classes):
+                tensor_type = TensorProto.INT64
+                classes_array = np.array(classes, dtype=np.int64)
+            elif all(isinstance(c, (float, np.floating)) for c in classes):
+                tensor_type = TensorProto.FLOAT
+                classes_array = np.array(classes, dtype=np.float32)
+            else:
+                tensor_type = TensorProto.STRING
+                classes_array = np.array([str(c) for c in classes], dtype=object)
+
+            classes_name = "class_labels"
+            if tensor_type == TensorProto.STRING:
+                classes_const_node = helper.make_node(
+                    "Constant",
+                    [],
+                    [classes_name],
+                    value=helper.make_tensor(
+                        name="classes_tensor",
+                        data_type=tensor_type,
+                        dims=[len(classes)],
+                        vals=[s.encode("utf-8") for s in classes_array],
+                    ),
+                )
+            else:
+                classes_const_node = helper.make_node(
+                    "Constant",
+                    [],
+                    [classes_name],
+                    value=helper.make_tensor(
+                        name="classes_tensor",
+                        data_type=tensor_type,
+                        dims=[len(classes)],
+                        vals=classes_array.flatten().tolist(),
+                    ),
+                )
+            ops.append(classes_const_node)
+
+            if is_multi:
+                prob_name = "probabilities"
+                softmax_node = helper.make_node(
+                    "Softmax", [raw_scores_name], [prob_name], axis=1
+                )
+                label_idx_name = "label_idx"
+                argmax_node = helper.make_node(
+                    "ArgMax", [prob_name], [label_idx_name], axis=1, keepdims=0
+                )
+                label_name = "label"
+                gather_node = helper.make_node(
+                    "Gather", [classes_name, label_idx_name], [label_name], axis=0
+                )
+                ops.extend([softmax_node, argmax_node, gather_node])
+                outputs = [
+                    helper.make_tensor_value_info(label_name, tensor_type, [None]),
+                    helper.make_tensor_value_info(
+                        prob_name, TensorProto.FLOAT, [None, n_classes]
+                    ),
+                ]
+            else:
+                p_name = "p"
+                sigmoid_node = helper.make_node("Sigmoid", [raw_scores_name], [p_name])
+                one_name = "one"
+                one_node = helper.make_node(
+                    "Constant",
+                    [],
+                    [one_name],
+                    value=helper.make_tensor("one_v", TensorProto.FLOAT, [1, 1], [1.0]),
+                )
+                one_minus_p_name = "one_minus_p"
+                sub_node = helper.make_node(
+                    "Sub", [one_name, p_name], [one_minus_p_name]
+                )
+                prob_name = "probabilities"
+                concat_node = helper.make_node(
+                    "Concat", [one_minus_p_name, p_name], [prob_name], axis=1
+                )
+                label_idx_name = "label_idx"
+                argmax_node = helper.make_node(
+                    "ArgMax", [prob_name], [label_idx_name], axis=1, keepdims=0
+                )
+                label_name = "label"
+                gather_node = helper.make_node(
+                    "Gather", [classes_name, label_idx_name], [label_name], axis=0
+                )
+                ops.extend(
+                    [
+                        sigmoid_node,
+                        one_node,
+                        sub_node,
+                        concat_node,
+                        argmax_node,
+                        gather_node,
+                    ]
+                )
+                outputs = [
+                    helper.make_tensor_value_info(label_name, tensor_type, [None]),
+                    helper.make_tensor_value_info(
+                        prob_name, TensorProto.FLOAT, [None, 2]
+                    ),
+                ]
+        else:
+            prediction_name = "prediction"
+            reg_node.output[0] = prediction_name
+            outputs = [
+                helper.make_tensor_value_info(
+                    prediction_name, TensorProto.FLOAT, [None, 1]
+                )
+            ]
+
+        graph_def = helper.make_graph(ops, name, [input_type], outputs)
+        model_def = helper.make_model(
+            graph_def,
+            producer_name="perpetual",
+            opset_imports=[
+                helper.make_opsetid("", 13),
+                helper.make_opsetid("ai.onnx.ml", 2),
+            ],
+        )
+        model_def.ir_version = 6
+        onnx.save(model_def, path)