Torchquad
Numerical integration in arbitrary dimensions on the GPU using PyTorch / TF / JAX
Install / Use
/learn @esa/TorchquadREADME
torchquad
<!-- *** Based on https://github.com/othneildrew/Best-README-Template -->
About The Project
The torchquad module allows utilizing GPUs for efficient numerical integration with PyTorch and other numerical Python3 modules. The software is free to use and is designed for the machine learning community and research groups focusing on topics requiring high-dimensional integration.
Built With
This project is built with the following packages:
- autoray, which means the implemented quadrature supports NumPy and can be used for machine learning with modules such as PyTorch, JAX and Tensorflow, where it is fully differentiable
- conda, which will take care of all requirements for you
If torchquad proves useful to you, please consider citing the accompanying paper.
<!-- GOALS -->Goals
- Supporting science: Multidimensional numerical integration is needed in many fields, such as physics (from particle physics to astrophysics), in applied finance, in medical statistics, and others. torchquad aims to assist research groups in such fields, as well as the general machine learning community.
- Withstanding the curse of dimensionality: The curse of dimensionality makes deterministic methods in particular, but also stochastic ones, computationally expensive when the dimensionality increases. However, many integration methods are embarrassingly parallel, which means they can strongly benefit from GPU parallelization. The curse of dimensionality still applies but the improved scaling alleviates the computational impact.
- Delivering a convenient and functional tool: torchquad is built with autoray, which means it is fully differentiable if the user chooses, for example, PyTorch as the numerical backend. Furthermore, the library of available and upcoming methods in torchquad offers high-effeciency integration for any need.
Getting Started
This is a brief guide for how to set up torchquad.
Prerequisites
We recommend using conda, especially if you want to utilize the GPU. With PyTorch it will automatically set up CUDA and the cudatoolkit for you, for example. Note that torchquad also works on the CPU; however, it is optimized for GPU usage. torchquad's GPU support is tested only on NVIDIA cards with CUDA. We are investigating future support for AMD cards through ROCm.
For a detailed list of required packages and packages for numerical backends, please refer to the conda environment files environment.yml and environment_all_backends.yml. torchquad has been tested with JAX 0.2.25, NumPy 1.19.5, PyTorch 1.10.0 and Tensorflow 2.7.0 on Linux; other versions of the backends should work as well but some may require additional setup on other platforms such as Windows.
Installation
The easiest way to install torchquad is simply to
conda install torchquad -c conda-forge
Alternatively, it is also possible to use
pip install torchquad
The PyTorch backend with CUDA support can be installed with
conda install "cudatoolkit>=11.1" "pytorch>=1.9=*cuda*" -c conda-forge -c pytorch
Note that since PyTorch is not yet on conda-forge for Windows, we have explicitly included it here using -c pytorch.
Note also that installing PyTorch with pip may not set it up with CUDA support. Therefore, we recommend to use conda.
Here are installation instructions for other numerical backends:
conda install "tensorflow>=2.6.0=cuda*" -c conda-forge
pip install "jax[cuda]>=0.4.17" --find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html # linux only
conda install "numpy>=1.19.5" -c conda-forge
More installation instructions for numerical backends can be found in environment_all_backends.yml and at the backend documentations, for example https://pytorch.org/get-started/locally/, https://github.com/google/jax/#installation and https://www.tensorflow.org/install/gpu, and often there are multiple ways to install them.
Test
After installing torchquad and PyTorch through conda or pip,
users can test torchquad's correct installation with:
import torchquad
torchquad._deployment_test()
After cloning the repository, developers can check the functionality of torchquad by running
pip install -e .
pytest
Usage
This is a brief example how torchquad can be used to compute a simple integral with PyTorch. For a more thorough introduction please refer to the tutorial section in the documentation.
The full documentation can be found on readthedocs.
# To avoid copying things to GPU memory,
# ideally allocate everything in torch on the GPU
# and avoid non-torch function calls
import torch
from torchquad import MonteCarlo, set_up_backend
# Enable GPU support if available and set the floating point precision
set_up_backend("torch", data_type="float32")
# The function we want to integrate, in this example
# f(x0,x1) = sin(x0) + e^x1 for x0=[0,1] and x1=[-1,1]
# Note that the function needs to support multiple evaluations at once (first
# dimension of x here)
# Expected result here is ~3.2698
def some_function(x):
return torch.sin(x[:, 0]) + torch.exp(x[:, 1])
# Declare an integrator;
# here we use the simple, stochastic Monte Carlo integration method
mc = MonteCarlo()
# Compute the function integral by sampling 10000 points over domain
integral_value = mc.integrate(
some_function,
dim=2,
N=10000,
integration_domain=[[0, 1], [-1, 1]],
backend="torch",
)
Logging Configuration
By default, torchquad disables its internal logging when installed from PyPI to avoid interfering with other loggers in your application. To enable logging change TORCHQUAD_DISABLE_LOGGING in __init__.py:
-
Set the log level: Use the
TORCHQUAD_LOG_LEVELenvironment variable:export TORCHQUAD_LOG_LEVEL=DEBUG # For detailed debugging export TORCHQUAD_LOG_LEVEL=INFO # For general information export TORCHQUAD_LOG_LEVEL=WARNING # For warnings only (default when enabled) -
Enable logging programmatically:
import torchquad torchquad.set_log_level("DEBUG") # This will enable and configure logging
Multi-GPU Usage
torchquad supports multi-GPU systems through standard PyTorch practices. The recommended approach is to use the CUDA_VISIBLE_DEVICES environment variable to control GPU selection:
``
