first commit

64e892b verified 10 months ago

12.3 kB

"""
Script 05: Model Training

This script trains the ordinal classification model:
- Uses LightGBM for multi-class ordinal classification
- Implements class weighting for imbalanced data
- Performs cross-validation
- Includes hyperparameter tuning with Optuna
- Saves the trained model

Ordinal Classification Approach:
Since fire size classes have a natural order (Small < Medium < Large),
we use ordinal-aware training with cumulative link model concepts.

Usage:
    python scripts/05_train_model.py [--tune]
"""

import argparse
import json
import sys
from pathlib import Path

import joblib
import lightgbm as lgb
import numpy as np
import optuna
import pandas as pd
from sklearn.metrics import (
    accuracy_score,
    balanced_accuracy_score,
    classification_report,
    cohen_kappa_score,
    f1_score,
)
from sklearn.model_selection import StratifiedKFold
from sklearn.utils.class_weight import compute_class_weight

# Add project root to path
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from config.config import (
    TRAIN_PARQUET,
    TEST_PARQUET,
    MODELS_DIR,
    TARGET_COLUMN,
    TARGET_CLASS_NAMES,
    LIGHTGBM_PARAMS,
    OPTUNA_SEARCH_SPACE,
    N_OPTUNA_TRIALS,
    N_FOLDS,
    RANDOM_STATE,
    USE_CLASS_WEIGHTS,
    PRIMARY_METRIC
)


def load_data() -> tuple[pd.DataFrame, pd.DataFrame]:
    """Load train and test data."""
    print("Loading data...")
    train_df = pd.read_parquet(TRAIN_PARQUET)
    test_df = pd.read_parquet(TEST_PARQUET)
    print(f"  Train: {len(train_df):,} rows")
    print(f"  Test: {len(test_df):,} rows")
    return train_df, test_df


def get_feature_columns(df: pd.DataFrame) -> list:
    """Get list of feature columns for modeling."""
    exclude_cols = [
        TARGET_COLUMN,
        'NWCG_REPORTING_AGENCY', 'STAT_CAUSE_DESCR', 'STATE', 'OWNER_DESCR',
        'COUNTY'
    ]
    feature_cols = [col for col in df.columns if col not in exclude_cols]
    return feature_cols


def prepare_data(train_df: pd.DataFrame, test_df: pd.DataFrame) -> tuple:
    """Prepare features and targets for training."""
    print("\nPreparing data...")
    
    feature_cols = get_feature_columns(train_df)
    
    X_train = train_df[feature_cols].values
    y_train = train_df[TARGET_COLUMN].values
    X_test = test_df[feature_cols].values
    y_test = test_df[TARGET_COLUMN].values
    
    print(f"  Features: {len(feature_cols)}")
    print(f"  Feature columns: {feature_cols}")
    
    return X_train, y_train, X_test, y_test, feature_cols


def compute_weights(y_train: np.ndarray) -> np.ndarray:
    """Compute sample weights for class imbalance."""
    print("\nComputing class weights...")
    
    classes = np.unique(y_train)
    class_weights = compute_class_weight(
        class_weight='balanced',
        classes=classes,
        y=y_train
    )
    
    weight_dict = dict(zip(classes, class_weights))
    print(f"  Class weights: {weight_dict}")
    
    # Create sample weights array
    sample_weights = np.array([weight_dict[y] for y in y_train])
    
    return sample_weights


def evaluate_model(y_true: np.ndarray, y_pred: np.ndarray, prefix: str = "") -> dict:
    """Evaluate model predictions."""
    metrics = {
        'accuracy': accuracy_score(y_true, y_pred),
        'balanced_accuracy': balanced_accuracy_score(y_true, y_pred),
        'macro_f1': f1_score(y_true, y_pred, average='macro'),
        'weighted_f1': f1_score(y_true, y_pred, average='weighted'),
        'cohen_kappa': cohen_kappa_score(y_true, y_pred, weights='linear')  # Linear weights for ordinal
    }
    
    if prefix:
        print(f"\n{prefix} Metrics:")
    else:
        print("\nMetrics:")
    
    for name, value in metrics.items():
        print(f"  {name}: {value:.4f}")
    
    return metrics


