Train

Supervised training with Optuna-backed hyperparameter search, feature selection, sample weight calculation, threshold optimisation, and structured reports. OptunaTrainer is the primary entry point for both classification and regression.

Configuration & State

bitbullet.train.core.config.TrainConfig dataclass

Configuration for model training.

Provides comprehensive control over the training process including hyperparameter optimization, feature selection, cross-validation, and evaluation strategies.

Example

config = TrainConfig(
    name="fraud_detector",
    model_type="lgbm",
    task="binary_classification",
    n_trials=100,
    cv_folds=5,
    optimization_metric="roc_auc",
    feature_selection="importance",
    threshold_optimization_method="youden"
)

__post_init__()

Validate configuration after initialization.

to_dict()

Convert config to dictionary.

Returns:

Type	Description
Dict[str, Any]	Dictionary representation of config

bitbullet.train.core.state.TrainState dataclass

Stores the complete state of a training run.

Captures everything needed to understand, reproduce, and analyze a training session, including SHAP-based explanations.

Attributes:

Name	Type	Description
model	Optional[BaseModel]	The trained model instance
best_params	Dict[str, Any]	Best hyperparameters found during optimization
best_score	float	Best CV score achieved
cv_scores	List[float]	Cross-validation scores for each fold
cv_predictions	Optional[DataFrame]	Out-of-fold predictions (if saved)
fold_metadata	List[Dict[str, Any]]	Auditable splitter boundaries for each CV fold
optimal_threshold	Optional[float]	Optimal classification threshold
feature_importance	Optional[DataFrame]	Feature importance scores (tree-based)
feature_names	List[str]	List of features used
n_features_selected	int	Number of features after selection
shap_values	Optional[ndarray]	SHAP values matrix for model explanations
shap_expected_value	Optional[Union[float, ndarray]]	Base value(s) for SHAP explanations
shap_feature_importance	Optional[DataFrame]	Feature importance from SHAP analysis
shap_background_data	Optional[DataFrame]	Background dataset used for SHAP
shap_interaction_values	Optional[ndarray]	SHAP interaction values (if calculated)
training_time_seconds	float	Total training time
optimization_time_seconds	float	Time spent on hyperparameter optimization
study	Optional[Any]	Optuna study object (if using Optuna)
training_history	Dict[str, List[float]]	Training metrics over time
fold_metrics	List[Dict[str, float]]	Metrics for each CV fold
started_at	Optional[datetime]	Training start timestamp
completed_at	Optional[datetime]	Training completion timestamp

cv_score_mean property

Get mean CV score.

cv_score_std property

Get CV score standard deviation.

is_complete property

Check if training is complete.

summary()

Generate a human-readable summary.

Returns:

Type	Description
str	Formatted summary string

Base Trainer

bitbullet.train.core.base.BaseTrainer

Bases: ABC

Abstract base class for all model trainers.

Implements the template method pattern, providing a consistent training workflow while allowing subclasses to customize specific steps like hyperparameter optimization.

The training workflow: 1. Validate input data 2. Feature selection (optional) 3. Sample weight calculation (optional) 4. Hyperparameter optimization 5. Train final model with best parameters 6. Threshold optimization (classification only) 7. Generate reports

Example

from bitbullet.train.trainers.optuna_trainer import OptunaTrainer
from bitbullet.train.core.config import TrainConfig

config = TrainConfig(
    name="my_model",
    model_type="lgbm",
    n_trials=50
)

trainer = OptunaTrainer(config)
model = trainer.fit(X_train, y_train, X_val, y_val)

best_params property

Get the best hyperparameters.

best_score property

Get the best CV score.

config property

Get the training configuration.

cv_scores property

Get CV scores for all folds.

feature_importance property

Get feature importance DataFrame.

optimal_threshold property

Get the optimal classification threshold.

state property

Get the training state.

__init__(config, prepare_final_features=None)

Initialize trainer.

Parameters:

Name	Type	Description	Default
config	TrainConfig	Training configuration	required
prepare_final_features	Optional[Callable[[], PreparedFeatureProvider[tuple[DataFrame, Optional[DataFrame]]]]]	Optional runtime preparer called after hyperparameter optimization. It returns a provider for the final training features and optional validation features. The provider is closed after all final-model consumers finish.	None

