add logistic regression baseline PS and an azure job to run it

[kirilklein]

Adds a conventional (logistic regression) propensity score baseline to compare against the transformer PS, and wires train_baseline up as an Azure job. Requested by reviewers: a traditional PS model as a reference, plus balance/overlap diagnostics across methods.

What changed

• corebehrt/main_causal/helper/baseline_models.py (new): model construction, fitting, tuning ranges and CatBoost device handling for the two baselines.
• helper/train_baseline.py: delegates to the factory; nested CV, folds and prediction saving unchanged.
• configs/causal/finetune/baseline.yaml (new): default config, model: logistic.
• azure/components/train_baseline.py (new) + train_baseline registered in the job CLI.
• Existing CatBoost configs (simulated_bl.yaml, both experiments/*/base_configs/train_baseline.yaml) pin model: catboost so nothing silently switches.

The baseline writes combined_predictions.csv in the same format as finetune_exp_y, so calibrate_exp_y → estimate → get_stats run downstream unchanged. Running get_stats with the baseline PS gives the covariate balance / overlap / ESS comparison against the transformer.

Run

python -m corebehrt.azure job train_baseline CPU-20-LP -e ps_baseline \
  -c corebehrt/configs/causal/finetune/baseline.yaml -o model=BaselinePSLogistic

Input is the same prepared_data asset the transformer used, so both models see identical folds.

Behaviour fix

Hyperparameters pinned in the config were logged as FIXED but only reached the final refit — the Optuna trial models ignored them. They are now applied to both.

Testing

• 6 new unit tests in tests/test_main_causal/test_helper/test_baseline_models.py; 11 tests in tests/test_main_causal pass.
• End-to-end smoke run of nested_cv_loop on synthetic patients completes for both logistic and catboost and writes combined_predictions.csv.
• ruff format --check and the CI lint selection are clean.

🤖 Generated with Claude Code

Summary by CodeRabbit

• New Features

  • Added baseline training as a supported Azure job.
  • Added configurable logistic regression and CatBoost baseline models.
  • Added model-specific hyperparameter tuning and optional validation-based early stopping.
  • Added baseline fine-tuning configurations for causal modeling workflows.

• Bug Fixes

  • Improved baseline training compatibility across CPU and GPU environments.

• Tests

  • Added coverage for model selection, configuration handling, tuning ranges, model fitting, and predictions.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 9095c66a-34e1-487b-812d-f9dc37c83ce6

📥 Commits

Reviewing files that changed from the base of the PR and between 017afaf and 64f4ae9.

📒 Files selected for processing (3)

• corebehrt/main_causal/helper/baseline_models.py
• corebehrt/main_causal/helper/train_baseline.py
• tests/test_main_causal/test_helper/test_baseline_models.py

🚧 Files skipped from review as they are similar to previous changes (3)

• tests/test_main_causal/test_helper/test_baseline_models.py
• corebehrt/main_causal/helper/baseline_models.py
• corebehrt/main_causal/helper/train_baseline.py

---

📝 Walkthrough

Walkthrough

The PR adds shared logistic regression and CatBoost baseline-model handling, integrates model selection into causal training and tuning, adds baseline configurations, and exposes baseline training through an Azure job.

Changes

Baseline training workflow

Layer / File(s)	Summary
Baseline model abstraction corebehrt/main_causal/helper/baseline_models.py, tests/test_main_causal/test_helper/test_baseline_models.py	Adds model defaults, tuning ranges, estimator construction, fitting, device handling, CatBoost parameter preparation, and tests for both supported models.
Model selection in causal training corebehrt/main_causal/helper/train_baseline.py	Passes the configured model through parameter setup, hyperparameter tuning, fold training, and final model fitting.
Configuration and Azure job wiring corebehrt/configs/causal/finetune/baseline.yaml, corebehrt/configs/causal/finetune/simulated_bl.yaml, experiments/causal_pipeline/*/base_configs/train_baseline.yaml, corebehrt/azure/components/train_baseline.py, corebehrt/azure/main/job.py	Adds baseline model configurations and exposes train_baseline through the Azure job CLI and wrapper.

Estimated code review effort: 4 (Complex) | ~45 minutes

Sequence Diagram(s)

sequenceDiagram
  participant CausalConfig
  participant main_baseline
  participant baseline_models
  participant Optuna
  CausalConfig->>main_baseline: provide baseline model selection
  main_baseline->>baseline_models: request model parameters and tuning ranges
  main_baseline->>Optuna: run model-specific hyperparameter tuning
  Optuna->>baseline_models: build and fit trial model
  main_baseline->>baseline_models: build and fit final fold model

Loading

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 42.86% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly identifies the new logistic regression propensity-score baseline and the Azure job that runs it.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches 💡 1

📝 Generate docstrings 💡

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch feat/logistic-baseline-ps

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@corebehrt/main_causal/helper/baseline_models.py`:
- Around line 101-110: Update build_model() to merge configured
task_type/devices with detected device settings before calling
_prepare_catboost_params(), so parameter filtering uses the effective device
mode; preserve explicit configuration precedence in the final CatBoostClassifier
arguments. Apply the same effective device-mode construction in
get_tuning_ranges() before its parameter preparation.

🪄 Autofix

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 5ff38540-7ad5-4b86-9fad-c214842e2b37

📥 Commits

Reviewing files that changed from the base of the PR and between 545e880 and 017afaf.

📒 Files selected for processing (9)

• corebehrt/azure/components/train_baseline.py
• corebehrt/azure/main/job.py
• corebehrt/configs/causal/finetune/baseline.yaml
• corebehrt/configs/causal/finetune/simulated_bl.yaml
• corebehrt/main_causal/helper/baseline_models.py
• corebehrt/main_causal/helper/train_baseline.py
• experiments/causal_pipeline/base_configs/train_baseline.yaml
• experiments/causal_pipeline_resample/base_configs/train_baseline.yaml
• tests/test_main_causal/test_helper/test_baseline_models.py