# sennadir/ts_jepa architecture

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

- Interactive diagram: https://gitdiagram.com/sennadir/ts_jepa
- Repository: https://github.com/sennadir/ts_jepa

## Overview

TS-JEPA is a PyTorch time-series representation-learning project: it pretrains an encoder to predict masked patch embeddings, then reuses a saved encoder for downstream forecasting of the next patch. The principal path runs from CSV preparation and masking through encoder/predictor training and checkpointing, then through a frozen encoder and linear decoder to forecast metrics. The graph also retains the separate autoregressive evaluation entry point documented by the tree. Some evaluation wiring is inferred only from README descriptions and filenames; unsampled implementations are not expanded.

Main entry points: [`pretrain.py`](https://github.com/sennadir/ts_jepa/blob/main/pretrain.py) (Pretrain loop), [`eval_forecast_last_pred.py`](https://github.com/sennadir/ts_jepa/blob/main/eval_forecast_last_pred.py) (Last-value forecast), [`src/models/encoder.py`](https://github.com/sennadir/ts_jepa/blob/main/src/models/encoder.py) (Time-series encoder), [`src/models/predictor.py`](https://github.com/sennadir/ts_jepa/blob/main/src/models/predictor.py) (JEPA predictor)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_data["Data preparation"]
  node_csv["Weather CSV<br/>[weather.csv]"]
  node_jepa_loader["Pretrain loader<br/>[data_loader.py]"]
  node_csv_dataset["Masked dataset<br/>[data_class.py]"]
  node_eval_dataset["Forecast dataset<br/>[data_class.py]"]
  node_eval_loader["Evaluation loader<br/>[data_loader.py]"]
end

subgraph group_learning["JEPA pretraining"]
  node_pretrain["Pretrain loop<br/>[pretrain.py]"]
  node_ema["EMA target encoder<br/>[encoder.py]"]
  node_predictor["JEPA predictor<br/>[predictor.py]"]
  node_loss_update["Loss and EMA update<br/>[pretrain.py]"]
  node_checkpoint["Encoder checkpoint"]
end

subgraph group_forecast["Forecast evaluation"]
  node_forecast["Last-value forecast"]
  node_autoregressive["Autoregressive evaluation"]
  node_decoder["Linear decoder<br/>[decoder.py]"]
  node_metrics["Forecast metrics<br/>[utils.py]"]
end

subgraph group_shared["Model components"]
  node_encoder["Time-series encoder<br/>[encoder.py]"]
  node_tokenizer["Patch tokenizer<br/>[tokenizer.py]"]
  node_transformer["Transformer blocks<br/>[modules.py]"]
  node_mask["Mask operations<br/>[mask_utils.py]"]
end

node_user(("Researcher"))

node_user -->|"starts pretraining"| node_pretrain
node_user -->|"starts evaluation"| node_forecast
node_user -->|"starts evaluation"| node_autoregressive
node_pretrain -->|"requests batches"| node_jepa_loader
node_jepa_loader -->|"wraps dataset"| node_csv_dataset
node_csv_dataset -->|"reads CSV"| node_csv
node_pretrain -->|"encodes patches"| node_encoder
node_encoder -->|"tokenizes patches"| node_tokenizer
node_encoder -->|"applies blocks"| node_transformer
node_pretrain -->|"copies encoder"| node_ema
node_pretrain -->|"predicts masked targets"| node_predictor
node_predictor -->|"applies blocks"| node_transformer
node_predictor -->|"selects tokens"| node_mask
node_pretrain -->|"selects targets"| node_mask
node_pretrain -->|"computes loss and updates EMA"| node_loss_update
node_loss_update -->|"saves encoder"| node_checkpoint
node_forecast -->|"requests windows"| node_eval_loader
node_eval_loader -->|"wraps dataset"| node_eval_dataset
node_eval_dataset -->|"reads CSV"| node_csv
node_forecast -->|"loads weights"| node_checkpoint
node_forecast -->|"encodes context"| node_encoder
node_forecast -->|"predicts next patch"| node_decoder
node_forecast -->|"computes errors"| node_metrics

click node_csv "https://github.com/sennadir/ts_jepa/blob/main/data/weather/weather.csv"
click node_jepa_loader "https://github.com/sennadir/ts_jepa/blob/main/src/data_loaders/data_loader.py"
click node_csv_dataset "https://github.com/sennadir/ts_jepa/blob/main/src/data_loaders/data_class.py"
click node_eval_dataset "https://github.com/sennadir/ts_jepa/blob/main/src/data_loaders/data_class.py"
click node_pretrain "https://github.com/sennadir/ts_jepa/blob/main/pretrain.py"
click node_ema "https://github.com/sennadir/ts_jepa/blob/main/src/models/encoder.py"
click node_predictor "https://github.com/sennadir/ts_jepa/blob/main/src/models/predictor.py"
click node_loss_update "https://github.com/sennadir/ts_jepa/blob/main/pretrain.py"
click node_checkpoint "https://github.com/sennadir/ts_jepa/blob/main/logs/output_model/weather/lr_1e-07_ema_momentum_0.998_mask_ratio_0.7_ratio_patches_10_encoder_128_2_1_predictor_128_2_1_epoch_5000.pt"
click node_forecast "https://github.com/sennadir/ts_jepa/blob/main/eval_forecast_last_pred.py"
click node_autoregressive "https://github.com/sennadir/ts_jepa/blob/main/eval_autoregressive.py"
click node_eval_loader "https://github.com/sennadir/ts_jepa/blob/main/src/data_loaders/data_loader.py"
click node_encoder "https://github.com/sennadir/ts_jepa/blob/main/src/models/encoder.py"
click node_tokenizer "https://github.com/sennadir/ts_jepa/blob/main/src/models/tokenizer.py"
click node_transformer "https://github.com/sennadir/ts_jepa/blob/main/src/models/utils/modules.py"
click node_mask "https://github.com/sennadir/ts_jepa/blob/main/src/models/utils/mask_utils.py"
click node_decoder "https://github.com/sennadir/ts_jepa/blob/main/src/models/decoder.py"
click node_metrics "https://github.com/sennadir/ts_jepa/blob/main/main/utils.py"

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_csv,node_jepa_loader,node_csv_dataset,node_eval_dataset,node_eval_loader toneBlue
class node_pretrain,node_ema,node_predictor,node_loss_update,node_checkpoint toneAmber
class node_forecast,node_autoregressive,node_decoder,node_metrics toneMint
class node_encoder,node_tokenizer,node_transformer,node_mask toneRose
class node_user toneIndigo
```

## Components

### Data preparation

- [Weather CSV](https://github.com/sennadir/ts_jepa/blob/main/data/weather/weather.csv): `data/weather/weather.csv`
- [Pretrain loader](https://github.com/sennadir/ts_jepa/blob/main/src/data_loaders/data_loader.py): `src/data_loaders/data_loader.py`
- [Masked dataset](https://github.com/sennadir/ts_jepa/blob/main/src/data_loaders/data_class.py): `src/data_loaders/data_class.py`
- [Forecast dataset](https://github.com/sennadir/ts_jepa/blob/main/src/data_loaders/data_class.py): `src/data_loaders/data_class.py`
- [Evaluation loader](https://github.com/sennadir/ts_jepa/blob/main/src/data_loaders/data_loader.py): `src/data_loaders/data_loader.py`

### JEPA pretraining

- [Pretrain loop](https://github.com/sennadir/ts_jepa/blob/main/pretrain.py): `pretrain.py`
- [EMA target encoder](https://github.com/sennadir/ts_jepa/blob/main/src/models/encoder.py): `src/models/encoder.py`
- [JEPA predictor](https://github.com/sennadir/ts_jepa/blob/main/src/models/predictor.py): `src/models/predictor.py`
- [Loss and EMA update](https://github.com/sennadir/ts_jepa/blob/main/pretrain.py): `pretrain.py`
- [Encoder checkpoint](https://github.com/sennadir/ts_jepa/blob/main/logs/output_model/weather/lr_1e-07_ema_momentum_0.998_mask_ratio_0.7_ratio_patches_10_encoder_128_2_1_predictor_128_2_1_epoch_5000.pt): `logs/output_model/weather/lr_1e-07_ema_momentum_0.998_mask_ratio_0.7_ratio_patches_10_encoder_128_2_1_predictor_128_2_1_epoch_5000.pt`

### Forecast evaluation

- [Last-value forecast](https://github.com/sennadir/ts_jepa/blob/main/eval_forecast_last_pred.py): `eval_forecast_last_pred.py`
- [Autoregressive evaluation](https://github.com/sennadir/ts_jepa/blob/main/eval_autoregressive.py): `eval_autoregressive.py`
- [Linear decoder](https://github.com/sennadir/ts_jepa/blob/main/src/models/decoder.py): `src/models/decoder.py`
- [Forecast metrics](https://github.com/sennadir/ts_jepa/blob/main/main/utils.py): `main/utils.py`

### Model components

- [Time-series encoder](https://github.com/sennadir/ts_jepa/blob/main/src/models/encoder.py): `src/models/encoder.py`
- [Patch tokenizer](https://github.com/sennadir/ts_jepa/blob/main/src/models/tokenizer.py): `src/models/tokenizer.py`
- [Transformer blocks](https://github.com/sennadir/ts_jepa/blob/main/src/models/utils/modules.py): `src/models/utils/modules.py`
- [Mask operations](https://github.com/sennadir/ts_jepa/blob/main/src/models/utils/mask_utils.py): `src/models/utils/mask_utils.py`

### Other components

- **Researcher**

## Connections

- Researcher → Pretrain loop: starts pretraining
- Researcher → Last-value forecast: starts evaluation
- Researcher → Autoregressive evaluation: starts evaluation
- Pretrain loop → Pretrain loader: requests batches
- Pretrain loader → Masked dataset: wraps dataset
- Masked dataset → Weather CSV: reads CSV
- Pretrain loop → Time-series encoder: encodes patches
- Time-series encoder → Patch tokenizer: tokenizes patches
- Time-series encoder → Transformer blocks: applies blocks
- Pretrain loop → EMA target encoder: copies encoder
- Pretrain loop → JEPA predictor: predicts masked targets
- JEPA predictor → Transformer blocks: applies blocks
- JEPA predictor → Mask operations: selects tokens
- Pretrain loop → Mask operations: selects targets
- Pretrain loop → Loss and EMA update: computes loss and updates EMA
- Loss and EMA update → Encoder checkpoint: saves encoder
- Last-value forecast → Evaluation loader: requests windows
- Evaluation loader → Forecast dataset: wraps dataset
- Forecast dataset → Weather CSV: reads CSV
- Last-value forecast → Encoder checkpoint: loads weights
- Last-value forecast → Time-series encoder: encodes context
- Last-value forecast → Linear decoder: predicts next patch
- Last-value forecast → Forecast metrics: computes errors