Skip to main content

PyTorch-native reservoir computing library with GPU acceleration

Project description

ResDAG

Python 3.11+ PyTorch 2.0+ License: MIT

A modern, GPU-accelerated reservoir computing library for PyTorch.

resdag brings the power of Echo State Networks (ESNs) and reservoir computing to PyTorch with a clean, modular API. Built for researchers and practitioners who need fast, flexible, and production-ready reservoir computing models.


โœจ Key Features

  • ๐Ÿš€ GPU-Accelerated: Full GPU support for training and inference
  • ๐ŸŽฏ Pure PyTorch: Native nn.Module components, TorchScript compatible
  • ๐Ÿงฉ Modular Design: Build complex architectures with simple building blocks
  • ๐Ÿ“Š Multiple Topologies: 15+ graph topologies for reservoir initialization
  • ๐Ÿ”ฌ Algebraic Training: Efficient ridge regression via Conjugate Gradient
  • ๐ŸŽจ Flexible API: Compose models using pytorch_symbolic
  • ๐Ÿ“ˆ HPO Ready: Built-in Optuna integration for hyperparameter optimization
  • ๐Ÿ”ง Production Ready: Stateful layers, model persistence, GPU compilation

๐Ÿ“ฆ Installation

From pip (recommended)

pip install resdag

# With hyperparameter optimization support
pip install resdag[hpo]

From source

git clone https://github.com/El3ssar/resdag.git
cd resdag
pip install -e .

# Or using uv (faster)
uv sync

๐Ÿš€ Quick Start

Your First ESN in 30 Seconds

import torch
import pytorch_symbolic as ps
from resdag import ESNModel, ReservoirLayer, CGReadoutLayer, ESNTrainer

# 1. Define the model architecture
inp = ps.Input((100, 3))  # (seq_len, features)
reservoir = ReservoirLayer(
    reservoir_size=500,
    feedback_size=3,
    spectral_radius=0.9,
    topology="erdos_renyi"
)(inp)
readout = CGReadoutLayer(500, 3, alpha=1e-6, name="output")(reservoir)
model = ESNModel(inp, readout)

# 2. Train the model (algebraic, not SGD!)
trainer = ESNTrainer(model)
trainer.fit(
    warmup_inputs=(warmup_data,),
    train_inputs=(train_data,),
    targets={"output": train_targets}
)

# 3. Make predictions
predictions = model.forecast(forecast_warmup, horizon=1000)

Using Premade Models

from resdag.models import ott_esn

# Ott's ESN for chaotic systems (with state augmentation)
model = ott_esn(
    reservoir_size=500,
    feedback_size=3,
    output_size=3,
    spectral_radius=0.95,
)

# Train and forecast as above

๐Ÿ“– Core Concepts

Reservoir Layers

The heart of ESNs - stateful RNN layers with randomly initialized, fixed recurrent weights:

from resdag.layers import ReservoirLayer
from resdag.init.topology import get_topology

reservoir = ReservoirLayer(
    reservoir_size=500,        # Number of neurons
    feedback_size=3,           # Dimension of feedback input
    input_size=5,              # Optional: dimension of driving inputs
    spectral_radius=0.9,       # Controls memory/stability
    leak_rate=1.0,             # Leaky integration (1.0 = no leak)
    activation="tanh",         # Activation function
    topology=get_topology("watts_strogatz", k=4, p=0.3),
)

# Forward pass
states = reservoir(feedback)                    # Feedback only
states = reservoir(feedback, driving_input)     # With driving input

Readout Layers

Linear layers trained via ridge regression (not gradient descent):

from resdag.layers.readouts import CGReadoutLayer

readout = CGReadoutLayer(
    in_features=500,           # Reservoir size
    out_features=3,            # Output dimension
    alpha=1e-6,                # Ridge regularization
    name="output",             # Name for multi-readout models
)

# Fit using conjugate gradient
readout.fit(reservoir_states, targets)
output = readout(reservoir_states)

Model Composition

Build models using pytorch_symbolic for clean, functional composition:

