Guide: Model Training
Once the feature set has been extracted, the SpectralClassifier class orchestrates model training, optimization, and evaluation. It encapsulates a complete scikit-learn / imblearn pipeline, supports about a dozen models, and automatically logs results to Weights & Biases.
All training logic is in src/pipeline/classifier.py. The orchestration notebook is 00_master_pipeline.ipynb.
Quick start
from src.pipeline.classifier import SpectralClassifier
# XGBoost training in pure spectroscopic mode
clf = SpectralClassifier(model_type="XGBoost", prediction_target="main_class")
clf, feature_cols, X, y = clf.train_and_evaluate(
features_df,
spectro_only=True, # excludes ra, dec, redshift and most Gaia columns (see §3)
)
# Save artifacts
clf.save_model("data/models/spectral_classifier_xgboost.pkl")
Reference result achieved in spectro_only=True mode:
1. Available models
SpectralClassifier supports 11 classifiers through the model_type parameter. Switching models is trivial — the pipeline, tuning, and W&B logging adapt automatically.
"XGBoost""LightGBM""CatBoost""ExtraTrees""RandomForest""Ensemble""SVM""MLP""LogisticRegression""KNN""LDA"SpectralClassifier(model_type="LightGBM") — the pipeline, tuning and W&B logging adapt automatically.2. The training process
Data preparation
SpectralClassifier automatically prepares the data before training:
- Cleaning: Spectra without a valid
subclassand classes with < 10 samples are excluded. - Feature selection: In
spectro_only=Truemode, positional, observing-condition and most Gaia columns are filtered out (see §3 for the Gaia columns that remain). - Train/test split: 80/20 stratified by class.
Pipeline and W&B tracking
.pkl includes the scaler and imputer — no additional data preparation is required for prediction.Hyperparameter tuning
Option A — GridSearchCV (default):
# Stratified cross-validation k=5 · metric: balanced_accuracy
clf.train_and_evaluate(features_df, spectro_only=True)
Option B — Optuna (Bayesian optimization):
The Optuna infrastructure is deployed for large-scale exploration. It is especially useful for exploring n_estimators, max_depth, learning_rate, colsample, subsample, reg_alpha/lambda on the ~139k available spectra.
clf.train_and_evaluate(
features_df,
spectro_only=True,
tuner="optuna",
n_trials=200,
)
3. spectro_only=True mode
Enable this mode for all rigorous scientific results. It ensures that the model learns to classify stars from their intrinsic physics only, without any observational bias.
clf.train_and_evaluate(features_df, spectro_only=True)
Automatically excluded features:
- Angular coordinates:
ra,dec - Observing conditions:
redshift,seeing,snr_*,jd,fiber_id,magnitude_* - Gaia columns:
parallax,bp_rp,teff_gspphot,logg_gspphot,v_tan_kms, etc.
ra, dec and redshift are excluded because they can correlate with spectral type through LAMOST observational biases (observation programs targeted by type) rather than through stellar physics. The exported run reports do not show a measurable accuracy gain from this exclusion: in the 100k runs, v3 (before spectro_only) and v4 (spectro_only) reach a balanced accuracy of 86.4% and 86.4% respectively, and v5 (spectro_only plus isotonic calibration) 87.3%.
In practice, spectro_only does not remove every Gaia column: the reference run still selected 4 Gaia/photometric features, among them M_G (also M_BP, M_RP and ag_gspphot), so M_G is not excluded in that run.
4. Saving and artifacts
At the end of each run, the pipeline automatically generates the following files:
data/models/reports/<ISO>/Load and predict
clf_loaded = SpectralClassifier.load_model(
"data/models/spectral_classifier_xgboost_20260213T225019Z.pkl"
)
# The loaded pipeline includes the scaler and imputer
# No additional data preparation is required
predictions = clf_loaded.model_pipeline.predict(X_new)
probabilities = clf_loaded.model_pipeline.predict_proba(X_new)
5. Complete evaluation
After training, clf.evaluate() automatically generates:
clf.evaluate(X_test, y_test, plot=True)
| Output | Description |
|---|---|
| Classification report | Precision / Recall / F1 by class (A, F, G, K, M) |
| Confusion matrix | Normalized — shows the expected F/G confusion |
| ROC and PR curves | One-vs-rest by class |
| ECE score | Expected Calibration Error — probability quality |
| Feature importance | XGBoost gain or permutation importance |
6. Reproducibility
Each run produces a unique ISO timestamp identifier. To reproduce an exact run:
# Retrieve the training source file
import json
with open("data/reports/20260213T225019Z/spectral_classifier_xgboost_20260213T225019Z_meta.json") as f:
meta = json.load(f)
print(meta["trained_on_file"]) # → the features CSV used for training
print(meta["best_params_"]) # → final hyperparameters
print(meta["n_features_used"]) # → number of features used by this run
print(meta["balanced_accuracy"]) # → balanced accuracy of this run