Python install

PyKeOps is a Python 3 wrapper around the low-level KeOpsCore library which is written in C++/CUDA. It provides functions that can be used in any NumPy or PyTorch script.

Requirements

Python 3 with the numpy package.
A C++ compiler compatible with std=c++11: g++ version >=7 or clang++ version >=8.
The Cuda toolkit: version >=10.0 is recommended.
PyTorch (optional): version >= 1.5.

Using pip (recommended)

Just in case: in a terminal, check the consistency of the outputs of the commands which python, python --version, which pip and pip --version.
In a terminal, type:

pip install pykeops
Note that compiled shared objects (.so files on Unix, .dylib on macOS) will be stored in the folder ~/.cache/keops/, where ~ is the path to your home folder.

Test your installation, as described in the next section.

On Google Colab

Google provides free virtual machines where KeOps runs out-of-the-box. In a new Colab notebook, typing:

!pip install pykeops > install.log

should allow you to get a working version of KeOps in less than twenty seconds.

Using Docker or Singularity

We provide a reference Dockerfile and publish full containers on our DockerHub channel using the docker-images.sh script. These environments contain a full installation of CUDA, NumPy, PyTorch, R, KeOps and GeomLoss. Their PYTHONPATH are configured to ensure that git installations of KeOps or GeomLoss mounted in /opt/keops or /opt/geomloss take precedence over the pre-installed Pip versions.

As an example, here are the steps that we follow to render this website on the Jean Zay scientific cluster:

# First, clone the latest release of the KeOps repository in ~/code/keops:
mkdir ~/code
cd ~/code
git clone git@github.com:getkeops/keops.git

# Load singularity in our environment:
module load singularity

# Create a folder to store our Singularity files:
mkdir -p ~/scratch/containers
cd ~/scratch/containers

# Download the Docker image and store it as an immutable Singularity Image File:
# N.B.: Our image is pretty heavy (~7 Gb), so it is safer to create
#       cache folders on the hard drive instead of relying on the RAM-only tmpfs:
# N.B.: This step may take 15mn to 60mn, so you may prefer to execute it on your
#       local computer and then copy the resulting file `keops-full.sif` to the cluster.
#       Alternatively, on the Jean Zay cluster, you may use the `prepost` partition
#       to have access to both a large RAM and an internet connection.
mkdir cache
mkdir tmp
mkdir tmp2
SINGULARITY_TMPDIR=`pwd`/tmp SINGULARITY_CACHEDIR=`pwd`/cache \
singularity build --tmpdir `pwd`/tmp2 keops-full.sif docker://getkeops/keops-full:latest

# At this point, on the Jean Zay cluster, you should use a command like:
# idrcontmgr cp keops-full.sif
# to add our new environment to the cluster's container registry as explained here:
# http://www.idris.fr/jean-zay/cpu/jean-zay-utilisation-singularity.html

# Then, create a separate home folder for this image. This is to ensure
# that we won't see any conflict between different versions of the KeOps binaries,
# stored in the ~/.cache folder of the virtual machine:
mkdir -p ~/containers/singularity_homes/keops-full

# Ask the slurm scheduler to render our documentation.
sbatch keops-doc.batch

Where keops-doc.batch is an executable file that contains:

#!/bin/bash

#SBATCH -A dvd@a100  # Use a A100 GPU - dvd@v100 is also available
#SBATCH -C a100
#SBATCH --partition=gpu_p5
#SBATCH --job-name=keops_doc    # create a short name for your job
#SBATCH --mail-type=ALL         # Mail events (NONE, BEGIN, END, FAIL, ALL)
#SBATCH --mail-user=your.name@inria.fr   # Where to send mail
#SBATCH --nodes=1               # node count
#SBATCH --ntasks=1              # total number of tasks across all nodes
#SBATCH --cpus-per-task=8       # cpu-cores per task (>1 if multi-threaded tasks)
#SBATCH --gres=gpu:1     # GPU nodes are only available in gpu partition
#SBATCH --time=03:00:00          # total run time limit (HH:MM:SS)
#SBATCH --output=logs/keops_doc.out   # output file name
#SBATCH --error=logs/keops_doc.err    # error file name

echo "### Running $SLURM_JOB_NAME ###"

set -x
cd ${SLURM_SUBMIT_DIR}

module purge
module load singularity

# The Jean Zay compute nodes don't have access to the internet,
# which means that they cannot fetch data as required by e.g. the MNIST tutorial.
# A workaround is to run:
# from sklearn.datasets import fetch_openml
# fetch_openml("mnist_784", cache=True, as_frame=False)
# on the front-end node or on your laptop, copy
# ~/scikit_learn_data to $WORK/data/scikit_learn_data
# and then rely on the --bind option as detailed below:

