# kritanuscodingdelight/jax_genomics_predictor architecture

> Architecture diagram of the GitHub repository kritanuscodingdelight/jax_genomics_predictor, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-09-28). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/kritanuscodingdelight/jax_genomics_predictor
- Repository: https://github.com/kritanuscodingdelight/jax_genomics_predictor

## Overview

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`](https://github.com/kritanuscodingdelight/jax_genomics_predictor/blob/main/chromatin_accessibility_jax.ipynb) (Data downloads)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_data["Data preparation"]
  node_acquire["Data downloads"]
  node_peakparse["Peak parsing"]
  node_sequence["Sequence loading"]
  node_encoding["One-hot encoding"]
  node_dataset["Window dataset"]
  node_split["Train/val/test split"]
end

subgraph group_model["Model and training"]
  node_network["Residual CNN"]
  node_optimizer["Optimizer and state"]
  node_training["Training loop"]
  node_checkpoint["Best-val checkpoint"]
end

subgraph group_analysis["Analysis"]
  node_eda["Exploratory analysis"]
  node_evaluation["Test evaluation"]
  node_saliency["Gradient saliency"]
end

node_researcher(("Researcher"))
node_encode["ENCODE API"]
node_ucsc["UCSC genome"]

node_researcher -->|"runs notebook"| node_acquire
node_acquire -->|"queries and downloads"| node_encode
node_acquire -->|"downloads FASTA"| node_ucsc
node_acquire -->|"provides peak file"| node_peakparse
node_acquire -->|"provides FASTA"| node_sequence
node_sequence -->|"provides DNA"| node_encoding
node_peakparse -->|"provides peaks"| node_dataset
node_encoding -->|"encodes windows"| node_dataset
node_dataset -->|"creates samples"| node_split
node_split -->|"provides samples"| node_eda
node_split -->|"provides training data"| node_training
node_split -->|"provides test data"| node_evaluation
node_network -->|"provides model"| node_training
node_optimizer -->|"provides train state"| node_training
node_training -->|"selects best validation"| node_checkpoint
node_checkpoint -->|"restores best parameters"| node_evaluation
node_checkpoint -->|"provides trained model"| node_saliency
node_evaluation -->|"reports metrics"| node_researcher
node_saliency -->|"shows importance profile"| node_researcher

click node_acquire "https://github.com/kritanuscodingdelight/jax_genomics_predictor/blob/main/chromatin_accessibility_jax.ipynb"
click node_peakparse "https://github.com/kritanuscodingdelight/jax_genomics_predictor/blob/main/chromatin_accessibility_jax.ipynb"
click node_sequence "https://github.com/kritanuscodingdelight/jax_genomics_predictor/blob/main/chromatin_accessibility_jax.ipynb"
click node_encoding "https://github.com/kritanuscodingdelight/jax_genomics_predictor/blob/main/chromatin_accessibility_jax.ipynb"
click node_dataset "https://github.com/kritanuscodingdelight/jax_genomics_predictor/blob/main/chromatin_accessibility_jax.ipynb"
click node_split "https://github.com/kritanuscodingdelight/jax_genomics_predictor/blob/main/chromatin_accessibility_jax.ipynb"
click node_eda "https://github.com/kritanuscodingdelight/jax_genomics_predictor/blob/main/chromatin_accessibility_jax.ipynb"
click node_network "https://github.com/kritanuscodingdelight/jax_genomics_predictor/blob/main/chromatin_accessibility_jax.ipynb"
click node_optimizer "https://github.com/kritanuscodingdelight/jax_genomics_predictor/blob/main/chromatin_accessibility_jax.ipynb"
click node_training "https://github.com/kritanuscodingdelight/jax_genomics_predictor/blob/main/chromatin_accessibility_jax.ipynb"
click node_checkpoint "https://github.com/kritanuscodingdelight/jax_genomics_predictor/blob/main/chromatin_accessibility_jax.ipynb"
click node_evaluation "https://github.com/kritanuscodingdelight/jax_genomics_predictor/blob/main/chromatin_accessibility_jax.ipynb"
click node_saliency "https://github.com/kritanuscodingdelight/jax_genomics_predictor/blob/main/chromatin_accessibility_jax.ipynb"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_acquire,node_peakparse,node_sequence,node_encoding,node_dataset,node_split toneBlue
class node_network,node_optimizer,node_training,node_checkpoint toneAmber
class node_eda,node_evaluation,node_saliency,node_encode toneMint
class node_researcher,node_ucsc toneIndigo
```

## Components

### Data preparation

- [Data downloads](https://github.com/kritanuscodingdelight/jax_genomics_predictor/blob/main/chromatin_accessibility_jax.ipynb): `chromatin_accessibility_jax.ipynb`
- [Peak parsing](https://github.com/kritanuscodingdelight/jax_genomics_predictor/blob/main/chromatin_accessibility_jax.ipynb): `chromatin_accessibility_jax.ipynb`
- [Sequence loading](https://github.com/kritanuscodingdelight/jax_genomics_predictor/blob/main/chromatin_accessibility_jax.ipynb): `chromatin_accessibility_jax.ipynb`
- [One-hot encoding](https://github.com/kritanuscodingdelight/jax_genomics_predictor/blob/main/chromatin_accessibility_jax.ipynb): `chromatin_accessibility_jax.ipynb`
- [Window dataset](https://github.com/kritanuscodingdelight/jax_genomics_predictor/blob/main/chromatin_accessibility_jax.ipynb): `chromatin_accessibility_jax.ipynb`
- [Train/val/test split](https://github.com/kritanuscodingdelight/jax_genomics_predictor/blob/main/chromatin_accessibility_jax.ipynb): `chromatin_accessibility_jax.ipynb`

### Model and training

- [Residual CNN](https://github.com/kritanuscodingdelight/jax_genomics_predictor/blob/main/chromatin_accessibility_jax.ipynb): `chromatin_accessibility_jax.ipynb`
- [Optimizer and state](https://github.com/kritanuscodingdelight/jax_genomics_predictor/blob/main/chromatin_accessibility_jax.ipynb): `chromatin_accessibility_jax.ipynb`
- [Training loop](https://github.com/kritanuscodingdelight/jax_genomics_predictor/blob/main/chromatin_accessibility_jax.ipynb): `chromatin_accessibility_jax.ipynb`
- [Best-val checkpoint](https://github.com/kritanuscodingdelight/jax_genomics_predictor/blob/main/chromatin_accessibility_jax.ipynb): `chromatin_accessibility_jax.ipynb`

### Analysis

- [Exploratory analysis](https://github.com/kritanuscodingdelight/jax_genomics_predictor/blob/main/chromatin_accessibility_jax.ipynb): `chromatin_accessibility_jax.ipynb`
- [Test evaluation](https://github.com/kritanuscodingdelight/jax_genomics_predictor/blob/main/chromatin_accessibility_jax.ipynb): `chromatin_accessibility_jax.ipynb`
- [Gradient saliency](https://github.com/kritanuscodingdelight/jax_genomics_predictor/blob/main/chromatin_accessibility_jax.ipynb): `chromatin_accessibility_jax.ipynb`

### Other components

- **Researcher**
- **ENCODE API**
- **UCSC genome**

## Connections

- 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