import pytorch_symbolic as ps
from resdag import ESNModel
from resdag.layers import ReservoirLayer, Concatenate
from resdag.layers.readouts import CGReadoutLayer

# Multi-input model with driving signal
feedback = ps.Input((100, 3))
driver = ps.Input((100, 5))

reservoir = ReservoirLayer(500, feedback_size=3, input_size=5)(feedback, driver)
readout = CGReadoutLayer(500, 3, name="output")(reservoir)

model = ESNModel([feedback, driver], readout)

Training

Efficient algebraic training via ESNTrainer:

from resdag.training import ESNTrainer

trainer = ESNTrainer(model)

# Two-phase training: warmup + fitting
trainer.fit(
    warmup_inputs=(warmup_feedback, warmup_driver),  # Synchronize states
    train_inputs=(train_feedback, train_driver),      # Fit readout
    targets={"output": targets},                      # One target per readout
)

Forecasting

Two-phase forecasting: teacher-forced warmup + autoregressive generation:

# Simple forecast (feedback only)
predictions = model.forecast(warmup_data, horizon=1000)

# Input-driven forecast (with external signals)
predictions = model.forecast(
    warmup_feedback,
    warmup_driver,
    horizon=1000,
    forecast_drivers=(future_driver,),  # Provide future driving inputs
)

# Include warmup in output
full_output = model.forecast(
    warmup_data,
    horizon=1000,
    return_warmup=True,
)

๐ŸŽฏ Advanced Usage

Graph Topologies

resdag supports 15+ graph topologies for reservoir initialization:

from resdag.init.topology import get_topology, show_topologies

# List all available topologies
show_topologies()

# Get details for a specific topology
show_topologies("erdos_renyi")

# Create topology initializer
topology = get_topology("watts_strogatz", k=6, p=0.3, seed=42)

# Use in reservoir
reservoir = ReservoirLayer(
    reservoir_size=500,
    feedback_size=3,
    topology=topology,
    spectral_radius=0.95,
)

Available topologies:

  • erdos_renyi - Random graphs with edge probability
  • watts_strogatz - Small-world networks
  • barabasi_albert - Scale-free networks
  • complete - Fully connected
  • ring_chord - Ring with chords
  • dendrocycle - Dendritic cycles
  • And many more!

Input/Feedback Initializers

Custom initialization strategies for input/feedback weights:

from resdag.init.input_feedback import get_input_feedback

# List available initializers
from resdag.init.input_feedback import show_input_initializers
show_input_initializers()

# Use custom initializer
reservoir = ReservoirLayer(
    reservoir_size=500,
    feedback_size=3,
    feedback_initializer=get_input_feedback("chebyshev"),
    input_initializer=get_input_feedback("random", input_scaling=0.5),
)

Multi-Readout Models

Train models with multiple outputs:

# Define multiple readouts
reservoir = ReservoirLayer(500, feedback_size=3)(inp)

readout1 = CGReadoutLayer(500, 3, name="position")(reservoir)
readout2 = CGReadoutLayer(500, 3, name="velocity")(reservoir)

model = ESNModel(inp, [readout1, readout2])

# Train both readouts
trainer.fit(
    warmup_inputs=(warmup,),
    train_inputs=(train,),
    targets={
        "position": position_targets,
        "velocity": velocity_targets,
    },
)

Data Utilities

Built-in utilities for data loading and preparation:

from resdag.utils.data import load_file, prepare_esn_data

# Load time series
data = load_file("lorenz.csv")  # Auto-detects format

# Split into ESN training phases
warmup, train, target, f_warmup, val = prepare_esn_data(
    data,
    warmup_steps=100,     # Reservoir synchronization
    train_steps=500,      # Readout training
    val_steps=200,        # Validation
    normalize="minmax",   # Normalization method
)

Hyperparameter Optimization

Built-in Optuna integration for HPO:

from resdag.hpo import run_hpo
from resdag.models import ott_esn

def model_creator(reservoir_size, spectral_radius):
    return ott_esn(
        reservoir_size=reservoir_size,
        feedback_size=3,
        output_size=3,
        spectral_radius=spectral_radius,
    )