def cross_validate(X: np.ndarray, y: np.ndarray, params: dict, 
                   sample_weights: np.ndarray = None) -> tuple[float, float]:
    """Perform cross-validation and return mean and std of primary metric."""
    
    skf = StratifiedKFold(n_splits=N_FOLDS, shuffle=True, random_state=RANDOM_STATE)
    scores = []
    
    for fold, (train_idx, val_idx) in enumerate(skf.split(X, y)):
        X_fold_train, X_fold_val = X[train_idx], X[val_idx]
        y_fold_train, y_fold_val = y[train_idx], y[val_idx]
        
        if sample_weights is not None:
            weights_fold = sample_weights[train_idx]
        else:
            weights_fold = None
        
        # Create LightGBM datasets
        train_data = lgb.Dataset(X_fold_train, label=y_fold_train, weight=weights_fold)
        val_data = lgb.Dataset(X_fold_val, label=y_fold_val, reference=train_data)
        
        # Train model
        model = lgb.train(
            params,
            train_data,
            num_boost_round=params.get('n_estimators', 500),
            valid_sets=[val_data],
            callbacks=[lgb.early_stopping(50, verbose=False)]
        )
        
        # Predict
        y_pred = model.predict(X_fold_val)
        y_pred_class = np.argmax(y_pred, axis=1)
        
        # Score
        score = f1_score(y_fold_val, y_pred_class, average='macro')
        scores.append(score)
    
    return np.mean(scores), np.std(scores)


def objective(trial: optuna.Trial, X: np.ndarray, y: np.ndarray, 
              sample_weights: np.ndarray) -> float:
    """Optuna objective function for hyperparameter tuning."""
    
    params = LIGHTGBM_PARAMS.copy()
    
    # Sample hyperparameters
    params['n_estimators'] = trial.suggest_int('n_estimators', *OPTUNA_SEARCH_SPACE['n_estimators'])
    params['max_depth'] = trial.suggest_int('max_depth', *OPTUNA_SEARCH_SPACE['max_depth'])
    params['learning_rate'] = trial.suggest_float('learning_rate', *OPTUNA_SEARCH_SPACE['learning_rate'], log=True)
    params['num_leaves'] = trial.suggest_int('num_leaves', *OPTUNA_SEARCH_SPACE['num_leaves'])
    params['min_child_samples'] = trial.suggest_int('min_child_samples', *OPTUNA_SEARCH_SPACE['min_child_samples'])
    params['subsample'] = trial.suggest_float('subsample', *OPTUNA_SEARCH_SPACE['subsample'])
    params['colsample_bytree'] = trial.suggest_float('colsample_bytree', *OPTUNA_SEARCH_SPACE['colsample_bytree'])
    params['reg_alpha'] = trial.suggest_float('reg_alpha', *OPTUNA_SEARCH_SPACE['reg_alpha'])
    params['reg_lambda'] = trial.suggest_float('reg_lambda', *OPTUNA_SEARCH_SPACE['reg_lambda'])
    
    # Cross-validate
    mean_score, _ = cross_validate(X, y, params, sample_weights)
    
    return mean_score


def tune_hyperparameters(X: np.ndarray, y: np.ndarray, 
                         sample_weights: np.ndarray) -> dict:
    """Tune hyperparameters using Optuna."""
    print("\n" + "="*60)
    print("HYPERPARAMETER TUNING")
    print("="*60)
    
    print(f"\nRunning {N_OPTUNA_TRIALS} Optuna trials...")
    
    # Create study
    study = optuna.create_study(
        direction='maximize',
        sampler=optuna.samplers.TPESampler(seed=RANDOM_STATE)
    )
    
    # Optimize
    study.optimize(
        lambda trial: objective(trial, X, y, sample_weights),
        n_trials=N_OPTUNA_TRIALS,
        show_progress_bar=True
    )
    
    print(f"\nBest trial:")
    print(f"  Value (macro F1): {study.best_trial.value:.4f}")
    print(f"  Params: {study.best_trial.params}")
    
    # Merge best params with base params
    best_params = LIGHTGBM_PARAMS.copy()
    best_params.update(study.best_trial.params)
    
    best_params_path = MODELS_DIR / 'best_params.json'
    with open(best_params_path, 'w') as f:
        json.dump(study.best_trial.params, f)
    print(f"  Best params saved: {best_params_path}")

    return best_params