fit(X, y, X_val=None, y_val=None, sample_weight=None)

Train a model using the complete pipeline.

Parameters:

Name	Type	Description	Default
X	Union[DataFrame, ndarray]	Training features, or the row-identity carrier used by runtime feature providers.	required
y	Union[Series, ndarray]	Training targets.	required
X_val	Optional[Union[DataFrame, ndarray]]	Validation features or matching row-identity carrier.	None
y_val	Optional[Union[Series, ndarray]]	Validation targets.	None
sample_weight	Optional[Union[Series, ndarray]]	Optional row-aligned sample weights.	None

Returns:

Type	Description
BaseModel	Trained model.

Any prepared final feature provider is released before this method returns and is also released when training raises.

Trainers

Fold-fitted preprocessing

OptunaTrainer(..., fold_features=build_fold) accepts an optional runtime-only callback. It receives positional training and validation indices relative to the X/y supplied to fit, and returns two pandas feature frames. Fit a fresh preprocessor and any selector using only the training membership, then transform validation with that fitted preprocessor. Return the exact original row indices in their validated order and matching, unique columns. Folds may have different feature widths, for example because training-only category vocabularies differ.

The callback is called for every fold and trial. It is not a serialized inference component and is not used for the final fit: that fit uses the caller-supplied X and the configured final feature selection. The caller must retain the final fitted preprocessor separately for inference. Do not build validation features with a preprocessor fitted on the whole outer training partition. Omitting this callback preserves the existing trainer behaviour; it does not automatically make an externally preprocessed matrix fold-local.

For deterministic preprocessing that should be reused across trials, use prepare_fold_features instead. The preparer receives each validated fold's row positions exactly once and returns a PreparedFeatureProvider. Its load callback supplies the fixed training and validation frames whenever a trial uses that fold. Its close callback releases caller-owned resources as soon as the search finishes, including exceptional exits. A provider's loader must be safe for concurrent reads when n_jobs permits parallel trials. The loaded frames are treated as immutable. fold_features and prepare_fold_features cannot be combined.

prepare_final_features defers construction of the final training matrix until hyperparameter search has finished and every prepared fold provider has been closed. It returns a provider whose loader supplies (X_train, X_val_or_none). The frames must preserve the row identity supplied to fit, while their columns may reflect caller-owned preprocessing. Final feature selection, fitting, threshold analysis, importance, explanations, and reports complete before this provider is closed. Search with deferred final features requires a fold-local factory or preparer, because fitting trials against a different eager matrix would change the model-selection question. Omitting both new hooks preserves the existing eager training path.

bitbullet.train.core.features.PreparedFeatureProvider

Bases: Generic[PreparedValue]

Load a prepared feature value and release its caller-owned resources.

load may be called more than once. When a trainer runs trials in parallel, it may also be called concurrently, so callers choosing a lazy loader must make those reads safe. close is idempotent and the trainer calls it once the value can no longer be requested, including error paths.

close()

Release the provider exactly once.

fixed(value) classmethod

Create a provider for an already prepared in-memory value.

load()

Return the prepared value while this provider is active.

TrainConfig(protected_features=["structural_input"]) requires those inputs in the final and fold matrices and retains them through feature selection. The maximum selected feature count must accommodate every protected input. This is explicit caller configuration, not automatic feature inference.

bitbullet.train.trainers.optuna_trainer.OptunaTrainer

Bases: BaseTrainer

Trainer using Optuna for hyperparameter optimization.

Features: - Optuna TPE sampler for efficient search - Stratified K-Fold cross-validation - Early stopping with best iteration tracking - Automatic handling of sample weights - Comprehensive metric tracking

Example

from bitbullet.train.trainers.optuna_trainer import OptunaTrainer
from bitbullet.train.core.config import TrainConfig

config = TrainConfig(
    name="fraud_detector",
    model_type="lgbm",
    task="binary_classification",
    n_trials=100,
    cv_folds=5,
    optimize_threshold=True
)

trainer = OptunaTrainer(config)
model = trainer.fit(X_train, y_train, X_val, y_val)

# Access results
print(f"Best CV score: {trainer.best_score:.4f}")
print(f"Best params: {trainer.best_params}")
print(f"Optimal threshold: {trainer.optimal_threshold:.3f}")

