140 lines
4.3 KiB
Python
140 lines
4.3 KiB
Python
"""
|
|
XGBoost model wrapper for candlestick pattern classification.
|
|
|
|
Provides a wrapper around XGBoost's XGBClassifier with support for
|
|
class weight balancing.
|
|
"""
|
|
|
|
from typing import Any, Dict, Optional
|
|
|
|
import numpy as np
|
|
from xgboost import XGBClassifier
|
|
from sklearn.utils.class_weight import compute_class_weight
|
|
from sklearn.preprocessing import LabelEncoder
|
|
|
|
|
|
class XGBoostModel:
|
|
"""
|
|
XGBoost classifier wrapper for candlestick patterns.
|
|
|
|
Attributes:
|
|
model: The underlying XGBClassifier instance
|
|
classes_: Fitted class labels
|
|
feature_importances_: Feature importance scores (after fitting)
|
|
"""
|
|
|
|
def __init__(self, hyperparameters: Dict[str, Any], class_weights: Optional[str] = None):
|
|
"""
|
|
Initialize XGBoost model.
|
|
|
|
Args:
|
|
hyperparameters: Model hyperparameters from config
|
|
class_weights: "balanced" for inverse-frequency weighting, None for no weighting
|
|
"""
|
|
self.hyperparameters = hyperparameters.copy()
|
|
self.class_weights = class_weights
|
|
self._sample_weights = None
|
|
self.label_encoder_ = None
|
|
|
|
# XGBoost doesn't have built-in class_weight parameter like sklearn
|
|
# We'll compute sample weights manually when class_weights is "balanced"
|
|
|
|
# Initialize XGBoost model
|
|
self.model = XGBClassifier(**self.hyperparameters)
|
|
|
|
def fit(self, X: np.ndarray, y: np.ndarray):
|
|
"""
|
|
Train the XGBoost model.
|
|
|
|
Args:
|
|
X: Training features (n_samples, n_features)
|
|
y: Training labels (n_samples,)
|
|
|
|
Returns:
|
|
self
|
|
"""
|
|
classes = np.unique(y)
|
|
if classes.size < 2:
|
|
raise ValueError(
|
|
f"XGBoost requires at least 2 classes for training; got {classes.size} ({classes})"
|
|
)
|
|
|
|
y_encoded = y
|
|
if not (
|
|
np.issubdtype(np.asarray(y).dtype, np.integer)
|
|
and np.array_equal(np.sort(classes), np.arange(classes.size))
|
|
):
|
|
self.label_encoder_ = LabelEncoder()
|
|
y_encoded = self.label_encoder_.fit_transform(y)
|
|
|
|
# Compute sample weights if class weighting is enabled
|
|
if self.class_weights == "balanced":
|
|
# Compute class weights
|
|
class_weights = compute_class_weight(
|
|
class_weight="balanced",
|
|
classes=np.unique(y_encoded),
|
|
y=y_encoded
|
|
)
|
|
|
|
# Map class weights to sample weights
|
|
class_weight_dict = dict(zip(np.unique(y_encoded), class_weights))
|
|
sample_weights = np.array([class_weight_dict[label] for label in y_encoded])
|
|
|
|
# Fit with sample weights
|
|
self.model.fit(X, y_encoded, sample_weight=sample_weights)
|
|
else:
|
|
# Fit without sample weights
|
|
self.model.fit(X, y_encoded)
|
|
|
|
return self
|
|
|
|
def predict(self, X: np.ndarray) -> np.ndarray:
|
|
"""
|
|
Predict class labels.
|
|
|
|
Args:
|
|
X: Features (n_samples, n_features)
|
|
|
|
Returns:
|
|
Predicted labels (n_samples,)
|
|
"""
|
|
preds = self.model.predict(X)
|
|
if self.label_encoder_ is not None:
|
|
return self.label_encoder_.inverse_transform(preds.astype(int))
|
|
return preds
|
|
|
|
def predict_proba(self, X: np.ndarray) -> np.ndarray:
|
|
"""
|
|
Predict class probabilities.
|
|
|
|
Args:
|
|
X: Features (n_samples, n_features)
|
|
|
|
Returns:
|
|
Class probabilities (n_samples, n_classes)
|
|
"""
|
|
return self.model.predict_proba(X)
|
|
|
|
@property
|
|
def classes_(self):
|
|
"""Get fitted class labels."""
|
|
if self.label_encoder_ is not None:
|
|
return self.label_encoder_.classes_
|
|
return self.model.classes_
|
|
|
|
@property
|
|
def feature_importances_(self):
|
|
"""Get feature importance scores."""
|
|
return self.model.feature_importances_
|
|
|
|
def get_params(self) -> Dict[str, Any]:
|
|
"""
|
|
Get model parameters.
|
|
|
|
Returns:
|
|
Dictionary of model hyperparameters
|
|
"""
|
|
return self.model.get_params()
|
|
|
|
def __repr__(self):
|
|
return f"XGBoostModel(n_estimators={self.hyperparameters.get('n_estimators', 100)})"
|