def train_final_model(X_train: np.ndarray, y_train: np.ndarray,
                      X_test: np.ndarray, y_test: np.ndarray,
                      params: dict, sample_weights: np.ndarray,
                      feature_names: list) -> lgb.Booster:
    """Train final model on full training data."""
    print("\n" + "="*60)
    print("TRAINING FINAL MODEL")
    print("="*60)
    
    # Create datasets
    train_data = lgb.Dataset(X_train, label=y_train, weight=sample_weights,
                             feature_name=feature_names)
    val_data = lgb.Dataset(X_test, label=y_test, reference=train_data,
                           feature_name=feature_names)
    
    # Train
    print("\nTraining...")
    model = lgb.train(
        params,
        train_data,
        num_boost_round=params.get('n_estimators', 2000),
        valid_sets=[train_data, val_data],
        valid_names=['train', 'test'],
        callbacks=[
            lgb.early_stopping(50, verbose=True),
            lgb.log_evaluation(period=50)
        ]
    )
    
    # Evaluate
    print("\n" + "-"*40)
    
    # Train predictions
    y_train_pred = np.argmax(model.predict(X_train), axis=1)
    evaluate_model(y_train, y_train_pred, "Train")
    
    # Test predictions
    y_test_pred = np.argmax(model.predict(X_test), axis=1)
    test_metrics = evaluate_model(y_test, y_test_pred, "Test")
    
    # Classification report
    print("\nClassification Report (Test):")
    print(classification_report(y_test, y_test_pred, target_names=TARGET_CLASS_NAMES))
    
    return model, test_metrics


def save_model(model: lgb.Booster, params: dict, feature_names: list, metrics: dict) -> None:
    """Save trained model and metadata."""
    print("\nSaving model...")
    
    MODELS_DIR.mkdir(parents=True, exist_ok=True)
    
    # Save LightGBM model
    model_path = MODELS_DIR / 'wildfire_model.txt'
    model.save_model(str(model_path))
    print(f"  Model: {model_path}")
    
    # Save metadata
    metadata = {
        'params': params,
        'feature_names': feature_names,
        'metrics': metrics,
        'target_classes': TARGET_CLASS_NAMES
    }
    metadata_path = MODELS_DIR / 'model_metadata.joblib'
    joblib.dump(metadata, metadata_path)
    print(f"  Metadata: {metadata_path}")


def main():
    """Main training pipeline."""
    # Parse arguments
    parser = argparse.ArgumentParser(description='Train wildfire classification model')
    parser.add_argument('--tune', action='store_true', help='Run hyperparameter tuning')
    args = parser.parse_args()
    
    print("\n" + "="*60)
    print("MODEL TRAINING")
    print("="*60)
    
    # Load data
    train_df, test_df = load_data()
    
    # Prepare data
    X_train, y_train, X_test, y_test, feature_cols = prepare_data(train_df, test_df)
    
    # Compute class weights
    sample_weights = None
    if USE_CLASS_WEIGHTS:
        sample_weights = compute_weights(y_train)
    
    # Get parameters
    if args.tune:
        params = tune_hyperparameters(X_train, y_train, sample_weights)
    else:
        best_params_path = MODELS_DIR / 'best_params.json'
        if best_params_path.exists():
            # Load saved best params
            with open(best_params_path, 'r') as f:
                tuned_params = json.load(f)
            params = LIGHTGBM_PARAMS.copy()
            params.update(tuned_params)
            print(f"Loaded best params from {best_params_path}")
        else:
            # Fallback to defaults
            params = LIGHTGBM_PARAMS.copy()
            params['n_estimators'] = 500
            params['max_depth'] = 8
            params['learning_rate'] = 0.05
            params['num_leaves'] = 64
            params['min_child_samples'] = 50
            params['subsample'] = 0.8
            params['colsample_bytree'] = 0.8
            print("No saved params found; using defaults")
    
    # Train final model
    model, metrics = train_final_model(
        X_train, y_train, X_test, y_test,
        params, sample_weights, feature_cols
    )
    
    # Save model
    save_model(model, params, feature_cols, metrics)
    
    print("\n" + "="*60)
    print("✓ Training Complete!")
    print("="*60 + "\n")


