How kritanuscodingdelight/jax_genomics_predictor works
This notebook is an end-to-end JAX/Flax chromatin-accessibility prediction workflow: it obtains peak and reference-genome data, constructs labelled sequence windows, trains a convolutional model, evaluates predictions, and interprets input saliency. The README documents these stages and their order; no notebook source was available to verify implementation details beyond that documentation.
Main entry points: chromatin_accessibility_jax.ipynb Data downloads
Components 16
Data preparation
- Data downloads
chromatin_accessibility_jax.ipynb - Peak parsing
chromatin_accessibility_jax.ipynb - Sequence loading
chromatin_accessibility_jax.ipynb - One-hot encoding
chromatin_accessibility_jax.ipynb - Window dataset
chromatin_accessibility_jax.ipynb - Train/val/test split
chromatin_accessibility_jax.ipynb
Model and training
- Residual CNN
chromatin_accessibility_jax.ipynb - Optimizer and state
chromatin_accessibility_jax.ipynb - Training loop
chromatin_accessibility_jax.ipynb - Best-val checkpoint
chromatin_accessibility_jax.ipynb
Analysis
- Exploratory analysis
chromatin_accessibility_jax.ipynb - Test evaluation
chromatin_accessibility_jax.ipynb - Gradient saliency
chromatin_accessibility_jax.ipynb
Other components
- Researcher
- ENCODE API
- UCSC genome
Connections 19
- Researcher → Data downloads: runs notebook
- Data downloads → ENCODE API: queries and downloads
- Data downloads → UCSC genome: downloads FASTA
- Data downloads → Peak parsing: provides peak file
- Data downloads → Sequence loading: provides FASTA
- Sequence loading → One-hot encoding: provides DNA
- Peak parsing → Window dataset: provides peaks
- One-hot encoding → Window dataset: encodes windows
- Window dataset → Train/val/test split: creates samples
- Train/val/test split → Exploratory analysis: provides samples
- Train/val/test split → Training loop: provides training data
- Train/val/test split → Test evaluation: provides test data
- Residual CNN → Training loop: provides model
- Optimizer and state → Training loop: provides train state
- Training loop → Best-val checkpoint: selects best validation
- Best-val checkpoint → Test evaluation: restores best parameters
- Best-val checkpoint → Gradient saliency: provides trained model
- Test evaluation → Researcher: reports metrics
- Gradient saliency → Researcher: shows importance profile