singularity exec \
-H $WORK/containers/singularity_homes/keops-full/:/home \
--bind ~/keops-doc.sh:/home/keops-doc.sh \
--bind $WORK/code:/home/code \
--bind $WORK/code/keops:/opt/keops \
--bind $WORK/data/scikit_learn_data:/home/scikit_learn_data \
--nv \
$SINGULARITY_ALLOWED_DIR/keops-full.sif \
/home/keops-doc.sh

And keops-doc.sh is an executable file that contains:

#!/bin/bash

echo "Rendering the KeOps documentation"

# Clean the cache folder of binaries:
python -c "import pykeops; pykeops.clean_pykeops()"

# First of all, make sure that all unit tests pass:
cd /home/code/keops
pytest -v

# Then, render the doc properly:
cd doc
# Remove the previous built pages:
make clean
# Render the website:
make html

# Re-render the doc to remove compilation messages:
make clean
make html

zip -r keops_doc.zip _build/

From source using git

The simplest way of installing a specific version of KeOps is to use some advanced pip syntax:

pip install git+https://github.com/getkeops/keops.git@main#subdirectory=keopscore
pip install git+https://github.com/getkeops/keops.git@main#subdirectory=pykeops

Alternatively, you may:

Clone the KeOps repo at a location of your choice (denoted here as /path/to):

git clone --recursive https://github.com/getkeops/keops.git /path/to/libkeops
Note that compiled .so routines will be stored in the folder /path/to/libkeops/pykeops/build: this directory must have write permission.

Install via pip in editable mode as follows :

pip install -e /path/to/libkeops/keopscore -e /path/to/libkeops/pykeops

Otherwise you may add the directories /path/to/libkeops/keopscore and /path/to/libkeops/pykeops to your python path. This can be done once and for all, by adding the path to to your ~/.bashrc. In a terminal, type:
echo "export PYTHONPATH=$PYTHONPATH:/path/to/libkeops/keopscore:/path/to/libkeops/pykeops" >> ~/.bashrc
Alternatively, you may add the following line to the beginning of your python scripts:
import os.path
import sys
sys.path.append('/path/to/libkeops/keopscore')
      sys.path.append('/path/to/libkeops/pykeops')

Test your installation, as described in the next section.

Testing your installation

You can use the following test functions to compile and run simple KeOps formulas. If the compilation fails, it returns the full log.

In a python terminal,

import pykeops
pykeops.test_numpy_bindings()    # perform the compilation

should return:

pyKeOps with numpy bindings is working!

If you use PyTorch, the following code:

import pykeops
pykeops.test_torch_bindings()    # perform the compilation

should return:

pyKeOps with torch bindings is working!

Please note that running pytest -v in a copy of our git repository will also let you perform an in-depth test of the entire KeOps codebase.

Troubleshooting

Compilation issues

First of all, make sure that you are using a C++ compiler which is compatible with the C++11 revision. Otherwise, compilation of formulas may fail in unexpected ways. Depending on your system, you can:

Install a compiler system-wide: for instance, on Debian-based Linux distributions, you can install g++ with apt and then use update-alternatives to choose a suitable compiler as default. Don’t forget to pick compatible versions for both gcc and g++.
Install a compiler locally: if you are using a conda environment, you can install a new instance of gcc and g++ by following the documentation of conda.
If you have a conda environment with CUDA toolkit and pyKeOps, the compiling test with pykeops.test_numpy_bindings()``will fail unless you also have a system-wide CUDA toolkit installation, due to missing ``cuda.h file. See <https://conda-forge.org/docs/user/faq.html?highlight=cuda>`_, question “How can I compile CUDA (host or device) codes in my environment?”

Cache directory

If you experience problems with compilation, it may be a good idea to flush the build folder that KeOps uses as a cache for already-compiled formulas. To do this, just type:

import pykeops
pykeops.clean_pykeops()

You can change the build folder by using the set_build_folder() function:

import pykeops
print(pykeops.get_build_folder())  # display current build_folder
pykeops.set_build_folder("/my/new/location")  # change the build folder
print(pykeops.get_build_folder())  # display new build_folder

Note that the command set_build_folder() without any argument will reset the location to the default one (~/.keops/build on unix-like systems)

Verbosity level

You can deactivate all messages and warnings by setting the environment variable PYKEOPS_VERBOSE to 0. In a terminal, type:

export PYKEOPS_VERBOSE=0
python my_script_calling_pykeops.py

Alternatively, you can disable verbose compilation from your python script using the function pykeops.set_verbose. In a python shell, type:

import pykeops
pykeops.set_verbose(False)