Neural spatiotemporal forecasting with PyTorch

---

🚀 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.

• Create Custom Models and Datasets   Easily build your own custom models and datasets for spatiotemporal data analysis. Whether you're working with sensor networks, environmental data, or any other spatiotemporal domain, 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.

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.

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:

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:

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:

uv pip install torch-spatiotemporal --torch-backend=auto

For development tools and tests, install the dev tools into the same environment:

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:

uv run pytest -m "not slow"

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:

model.compile(mode="reduce-overhead")

The forecasting and imputation examples expose the same feature through Hydra:

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.

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. Their wheels must match the exact PyTorch and CUDA build.

The repository still includes conda_env.yml for users who prefer conda:

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.

The best way to start using tsl is by following the tutorial notebook in examples/notebooks/a_gentle_introduction_to_tsl.ipynb.

Visit the documentation to learn more about the library, including detailed API references, examples, and tutorials.

The documentation is hosted on readthedocs. For local access, you can build it from the docs directory.

Contributions are welcome! For major changes or new features, please open an issue first to discuss your ideas. See the Contributing guidelines for more details on how to get involved. Help us build a better tsl!

Thanks to all contributors! 🧡

If you use Torch Spatiotemporal for your research, please consider citing the library

@software{Cini_Torch_Spatiotemporal_2022,
    author = {Cini, Andrea and Marisca, Ivan},
    license = {MIT},
    month = {3},
    title = {{Torch Spatiotemporal}},
    url = {https://github.com/TorchSpatiotemporal/tsl},
    year = {2022}
}

By Andrea Cini and Ivan Marisca.

This project is licensed under the terms of the MIT license. See the LICENSE file for details.