def search_space(trial):
    return {
        "reservoir_size": trial.suggest_int("reservoir_size", 100, 1000, step=100),
        "spectral_radius": trial.suggest_float("spectral_radius", 0.5, 1.5),
    }

def data_loader(trial):
    return {
        "warmup": warmup,
        "train": train,
        "target": target,
        "f_warmup": f_warmup,
        "val": val,
    }

# Run optimization
study = run_hpo(
    model_creator=model_creator,
    search_space=search_space,
    data_loader=data_loader,
    n_trials=100,
    loss="efh",           # Expected Forecast Horizon (best for chaotic systems)
    n_workers=4,          # Parallel optimization
)

print(f"Best params: {study.best_params}")

Available loss functions:

  • efh - Expected Forecast Horizon (recommended for chaotic systems)
  • horizon - Contiguous valid forecast steps
  • lyap - Lyapunov-weighted loss
  • standard - Mean geometric error
  • discounted - Time-discounted RMSE

๐Ÿ“š Examples

Comprehensive examples are available in the examples/ directory:

  • 00_registry_system.py - Working with topology and initializer registries
  • 01_reservoir_with_topology.py - Using different graph topologies
  • 02_input_feedback_initializers.py - Custom weight initialization
  • 06_premade_models.py - Using premade ESN architectures
  • 07_save_load_models.py - Model persistence and checkpointing
  • 08_forecasting.py - Time series forecasting examples
  • 09_training.py - Training workflows with ESNTrainer
  • 10_hpo.py - Hyperparameter optimization examples

Run any example:

python examples/08_forecasting.py

๐Ÿ”ฌ Use Cases

Chaotic System Prediction

from resdag.models import ott_esn
from resdag.training import ESNTrainer

# Lorenz attractor prediction
model = ott_esn(reservoir_size=500, feedback_size=3, output_size=3)

trainer = ESNTrainer(model)
trainer.fit(warmup_inputs=(warmup,), train_inputs=(train,), targets={"output": target})

# Long-term prediction
predictions = model.forecast(forecast_warmup, horizon=5000)

Input-Driven Systems

# System with external forcing
feedback = ps.Input((100, 3))
driver = ps.Input((100, 2))

reservoir = ReservoirLayer(500, feedback_size=3, input_size=2)(feedback, driver)
readout = CGReadoutLayer(500, 3, name="output")(reservoir)
model = ESNModel([feedback, driver], readout)

# Forecast with future driving signals
predictions = model.forecast(
    warmup_feedback,
    warmup_driver,
    horizon=1000,
    forecast_drivers=(future_driver,),
)

Multi-Scale Predictions

# Predict at multiple timescales
reservoir = ReservoirLayer(1000, feedback_size=3)(inp)

short_term = CGReadoutLayer(1000, 3, name="1step")(reservoir)
medium_term = CGReadoutLayer(1000, 3, name="10step")(reservoir)
long_term = CGReadoutLayer(1000, 3, name="100step")(reservoir)

model = ESNModel(inp, [short_term, medium_term, long_term])

๐ŸŽ“ Documentation

Full documentation is available in the docs/ directory:

API Reference

Generate API documentation using Sphinx:

cd docs/
sphinx-apidoc -o api/ ../src/resdag
make html

๐Ÿงช Testing

Run the test suite:

# Run all tests
pytest

# Run with coverage
pytest --cov=resdag --cov-report=html

# Run specific test module
pytest tests/test_layers/test_reservoir.py

Current test coverage: 57% (240 tests passing)


๐Ÿ› ๏ธ Development

Setting up Development Environment

# Clone repository
git clone https://github.com/El3ssar/resdag.git
cd resdag

# Install with development dependencies
uv sync --dev

# Or with pip
pip install -e ".[dev]"

Code Quality

We use ruff for linting and black for formatting:

# Format code
black src/ tests/

# Lint code
ruff check src/ tests/

# Type checking
mypy src/

Project Structure

