antigen-ml 1.0.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Amit Sah, SignorLab, North Dakota State University
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,413 @@
1
+ Metadata-Version: 2.4
2
+ Name: antigen-ml
3
+ Version: 1.0.0
4
+ Summary: Sequence- and structure-based antigenicity prediction using classical and ESM-2 features
5
+ Author-email: Amit Sah <amit.sah@ndsu.edu>
6
+ License-Expression: MIT
7
+ Project-URL: Homepage, https://github.com/asah15/antigen-ml
8
+ Project-URL: Repository, https://github.com/asah15/antigen-ml
9
+ Project-URL: Bug Tracker, https://github.com/asah15/antigen-ml/issues
10
+ Keywords: antigenicity,protein,machine learning,ESM-2,bioinformatics,immunoinformatics
11
+ Classifier: Programming Language :: Python :: 3
12
+ Classifier: Programming Language :: Python :: 3.10
13
+ Classifier: Programming Language :: Python :: 3.11
14
+ Classifier: Programming Language :: Python :: 3.12
15
+ Classifier: Operating System :: OS Independent
16
+ Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
17
+ Classifier: Intended Audience :: Science/Research
18
+ Requires-Python: >=3.10
19
+ Description-Content-Type: text/markdown
20
+ License-File: LICENSE
21
+ Requires-Dist: numpy>=1.26
22
+ Requires-Dist: pandas>=2.0
23
+ Requires-Dist: scikit-learn>=1.4
24
+ Requires-Dist: xgboost>=2.0
25
+ Requires-Dist: biopython>=1.83
26
+ Requires-Dist: matplotlib>=3.8
27
+ Provides-Extra: esm
28
+ Requires-Dist: torch>=2.1; extra == "esm"
29
+ Requires-Dist: transformers>=4.40; extra == "esm"
30
+ Provides-Extra: gui
31
+ Requires-Dist: Pillow>=10.0; extra == "gui"
32
+ Provides-Extra: all
33
+ Requires-Dist: torch>=2.1; extra == "all"
34
+ Requires-Dist: transformers>=4.40; extra == "all"
35
+ Requires-Dist: Pillow>=10.0; extra == "all"
36
+ Dynamic: license-file
37
+
38
+ # antigen-ml
39
+
40
+ A machine-learning pipeline for predicting protein antigenicity from sequence and structural features, with optional ESM-2 protein language model embeddings.
41
+
42
+ **Author:** Amit Sah — SignorLab, North Dakota State University
43
+ **Language:** Python 3.10+
44
+ **License:** MIT
45
+
46
+ ---
47
+
48
+ ## Overview
49
+
50
+ antigen-ml trains classifiers to discriminate antigenic proteins (recognized by the immune system) from non-antigenic proteins. It supports two complementary feature tracks:
51
+
52
+ | Track | Features | Dimensionality | Notes |
53
+ |---|---|---|---|
54
+ | Classical | AAC, DPC, CTD, physicochemical, structural | 503-dim | No GPU needed |
55
+ | ESM-2 | Pretrained protein language model embeddings | 480–1280-dim | Requires PyTorch + internet |
56
+
57
+ Both tracks can be combined. The pipeline fetches labeled training data from [IEDB](https://www.iedb.org/) and [UniProt](https://www.uniprot.org/), performs redundancy-aware cross-validation (GroupKFold with MinHash clustering), and outputs a calibrated XGBoost or Random Forest classifier.
58
+
59
+ ### Results on the full IEDB + UniProt dataset (n = 3,763)
60
+
61
+ | Model | Features | CV AUC | Test AUC | Test MCC | Best threshold |
62
+ |---|---|---|---|---|---|
63
+ | XGBoost | Classical (503-dim) | 0.858 | 0.858 | 0.571 | 0.40 |
64
+ | XGBoost | ESM-2 35M + Classical | ≥0.95 | — | — | — |
65
+
66
+ ---
67
+
68
+ ## Repository layout
69
+
70
+ ```
71
+ antigen-ml/
72
+ ├── train.py # Main training script
73
+ ├── predict.py # Inference on new FASTA sequences
74
+ ├── requirements.txt # Python dependencies
75
+ ├── LICENSE
76
+ ├── README.md
77
+ │
78
+ ├── src/
79
+ │ ├── download_data.py # Fetch labeled data from IEDB + UniProt
80
+ │ ├── data.py # FASTA/CSV I/O, cleaning, train/test split
81
+ │ ├── features.py # AAC, DPC, CTD, physicochemical features
82
+ │ ├── struct_features.py # Structural propensity features (28-dim)
83
+ │ ├── models_classical.py # Logistic regression, SVM, RF, XGBoost
84
+ │ ├── models_deep.py # ESM-2 embedding extractor + ProteinCNN
85
+ │ └── evaluate.py # AUC, MCC, PR, threshold search, plots
86
+ │
87
+ ├── data/
88
+ │ ├── raw/ # Downloaded IEDB CSVs and UniProt FASTAs (not tracked)
89
+ │ └── processed/ # dataset.csv built by download_data.py (not tracked)
90
+ │
91
+ └── results/ # Model weights, plots, summary (best_model.pkl not tracked)
92
+ ```
93
+
94
+ ---
95
+
96
+ ## Installation
97
+
98
+ **1. Clone the repository**
99
+
100
+ ```bash
101
+ git clone https://github.com/asah15/antigen-ml.git
102
+ cd antigen-ml
103
+ ```
104
+
105
+ **2. Create a virtual environment (recommended)**
106
+
107
+ ```bash
108
+ python -m venv .venv
109
+ # Windows
110
+ .venv\Scripts\activate
111
+ # macOS / Linux
112
+ source .venv/bin/activate
113
+ ```
114
+
115
+ **3. Install core dependencies**
116
+
117
+ ```bash
118
+ pip install -r requirements.txt
119
+ ```
120
+
121
+ **4. (Optional) Install ESM-2 dependencies**
122
+
123
+ Required only if you use `--features esm` during training.
124
+
125
+ ```bash
126
+ # CPU-only PyTorch (smaller download):
127
+ pip install torch --index-url https://download.pytorch.org/whl/cpu
128
+ pip install transformers>=4.40
129
+
130
+ # GPU (CUDA 12.x):
131
+ pip install torch --index-url https://download.pytorch.org/whl/cu121
132
+ pip install transformers>=4.40
133
+ ```
134
+
135
+ ---
136
+
137
+ ## Quickstart
138
+
139
+ ### Step 1 — Download labeled training data
140
+
141
+ The pipeline supports three data modes. **Mode A** (IEDB bulk export) gives the most sequences and is recommended.
142
+
143
+ #### Mode A: IEDB bulk export (recommended, ~1,900 antigens)
144
+
145
+ 1. Go to [https://www.iedb.org/database_export_v3.php](https://www.iedb.org/database_export_v3.php)
146
+ 2. Download **B Cell Assays** → `bcell_full_v3.zip`
147
+ 3. Extract so that `bcell_full_v3_single_file/bcell_full_v3.csv` is accessible
148
+
149
+ ```bash
150
+ python src/download_data.py \
151
+ --mode A \
152
+ --iedb_csv path/to/bcell_full_v3_single_file/bcell_full_v3.csv \
153
+ --out_csv data/processed/dataset.csv
154
+ ```
155
+
156
+ #### Mode B: UniProt KW-0045 keyword query (~600 antigens, internet only)
157
+
158
+ ```bash
159
+ python src/download_data.py \
160
+ --mode B \
161
+ --out_csv data/processed/dataset.csv
162
+ ```
163
+
164
+ #### Mode C: Synthetic demo (no data needed, for testing the pipeline)
165
+
166
+ ```bash
167
+ python src/download_data.py \
168
+ --mode C \
169
+ --out_csv data/processed/dataset.csv
170
+ ```
171
+
172
+ ---
173
+
174
+ ### Step 2 — Train a classifier
175
+
176
+ #### Classical features (fast, no GPU required)
177
+
178
+ ```bash
179
+ python train.py \
180
+ --csv data/processed/dataset.csv \
181
+ --features aac,dpc,ctd,phys,struct \
182
+ --models rf,xgb \
183
+ --cv 5
184
+ ```
185
+
186
+ #### ESM-2 embeddings only
187
+
188
+ ```bash
189
+ python train.py \
190
+ --csv data/processed/dataset.csv \
191
+ --features esm \
192
+ --esm_model facebook/esm2_t12_35M_UR50D \
193
+ --models xgb
194
+ ```
195
+
196
+ Embeddings are automatically cached to `results/esm_cache.npy` after the first run. Subsequent training runs load from cache and skip the slow embedding step.
197
+
198
+ #### Combined: ESM-2 + classical features (highest accuracy)
199
+
200
+ ```bash
201
+ python train.py \
202
+ --csv data/processed/dataset.csv \
203
+ --features aac,dpc,ctd,phys,struct,esm \
204
+ --esm_model facebook/esm2_t12_35M_UR50D \
205
+ --models xgb
206
+ ```
207
+
208
+ #### All options
209
+
210
+ ```
211
+ --csv PATH Labeled CSV with 'sequence' and 'label' columns
212
+ --pos PATH FASTA of antigen proteins (use with --neg)
213
+ --neg PATH FASTA of non-antigen proteins (use with --pos)
214
+ --demo Use 300 synthetic sequences per class (no files needed)
215
+
216
+ --features STR Comma-separated: aac,dpc,ctd,phys,struct,esm
217
+ Default: aac,dpc,ctd,phys,struct
218
+ --esm_model STR ESM-2 HuggingFace model ID (default: facebook/esm2_t12_35M_UR50D)
219
+ Options:
220
+ facebook/esm2_t6_8M_UR50D 320-dim (fastest)
221
+ facebook/esm2_t12_35M_UR50D 480-dim (default)
222
+ facebook/esm2_t30_150M_UR50D 640-dim
223
+ facebook/esm2_t33_650M_UR50D 1280-dim (best quality)
224
+ --esm_cache PATH Path to save/load ESM embeddings (default: results/esm_cache.npy)
225
+ --models STR Comma-separated: logreg,svm,rf,xgb (default: rf,xgb)
226
+ --cv INT Cross-validation folds (default: 5)
227
+ --out DIR Output directory (default: results/)
228
+ --seed INT Random seed (default: 42)
229
+ ```
230
+
231
+ **Outputs** written to `results/`:
232
+
233
+ | File | Description |
234
+ |---|---|
235
+ | `best_model.pkl` | Trained classifier (best CV AUC) |
236
+ | `summary.csv` | Per-model CV and test metrics |
237
+ | `roc_pr.png` | ROC and Precision-Recall curves |
238
+ | `metric_bars.png` | AUC bar chart across models |
239
+ | `feature_importance.png` | Top-30 feature importances (RF/XGB) |
240
+ | `esm_cache.npy` | Cached ESM-2 embeddings (if applicable) |
241
+
242
+ ---
243
+
244
+ ### Step 3 — Predict antigenicity for new proteins
245
+
246
+ ```bash
247
+ python predict.py \
248
+ --fasta my_proteins.fasta \
249
+ --model results/best_model.pkl \
250
+ --threshold 0.40 \
251
+ --out predictions.tsv
252
+ ```
253
+
254
+ #### Threshold guidance
255
+
256
+ The default threshold is 0.50. For the XGBoost + classical model, threshold **0.40** maximizes MCC on the held-out test set. Use the value printed during training (`best-thr=X.XX`) for the model you saved.
257
+
258
+ #### Output columns
259
+
260
+ | Column | Description |
261
+ |---|---|
262
+ | `protein_id` | Sequence ID from FASTA header (first word) |
263
+ | `length` | Amino-acid length after cleaning |
264
+ | `antigenicity_score` | Model probability (0–1); higher = more likely antigen |
265
+ | `prediction` | `antigen` or `non-antigen` at the given threshold |
266
+ | `flag` | `low_confidence` if score is between 0.40 and 0.60 |
267
+
268
+ #### All options
269
+
270
+ ```
271
+ --fasta PATH Input FASTA file (required)
272
+ --model PATH Trained model pickle (default: results/best_model.pkl)
273
+ --threshold FLOAT Decision threshold (default: 0.5)
274
+ --out PATH Output TSV path (default: stdout)
275
+ --min_len INT Skip sequences shorter than N aa (default: 30)
276
+ --max_len INT Truncate sequences longer than N aa (default: 6000)
277
+ ```
278
+
279
+ ---
280
+
281
+ ## Feature sets
282
+
283
+ | Set | Flag | Dim | Description |
284
+ |---|---|---|---|
285
+ | Amino acid composition | `aac` | 20 | Fraction of each of the 20 standard amino acids |
286
+ | Dipeptide composition | `dpc` | 400 | Fraction of each of the 400 dipeptide pairs |
287
+ | CTD | `ctd` | 42 | Composition, transition, distribution descriptors (Dubchak 1995) |
288
+ | Physicochemical | `phys` | 13 | MW, charge, pI, GRAVY, aromaticity, instability index, etc. |
289
+ | Structural propensity | `struct` | 28 | Chou-Fasman, Eisenberg hydrophobicity, Vihinen flexibility, Janin surface exposure |
290
+ | ESM-2 embeddings | `esm` | 320–1280 | Mean-pooled last-hidden-state from ESM-2 (Meta AI) |
291
+
292
+ ---
293
+
294
+ ## Data sources
295
+
296
+ - **IEDB** (Immune Epitope Database): B-cell assay bulk export v3. Positive labels are proteins with experimentally confirmed B-cell epitopes. [https://www.iedb.org](https://www.iedb.org)
297
+ - **UniProt/Swiss-Prot**: Reviewed entries. Positives are annotated with keyword KW-0045 (Antigen) in viral and bacterial organisms. Negatives are human, yeast, and plant entries without that keyword. [https://www.uniprot.org](https://www.uniprot.org)
298
+
299
+ ---
300
+
301
+ ## Reproducing the benchmark
302
+
303
+ ```bash
304
+ # 1. Download IEDB bulk export manually (see Step 1 above), then:
305
+ python src/download_data.py \
306
+ --mode A \
307
+ --iedb_csv path/to/bcell_full_v3.csv \
308
+ --out_csv data/processed/dataset.csv
309
+
310
+ # 2. Train
311
+ python train.py \
312
+ --csv data/processed/dataset.csv \
313
+ --features aac,dpc,ctd,phys,struct \
314
+ --models rf,xgb \
315
+ --seed 42
316
+
317
+ # 3. Inspect results
318
+ cat results/summary.csv
319
+ ```
320
+
321
+ ---
322
+
323
+ ## Windows GUI
324
+
325
+ A point-and-click interface is included for running predictions and viewing training results without touching the command line.
326
+
327
+ ---
328
+
329
+ ### Option A — pip install (cleanest, no install wizard needed)
330
+
331
+ Works on Windows, macOS, and Linux. No `.exe`, no `.bat`, no security warnings.
332
+
333
+ ```bash
334
+ pip install antigen-ml # core dependencies
335
+ pip install antigen-ml[gui] # + Pillow for plot display in the GUI
336
+ ```
337
+
338
+ Then launch from any terminal:
339
+
340
+ ```bash
341
+ antigenml-gui # opens the GUI window
342
+ antigenml-train --csv ... # training pipeline
343
+ antigenml-predict --fasta ... # inference
344
+ ```
345
+
346
+ > PyPI page: [pypi.org/project/antigen-ml](https://pypi.org/project/antigen-ml)
347
+
348
+ ---
349
+
350
+ ### Option B — Python installer script (Windows, no security warnings)
351
+
352
+ Windows Smart App Control blocks unsigned `.bat` and `.exe` files, but
353
+ `.py` files run through Python are never blocked.
354
+
355
+ **Requirements:** Python 3.10+ in PATH ([python.org/downloads](https://www.python.org/downloads/)).
356
+
357
+ ```
358
+ python windows\install.py
359
+ ```
360
+
361
+ This creates a `.venv`, installs dependencies, and places an **antigen-ml GUI** shortcut on your Desktop.
362
+
363
+ ---
364
+
365
+ ### Option C — Standalone .exe (no Python required)
366
+
367
+ Build a fully self-contained executable using PyInstaller and optionally wrap it into a one-click Setup wizard with Inno Setup.
368
+
369
+ **Step 1 — Install build tools** (once):
370
+
371
+ ```
372
+ pip install pyinstaller
373
+ ```
374
+
375
+ Inno Setup (optional, for the Setup wizard): [jrsoftware.org/isinfo.php](https://jrsoftware.org/isinfo.php)
376
+
377
+ **Step 2 — Build:**
378
+
379
+ ```
380
+ windows\build_exe.bat
381
+ ```
382
+
383
+ If Inno Setup is installed this produces `windows\Output\antigenml-setup.exe`, a standard Windows installer wizard. Without Inno Setup it produces the bundle at `dist\antigenml-gui\antigenml-gui.exe`; distribute the entire `antigenml-gui\` folder.
384
+
385
+ > **Note:** The standalone .exe excludes ESM-2 (PyTorch) because it would add ~2 GB. Use `train.py` with `--features aac,dpc,ctd,phys,struct` (classical features only) when distributing models that run with the standalone build.
386
+
387
+ ---
388
+
389
+ ### GUI features
390
+
391
+ **Predict tab** — Select a FASTA file and a trained model (`.pkl`), adjust the threshold and sequence-length filters, then click **Run**. Progress streams into the log panel in real time. Results appear in a table when the run completes: antigenic proteins are highlighted in orange, non-antigenic in green, and low-confidence calls (score 0.40–0.60) in yellow.
392
+
393
+ **Results tab** — Browse the plots saved by `train.py` (ROC & PR curves, AUC comparison, feature importances) and the per-model metrics from `results/summary.csv`. Click **Refresh** after a new training run to reload everything.
394
+
395
+ ---
396
+
397
+ ## Citation
398
+
399
+ If you use this tool in published work, please cite:
400
+
401
+ ```
402
+ Sah, A. (2026). antigen-ml: sequence- and structure-based antigenicity prediction
403
+ using classical and protein language model features. GitHub.
404
+ https://github.com/ash15/antigen-ml
405
+ ```
406
+
407
+ ---
408
+
409
+ ## References
410
+
411
+ - Dubchak, I. et al. (1995). Prediction of protein folding class using global description of amino acid sequence. *PNAS* 92(19), 8700–8704.
412
+ - Lin, Z. et al. (2023). Evolutionary-scale prediction of atomic-level protein structure with a language model. *Science* 379(6637), 1123–1130. (ESM-2)
413
+ - Vita, R. et al. (2019). The Immune Epitope Database (IEDB): 2018 update. *Nucleic Acids Research* 47(D1), D339–D343.