1. ML Model Comparison on a Single Dataset
Starting from the optimized physical model of section 3, you will now systematically compare several ML models for residual correction, and track all experiments with MLflow.
- Dataset:
rgs000_0.csv
- Input:
(y_nisp, z_nisp, wavelength)
- Target: residuals
r_y, r_z from your fitted physical model
- Metric: RMSE on the test set (80/20 split)
Every model + hyperparameter combination must be logged as a separate MLflow run.
MLflow: Experiment Tracking
MLflow records parameters, metrics, and artifacts for each run, enabling reproducible comparison.
Key concepts
— Experiment: a named group of runs
— Run: one training trial (model + hyperparams)
— Parameters: logged with mlflow.log_param
— Metrics: logged with mlflow.log_metric
Minimal usage pattern
- Create or set an experiment
- Start a run with
mlflow.start_run()
- Log params and metrics inside the run
- End the run
- Compare runs in the UI:
mlflow ui
Docs: https://mlflow.org/docs/latest/tracking.html
Models to Compare
| Model | Key hyperparameters to vary | sklearn class |
|---|
| Linear Regression | None (baseline) | LinearRegression |
| Random Forest | n_estimators, max_depth | RandomForestRegressor |
| Gradient Boosting | max_iter, max_depth, learning_rate | HistGradientBoostingRegressor |
| SVM | C, kernel, epsilon | SVR |
| MLP (neural net) | hidden_layer_sizes, activation, alpha | MLPRegressor |
For each model, test at least 3 hyperparameter configurations. Log each as a separate MLflow run.
Docs: https://scikit-learn.org/stable/supervised_learning.html
Systematic Hyperparameter Search
Instead of trying hyperparameters manually, use GridSearchCV or RandomizedSearchCV to automate the search.
1. Define a parameter grid (dict of lists)
2. Pass it to GridSearchCV(estimator, param_grid, cv=5, scoring="neg_root_mean_squared_error")
3. Fit on the training set
4. Read .best_params_ and .best_score_
5. Log the best result to MLflow
Use RandomizedSearchCV for large grids (MLP, SVM). It samples randomly instead of exhaustively.
Docs: https://scikit-learn.org/stable/modules/grid_search.html
Task 1 — Systematic ML Comparison
- Install and import MLflow. Create an experiment named
"residual_correction".
- For each model: run a hyperparameter search. Log best params, train RMSE, test RMSE, and train/test ratio per run.
- Run
mlflow ui and compare all runs. Identify the best model for r_y and for r_z.
- Plot: predicted vs actual residuals, and histogram of residuals before/after ML correction.
Expected output: one MLflow experiment with 15–30 runs, a results table, and the best model saved to disk.
MLflow sklearn: https://mlflow.org/docs/latest/python_api/mlflow.sklearn.html
Analyzing the Results
After the comparison, answer the following questions in your report:
- Which model achieves the lowest test RMSE? By how much does it improve over the physical model alone?
- Is the improvement in
r_y and r_z symmetric? Why or why not?
- Which model overfits the most (largest train/test RMSE ratio)? Is this expected?
- Does the best model generalize? Test it on a different dataset without retraining.
- Are the residuals after ML correction approximately random, or is there remaining structure?
2. Applying the Pipeline to All Datasets
You have a working pipeline: physical model optimization + ML residual correction. Apply it to each of the 6 calibration configurations.
| Dataset | GWA angle | Tilt |
|---|
rgs000_0 | 0° | 0° |
rgs000_m4 | 0° | −4° |
rgs000_p4 | 0° | +4° |
rgs180_0 | 180° | 0° |
rgs180_m4 | 180° | −4° |
rgs180_p4 | 180° | +4° |
For each dataset: (1) optimize physical parameters, (2) apply best ML correction from task 1, (3) log results to a new MLflow experiment "per_dataset".
Task 2 — Per-Dataset Results
Steps to follow:
- Wrap your pipeline into a single reusable function that takes a dataset path and returns optimal parameters + RMSE.
- Run it on all 6 datasets. Log physical RMSE and hybrid RMSE per dataset.
- Collect the fitted physical parameters in a table (one row per dataset).
- Identify which parameters are consistent across datasets (candidates for global fitting in step 3).
- Compare RMSE across datasets: which configuration is hardest to fit? Propose a physical explanation.
Key question: are the collimator parameters consistent across all 6 datasets? If yes, they can be treated as global parameters in step 3.
Parameter Consistency Analysis
After fitting each dataset independently, compare the optimal parameters.
Physically shared
These describe the hardware — constant across all configurations:
- Collimator focal length
- Camera focal length
- Prism refractive index
- Grating period
Configuration-specific
These change per dataset:
- GWA rotation angle
- Detector tilt angle
- Alignment offsets
If a nominally shared parameter converges to very different values across datasets, this signals a model deficiency or a wrong assumption about which parameters are truly global.
3. Multi-Dataset Fitting
The 6 calibration datasets share the same physical instrument.
Some parameters are shared, others are dataset-specific.
Shared parameters
— Collimator focal length
— Camera focal length
— Prism angles / index
— Grating period
Dataset-specific
— GWA angle (per configuration)
— Detector alignment offsets
— Grism rotation
Strategy: build a combined cost function that sums the RMSE across all datasets, with shared parameters appearing in every term.
# Schematic only
def cost_combined(theta):
shared, specific = unpack(theta)
total = 0
for dataset, params in zip(datasets, specific):
pred = model.forward(dataset.inputs, shared, params)
total += mse(pred, dataset.targets)
return total
3. Multi-Dataset Optimization Strategy
Fitting each dataset independently does not enforce consistency of shared parameters. A structured alternating strategy does.
Stage 1. Fit physical model independently on each dataset (no ML correction)
Stage 2. Fix dataset-specific parameters; optimize shared parameters jointly on all 6 datasets
Stage 3. Fix shared parameters; re-optimize specific parameters on each dataset
Stage 4. Fix all physical parameters; add ML residual correction per dataset
This is an alternating optimization strategy. Stages 2 and 3 can be iterated until convergence before applying stage 4.
Stage 2: Joint Optimization of Shared Parameters
The combined cost sums over all datasets, with shared parameters:
\[\mathcal{L}_\text{joint}(\theta_\text{shared}) = \sum_{d=1}^{6} \mathcal{L}_d(\theta_\text{shared},\, \theta_d^*)\]
where θ_d* are the dataset-specific parameters fixed from stage 1.
Steps:
- Load all 6 datasets and their stage-1 specific parameters.
- Build a combined cost function over all datasets.
- Optimize only the shared parameters.
- Compare resulting shared parameters to per-dataset stage-1 values.
Stages 3 and 4
Stage 3 — Re-optimize specific parameters
Fix the shared parameters from stage 2. Re-fit dataset-specific parameters independently on each dataset. This is a refinement step — RMSE should improve slightly.
Stage 4 — ML correction
With all physical parameters fixed, compute residuals on each dataset. Train the best ML model (from task 1) per dataset to correct them.
After all 4 stages, compare the final RMSE per dataset against the independent results from task 2.
Expected: the joint strategy produces more consistent and slightly lower RMSE.
Task 3 — Implement the Multi-Dataset Strategy
Steps to follow:
- Stage 1: use results from task 2 (already done).
- Stage 2: build a combined cost function. Optimize shared parameters. Log to MLflow experiment
"stage2".
- Stage 3: re-fit specific parameters with shared ones fixed. Log to
"stage3".
- Stage 4: add ML correction per dataset. Log hybrid RMSE to
"stage4".
- Build a comparison table: RMSE per dataset × strategy (independent / joint / joint+ML).
- Which datasets benefit most from joint fitting? Which parameters changed most between stages?
4. Joint Refinement of Physics and ML
This section is more advanced. Attempt it only after completing tasks 1–3.
In the sequential approach, the physical model and the ML corrector are trained one after the other. In joint refinement, both are optimized together, with a penalty to prevent the ML from dominating.
\[\mathcal{L}_\text{total}(\theta, \phi) = \underbrace{\mathcal{L}_\text{data}(\theta, \phi)}_{\text{prediction error}} + \lambda \underbrace{\Omega(\phi)}_{\text{ML complexity penalty}}\]
θ = physical parameters, φ = ML model, λ = regularization weight
The Regularization Penalty
Without a penalty, the optimizer may delegate all prediction error to the ML component, making the physical parameters meaningless.
Possible penalty choices for Ω(φ):
| Penalty | Effect |
|---|
| \(\|\delta y_\text{ML}\|^2\) | Forces ML correction to be small overall |
| \(\text{Var}(\delta y_\text{ML})\) | Forces correction to be spatially uniform |
| Fraction of variance explained by ML | Directly caps ML contribution |
Start with the L2 norm penalty. Tune λ by monitoring the ratio of ML-explained variance to total variance.
Task 4 — Joint Refinement
Steps to follow:
- Define a combined loss: prediction MSE + λ × L2 norm of ML corrections.
- Implement an alternating optimization loop: (a) optimize θ with φ fixed; (b) retrain φ with the penalized objective; repeat.
- Run for several values of λ (e.g. 0, 0.01, 0.1, 1.0). Log each to MLflow.
- For each λ: report physical RMSE, ML correction RMS magnitude, and hybrid RMSE.
- Plot ML correction magnitude vs λ. It should decrease monotonically.
- Identify the λ that gives the best trade-off between accuracy and physical interpretability.
If λ = 0 gives no improvement over the sequential approach, the joint refinement does not add value. That is also a valid scientific result.
5. Comparison with Published Results
The reference paper (arXiv 2506.08378) reports calibration accuracy for NISP dispersion. Compare your results against theirs.
Before comparing, verify that you are using the same:
- Definition of the residual (predicted − observed, or opposite?)
- Metric (RMSE, MAE, or other?)
- Unit (mm, pixels, arcsec?)
- Dataset (same configurations? same wavelength range?)
If your model differs from the reference (fewer parameters, different optics), the comparison is approximate. State this explicitly in your report.
Reference: https://arxiv.org/abs/2506.08378
Task 5 — Read and Compare
Steps to follow:
- Read the results section of arXiv 2506.08378. Note all reported RMSE values and their units.
- Identify which calibration configurations and wavelength ranges they use.
- Convert your RMSE to the same unit if necessary (1 NISP pixel ≈ 0.3 mm).
- Fill the comparison table on the next slide.
- Identify the main sources of discrepancy between your approach and the reference.
- Propose at least one concrete improvement that could close the gap.
Comparison Table
Fill this table. Use test-set RMSE only. Units: mm.
| Model | RMSE y (mm) | RMSE z (mm) | Source |
|---|
| Physical model (after fit) | — | — | Task 2 |
| Hybrid (best ML, single dataset) | — | — | Task 1 |
| Hybrid (joint fitting) | — | — | Task 3 |
| Hybrid (joint refinement) | — | — | Task 4 |
| Reference (arXiv 2506.08378) | — | — | Paper |
If your result is significantly worse than the reference: identify which model component is likely responsible and propose an improvement.