Major Refactor: CLI Pipeline, Model Packaging, and PLINK-Native Workflow - #40
Major Refactor: CLI Pipeline, Model Packaging, and PLINK-Native Workflow#40MattScicluna wants to merge 41 commits into
Conversation
Efficient batch loading
code cleanup
Clean outputs
- partition_data.py: save sample_ids into partition NPZ so check
--partition-file is self-contained (no extra args needed)
- log_utils.py: add seed param to create_out_dir (fixes multi-seed
overwrite bug; new dir naming: {exp_name}_seed{seed}_fold{fold});
add save_predictions_tsv() helper
- train.py: pass seed_override to create_out_dir; write predictions.tsv
after test step (same format as dietnet predict) for both PLINK and
HDF5 label_names paths
- cli.py: fix out_dir lookup in _train_with_plink_packages to match new
seed-aware naming; add --partition-file/--fold to dietnet check for
test-fold filtering; extract _print_accuracy_table() helper
- scripts: replace inflated-accuracy predict+check-all pattern with
per-fold predictions.tsv check loop in all three sbatch scripts;
sbatch_train_hgdp_ukbb.sh grows to 7 steps (external predict kept)
Co-Authored-By: Claude Sonnet 4.6 <[email protected]>
Co-authored-by: Copilot <[email protected]>
Co-authored-by: MattScicluna <[email protected]>
Fix unreachable Comet best-validation logging in early-stopping block
Integrate with main version
There was a problem hiding this comment.
Pull request overview
This PR refactors DietNetwork from a set of research scripts into an installable, Click-based CLI application with PLINK-native training/inference, portable model packaging, and smoke-testable workflows.
Changes:
- Introduces a packaged
dietnetCLI workflow for partitioning, embedding generation, training, preprocessing, prediction, and checking. - Adds model packaging + pretrained model/test-data registries, with download/cache helpers and PLINK preprocessing/alignment utilities.
- Adds reproducible scripts (setup, packaging, smoke tests) and removes legacy/experimental code and generated Sphinx build artifacts.
Reviewed changes
Copilot reviewed 68 out of 78 changed files in this pull request and generated 7 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/pop_cmap_1kgp.json | Adds a population→color mapping used by analysis/plotting utilities. |
| tests/kgp_precomputed/smoke_test_config.sh | Centralizes smoke-test environment configuration (cache/data paths, thresholds). |
| tests/kgp_precomputed/run_train_smoke.sh | Adds an end-to-end train→predict→check smoke test using the CLI. |
| tests/kgp_precomputed/run_smoke_test_single.sh | Adds a single-model inference smoke test using the CLI. |
| tests/kgp_precomputed/run_smoke_test.sh | Adds an ensemble inference smoke test using the CLI. |
| tests/kgp_precomputed/download_test_data.sh | Adds a helper to fetch/link bundled PLINK test data for smoke tests. |
| tests/kgp_precomputed/download_model.sh | Adds a helper to pre-download the pretrained model into the local cache. |
| tests/kgp_precomputed/data/config.yaml | Adds a small config used by smoke-test training. |
| setup.sh | Adds a setup script for venv install and downloading PLINK2 into bin/. |
| scripts/train_predict.sh | Provides a template script to run the full train+predict CLI pipeline. |
| scripts/predict_external.sh | Provides a template script for inference-only runs via the CLI. |
| scripts/package_test_data.sh | Adds a packaging helper to tar test PLINK data and compute SHA256. |
| scripts/package_models.sh | Adds a packaging helper to tar pretrained model packages and compute SHA256. |
| scripts/UPLOAD_INSTRUCTIONS.md | Documents how to package and publish model/test-data tarballs and update registries. |
| pyproject.toml | Migrates to modern packaging (hatchling) and defines the dietnet console entry point. |
| docs/source/index.rst | Removes legacy Sphinx docs source content. |
| docs/source/conf.py | Removes legacy Sphinx configuration. |
| docs/make.bat | Removes legacy Sphinx build script (Windows). |
| docs/Makefile | Removes legacy Sphinx build Makefile. |
| docs/build/html/searchindex.js | Removes generated Sphinx HTML artifact. |
| docs/build/html/search.html | Removes generated Sphinx HTML artifact. |
| docs/build/html/index.html | Removes generated Sphinx HTML artifact. |
| docs/build/html/genindex.html | Removes generated Sphinx HTML artifact. |
| docs/build/html/_static/pygments.css | Removes generated Sphinx HTML artifact. |
| docs/build/html/_static/documentation_options.js | Removes generated Sphinx HTML artifact. |
| docs/build/html/_static/doctools.js | Removes generated Sphinx HTML artifact. |
| docs/build/html/_static/custom.css | Removes generated Sphinx HTML artifact. |
| docs/build/html/_sources/index.rst.txt | Removes generated Sphinx HTML artifact. |
| docs/build/html/.buildinfo | Removes generated Sphinx HTML artifact. |
| README.md | Rewrites docs around installation, CLI usage, pretrained presets, and smoke tests. |
| Dietnet/vae_train.py | Removes legacy experimental VAE training code. |
| Dietnet/train_multiGPU.py | Removes legacy multi-GPU training script. |
| Dietnet/optimize.py | Removes legacy hyperparameter optimization script. |
| Dietnet/make_attributions.py | Removes legacy attribution pipeline entrypoint. |
| Dietnet/get_preprocessing_params.py | Removes legacy preprocessing-params script (replaced by CLI workflow). |
| Dietnet/get_last_layer.py | Removes legacy last-layer extraction script. |
| Dietnet/generate_embedding_by_fold.py | Removes legacy fold embedding helper (superseded by new pipeline). |
| Dietnet/evaluate.py | Removes legacy evaluation script. |
| Dietnet/Interpretability/utils.py | Removes legacy interpretability utilities. |
| Dietnet/Interpretability/graph_attribution_manager.py | Removes legacy interpretability/graph attribution manager. |
| Dietnet/Interpretability/experiment_scripts/create_dataset_less_snps.py | Removes legacy interpretability experiment script. |
| Dietnet/Interpretability/custom_engine.py | Removes legacy custom attribution engine. |
| Dietnet/Interpretability/attribution_manager.py | Removes legacy attribution manager implementation. |
| Dietnet/test_external_dataset.py | Updates imports and adds an argument-driven entrypoint helper. |
| Dietnet/pretrained_models.py | Adds registries for pretrained model presets and test data (URLs, SHA256, metadata). |
| Dietnet/partition_data.py | Adds PLINK-aware partitioning and optional stratification support. |
| Dietnet/model_manager.py | Adds download/cache/validation helpers for pretrained model presets. |
| Dietnet/helpers/snp_alignment.py | Adds BIM parsing + SNP mapping/alignment helpers for inference. |
| Dietnet/helpers/plink_utils.py | Adds PLINK2-based preprocessing to enforce allele coding and SNP extraction. |
| Dietnet/helpers/model_package.py | Adds a portable on-disk model package format + conversion utilities. |
| Dietnet/helpers/model.py | Renames internal submodules while keeping backward-compatible attributes. |
| Dietnet/helpers/mainloop_utils.py | Updates imports and removes legacy/disused model-loading helpers. |
| Dietnet/helpers/log_utils.py | Updates output-dir naming (seed+fold) and adjusts summary/prediction utilities. |
| Dietnet/generate_embedding.py | Adds PLINK-aware embedding generation and fold-aware embedding computation. |
| Dietnet/create_dataset.py | Updates imports and adds an argument-driven entrypoint helper. |
| Dietnet/compute_input_features_mean.py | Refactors to compute per-fold mean/std for both HDF5 and PLINK datasets. |
| Dietnet/analysis/prediction_plots.py | Adds plotting utility for compact prediction vote breakdowns. |
| .gitignore | Adds ignores for caches, generated artifacts, tarballs, and local bins/logs. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
You can also share your feedback on Copilot code review. Take the survey.
| print(f"Extracting to {dest}...") | ||
| dest.mkdir(parents=True, exist_ok=True) | ||
| try: | ||
| with tarfile.open(temp_file, "r:gz") as tar: | ||
| tar.extractall(dest) | ||
| except Exception as e: |
There was a problem hiding this comment.
tar.extractall(dest) is vulnerable to path traversal if the archive contains malicious member paths. Even with checksums, it’s safer to validate that all extracted paths stay within dest (or use a safe extraction helper) before extracting.
| def get_cache_dir() -> Path: | ||
| """ | ||
| Get cache directory for models. | ||
|
|
||
| Returns: | ||
| Path to cache directory (~/.cache/dietnet/) | ||
| """ | ||
| cache_home = Path.home() / '.cache' / 'dietnet' | ||
| cache_home.mkdir(parents=True, exist_ok=True) | ||
| return cache_home |
There was a problem hiding this comment.
get_cache_dir() hardcodes ~/.cache/dietnet and ignores XDG_CACHE_HOME, while the shell scripts and ecosystem conventions use XDG. Consider using XDG_CACHE_HOME when set (falling back to ~/.cache) to keep cache location consistent across the project.
| try: | ||
| with tarfile.open(temp_file, 'r:gz') as tar: | ||
| tar.extractall(dest) |
There was a problem hiding this comment.
tar.extractall(dest) can write files outside dest if the tarball contains path-traversal entries. Use a safe extraction routine that validates member paths (or tarfile’s safe-filtering where available) before extracting downloaded archives.
| cmd = [ | ||
| plink_bin, | ||
| '--bfile', test_plink_prefix, | ||
| '--extract', model_bim_file, | ||
| '--alt1-allele', 'force', model_bim_file, '5', '2', | ||
| '--make-bed', | ||
| '--out', output_prefix | ||
| ] |
There was a problem hiding this comment.
--extract expects a list of variant IDs (one per line). Passing a .bim file will typically feed the chromosome column as the ID list, so the extract step will be wrong or fail. Consider generating an explicit variant-ID list (e.g., from column 2 of the BIM, or using the packaged snps.txt if it contains variant IDs) and use that for --extract.
| bim = pd.read_csv( | ||
| bim_path, | ||
| sep='\t', | ||
| header=None, | ||
| names=['chr', 'snp_id', 'cm', 'pos', 'a1', 'a2'], | ||
| dtype={'chr': str, 'snp_id': str, 'cm': float, 'pos': int, 'a1': str, 'a2': str} | ||
| ) |
There was a problem hiding this comment.
PLINK .bim files are whitespace-delimited (spaces or tabs). Using sep='\t' will fail to parse common BIM files that are space-separated. Consider sep=r'\s+' (or delim_whitespace=True) to robustly handle both.
| # Add project root to path | ||
| project_root = Path(__file__).resolve().parent.parent | ||
| sys.path.insert(0, str(project_root)) |
There was a problem hiding this comment.
The embedded Python snippet is executed via python3 -, where __file__ is not defined; Path(__file__).resolve() will raise NameError and the download script will fail. Compute the project root from Path.cwd() (or pass the repo root as an argument/environment variable) instead of relying on __file__.
| print("Extracting...") | ||
| try: | ||
| with tarfile.open(temp_file, 'r:gz') as tar: | ||
| tar.extractall(test_data_dir.parent) | ||
| except Exception as e: |
There was a problem hiding this comment.
tar.extractall(...) is vulnerable to path traversal if the tarball contains entries with ../ or absolute paths. Since this archive is downloaded, extract with a safe-filter (e.g., validate all member paths are within the destination dir) before calling extractall.
Major Refactor: CLI Pipeline, Model Packaging, and PLINK-Native Workflow
This PR substantially refactors the repository to transform it from a collection of research scripts into a packaged, installable CLI tool with a reproducible training and inference pipeline.
Key improvements include a unified CLI interface, PLINK-native I/O, portable model packaging, and improved dataset handling. Legacy experimental code and build artifacts have been removed to simplify the codebase.
Major Additions
Unified CLI & Pipeline
Dietnet/cli.py(+1462 lines)Introduces a new Click-based CLI that unifies all commands under the
dietnetentry point.Benefits:
Model Packaging & Inference
Files
Dietnet/helpers/model_package.pyDietnet/predict_with_plink.pyDietnet/pretrained_models.pyFeatures:
This enables trained models to be easily distributed and applied to new genotype datasets.
PLINK & SNP Utilities
Files
Dietnet/helpers/plink_utils.pyDietnet/helpers/snp_alignment.pyAdds:
Model Manager
Dietnet/model_manager.pyCentralized management of model state and metadata.
Scripts & Tests
Added infrastructure for reproducible workflows:
scripts/)tests/kgp_precomputed/)setup.sh)Analysis Utilities
Dietnet/analysis/prediction_plots.pyAdds tools for generating population-level prediction plots.
Modern Python Packaging
pyproject.tomlReplaces legacy setup configuration with modern packaging using hatchling.
Major Deletions
Removed legacy and experimental code that is no longer maintained.
Legacy Training Scripts
train_multiGPU.pyvae_train.pyoptimize.pyevaluate.pyLegacy Utilities
get_last_layer.pyget_preprocessing_params.pygenerate_embedding_by_fold.pyInterpretability Module
Removed entire module:
Attribution Pipeline
make_attributions.pyBuild Artifacts
Removed generated documentation:
(~17k lines of generated HTML/JS)
Assets
Images/folder removedSignificant Rewrites
Dietnet/train.pyDietnet/helpers/dataset_utils.pyDietnet/helpers/mainloop_utils.pyDietnet/generate_embedding.pyDietnet/compute_input_features_mean.pyREADME.mdSummary
This branch modernizes the repository by:
Overall, this refactor improves maintainability, usability, and reproducibility of the DietNetwork pipeline.