if __name__ == "__main__":
    main()

	"""
	Script 05: Model Training

	This script trains the ordinal classification model:
	- Uses LightGBM for multi-class ordinal classification
	- Implements class weighting for imbalanced data
	- Performs cross-validation
	- Includes hyperparameter tuning with Optuna
	- Saves the trained model

	Ordinal Classification Approach:
	Since fire size classes have a natural order (Small < Medium < Large),
	we use ordinal-aware training with cumulative link model concepts.

	Usage:
	python scripts/05_train_model.py [--tune]
	"""

	import argparse
	import json
	import sys
	from pathlib import Path

	import joblib
	import lightgbm as lgb
	import numpy as np
	import optuna
	import pandas as pd
	from sklearn.metrics import (
	accuracy_score,
	balanced_accuracy_score,
	classification_report,
	cohen_kappa_score,
	f1_score,
	)
	from sklearn.model_selection import StratifiedKFold
	from sklearn.utils.class_weight import compute_class_weight

	# Add project root to path
	project_root = Path(__file__).parent.parent
	sys.path.insert(0, str(project_root))

	from config.config import (
	TRAIN_PARQUET,
	TEST_PARQUET,
	MODELS_DIR,
	TARGET_COLUMN,
	TARGET_CLASS_NAMES,
	LIGHTGBM_PARAMS,
	OPTUNA_SEARCH_SPACE,
	N_OPTUNA_TRIALS,
	N_FOLDS,
	RANDOM_STATE,
	USE_CLASS_WEIGHTS,
	PRIMARY_METRIC
	)


	def load_data() -> tuple[pd.DataFrame, pd.DataFrame]:
	"""Load train and test data."""
	print("Loading data...")
	train_df = pd.read_parquet(TRAIN_PARQUET)
	test_df = pd.read_parquet(TEST_PARQUET)
	print(f" Train: {len(train_df):,} rows")
	print(f" Test: {len(test_df):,} rows")
	return train_df, test_df


	def get_feature_columns(df: pd.DataFrame) -> list:
	"""Get list of feature columns for modeling."""
	exclude_cols = [
	TARGET_COLUMN,
	'NWCG_REPORTING_AGENCY', 'STAT_CAUSE_DESCR', 'STATE', 'OWNER_DESCR',
	'COUNTY'
	]
	feature_cols = [col for col in df.columns if col not in exclude_cols]
	return feature_cols


	def prepare_data(train_df: pd.DataFrame, test_df: pd.DataFrame) -> tuple:
	"""Prepare features and targets for training."""
	print("\nPreparing data...")

	feature_cols = get_feature_columns(train_df)

	X_train = train_df[feature_cols].values
	y_train = train_df[TARGET_COLUMN].values
	X_test = test_df[feature_cols].values
	y_test = test_df[TARGET_COLUMN].values

	print(f" Features: {len(feature_cols)}")
	print(f" Feature columns: {feature_cols}")

	return X_train, y_train, X_test, y_test, feature_cols


	def compute_weights(y_train: np.ndarray) -> np.ndarray:
	"""Compute sample weights for class imbalance."""
	print("\nComputing class weights...")

	classes = np.unique(y_train)
	class_weights = compute_class_weight(
	class_weight='balanced',
	classes=classes,
	y=y_train
	)

	weight_dict = dict(zip(classes, class_weights))
	print(f" Class weights: {weight_dict}")

	# Create sample weights array
	sample_weights = np.array([weight_dict[y] for y in y_train])

	return sample_weights


	def evaluate_model(y_true: np.ndarray, y_pred: np.ndarray, prefix: str = "") -> dict:
	"""Evaluate model predictions."""
	metrics = {
	'accuracy': accuracy_score(y_true, y_pred),
	'balanced_accuracy': balanced_accuracy_score(y_true, y_pred),
	'macro_f1': f1_score(y_true, y_pred, average='macro'),
	'weighted_f1': f1_score(y_true, y_pred, average='weighted'),
	'cohen_kappa': cohen_kappa_score(y_true, y_pred, weights='linear') # Linear weights for ordinal
	}

	if prefix:
	print(f"\n{prefix} Metrics:")
	else:
	print("\nMetrics:")

	for name, value in metrics.items():
	print(f" {name}: {value:.4f}")

	return metrics


	def cross_validate(X: np.ndarray, y: np.ndarray, params: dict,
	sample_weights: np.ndarray = None) -> tuple[float, float]:
	"""Perform cross-validation and return mean and std of primary metric."""

	skf = StratifiedKFold(n_splits=N_FOLDS, shuffle=True, random_state=RANDOM_STATE)
	scores = []

	for fold, (train_idx, val_idx) in enumerate(skf.split(X, y)):
	X_fold_train, X_fold_val = X[train_idx], X[val_idx]
	y_fold_train, y_fold_val = y[train_idx], y[val_idx]

	if sample_weights is not None:
	weights_fold = sample_weights[train_idx]
	else:
	weights_fold = None

	# Create LightGBM datasets
	train_data = lgb.Dataset(X_fold_train, label=y_fold_train, weight=weights_fold)
	val_data = lgb.Dataset(X_fold_val, label=y_fold_val, reference=train_data)

	# Train model
	model = lgb.train(
	params,
	train_data,
	num_boost_round=params.get('n_estimators', 500),
	valid_sets=[val_data],
	callbacks=[lgb.early_stopping(50, verbose=False)]
	)

	# Predict
	y_pred = model.predict(X_fold_val)
	y_pred_class = np.argmax(y_pred, axis=1)

	# Score
	score = f1_score(y_fold_val, y_pred_class, average='macro')
	scores.append(score)

	return np.mean(scores), np.std(scores)


	def objective(trial: optuna.Trial, X: np.ndarray, y: np.ndarray,
	sample_weights: np.ndarray) -> float:
	"""Optuna objective function for hyperparameter tuning."""

	params = LIGHTGBM_PARAMS.copy()

	# Sample hyperparameters
	params['n_estimators'] = trial.suggest_int('n_estimators', *OPTUNA_SEARCH_SPACE['n_estimators'])
	params['max_depth'] = trial.suggest_int('max_depth', *OPTUNA_SEARCH_SPACE['max_depth'])
	params['learning_rate'] = trial.suggest_float('learning_rate', *OPTUNA_SEARCH_SPACE['learning_rate'], log=True)
	params['num_leaves'] = trial.suggest_int('num_leaves', *OPTUNA_SEARCH_SPACE['num_leaves'])
	params['min_child_samples'] = trial.suggest_int('min_child_samples', *OPTUNA_SEARCH_SPACE['min_child_samples'])
	params['subsample'] = trial.suggest_float('subsample', *OPTUNA_SEARCH_SPACE['subsample'])
	params['colsample_bytree'] = trial.suggest_float('colsample_bytree', *OPTUNA_SEARCH_SPACE['colsample_bytree'])
	params['reg_alpha'] = trial.suggest_float('reg_alpha', *OPTUNA_SEARCH_SPACE['reg_alpha'])
	params['reg_lambda'] = trial.suggest_float('reg_lambda', *OPTUNA_SEARCH_SPACE['reg_lambda'])

	# Cross-validate
	mean_score, _ = cross_validate(X, y, params, sample_weights)

	return mean_score


	def tune_hyperparameters(X: np.ndarray, y: np.ndarray,
	sample_weights: np.ndarray) -> dict:
	"""Tune hyperparameters using Optuna."""
	print("\n" + "="*60)
	print("HYPERPARAMETER TUNING")
	print("="*60)

	print(f"\nRunning {N_OPTUNA_TRIALS} Optuna trials...")

	# Create study
	study = optuna.create_study(
	direction='maximize',
	sampler=optuna.samplers.TPESampler(seed=RANDOM_STATE)
	)

	# Optimize
	study.optimize(
	lambda trial: objective(trial, X, y, sample_weights),
	n_trials=N_OPTUNA_TRIALS,
	show_progress_bar=True
	)

	print(f"\nBest trial:")
	print(f" Value (macro F1): {study.best_trial.value:.4f}")
	print(f" Params: {study.best_trial.params}")

	# Merge best params with base params
	best_params = LIGHTGBM_PARAMS.copy()
	best_params.update(study.best_trial.params)

	best_params_path = MODELS_DIR / 'best_params.json'
	with open(best_params_path, 'w') as f:
	json.dump(study.best_trial.params, f)
	print(f" Best params saved: {best_params_path}")

	return best_params


	def train_final_model(X_train: np.ndarray, y_train: np.ndarray,
	X_test: np.ndarray, y_test: np.ndarray,
	params: dict, sample_weights: np.ndarray,
	feature_names: list) -> lgb.Booster:
	"""Train final model on full training data."""
	print("\n" + "="*60)
	print("TRAINING FINAL MODEL")
	print("="*60)

	# Create datasets
	train_data = lgb.Dataset(X_train, label=y_train, weight=sample_weights,
	feature_name=feature_names)
	val_data = lgb.Dataset(X_test, label=y_test, reference=train_data,
	feature_name=feature_names)

	# Train
	print("\nTraining...")
	model = lgb.train(
	params,
	train_data,
	num_boost_round=params.get('n_estimators', 2000),
	valid_sets=[train_data, val_data],
	valid_names=['train', 'test'],
	callbacks=[
	lgb.early_stopping(50, verbose=True),
	lgb.log_evaluation(period=50)
	]
	)

	# Evaluate
	print("\n" + "-"*40)

	# Train predictions
	y_train_pred = np.argmax(model.predict(X_train), axis=1)
	evaluate_model(y_train, y_train_pred, "Train")

	# Test predictions
	y_test_pred = np.argmax(model.predict(X_test), axis=1)
	test_metrics = evaluate_model(y_test, y_test_pred, "Test")

	# Classification report
	print("\nClassification Report (Test):")
	print(classification_report(y_test, y_test_pred, target_names=TARGET_CLASS_NAMES))

	return model, test_metrics


	def save_model(model: lgb.Booster, params: dict, feature_names: list, metrics: dict) -> None:
	"""Save trained model and metadata."""
	print("\nSaving model...")

	MODELS_DIR.mkdir(parents=True, exist_ok=True)

	# Save LightGBM model
	model_path = MODELS_DIR / 'wildfire_model.txt'
	model.save_model(str(model_path))
	print(f" Model: {model_path}")

	# Save metadata
	metadata = {
	'params': params,
	'feature_names': feature_names,
	'metrics': metrics,
	'target_classes': TARGET_CLASS_NAMES
	}
	metadata_path = MODELS_DIR / 'model_metadata.joblib'
	joblib.dump(metadata, metadata_path)
	print(f" Metadata: {metadata_path}")


	def main():
	"""Main training pipeline."""
	# Parse arguments
	parser = argparse.ArgumentParser(description='Train wildfire classification model')
	parser.add_argument('--tune', action='store_true', help='Run hyperparameter tuning')
	args = parser.parse_args()

	print("\n" + "="*60)
	print("MODEL TRAINING")
	print("="*60)

	# Load data
	train_df, test_df = load_data()

	# Prepare data
	X_train, y_train, X_test, y_test, feature_cols = prepare_data(train_df, test_df)

	# Compute class weights
	sample_weights = None
	if USE_CLASS_WEIGHTS:
	sample_weights = compute_weights(y_train)

	# Get parameters
	if args.tune:
	params = tune_hyperparameters(X_train, y_train, sample_weights)
	else:
	best_params_path = MODELS_DIR / 'best_params.json'
	if best_params_path.exists():
	# Load saved best params
	with open(best_params_path, 'r') as f:
	tuned_params = json.load(f)
	params = LIGHTGBM_PARAMS.copy()
	params.update(tuned_params)
	print(f"Loaded best params from {best_params_path}")
	else:
	# Fallback to defaults
	params = LIGHTGBM_PARAMS.copy()
	params['n_estimators'] = 500
	params['max_depth'] = 8
	params['learning_rate'] = 0.05
	params['num_leaves'] = 64
	params['min_child_samples'] = 50
	params['subsample'] = 0.8
	params['colsample_bytree'] = 0.8
	print("No saved params found; using defaults")

	# Train final model
	model, metrics = train_final_model(
	X_train, y_train, X_test, y_test,
	params, sample_weights, feature_cols
	)

	# Save model
	save_model(model, params, feature_cols, metrics)

	print("\n" + "="*60)
	print("✓ Training Complete!")
	print("="*60 + "\n")


	if __name__ == "__main__":
	main()