# elmonstrodelbrest/chaosai architecture

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

- Interactive diagram: https://gitdiagram.com/elmonstrodelbrest/chaosai
- Repository: https://github.com/elmonstrodelbrest/chaosai

## Overview

ChaosAI is a research pipeline for learning financial time-series representations and evaluating cross-sectional portfolio policies, not a live trading system. The principal workflow tokenizes market data, pretrains a JEPA foundation model with a CFM auxiliary objective, precomputes embeddings and an RL buffer, trains a cross-sectional DQN, then evaluates it on held-out data, optionally applying Multiverse Crossing confidence filtering. The graph preserves the README’s tokenizer, macro/on-chain context, and TD-MPC2 capabilities. Several links follow documented scripts and method stages; unsampled implementations are not inferred beyond those documented relationships.

Main entry points: [`src/jax_v6/jepa.py`](https://github.com/elmonstrodelbrest/chaosai/blob/main/src/jax_v6/jepa.py) (JEPA pretraining), [`src/jax_v6/strate_iv/dqn_agent.py`](https://github.com/elmonstrodelbrest/chaosai/blob/main/src/jax_v6/strate_iv/dqn_agent.py) (Cross-sectional DQN), [`src/jax_v6/strate_i/tokenizer.py`](https://github.com/elmonstrodelbrest/chaosai/blob/main/src/jax_v6/strate_i/tokenizer.py) (FSQ tokenizer), [`src/jax_v6/strate_iv/env_cross_sectional.py`](https://github.com/elmonstrodelbrest/chaosai/blob/main/src/jax_v6/strate_iv/env_cross_sectional.py) (Portfolio environment)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_data["Data and context"]
  node_market_data["Market data"]
  node_tokenizer["FSQ tokenizer<br/>[tokenizer.py]"]
  node_tokenized["Token sequences"]
  node_grain["Grain loader<br/>[grain_loader.py]"]
  node_macro["Macro context<br/>[macro_loader.py]"]
end

subgraph group_foundation["Representation learning"]
  node_jepa["JEPA pretraining<br/>[jepa.py]"]
  node_encoder["Mamba encoder<br/>[mamba2_encoder.py]"]
  node_predictor["JEPA predictors<br/>[predictor.py]"]
  node_flow["Flow predictor<br/>[flow_predictor.py]"]
  node_losses["Representation losses<br/>[vicreg.py]"]
  node_training["Training runtime<br/>[train_step.py]"]
  node_embeddings["JEPA embeddings"]
end

subgraph group_policy["Portfolio learning"]
  node_rlbuffer[("RL buffer")]
  node_environment["Portfolio environment"]
  node_dqn["Cross-sectional DQN<br/>[dqn_agent.py]"]
  node_replay[("Replay buffer<br/>[replay_buffer.py]")]
end

subgraph group_evaluation["Evaluation and robustness"]
  node_mvx["Multiverse Crossing"]
  node_evaluation_run["Held-out evaluation"]
  node_mvx_eval["MVX evaluation<br/>[eval_oos_mvx.py]"]
end

subgraph group_alternate["Additional modeling"]
  node_onchain["On-chain graphs<br/>[graph_builder.py]"]
  node_gnn["On-chain GNN<br/>[gnn_model.py]"]
  node_tdmpc["TD-MPC2 agent<br/>[tdmpc2.py]"]
end

node_researcher(("Researcher"))

node_researcher -->|"runs tokenizer"| node_tokenizer
node_market_data -->|"tokenizes OHLCV"| node_tokenizer
node_tokenizer -->|"produces codes"| node_tokenized
node_tokenized -->|"reads shards"| node_grain
node_grain -->|"provides batches"| node_jepa
node_macro -.->|"provides context"| node_jepa
node_onchain -->|"builds graph input"| node_gnn
node_gnn -.->|"provides embeddings"| node_jepa
node_jepa -->|"invokes encoder"| node_encoder
node_jepa -->|"invokes predictor"| node_predictor
node_jepa -->|"uses CFM objective"| node_flow
node_jepa -->|"computes losses"| node_losses
node_training -->|"updates model"| node_jepa
node_jepa -->|"encodes sequences"| node_embeddings
node_embeddings -->|"precomputes transitions"| node_rlbuffer
node_rlbuffer -->|"supplies experience"| node_replay
node_replay -->|"samples batches"| node_dqn
node_dqn -->|"selects portfolio actions"| node_environment
node_environment -->|"returns rewards"| node_dqn
node_embeddings -->|"perturbs representations"| node_mvx
node_dqn -->|"evaluates policy"| node_evaluation_run
node_rlbuffer -->|"provides held-out data"| node_evaluation_run
node_evaluation_run -->|"reports results"| node_researcher
node_dqn -->|"evaluates with filter"| node_mvx_eval
node_mvx -->|"provides confidence"| node_mvx_eval
node_rlbuffer -->|"provides held-out data"| node_mvx_eval
node_tdmpc -.->|"uses portfolio environment"| node_environment

click node_tokenizer "https://github.com/elmonstrodelbrest/chaosai/blob/main/src/jax_v6/strate_i/tokenizer.py"
click node_grain "https://github.com/elmonstrodelbrest/chaosai/blob/main/src/jax_v6/data/grain_loader.py"
click node_macro "https://github.com/elmonstrodelbrest/chaosai/blob/main/src/jax_v6/data/macro_loader.py"
click node_onchain "https://github.com/elmonstrodelbrest/chaosai/blob/main/src/strate_v/graph_builder.py"
click node_gnn "https://github.com/elmonstrodelbrest/chaosai/blob/main/src/strate_v/gnn_model.py"
click node_jepa "https://github.com/elmonstrodelbrest/chaosai/blob/main/src/jax_v6/jepa.py"
click node_encoder "https://github.com/elmonstrodelbrest/chaosai/blob/main/src/jax_v6/encoders/mamba2_encoder.py"
click node_predictor "https://github.com/elmonstrodelbrest/chaosai/blob/main/src/jax_v6/predictors/predictor.py"
click node_flow "https://github.com/elmonstrodelbrest/chaosai/blob/main/src/jax_v6/predictors/flow_predictor.py"
click node_losses "https://github.com/elmonstrodelbrest/chaosai/blob/main/src/jax_v6/losses/vicreg.py"
click node_training "https://github.com/elmonstrodelbrest/chaosai/blob/main/src/jax_v6/training/train_step.py"
click node_environment "https://github.com/elmonstrodelbrest/chaosai/blob/main/src/jax_v6/strate_iv/env_cross_sectional.py"
click node_dqn "https://github.com/elmonstrodelbrest/chaosai/blob/main/src/jax_v6/strate_iv/dqn_agent.py"
click node_replay "https://github.com/elmonstrodelbrest/chaosai/blob/main/src/jax_v6/strate_iv/replay_buffer.py"
click node_mvx "https://github.com/elmonstrodelbrest/chaosai/blob/main/src/jax_v6/strate_iv/multiverse_crossing.py"
click node_evaluation_run "https://github.com/elmonstrodelbrest/chaosai/blob/main/scripts/eval_oos_temporal.py"
click node_mvx_eval "https://github.com/elmonstrodelbrest/chaosai/blob/main/scripts/eval_oos_mvx.py"
click node_tdmpc "https://github.com/elmonstrodelbrest/chaosai/blob/main/src/jax_v6/strate_iv/tdmpc2.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_market_data,node_tokenizer,node_tokenized,node_grain,node_macro toneBlue
class node_jepa,node_encoder,node_predictor,node_flow,node_losses,node_training,node_embeddings toneAmber
class node_rlbuffer,node_environment,node_dqn,node_replay toneMint
class node_mvx,node_evaluation_run,node_mvx_eval toneRose
class node_onchain,node_gnn,node_tdmpc,node_researcher toneIndigo
```

## Components

### Data and context

- **Market data**
- [FSQ tokenizer](https://github.com/elmonstrodelbrest/chaosai/blob/main/src/jax_v6/strate_i/tokenizer.py): `src/jax_v6/strate_i/tokenizer.py`
- **Token sequences**
- [Grain loader](https://github.com/elmonstrodelbrest/chaosai/blob/main/src/jax_v6/data/grain_loader.py): `src/jax_v6/data/grain_loader.py`
- [Macro context](https://github.com/elmonstrodelbrest/chaosai/blob/main/src/jax_v6/data/macro_loader.py): `src/jax_v6/data/macro_loader.py`

### Representation learning

- [JEPA pretraining](https://github.com/elmonstrodelbrest/chaosai/blob/main/src/jax_v6/jepa.py): `src/jax_v6/jepa.py`
- [Mamba encoder](https://github.com/elmonstrodelbrest/chaosai/blob/main/src/jax_v6/encoders/mamba2_encoder.py): `src/jax_v6/encoders/mamba2_encoder.py`
- [JEPA predictors](https://github.com/elmonstrodelbrest/chaosai/blob/main/src/jax_v6/predictors/predictor.py): `src/jax_v6/predictors/predictor.py`
- [Flow predictor](https://github.com/elmonstrodelbrest/chaosai/blob/main/src/jax_v6/predictors/flow_predictor.py): `src/jax_v6/predictors/flow_predictor.py`
- [Representation losses](https://github.com/elmonstrodelbrest/chaosai/blob/main/src/jax_v6/losses/vicreg.py): `src/jax_v6/losses/vicreg.py`
- [Training runtime](https://github.com/elmonstrodelbrest/chaosai/blob/main/src/jax_v6/training/train_step.py): `src/jax_v6/training/train_step.py`
- **JEPA embeddings**

### Portfolio learning

- **RL buffer**
- [Portfolio environment](https://github.com/elmonstrodelbrest/chaosai/blob/main/src/jax_v6/strate_iv/env_cross_sectional.py): `src/jax_v6/strate_iv/env_cross_sectional.py`
- [Cross-sectional DQN](https://github.com/elmonstrodelbrest/chaosai/blob/main/src/jax_v6/strate_iv/dqn_agent.py): `src/jax_v6/strate_iv/dqn_agent.py`
- [Replay buffer](https://github.com/elmonstrodelbrest/chaosai/blob/main/src/jax_v6/strate_iv/replay_buffer.py): `src/jax_v6/strate_iv/replay_buffer.py`

### Evaluation and robustness

- [Multiverse Crossing](https://github.com/elmonstrodelbrest/chaosai/blob/main/src/jax_v6/strate_iv/multiverse_crossing.py): `src/jax_v6/strate_iv/multiverse_crossing.py`
- [Held-out evaluation](https://github.com/elmonstrodelbrest/chaosai/blob/main/scripts/eval_oos_temporal.py): `scripts/eval_oos_temporal.py`
- [MVX evaluation](https://github.com/elmonstrodelbrest/chaosai/blob/main/scripts/eval_oos_mvx.py): `scripts/eval_oos_mvx.py`

### Additional modeling

- [On-chain graphs](https://github.com/elmonstrodelbrest/chaosai/blob/main/src/strate_v/graph_builder.py): `src/strate_v/graph_builder.py`
- [On-chain GNN](https://github.com/elmonstrodelbrest/chaosai/blob/main/src/strate_v/gnn_model.py): `src/strate_v/gnn_model.py`
- [TD-MPC2 agent](https://github.com/elmonstrodelbrest/chaosai/blob/main/src/jax_v6/strate_iv/tdmpc2.py): `src/jax_v6/strate_iv/tdmpc2.py`

### Other components

- **Researcher**

## Connections

- Researcher → FSQ tokenizer: runs tokenizer
- Market data → FSQ tokenizer: tokenizes OHLCV
- FSQ tokenizer → Token sequences: produces codes
- Token sequences → Grain loader: reads shards
- Grain loader → JEPA pretraining: provides batches
- Macro context → JEPA pretraining: provides context
- On-chain graphs → On-chain GNN: builds graph input
- On-chain GNN → JEPA pretraining: provides embeddings
- JEPA pretraining → Mamba encoder: invokes encoder
- JEPA pretraining → JEPA predictors: invokes predictor
- JEPA pretraining → Flow predictor: uses CFM objective
- JEPA pretraining → Representation losses: computes losses
- Training runtime → JEPA pretraining: updates model
- JEPA pretraining → JEPA embeddings: encodes sequences
- JEPA embeddings → RL buffer: precomputes transitions
- RL buffer → Replay buffer: supplies experience
- Replay buffer → Cross-sectional DQN: samples batches
- Cross-sectional DQN → Portfolio environment: selects portfolio actions
- Portfolio environment → Cross-sectional DQN: returns rewards
- JEPA embeddings → Multiverse Crossing: perturbs representations
- Cross-sectional DQN → Held-out evaluation: evaluates policy
- RL buffer → Held-out evaluation: provides held-out data
- Held-out evaluation → Researcher: reports results
- Cross-sectional DQN → MVX evaluation: evaluates with filter
- Multiverse Crossing → MVX evaluation: provides confidence
- RL buffer → MVX evaluation: provides held-out data
- TD-MPC2 agent → Portfolio environment: uses portfolio environment