# Feature importance
print(trainer.feature_importance.head(10))

__init__(config, progress_callback=None, cv_splitter=None, fold_features=None, prepare_fold_features=None, prepare_final_features=None)

Initialize Optuna trainer.

Parameters:

Name	Type	Description	Default
config	TrainConfig	Training configuration	required
progress_callback	Optional[Callable]	Optional callback function called after each trial. Receives (trial_number: int, score: float, params: dict)	None
cv_splitter	Optional[Any]	Optional sklearn-compatible splitter. When supplied, it replaces the strategy constructed from TrainConfig. Splitters exposing split_with_metadata also populate state.fold_metadata for audit and reporting.	None
fold_features	Optional[Callable[[ndarray, ndarray], tuple[DataFrame, DataFrame]]]	Optional runtime factory receiving validated train and validation row positions. It must fit preprocessing and any feature selection using only those training rows, returning aligned train/validation feature frames. It is called afresh for each fold/trial and is never used for the final fit, which uses the X explicitly supplied to fit(). This lets callers keep raw covariates outside the trainer without leaking transforms fitted on the complete training partition into inner CV.	None
prepare_fold_features	Optional[Callable[[ndarray, ndarray], PreparedFeatureProvider[tuple[DataFrame, DataFrame]]]]	Optional runtime preparer called once for each validated fold before the search starts. It returns a provider that supplies the same training/validation features to every trial and is closed as soon as search finishes. This is mutually exclusive with fold_features.	None
prepare_final_features	Optional[Callable[[], PreparedFeatureProvider[tuple[DataFrame, Optional[DataFrame]]]]]	Optional runtime preparer called only after search finishes. It returns a provider for the final training features and optional external-validation features. The row membership must match the inputs supplied to fit.	None

Optimizers

bitbullet.train.optimizers.optuna_optimizer.OptunaOptimizer

Optuna hyperparameter optimizer.

Wraps Optuna's study API with sensible defaults and enhanced functionality for ML model optimization.

Example

optimizer = OptunaOptimizer(
    direction="maximize",
    sampler="tpe",
    random_state=42
)

def objective(trial):
    # Define hyperparameter search space
    params = {
        'n_estimators': trial.suggest_int('n_estimators', 50, 500),
        'max_depth': trial.suggest_int('max_depth', 3, 10),
        'learning_rate': trial.suggest_float('learning_rate', 0.01, 0.3, log=True)
    }

    # Train model and return score
    model = LGBMClassifier(**params)
    score = cross_val_score(model, X, y, cv=5).mean()
    return score

study = optimizer.optimize(objective, n_trials=100)
best_params = study.best_params

__init__(direction='maximize', sampler='tpe', random_state=42, study_name=None, storage=None, search_space=None)

Initialize Optuna optimizer.

Parameters:

Name	Type	Description	Default
direction	str	'maximize' or 'minimize'	'maximize'
sampler	str	'tpe', 'random', or 'cmaes'	'tpe'
random_state	int	Random seed for reproducibility	42
study_name	Optional[str]	Name for the study (for persistence)	None
storage	Optional[str]	Storage URL for study persistence (e.g., 'sqlite:///optuna.db')	None

extract_best_params(study) staticmethod

Extract best parameters from completed study.

Includes special handling for iteration-based parameters (e.g., best_iteration from early stopping).

Parameters:

Name	Type	Description	Default
study	Study	Completed Optuna study	required

Returns:

Type	Description
Dict[str, Any]	Dictionary of best parameters

Example

study = optimizer.optimize(objective, n_trials=100)
best_params = OptunaOptimizer.extract_best_params(study)

# Train final model
model = LGBMClassifier(**best_params)
model.fit(X, y)

get_trial_history(study) staticmethod

Get history of all trials.

Parameters:

Name	Type	Description	Default
study	Study	Completed Optuna study	required

Returns:

Type	Description
Dict[str, list]	Dictionary with trial history

Example

history = OptunaOptimizer.get_trial_history(study)

import matplotlib.pyplot as plt
plt.plot(history['trial_numbers'], history['values'])
plt.xlabel('Trial')
plt.ylabel('Score')
plt.show()

