chaos-ml 0.1.3

Unified chaotic systems forecasting toolkit

ChaosML

ChaosML is a unified toolkit for forecasting classic chaotic systems (Duffing, Lorenz-63, Lorenz-96) using PyTorch. It replaces notebook duplication with a config-driven CLI, reusable modules, and an interactive UI.

Highlights

• Shared data generation + scaling + windowing pipeline
• LSTM, BiLSTM, Transformer, and Encoder-Decoder models
• Reproducible runs (seeded), checkpointing, metrics, and plots
• Optional hyperparameter tuning (Optuna)
• Custom model plug-ins
• UI for running, comparing, and exporting results

Use cases

• Benchmarking sequence models on chaotic dynamics
• Testing forecasting stability and multi-step error growth
• Prototyping data-driven surrogates for complex dynamical systems
• Exploring education and research demos in nonlinear dynamics
• Real-world analogs: weather/climate prototypes, mechanical oscillators, power/grid dynamics, and stress-testing forecasting workflows in finance

Quick start (Windows)

Recommended Python: 3.11

Create a venv and install deps:

py -3.11 -m venv .venv
.\.venv\Scripts\python -m pip install -r requirements.txt

Run a single experiment:

.\.venv\Scripts\python -m chaos_ml.cli --config configs/lorenz63_lstm.json

Run all example configs:

.
un_all.ps1

Outputs are saved to runs/... (metrics, predictions, plots, model checkpoint).

If PyTorch fails to load on Windows, install the Microsoft Visual C++ 2015-2022 Redistributable (x64) and retry.

Install as a package (editable)

.\.venv\Scripts\python -m pip install -e .

Streamlit UI

.\.venv\Scripts\python -m streamlit run app.py

UI features:

• Select and edit configs
• Save presets
• Queue multiple runs
• Inspect run history, metrics, and plots
• Compare runs across metrics
• Top-N run spotlight and CSV export
• Export trained models

Streamlit Community Cloud

This repo includes runtime.txt for Python 3.11 compatibility. For faster cloud demos, use the *_demo.json presets in configs/ to keep runs short.

Config structure

Each config is a JSON file with:

• system: which system, params, and time span
• data: window, horizon, splits, scaler
• model: architecture and hyperparameters
• training: epochs, batch size, learning rate, patience
• tuning (optional): Optuna search config

Example (minimal):

{
  "system": {"name": "lorenz63", "t_end": 100, "t_points": 1000},
  "data": {"window": 10, "horizon": 1, "stride": 1, "train_ratio": 0.2, "val_ratio": 0.1, "scaler": "minmax"},
  "model": {"name": "lstm", "units": 256, "depth": 3, "dropout": 0.1},
  "training": {"epochs": 200, "batch_size": 64, "patience": 20, "learning_rate": 0.001}
}

Hyperparameter tuning (Optuna)

Add a tuning section to your config:

"tuning": {
  "enabled": true,
  "n_trials": 25,
  "epochs": 50,
  "search_space": {
    "window": {"low": 5, "high": 30, "step": 5},
    "learning_rate": {"low": 1e-4, "high": 1e-2, "log": true},
    "units": {"low": 64, "high": 512, "step": 64},
    "depth": {"low": 2, "high": 4, "step": 1},
    "dropout": {"low": 0.0, "high": 0.3, "step": 0.05}
  }
}

The best params are saved to runs/<output>/tuning/best_params.json and applied to the final training run.

Custom models

Set model.name to custom and point to your builder:

"model": {
  "name": "custom",
  "custom_path": "custom_model.py",
  "custom_builder": "build_model"
}

See custom_model.py for a working template.

Your custom_model.py should define:

import torch
from torch import nn

def build_model(model_cfg, input_shape, output_dim, horizon):
    # input_shape is (window, features)
    # return a torch.nn.Module
    return nn.Sequential(...)

Plot options (Lorenz-96)

• plot_options.lorenz96_lines: include line plots alongside heatmap
• plot_options.lorenz96_heatmap_mode: pair or error

Model export

Every run saves a best checkpoint to runs/<output>/model.pt. You can export from the UI (History tab) in either:

• pt (state_dict)
• torchscript

Enterprise usage notes

• Reproducibility: set seed in the config for deterministic runs.
• Separation of concerns: configs are versionable artifacts; UI edits write to JSON.
• Extensibility: custom models can be checked into repo or loaded from external paths.
• Auditability: each run writes config.json, history.json, and metrics.json.

Troubleshooting

• No module named streamlit: install in venv using python -m pip install streamlit.
• DLL load failed (PyTorch): install the VC++ 2015-2022 Redistributable (x64).
• Empty validation set: reduce window or increase train_ratio/val_ratio.

FAQ

Q: Can I run tuning for only one model? A: Yes. Select the model in the config and enable tuning.enabled.

Q: Where are results stored? A: runs/<output>/ contains the metrics, history, plots, and model checkpoint.

Q: Can I share a trained model? A: Yes. Export from the UI or copy model.pt from the run directory.

Performance tips

• Start with small t_points or fewer epochs while iterating.
• Use GPU if available; the CLI auto-detects CUDA.
• For transformers, increase window cautiously to control memory use.