Skip to content

model

XGBoostClassifierModel(classifier, *, input_features, output_features, threshold=0.5)

Bases: Model

Wrapper for an XGBoost classifier model with threshold support.

Source code in src/flowcean/xgboost/model.py 
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
def __init__(
    self,
    classifier: XGBClassifier,
    *,
    input_features: list[str],
    output_features: list[str],
    threshold: float = 0.5,
) -> None:
    self.classifier = classifier
    self.input_features = input_features
    self.output_features = output_features
    self.threshold = threshold
    # Initialize Protocol attributes
    self.pre_transform = Identity()
    self.post_transform = Identity()

predict_proba(input_features)

Predict class probabilities, applying preprocessing transforms.

Parameters:

Name Type Description Default
input_features LazyFrame

The inputs for which to predict probabilities.

 required

Returns:

Type Description
LazyFrame

The predicted probabilities for the positive class.
Source code in src/flowcean/xgboost/model.py 
47
48
49
50
51
52
53
54
55
56
57
def predict_proba(self, input_features: pl.LazyFrame) -> pl.LazyFrame:
    """Predict class probabilities, applying preprocessing transforms.

    Args:
        input_features: The inputs for which to predict probabilities.

    Returns:
        The predicted probabilities for the positive class.
    """
    input_features = self.preprocess(input_features)
    return self._predict_proba(input_features)

__getstate__()

Remove callbacks when pickling (they contain unpickleable locks).

Source code in src/flowcean/xgboost/model.py 
59
60
61
62
63
64
65
66
67
68
69
def __getstate__(self) -> dict:
    """Remove callbacks when pickling (they contain unpickleable locks)."""
    state = self.__dict__.copy()
    if "classifier" in state:
        # Copy the full fitted state and null out callbacks only
        classifier_dict = state["classifier"].__dict__.copy()
        classifier_dict["callbacks"] = None
        new_classifier = XGBClassifier.__new__(XGBClassifier)
        new_classifier.__dict__.update(classifier_dict)
        state["classifier"] = new_classifier
    return state

__setstate__(state)

Restore state after unpickling.

Source code in src/flowcean/xgboost/model.py 
71
72
73
def __setstate__(self, state: dict) -> None:
    """Restore state after unpickling."""
    self.__dict__.update(state)

XGBoostRegressorModel(regressor, *, input_features, output_features)

Bases: Model

Wrapper for an XGBoost regressor model.

Source code in src/flowcean/xgboost/model.py 
110
111
112
113
114
115
116
117
118
119
120
def __init__(
    self,
    regressor: XGBRegressor,
    *,
    input_features: list[str],
    output_features: list[str],
) -> None:
    super().__init__()
    self.regressor = regressor
    self.input_features = input_features
    self.output_features = output_features

__getstate__()

Remove callbacks when pickling (they contain unpickleable locks).

Source code in src/flowcean/xgboost/model.py 
122
123
124
125
126
127
128
129
130
131
132
def __getstate__(self) -> dict:
    """Remove callbacks when pickling (they contain unpickleable locks)."""
    state = self.__dict__.copy()
    if "regressor" in state:
        # Copy the full fitted state and null out callbacks only
        regressor_dict = state["regressor"].__dict__.copy()
        regressor_dict["callbacks"] = None
        new_regressor = XGBRegressor.__new__(XGBRegressor)
        new_regressor.__dict__.update(regressor_dict)
        state["regressor"] = new_regressor
    return state

__setstate__(state)

Restore state after unpickling.

Source code in src/flowcean/xgboost/model.py 
134
135
136
def __setstate__(self, state: dict) -> None:
    """Restore state after unpickling."""
    self.__dict__.update(state)