optimize(objective, n_trials=100, timeout=None, n_jobs=1, show_progress_bar=True, callbacks=None)

Run hyperparameter optimization.

Parameters:

Name	Type	Description	Default
objective	Callable	Objective function that takes a trial and returns a score	required
n_trials	int	Number of trials to run	100
timeout	Optional[float]	Time limit in seconds (optional)	None
n_jobs	int	Number of parallel jobs (-1 for all CPUs)	1
show_progress_bar	bool	Whether to show progress bar	True
callbacks	Optional[list]	List of callback functions	None

Returns:

Type	Description
Study	Completed Optuna study

Example

def objective(trial):
    params = trial.suggest_int('n_estimators', 50, 500)
    # Train and evaluate model
    return score

study = optimizer.optimize(
    objective=objective,
    n_trials=100,
    n_jobs=-1,
    show_progress_bar=True
)

print(f"Best score: {study.best_value}")
print(f"Best params: {study.best_params}")

print_study_summary(study) staticmethod

Print a summary of the optimization study.

Parameters:

Name	Type	Description	Default
study	Study	Completed Optuna study	required

Feature Selection

bitbullet.train.feature_selection.selector_factory.FeatureSelector

Factory for creating feature selectors.

create(method, **kwargs) staticmethod

Create a feature selector.

Parameters:

Name	Type	Description	Default
method	str	Selection method - 'importance': Tree-based feature importance (fast) - 'correlation': Remove correlated features - 'mutual_info': Mutual information with target - 'shap': SHAP-based importance (most accurate, slower) - 'rfe': Recursive Feature Elimination (thorough, slow) - 'statistical': Mutual information with correlation redundancy filtering	required
**kwargs		Method-specific parameters	{}

Returns:

Type	Description
BaseFeatureSelector	Feature selector instance

Example

# Tree importance (fast, good baseline)
selector = FeatureSelector.create('importance', top_k=50)
selected = selector.select(X, y)

# SHAP (most accurate, recommended for production)
selector = FeatureSelector.create('shap', top_k=50, model_type='lgbm')
selected = selector.select(X, y)

# Correlation-based (removes redundant features)
selector = FeatureSelector.create('correlation', threshold=0.95)
selected = selector.select(X, y)

# RFE (thorough, slow)
selector = FeatureSelector.create('rfe', n_features=50)
selected = selector.select(X, y)

# Statistical selector (model-free redundancy filtering)
selector = FeatureSelector.create('statistical', top_k=50, threshold=0.8)
selected = selector.select(X, y)

Sample Weights

bitbullet.train.utils.sample_weights.SampleWeightCalculator

Factory for creating sample weight calculators.

create(strategy, **kwargs) staticmethod

Create a sample weight calculator.

Parameters:

Name	Type	Description	Default
strategy	str	Strategy name ('exponential_decay', 'class_balance', 'custom')	required
**kwargs		Strategy-specific parameters	{}

Returns:

Type	Description
BaseSampleWeightCalculator	Sample weight calculator instance

Example

# Exponential decay for time-series
calculator = SampleWeightCalculator.create(
    'exponential_decay',
    time_column='days_diff',
    threshold_days=180,
    decay_rate=0.01
)
weights = calculator.calculate(X, y)

# Class balancing
calculator = SampleWeightCalculator.create('class_balance')
weights = calculator.calculate(X, y)

Threshold Optimisation

bitbullet.train.evaluation.threshold_optimizer.ThresholdOptimizer

Optimizes classification thresholds.

Supports multiple optimization methods: - Youden's Index (TPR - FPR) - F1 Score maximization - Custom metric optimization

Example

optimizer = ThresholdOptimizer(method="youden")
optimal_threshold, metrics = optimizer.optimize(y_true, y_pred_proba)

print(f"Optimal threshold: {optimal_threshold:.3f}")
print(f"Metrics at threshold: {metrics}")

__init__(method='youden')

Initialize threshold optimizer.

Parameters:

Name	Type	Description	Default
method	str	Optimization method ('youden', 'f1', 'custom')	'youden'

optimize(y_true, y_pred_proba)

Find optimal threshold.

Parameters:

Name	Type	Description	Default
y_true	ndarray	True labels	required
y_pred_proba	ndarray	Predicted probabilities	required

Returns:

