Neural spatiotemporal forecasting with PyTorch
๐ Getting Started - ๐ Documentation - ๐ป Introductory notebook
๐ Getting Started - ๐ Documentation - ๐ป Introductory notebook
tsl (Torch Spatiotemporal) is a library built to accelerate research on neural spatiotemporal data processing
methods, with a focus on Graph Neural Networks.
Built upon popular libraries such as 
PyTorch, 
PyG (PyTorch Geometric), and 
PyTorch Lightning, tsl provides a unified and user-friendly framework for efficient neural spatiotemporal data processing, that goes from data preprocessing to model prototyping.
tsl's high-level APIs empower you to develop tailored solutions.
* **Access a Wealth of Existing Datasets and Models** Leverage a vast collection of datasets and models from the spatiotemporal data processing literature. Explore and benchmark against state-of-the-art baselines, and test your brand new model on widely used public datasets.
* **Handle Irregularities and Missing Data** Seamlessly manage irregularities in your spatiotemporal data streams, including missing data and variations in network structures. Ensure the robustness and reliability of your data processing pipelines.
* **Streamlined Preprocessing** Automate the preprocessing phase with tsl's methods for scaling, resampling and clustering time series. Spend less time on data preparation and focus on extracting meaningful patterns and insights.
* **Efficient Data Structures** Utilize tsl's straightforward data structures, seamlessly integrated with PyTorch and PyG, to accelerate your workflows. Benefit from the flexibility and compatibility of these widely adopted libraries.
* **Scalability with PyTorch Lightning** Scale your computations effortlessly, from a single CPU to clusters of GPUs, with tsl's integration with PyTorch Lightning. Accelerate training and inference across various hardware configurations.
* **Modular Neural Layers** Build powerful and modular neural spatiotemporal models using tsl's collection of specialized layers. Create architectures with ease, leveraging the flexibility and extensibility of the library.
* **Reproducible Experiments** Ensure experiment reproducibility using the 
Hydra framework, a standard in the field. Validate and compare results confidently, promoting rigorous research in spatiotemporal data mining.
## Getting Started
Before you start using tsl, please review the documentation to get an understanding of the library and its capabilities.
You can also explore the examples provided in the `examples` directory to see how train deep learning models working with spatiotemporal data.
## Installation
tsl supports NumPy >= 1.26 and < 3, PyTorch >= 2.2 and < 2.13, and PyG >= 2.4 and < 3. `torch-sparse` is optional for `torch_sparse.SparseTensor` adjacency matrices and requires the matching `torch-scatter` runtime package.
The recommended setup for local development is [uv](https://docs.astral.sh/uv/):
```bash
uv venv --python 3.10
uv pip install torch-spatiotemporal --torch-backend=auto
# Optional SparseTensor stack; select matching wheels from the PyG guide:
# https://pytorch-geometric.readthedocs.io/en/latest/install/installation.html
# uv pip install torch-scatter torch-sparse -f https://data.pyg.org/whl/torch-${TORCH}+${CUDA}.html
```
On Windows PowerShell:
```powershell
uv venv --python 3.10
uv pip install torch-spatiotemporal --torch-backend=auto
```
`uv` resolves PyTorch and PyG alongside tsl, so no separate framework
installation is required. `--torch-backend=auto` lets uv choose the most
compatible PyTorch backend: CUDA on NVIDIA systems with a compatible driver,
and CPU wheels otherwise. macOS uses PyTorch's macOS wheels; GPU acceleration
there is handled by PyTorch/MPS rather than CUDA.
For a regular user install from PyPI:
```bash
uv pip install torch-spatiotemporal --torch-backend=auto
```
For development tools and tests, install the dev tools into the same environment:
```bash
uv pip install -e ".[dev]" --torch-backend=auto
uv run pytest
```
The editable-install command above is also the recommended setup for a local
project environment.
The default test command uses the coverage settings in `pyproject.toml`. To skip slow tests explicitly:
```bash
uv run pytest -m "not slow"
```
### Compiling models
Models marked with `ModelClass.can_be_compiled = True` are covered by
full-graph compilation tests. Compile an instance in place using PyTorch's
standard module API:
```python
model.compile(mode="reduce-overhead")
```
The forecasting and imputation examples expose the same feature through Hydra:
```bash
python examples/forecasting/run_traffic_experiment.py \
model=dcrnn dataset=la compile.enabled=true compile.mode=reduce-overhead
```
Compilation has an initial warm-up cost and is most useful for repeated
training or inference with stable tensor shapes. See the migration audit for
the tested model matrix and current limitations.
### PyG compiled extension
`torch-sparse` is optional and only needed for `torch_sparse.SparseTensor` connectivity. It requires `torch-scatter` at runtime, so install both matching wheels from the [PyG installation guide](https://pytorch-geometric.readthedocs.io/en/latest/install/installation.html). Their wheels must match the exact PyTorch and CUDA build.
### Conda
The repository still includes `conda_env.yml` for users who prefer conda:
```bash
conda env create -f conda_env.yml
conda activate tsl
```
The conda environment remains CUDA-oriented by default. Remove the `nvidia` channel and `pytorch-cuda` dependency from `conda_env.yml` for a CPU-only conda environment.
## Tutorial
The best way to start using tsl is by following the tutorial notebook in `examples/notebooks/a_gentle_introduction_to_tsl.ipynb`.
## Documentation
Visit the [documentation](https://torch-spatiotemporal.readthedocs.io/en/latest/) to learn more about the library, including detailed API references, examples, and tutorials.
The documentation is hosted on [readthedocs](https://torch-spatiotemporal.readthedocs.io/en/latest/). For local access, you can build it from the `docs` directory.
## Contributing
Contributions are welcome! For major changes or new features, please open an issue first to discuss your ideas. See the [Contributing guidelines](https://github.com/TorchSpatiotemporal/tsl/blob/dev/.github/CONTRIBUTING.md) for more details on how to get involved. Help us build a better tsl!
Thanks to all contributors! ๐งก