Surgery Prediction

Machine-learning pipeline that predicts orthognathic surgical movements from pre-operative cephalometric / landmark measurements.

For each patient, a set of anatomical measurements (maxilla and mandible landmark positions and displacements) is used to predict a set of surgical-movement targets (every column whose name contains Pred). Several regression models are trained and compared, and a stacking ensemble combines the strongest ones.

This repository contains the code and the generated results/plots. The raw patient data is not included (see Data).

---

Table of contents

• Overview
• Data
• Installation
• Quick start
• Project structure
• Models
• Generated outputs
• Notes for the next contributor
• License

---

Overview

The workflow is:

1. Load & clean the data — model/data_loader.py reads the Excel file, drops metadata columns, keeps numeric features, removes highly correlated features (corr_threshold=0.95), scales, and clips target outliers.
2. Rank the features — model/model_lgbm.py trains one LightGBM regressor per target and writes detailedranking.xlsx, a per-target feature ranking.
3. Train the models — each model/model_*.py trains a regression model for every target and reports MAE / R² / RMSE. run_all_models.py runs them all and builds a comparison table.
4. Stacking ensemble — model/model_stacking_ensemble.py selects the top-N features per target (from the ranking file), then trains a StackingRegressor (ElasticNet + LightGBM + Gaussian Process → LinearRegression).

All metrics are three standard regression scores:

• MAE — mean absolute error (lower is better)
• RMSE — root mean squared error (lower is better)
• R² — coefficient of determination (higher is better)

---

Data

The patient data is private and intentionally excluded from this repository. The data/ folder is git-ignored.

To run the pipeline, place the input Excel file at:

data/surgical_mov_1_1496.xlsx

(or pass a different file_path to load_and_prepare_data / the model functions).

The loader expects, in that Excel file:

Column	Role
#	Patient ID
Skl_Cl_2	Skeletal class of the patient (used to filter by class)
Gender, Surg, Genio, 3piece, assymetry	Metadata — ignored as model inputs
any column containing Pred	Target to predict (a surgical movement)
all other numeric columns	Features (landmark measurements)

The number in the file name (1496) refers to the patient cohort size; an earlier 500-patient cohort is reflected in some of the committed result folders.

---

Installation

Requires Python 3.10+.

python -m venv .venv
source .venv/bin/activate        # Windows: .venv\Scripts\activate
pip install -r requirements.txt

---

Quick start

Run everything from the repository root so that the model package imports resolve correctly.

# 1. Generate the feature ranking (needed by the stacking ensemble)
python -m model.model_lgbm

# 2. Train + compare every model
python run_all_models.py

# 3. Train the stacking ensemble (per skeletal class)
python -m model.model_stacking_ensemble

# Optional: LightGBM tuned with Optuna
python -m model.model_optuna_optimized

# Optional: summarize the feature-selection sweep
python compute_mean.py

You can also run a single model, e.g.:

python -m model.model_random_forest

---

Project structure

SurgicalMovementPrediction-Training/
├── data/                       # (git-ignored) put the input .xlsx here
├── model/                      # model package
│   ├── __init__.py
│   ├── data_loader.py          # shared data loading & cleaning
│   ├── model_lgbm.py           # LightGBM feature-ranking -> detailedranking.xlsx
│   ├── model_stacking_ensemble.py   # main model (stacking ensemble)
│   └── model_*.py              # individual regression models (see table below)
├── run_all_models.py           # run + compare all models -> ensemble_comparison.csv
├── compute_mean.py             # summarize the feature-selection sweep
├── plot_feature_importance.py  # feature-importance plots
├── plots/                      # plotting scripts + generated PDFs/PNGs
├── results/                    # generated metric CSVs (500 & 1496 patients)
├── ranking/                    # generated feature rankings (.xlsx)
├── predictions/                # generated per-patient predictions
├── feature_importance/         # generated per-class feature importances
├── requirements.txt
├── LICENSE                     # Apache-2.0
└── README.md

---

Models

Every model/model_*.py exposes a train_*_model() function returning its results (a DataFrame or a dict of averaged metrics), and can also be run directly.

Script	Model
model_baseline_clean.py	StandardScaler → PCA → LinearRegression (baseline)
model_pls_robust.py	Partial Least Squares regression
model_random_forest.py	Random Forest regressor
model_mlp_modern.py	Multi-layer perceptron (neural net)
model_regressor_chain.py	Regressor chain (multi-output)
model_gpr.py	Gaussian Process regression
model_knn.py	K-nearest-neighbors regression
model_elasticnet.py	ElasticNet (Lasso + Ridge)
model_multitask_elasticnet.py	Multi-task ElasticNet (joint targets)
model_bayesian_ridge.py	Bayesian Ridge regression
model_lgbm_simple.py	LightGBM, single train/test split
model_lgbm_cv_simple.py	LightGBM with cross-validation
model_optuna_optimized.py	LightGBM tuned with Optuna
model_lgbm.py	LightGBM feature-ranking generator
model_stacking_ensemble.py	Stacking ensemble (main model)

---

Generated outputs

The repository keeps previously generated results so you can see prior work without re-running the whole pipeline:

• results/ — per-model metric CSVs, split by cohort (500 patients, 1496 patients), including a feature selection/ sweep of the stacking ensemble over different numbers of top features.
• ranking/ — feature rankings (detailedranking.xlsx) per cohort.
• predictions/ — per-patient actual-vs-predicted values, per skeletal class.
• feature_importance/ — per-class feature importances.
• plots/ — predicted-vs-actual PDFs and feature-importance images.

Re-running a model overwrites the corresponding CSV/plot in the working directory.

---

Notes for the next contributor

• Run from the repo root. Imports use the model package (from model.data_loader import ...), so scripts must be launched from the root (e.g. python -m model.model_stacking_ensemble).
• The ranking comes first. The stacking ensemble reads detailedranking.xlsx to pick the top-N features per target. Generate it with python -m model.model_lgbm before running the stacking model on a new dataset. Mind the ranking path used in model_stacking_ensemble.py (ranking/1496 patients/detailedranking.xlsx).
• Class filtering. model_stacking_ensemble.train_stacking_ensemble_model can filter patients by skeletal class via class_patient=....
• Optional heavy steps. SHAP plots, permutation importance, and per-patient prediction CSVs are present but commented out in model_stacking_ensemble.py — uncomment them if you need those artifacts (they are slow).

---

License

Licensed under the Apache License 2.0.