Type	Description
Tuple[float, Dict[str, float]]	Tuple of (optimal_threshold, metrics_dict)

Reports

bitbullet.train.reports.training_report.TrainingReport dataclass

Container for training report data.

__str__()

Generate human-readable report string.

to_dict()

Convert report to dictionary for serialization.

to_json(path)

Save report to JSON file.

bitbullet.train.reports.training_report.TrainingReportGenerator

Generate comprehensive training reports.

from_trainer_state(trainer_state, config, study=None) staticmethod

Generate report from trainer state.

Parameters:

Name	Type	Description	Default
trainer_state		Trainer state object	required
config		Training configuration	required
study	Optional[Study]	Optuna study (optional)	None

Returns:

Type	Description
TrainingReport	TrainingReport instance

generate_optuna_importance_report(study, top_n=10) staticmethod

Generate hyperparameter importance report.

Analyzes which hyperparameters had the most impact on performance.

Parameters:

Name	Type	Description	Default
study	Study	Completed Optuna study	required
top_n	int	Number of top parameters to show	10

Returns:

Type	Description
DataFrame	DataFrame with parameter importance

generate_trial_history(study) staticmethod

Generate trial history DataFrame.

Parameters:

Name	Type	Description	Default
study	Study	Completed Optuna study	required

Returns:

Type	Description
DataFrame	DataFrame with trial history

print_training_summary(report) staticmethod

Print training summary to console.

Parameters:

Name	Type	Description	Default
report	TrainingReport	TrainingReport instance	required

save_full_report(report, study, output_dir) staticmethod

Save complete training report with all artifacts.

Parameters:

Name	Type	Description	Default
report	TrainingReport	TrainingReport instance	required
study	Optional[Study]	Optuna study (optional)	required
output_dir	str	Directory to save reports	required

bitbullet.train.reports.feature_report.FeatureReport dataclass

Container for feature analysis report data.

__str__()

Generate human-readable report string.

to_dict()

Convert report to dictionary for serialization.

bitbullet.train.reports.feature_report.FeatureReportGenerator

Generate comprehensive feature analysis reports.

Inspired by the assess_features method from rfi codebase but enhanced with additional analysis capabilities.

assess_feature_stability(importance_dfs, top_k=20) staticmethod

Assess stability of feature importance across multiple runs.

Useful for understanding which features are consistently important across different CV folds or training runs.

Parameters:

Name	Type	Description	Default
importance_dfs	List[DataFrame]	List of importance DataFrames from different runs	required
top_k	int	Number of top features to analyze	20

Returns:

Type	Description
DataFrame	DataFrame with stability metrics

from_model(model, feature_names=None, X=None, correlation_threshold=0.95) staticmethod

Generate feature report from a trained model.

Parameters:

Name	Type	Description	Default
model		Trained model with feature_importances_ attribute	required
feature_names	Optional[List[str]]	List of feature names	None
X	Optional[DataFrame]	Training data (for correlation analysis)	None
correlation_threshold	float	Threshold for high correlation detection	0.95

Returns:

Type	Description
FeatureReport	FeatureReport instance

from_selection_results(importance_df, selected_features, original_features, selection_method, X=None, correlation_threshold=0.95) staticmethod

Generate feature report from feature selection results.

Parameters:

Name	Type	Description	Default
importance_df	DataFrame	DataFrame with feature importance	required
selected_features	List[str]	List of selected features	required
original_features	List[str]	List of original features	required
selection_method	str	Method used for selection	required
X	Optional[DataFrame]	Training data (for correlation analysis)	None
correlation_threshold	float	Threshold for high correlation detection	0.95

Returns:

Type	Description
FeatureReport	FeatureReport instance

generate_feature_summary_stats(importance_df) staticmethod

Generate summary statistics for feature importance.

Parameters:

Name	Type	Description	Default
importance_df	DataFrame	DataFrame with feature importance	required

Returns:

Type	Description
Dict[str, float]	Dictionary with summary statistics

save_feature_report(report, output_dir, include_csv=True) staticmethod

Save feature report to files.

Parameters:

Name	Type	Description	Default
report	FeatureReport	FeatureReport instance	required
output_dir	str	Directory to save reports	required
include_csv	bool	Whether to save CSV files	True