resdag/
โ”œโ”€โ”€ src/resdag/
โ”‚   โ”œโ”€โ”€ composition/       # Model composition (pytorch_symbolic)
โ”‚   โ”œโ”€โ”€ layers/            # Reservoir, Readout, custom layers
โ”‚   โ”œโ”€โ”€ init/              # Weight initialization
โ”‚   โ”‚   โ”œโ”€โ”€ topology/      # Graph topologies
โ”‚   โ”‚   โ”œโ”€โ”€ input_feedback/ # Input/feedback initializers
โ”‚   โ”‚   โ””โ”€โ”€ graphs/        # Graph generation functions
โ”‚   โ”œโ”€โ”€ training/          # ESNTrainer and training utilities
โ”‚   โ”œโ”€โ”€ models/            # Premade model architectures
โ”‚   โ”œโ”€โ”€ hpo/               # Hyperparameter optimization
โ”‚   โ””โ”€โ”€ utils/             # Data loading and utilities
โ”œโ”€โ”€ tests/                 # Comprehensive test suite
โ”œโ”€โ”€ examples/              # Example scripts
โ””โ”€โ”€ docs/                  # Documentation

๐Ÿค Contributing

Contributions are welcome! Please:

  1. Fork the repository
  2. Create a feature branch (git checkout -b feature/amazing-feature)
  3. Make your changes with tests
  4. Ensure tests pass (pytest)
  5. Format code (black src/ tests/)
  6. Submit a pull request

See CONTRIBUTING.md for detailed guidelines.


๐Ÿ“„ Citation

If you use resdag in your research, please cite:

@software{resdag2026,
  author = {Daniel Estevez-Moya},
  title = {resdag: A PyTorch Library for Reservoir Computing},
  year = {2026},
  url = {https://github.com/El3ssar/resdag}
}

๐Ÿ“œ License

This project is licensed under the MIT License - see the LICENSE file for details.


๐Ÿ™ Acknowledgments


๐Ÿ“ฌ Contact


๐Ÿ—บ๏ธ Roadmap

  • Additional premade architectures (Liquid State Machines, Next-Gen RC)
  • Online learning capabilities
  • TorchScript export for production
  • ONNX support
  • Distributed training for large reservoirs
  • Interactive visualization tools
  • Benchmarking suite against other ESN libraries

โญ Star us on GitHub if you find this useful!

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

resdag-0.2.0.tar.gz (188.2 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

resdag-0.2.0-py3-none-any.whl (118.4 kB view details)

Uploaded Python 3

File details

Details for the file resdag-0.2.0.tar.gz.

File metadata

  • Download URL: resdag-0.2.0.tar.gz
  • Upload date:
  • Size: 188.2 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/6.1.0 CPython/3.13.7

File hashes

Hashes for resdag-0.2.0.tar.gz
Algorithm Hash digest
SHA256 92699cec5208f45f7187bedb31eaf0ea7415ac0343380f06df67d7d1acbd7a45
MD5 ee7c0a9279096f8eda98f493442c9503
BLAKE2b-256 a2964f1310f402e10036c0b05636731b47328b454718f4c23ef550285e99cec3

See more details on using hashes here.

Provenance

The following attestation bundles were made for resdag-0.2.0.tar.gz:

Publisher: release.yml on El3ssar/ResDAG

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file resdag-0.2.0-py3-none-any.whl.

File metadata

  • Download URL: resdag-0.2.0-py3-none-any.whl
  • Upload date:
  • Size: 118.4 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/6.1.0 CPython/3.13.7

File hashes

Hashes for resdag-0.2.0-py3-none-any.whl
Algorithm Hash digest
SHA256 d075ba81a07d2f04f235c89656d635182e7d46393acc92f1ae5535521f4a7314
MD5 e9e444328a3e914ed991b34ca9535ed5
BLAKE2b-256 2ed15796f501f1dc25f2e8be0c5ca69e190112eb48e503e50a82aac5bdc1f7d7

See more details on using hashes here.

Provenance

The following attestation bundles were made for resdag-0.2.0-py3-none-any.whl:

Publisher: release.yml on El3ssar/ResDAG

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page