Repository: bacpop/pp-sketchlib (language: python)

Issue:
Change interface to python command line
Use docopt for the python CLI
Allow single kmer length comparisons

TODO:
- [x] Check tests run
- [x] Remove printing of options dict
- [x] Remove hard coded sys.path
- [x] Update README
- [x] Check single k-mer length
- ~[ ] Update sketching page of poppunk docs~ now https://github.com/johnlees/PopPUNK/issues/191

Closes #58
Add sparse query which sparisifies on the fly
This will have much lower & more predictable memory usage.

Merge after #66 

TODO:
- [x] Add to CLI for sparse
- [x] Support CUDA (use cub sort)
- [x] Check CPU progress % still works
- [x] Check dense GPU mode still works
- [x] Add to tests

Allow sketching & Jaccard distance at single k-mer length
Currently at least two required
Add sparse query which sparisifies on the fly
This will have much lower & more predictable memory usage.

Merge after #66 

TODO:
- [x] Add to CLI for sparse
- [x] Support CUDA (use cub sort)
- [x] Check CPU progress % still works
- [x] Check dense GPU mode still works
- [x] Add to tests



Discussion / hints:


Combine this with a move to https://github.com/docopt/docopt to make the command line for sketchlib a bit more like familiar tools. Probably turn off random matches by default


Fix patch:
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 6dac2080..c231c7ab 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -20,27 +20,33 @@ set(TARGET_NAME pp_sketchlib)
 add_compile_definitions(PYTHON_EXT)
 
 # gcc: Add openmp
-# gcc: Add -O0 to remove optimizations when using debug
-IF(CMAKE_COMPILER_IS_GNUCC)
+IF(CMAKE_COMPILER_IS_GNUCC OR "$ENV{SKETCHLIB_INSTALL}" STREQUAL "conda")
     set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fopenmp")
-    set(CMAKE_CXX_FLAGS_DEBUG "${CMAKE_CXX_FLAGS_DEBUG} -O0")
-    set(CMAKE_C_FLAGS_DEBUG "${CMAKE_C_FLAGS_DEBUG} -O0")
-ENDIF(CMAKE_COMPILER_IS_GNUCC)
+ENDIF()
+
+# Add -O0 to remove optimizations when using debug
+set(CMAKE_CXX_FLAGS_DEBUG "${CMAKE_CXX_FLAGS_DEBUG} -O0")
+set(CMAKE_C_FLAGS_DEBUG "${CMAKE_C_FLAGS_DEBUG} -O0")
 
 if(UNIX AND NOT APPLE)
     if(CMAKE_CXX_COMPILER STREQUAL "icpc")
-        set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fast -xCASCADELAKE -DMKL_ILP64 -m64 -static-intel")
+        set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fast -march=\"native\" -DMKL_ILP64 -m64 -static-intel")
     else()
         set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -D__STDC_LIMIT_MACROS -D__STDC_CONSTANT_MACROS")
         set(CMAKE_LD_FLAGS "${CMAKE_LDFLAGS} -Wl,--as-needed")
     endif()
 endif()
 
+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall")
+set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} -ffast-math -funroll-loops -m64")
+
 # Set paths for non standard lib/ and include/ locations
 if(DEFINED ENV{CONDA_PREFIX})
     include_directories($ENV{CONDA_PREFIX}/include)
     link_directories($ENV{CONDA_PREFIX}/lib)
     link_directories($ENV{CONDA_PREFIX}/lib/intel64)
+else()
+    find_package(OpenMP)
 endif()
 
 # Add libraries
@@ -50,12 +56,12 @@ include_directories(${HDF5_INCLUDE_DIRS})
 include_directories(${CMAKE_CURRENT_SOURCE_DIR}/vendor/highfive/include)
 include_directories(${CMAKE_CURRENT_SOURCE_DIR}/src)
 
+find_package(ZLIB)
 execute_process(COMMAND pybind11-config --cmakedir OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE pybind11_DIR)
 find_package(pybind11 2.6 CONFIG REQUIRED)
 find_package(Eigen3 3.3 REQUIRED NO_MODULE)
 find_package(Armadillo REQUIRED)
 include_directories(${ARMADILLO_INCLUDE_DIRS})
-#find_package(OpenMP) # This links system openmp if present - conda sorts out rpath but take care
 
 # Define python library target
 add_library("${TARGET_NAME}" MODULE)
@@ -64,42 +70,46 @@ add_library("${TARGET_NAME}" MODULE)
 include(CheckLanguage)
 check_language(CUDA)
 if(CMAKE_CUDA_COMPILER)
-    message(STATUS "CUDA found, compiling both GPU and CPU code")
     enable_language(CUDA)
-
-    # PIC/relocatable-device-code needed as this is linked by gcc later
-    # -Xptxas -dlcm=ca turns cache on, but not needed in recent nvcc versions
-    # --cudart static: static linking of the CUDA libraries
-    # -gencode arch=compute_35 etc compiles for each (minimum) device version listed (v3.5, v5.0, v7.5)
-    set(CUDA_OPTS "-Xcompiler -fPIC -Xptxas -dlcm=ca --relocatable-device-code=true --expt-relaxed-constexpr")
-    # Turn on link time optimisation if available
-    if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER 11.0 AND CMAKE_BUILD_TYPE MATCHES Release)
-        string(APPEND CUDA_OPTS " -dlto -arch=sm_86")
+    if (CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER 11.0)
+        message(STATUS "CUDA >11.0 found, compiling both GPU and CPU code")
+
+        # PIC/relocatable-device-code needed as this is linked by gcc later
+        # -Xptxas -dlcm=ca turns cache on, but not needed in recent nvcc versions
+        # --cudart static: static linking of the CUDA libraries
+        # -gencode arch=compute_35 etc compiles for each (minimum) device version listed (v3.5, v5.0, v7.5)
+        set(CUDA_OPTS "-Xcompiler -fPIC -Xptxas -dlcm=ca --relocatable-device-code=true --expt-relaxed-constexpr")
+        # Turn on link time optimisation if available
+        if(CMAKE_BUILD_TYPE MATCHES Release)
+            string(APPEND CUDA_OPTS " -dlto -arch=sm_86")
+        else()
+            string(APPEND CUDA_OPTS " -gencode arch=compute_70,code=sm_70 -gencode arch=compute_75,code=sm_75 -gencode arch=compute_80,code=sm_80 -gencode arch=compute_86,code=sm_86")
+        endif()
+        if(CMAKE_BUILD_TYPE MATCHES Debug)
+            string(APPEND CUDA_OPTS " -G")
+        endif()
+
+        set(CMAKE_CUDA_FLAGS "${CUDA_OPTS}")
+
+        add_compile_definitions(GPU_AVAILABLE)
+        add_library("${TARGET_NAME}_CUDA" OBJECT src/gpu/dist.cu
+                                                src/gpu/sketch.cu
+                                                src/gpu/device_memory.cu
+                                                src/gpu/gpu_countmin.cu
+                                                src/gpu/device_reads.cu)
+        target_include_directories("${TARGET_NAME}_CUDA" PRIVATE "${EIGEN3_INCLUDE_DIR}" "${pybind11_INCLUDE_DIRS}")
+        set_property(TARGET "${TARGET_NAME}_CUDA"
+                    PROPERTY POSITION_INDEPENDENT_CODE ON
+                    CUDA_SEPARABLE_COMPILATION ON
+                    CUDA_RESOLVE_DEVICE_SYMBOLS ON   # try and ensure device link with nvcc
+                    CUDA_VISIBILITY_PRESET "hidden"
+                CUDA_RUNTIME_LIBRARY Static)
+              #CUDA_ARCHITECTURES OFF) # set off as done explicitly above (due to dlto complexities)
+        # CPU code/gcc compiled code needed by cuda lib
+        target_sources("${TARGET_NAME}" PRIVATE src/gpu/gpu_api.cpp)
     else()
-        string(APPEND CUDA_OPTS " -gencode arch=compute_70,code=sm_70 -gencode arch=compute_75,code=sm_75 -gencode arch=compute_80,code=sm_80 -gencode arch=compute_86,code=sm_86")
-    endif()
-    if(CMAKE_BUILD_TYPE MATCHES Debug)
-        string(APPEND CUDA_OPTS " -G")
+        message(STATUS "CUDA >=11.0 required, compiling CPU code only")
     endif()
-
-    set(CMAKE_CUDA_FLAGS "${CUDA_OPTS}")
-
-    add_compile_definitions(GPU_AVAILABLE)
-    add_library("${TARGET_NAME}_CUDA" OBJECT src/gpu/dist.cu
-                                             src/gpu/sketch.cu
-                                             src/gpu/device_memory.cu
-                                             src/gpu/gpu_countmin.cu
-                                             src/gpu/device_reads.cu)
-    target_include_directories("${TARGET_NAME}_CUDA" PRIVATE "${EIGEN3_INCLUDE_DIR}" "${pybind11_INCLUDE_DIRS}")
-    set_property(TARGET "${TARGET_NAME}_CUDA"
-                 PROPERTY POSITION_INDEPENDENT_CODE ON
-                 CUDA_SEPARABLE_COMPILATION ON
-                 CUDA_RESOLVE_DEVICE_SYMBOLS ON   # try and ensure device link with nvcc
-                 CUDA_VISIBILITY_PRESET "hidden"
-		         CUDA_RUNTIME_LIBRARY Static)
-	         #CUDA_ARCHITECTURES OFF) # set off as done explicitly above (due to dlto complexities)
-    # CPU code/gcc compiled code needed by cuda lib
-    target_sources("${TARGET_NAME}" PRIVATE src/gpu/gpu_api.cpp)
 else()
     message(STATUS "CUDA not found, compiling CPU code only")
 endif()
@@ -133,7 +143,12 @@ if(CMAKE_CUDA_COMPILER)
     set_property(TARGET "${TARGET_NAME}" PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS ON)
     #set_property(TARGET "${TARGET_NAME}" PROPERTY CUDA_ARCHITECTURES OFF)
 endif()
-target_link_libraries("${TARGET_NAME}" PRIVATE pybind11::module Eigen3::Eigen z ${HDF5_LIBRARIES} gomp openblas lapack gfortran m dl)
-#if(OpenMP_CXX_FOUND)
-#    target_link_libraries("${TARGET_NAME}" PRIVATE OpenMP::OpenMP_CXX)
-#endif()
+target_link_libraries("${TARGET_NAME}" PRIVATE pybind11::module Eigen3::Eigen ${HDF5_LIBRARIES} openblas lapack gfortran m dl)
+if(DEFINED ENV{CONDA_PREFIX} AND (NOT APPLE OR CMAKE_COMPILER_IS_GNUCC OR ENV{SKETCHLIB_INSTALL} EQUAL "conda"))
+    target_link_libraries("${TARGET_NAME}" PRIVATE gomp z)
+else()
+    target_link_libraries("${TARGET_NAME}" PRIVATE ZLIB::ZLIB)
+    if(OpenMP_CXX_FOUND)
+        target_link_libraries("${TARGET_NAME}" PRIVATE OpenMP::OpenMP_CXX)
+    endif()
+endif()
diff --git a/README.md b/README.md
index eabf3c83..0feb62ab 100644
--- a/README.md
+++ b/README.md
@@ -1,12 +1,13 @@
 # pp-sketchlib <img src='sketchlib_logo.png' align="right" height="139" />
 
 <!-- badges: start -->
-[![Build status](https://dev.azure.com/jlees/pp-sketchlib/_apis/build/status/johnlees.pp-sketchlib?branchName=master)](https://dev.azure.com/jlees/pp-sketchlib/_build/latest?definitionId=1&branchName=master)
+[![Build Status](https://dev.azure.com/jlees/pp-sketchlib/_apis/build/status/bacpop.pp-sketchlib?branchName=master)](https://dev.azure.com/jlees/pp-sketchlib/_build/latest?definitionId=4&branchName=master)
+[![Build status](https://badge.buildkite.com/b1bc9ccd16211ca5a55846b95e297554e5aa3b544d8cb752b0.svg?branch=master;theme=github)](https://buildkite.com/mrc-ide/pp-sketchlib)
 [![Anaconda package](https://anaconda.org/conda-forge/pp-sketchlib/badges/version.svg)](https://anaconda.org/conda-forge/pp-sketchlib)
 <!-- badges: end -->
 
 
-Library of sketching functions used by [PopPUNK](https://www.poppunk.net>).
+Library of sketching functions used by [PopPUNK](https://www.poppunk.net>). See documentation at http://poppunk.readthedocs.io/en/latest/sketching.html
 
 ## Installation
 
@@ -67,7 +68,7 @@ installed (tested on 10.2 and 11.0).
 Create a set of sketches and save these as a database:
 
 ```
-poppunk_sketch --sketch --rfile rfiles.txt --ref-db listeria --sketch-size 10000 --cpus 4 --min-k 15 --k-step 2
+sketchlib sketch -l rfiles.txt -o listeria --cpus 4
 ```
 
 The input file `rfiles.txt` has one sequence per line. The first column is the sample name, subsequent tab-separated
@@ -79,32 +80,29 @@ sample2    sample2.fa
 sample3    sample3_1.fq.gz     sample3_2.fq.gz
 ```
 
-Calculate core and accessory distances between databases with `--query`. If all-vs-all, only the upper triangle is calculated,
+Calculate core and accessory distances between databases with `query dist`. If all-vs-all, only the upper triangle is calculated,
 for example:
 
 ```
-poppunk_sketch --query --ref_db listeria --query_db listeria --cpus 4
+poppunk_sketch query dist listeria --cpus 4
 ```
 
-This will save output files as a database for use with PopPUNK. If you wish to output the
-distances add the `--print` option:
-
-```
-poppunk_sketch --query --ref_db listeria --query_db listeria --cpus 4 --print > distances.txt
-```
+This will print the distances to STDOUT and can be captured with `>`. If you wish to output save output files as a database for use with PopPUNK.add the `-o` option.
 
 ### Other options
 
 Sketching:
 
-- `--strand` ignores reverse complement k-mers, if input is all in the same sense
+- `--single-strand` ignores reverse complement k-mers, if input is all in the same sense
 - `--min-count` minimum k-mer count to include when using reads
 - `--exact-counter` uses a hash table to count k-mers, which is recommended for non-bacterial datasets.
 
 Query:
 
 - To only use some of the samples in the sketch database, you can add the `--subset` option with a file which lists the required sample names.
-- `--jaccard` will output the Jaccard distances, rather than core and accessory distances.
+- `query jaccard` will output the Jaccard distances, rather than core and accessory distances.
+- `query sparse` will output a sparse distance matrix,
+using either a `--threshold` or the k-nearest (`-kNN`).
 
 ### Large datasets
 
@@ -112,19 +110,17 @@ When working with large datasets, you can increase the `--cpus` to high numbers
 a roughly proportional performance increase.
 
 For calculating sketches of read datasets, or large numbers of distances, and you have a CUDA compatible GPU,
-you can calculate distances on your graphics device even more quickly. Add the `--use-gpu` option:
+you can calculate distances on your graphics device even more quickly. Add the `--gpu` option with the desired
+device ID:
 
 ```
-poppunk_sketch --sketch --rfile rfiles.txt --ref-db listeria --cpus 4 --use-gpu
-poppunk_sketch --query --ref-db listeria --query-db listeria --use-gpu
+sketchlib sketch -l rfiles.txt -o listeria --cpus 4 --gpu 0
+sketchlib query dist listeria --gpu 0
 ```
 
 Both CPU parallelism and the GPU will be used, so be sure to add
-both `--cpus` and `--use-gpu` for maximum speed. This is particularly efficient
-when sketching.
-
-You can set the `--gpu-id` if you have more than one device, which may be necessary on
-cluster systems. This mode can also benefit from having multiple CPU cores available too.
+both `--cpus` and `--gpu` for maximum speed. This is particularly efficient
+when sketching reads.
 
 ### Benchmarks
 
@@ -193,12 +189,16 @@ contain `sketch` and may contain `random`. Run `h5dump` to see the full contents
 Contents are programmatically accessible with any HDF5 API. See `__main__.py` for an
 example in python.
 
+See `poppunk_db_info` from the [PopPUNK](https://github.com/johnlees/PopPUNK) package for pretty printing.
+
 #### sketch
 
 Attributes:
 
 - `sketch_version` - version of sketching code used to create the database.
   The SHA1 hash of relevant code files (doesn't change with every commit).
+- `codon_phased` - 1 if codon-phased seeds were used.
+- `reverse_complement` - 0 if `--single-strand`.
 
 Contains a group for each sample, within each has attributes:
 
@@ -230,78 +230,6 @@ Datasets:
 - `table_keys` - sample order of `table_values`.
 - `table_values` - centroid ID assigned to each sample.
 
-C++
----
-I have yet to set up a proper namespace for this, but you can include this
-code (`api.hpp` will do most functions) and use the parts you need. If you
-are interested in this becoming more functional, please raise an issue.
-
-See `main.cpp` for examples:
-
-```
-#include <fstream>
-#include <iostream>
-
-#include "reference.hpp"
-#include "database.hpp"
-#include "random_match.hpp"
-#include "api.hpp"
-
-// Set k-mer lengths
-std::vector<size_t> kmer_lengths {15, 17, 19, 21, 23, 25, 27, 29};
-
-// Create a two sketches
-Reference ref(argv[1], {argv[2]}, kmer_lengths, 156, true, 0, false);
-Reference query(argv[3], {argv[4]}, kmer_lengths, 156, true, 0, false);
-
-// Use default random match chances
-RandomMC random(true);
-
-// Output some distances at a single k-mer length
-std::cout << ref.jaccard_dist(query, 15, random) << std::endl;
-std::cout << ref.jaccard_dist(query, 29, random) << std::endl;
-
-// Calculate core and accessory distances between two sketches
-auto core_acc = ref.core_acc_dist<RandomMC>(query, random);
-std::cout << std::get<0>(core_acc) << "\t" << std::get<1>(core_acc) << std::endl;
-
-// Save sketches to file
-Database sketch_db("sketch.h5");
-sketch_db.add_sketch(ref);
-sketch_db.add_sketch(query);
-
-// Read sketches from file
-Reference ref_read = sketch_db.load_sketch(argv[1]);
-Reference query_read = sketch_db.load_sketch(argv[3]);
-// Create sketches using multiple threads, saving to file
-std::vector<Reference> ref_sketches = create_sketches("full",
-                           {argv[1], argv[3]},
-                           {{argv[2]}, {argv[4]}},
-                           kmer_lengths,
-                           156,
-                           true,
-                           0,
-                           false,
-                           2);
-// Calculate distances between sketches using multiple threads
-MatrixXf dists = query_db(ref_sketches,
-                          ref_sketches,
-                          kmer_lengths,
-                          random,
-                          false,
-                          2);
-std::cout << dists << std::endl;
-
-// Read sketches from an existing database, using random access
-HighFive::File h5_db("listeria.h5");
-Database listeria_db(h5_db);
-std::vector<Reference> listeria_sketches;
-for (auto name_it = names.cbegin(); name_it != names.cend(); name_it++)
-{
-    listeria_sketches.push_back(listeria_db.load_sketch(*name_it));
-}
-```
-
 ## Algorithms
 
 ### Sketching
@@ -363,7 +291,7 @@ Blais & Blanchette is used (formula 6 in the paper cited below).
   sketch each separately and join the databases.
 - GPU sketching filters out any read containing an N, which may give slightly
   different results from the CPU code.
-- GPU sketching with variable read lengths is untested, but theoretically supported.
+- GPU sketching with variable read lengths is unsupported. Illumina data only for now!
 - GPU distances use lower precision than the CPU code, so slightly different results
   are expected.
 
@@ -427,6 +355,9 @@ Modifiers:
 - `PROFILE=1` runs with profiler flags for `ncu` and `nsys`
 - `GPU=1` also build CUDA code (assumes `/usr/local/cuda-11.1/` and SM v8.6)
 
+### azure
+The repository key for the ubuntu CUDA install is periodically updated, which may cause build failures. See https://developer.nvidia.com/blog/updating-the-cuda-linux-gpg-repository-key/ and update in `azure-pipelines.yml`.
+
 ### Test that Python can build an installable package
 
 Build a python source package and install it into an empty docker container with vanilla python 3. If this works, then there's a good chance that the version uploaded to pypi will work
diff --git a/azure-pipelines.yml b/azure-pipelines.yml
index 1553c2df..4e6bf1e0 100644
--- a/azure-pipelines.yml
+++ b/azure-pipelines.yml
@@ -29,7 +29,7 @@ steps:
 - script: |
     wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/cuda-ubuntu1804.pin
     sudo mv cuda-ubuntu1804.pin /etc/apt/preferences.d/cuda-repository-pin-600
-    sudo apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/7fa2af80.pub
+    sudo apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/3bf863cc.pub
     sudo add-apt-repository "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/ /"
     sudo apt-get update
     sudo apt-get -y install cuda=11.2.2-1
@@ -54,5 +54,5 @@ steps:
     export CUDA_HOME=/usr/local/cuda-11.2
     export PATH=${CUDA_HOME}/bin${PATH:+:${PATH}}
     export LD_LIBRARY_PATH=${CUDA_HOME}/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
-    cd test && python run_test.py --no-cpp
+    cd test && python run_test.py
   displayName: 'Run tests (run_test.py)'
diff --git a/environment.yml b/environment.yml
index 26467868..1c15a5ca 100644
--- a/environment.yml
+++ b/environment.yml
@@ -9,6 +9,7 @@ dependencies:
   - pip
   - numpy
   - scipy
+  - docopt
   - cmake >= 3.12
   - pybind11
   - zlib
@@ -21,4 +22,4 @@ dependencies:
   - armadillo
   - libgfortran-ng
   - nvcc_linux-64
-  - cudatoolkit==11.2
+  - cudatoolkit==11.2 # This is pinned due to version install on azure, see azure-pipelines.yml
diff --git a/pp_sketch/__init__.py b/pp_sketch/__init__.py
index 37874766..97349f95 100644
--- a/pp_sketch/__init__.py
+++ b/pp_sketch/__init__.py
@@ -3,4 +3,4 @@
 
 '''PopPUNK sketching functions'''
 
-__version__ = '1.7.6.2'
+__version__ = '2.0.0'
diff --git a/pp_sketch/__main__.py b/pp_sketch/__main__.py
index 29da2553..b7fb779d 100644
--- a/pp_sketch/__main__.py
+++ b/pp_sketch/__main__.py
@@ -1,8 +1,45 @@
-# Copyright 2019-2020 John Lees
-
-'''Wrapper around sketch functions'''
+# Copyright 2019-2022 John Lees
+"""sketchlib: fast sketching and core/accessory distance estimation from assemblies and reads
+
+Usage:
+  sketchlib sketch <files>... -o <output> [-k <kseq>|--kmer <k>] [-s <size>] [--single-strand] [--codon-phased] [--min-count <count>] [--exact-counter] [--cpus <cpus>] [--gpu <gpu>]
+  sketchlib sketch -l <file-list> -o <output> [-k <kseq>|--kmer <k>] [-s <size>] [--single-strand] [--codon-phased] [--min-count <count>] [--exact-counter] [--cpus <cpus>] [--gpu <gpu>]
+  sketchlib query dist <db1> [<db2>] [-o <output>] [--adj-random] [--subset <file>] [--cpus <cpus>] [--gpu <gpu>]
+  sketchlib query jaccard <db1> [<db2>] [-o <output>] [--kmer <k>] [--adj-random] [--subset <file>] [--cpus <cpus>]
+  sketchlib query sparse <db1> (--kNN <k>|--threshold <max>) [-o <output>] [--accessory] [--adj-random] [--subset <file>] [--cpus <cpus>] [--gpu <gpu>]
+  sketchlib query sparse jaccard <db1> --kNN <k> --kmer <k> [-o <output>] [--adj-random] [--subset <file>] [--cpus <cpus>]
+  sketchlib join <db1> <db2> -o <output>
+  sketchlib (add|remove) random <db1> [--cpus <cpus>]
+  sketchlib (-h | --help)
+  sketchlib (--version)
+
+Options:
+  -h --help     Show this help.
+  --version     Show version.
+
+  -o <output>    Output prefix.
+  -l <file-list> File with a list of input files.
+  --cpus <cpus>  Number of CPU threads to use [default: 1].
+  --gpu <gpu>    Use GPU with specified device ID [default: -1].
+
+  -k <kseq>     Sequence of k-mers to sketch (min,max,step) [default: 15,31,4].
+  --kmer <k>    Sketch (or distance) at a single k-mer length k.
+  -s <size>     Sketch size [default: 10000].
+  --single-strand  Ignore the reverse complement (e.g. in RNA viruses).
+  --codon-phased  Use codon phased seeds X--X--X
+  --min-count <count>  Minimum coverage count for k-mers from reads to be sketched [default: 20].
+  --exact-counter  Use an exact k-mer count filter for reads (for genomes >10Mb)
+
+  --adj-random  Adjust query matches for their chance of occurring at random
+  --subset <file>  Only query samples matching names in file
+
+  --kNN <k>  Use k nearest neighbours to sparsify
+  --threshold <max>  Remove distances over max to sparsify
+  --accessory  Use accessory distances rather than core to sparsify
+"""
 
 import os, sys
+import re
 
 import numpy as np
 from scipy.sparse import save_npz
@@ -51,7 +88,7 @@ def getSampleNames(db_prefix):
     ref = h5py.File(db_prefix + ".h5", 'r')
     for sample_name in list(ref['sketches'].keys()):
         rList.append(sample_name)
-    return(rList)
+    return rList
 
 def storePickle(rlist, qlist, self, X, pklName):
     """Saves core and accessory distances in a .npy file, names in a .pkl
@@ -76,169 +113,97 @@ def storePickle(rlist, qlist, self, X, pklName):
         save_npz(pklName + ".npz", X)
 
 def get_options():
-    import argparse
-
-    description = 'Run poppunk sketching/distances'
-    parser = argparse.ArgumentParser(description=description,
-                                     prog='poppunk_sketch')
-
-    modeGroup = parser.add_argument_group('Mode of operation')
-    mode = modeGroup.add_mutually_exclusive_group(required=True)
-    mode.add_argument('--sketch',
-                        action='store_true',
-                        default=False,
-                        help='Create a database of sketches')
-    mode.add_argument('--join',
-                        action='store_true',
-                        default=False,
-                        help='Combine two sketch databases '
-                             '--ref-db and --query-db')
-    mode.add_argument('--query',
-                        action='store_true',
-                        default=False,
-                        help='Find distances between two sketch databases')
-    mode.add_argument('--add-random',
-                        action='store_true',
-                        default=False,
-                        help='Calculate random match chances and add these to a database')
-
-    io = parser.add_argument_group('Input/output')
-    io.add_argument('--rfile',
-                    help='Samples to sketch')
-    io.add_argument('--ref-db',
-                    help='Prefix of reference database file')
-    io.add_argument('--query-db',
-                    help='Prefix of query database file')
-    io.add_argument('--output',
-                    default='ppsketch',
-                    help="Output prefix [default = 'ppsketch']")
-
-    kmerGroup = parser.add_argument_group('Kmer options')
-    kmerGroup.add_argument('--read-k', default = False, action='store_true',
-                            help='Use k-mer lengths found in query databases (query mode only)')
-    kmerGroup.add_argument('--min-k', default = 13, type=int, help='Minimum kmer length [default = 13]')
-    kmerGroup.add_argument('--max-k', default = 29, type=int, help='Maximum kmer length [default = 29]')
-    kmerGroup.add_argument('--k-step', default = 4, type=int, help='K-mer step size [default = 4]')
-
-    sketchGroup = parser.add_argument_group('Sketch options')
-    sketchGroup.add_argument('--sketch-size', default=10000, type=int, help='K-mer sketch size [default = 10000]')
-    sketchGroup.add_argument('--strand', default=True, action='store_false', help='Set to ignore complementary strand sequence '
-                                                                                'e.g. for RNA viruses with preserved strand')
-    sketchGroup.add_argument('--codon-phased', default=False, action='store_true', help='Use codon-phased seeds [default = False]')
-    sketchGroup.add_argument('--min-count', default=20, type=int, help='Minimum k-mer count from reads [default = 20]')
-    sketchGroup.add_argument('--exact-counter', default=False, action='store_true',
-                            help='Use an exact rather than approximate k-mer counter '
-                                  'when using reads as input (increases memory use) '
-                                  '[default = False]')
-    sketchGroup.add_argument('--no-random', default = False, action = 'store_true',
-                             help = 'Do not calculate random match chances. Use this '
-                                    'if sketching input for querying against a reference.'
-                                    ' [default = False]')
-
-    query = parser.add_argument_group('Query options')
-    query.add_argument('--subset',
-                        help='List of names to include in query, default is to use '
-                             'everything in the database',
-                        default=None)
-    query.add_argument('--print',
-                        default=False,
-                        action='store_true',
-                        help='Print results to stdout instead of file')
-    query.add_argument('--sparse',
-                        default=False,
-                        action='store_true',
-                        help='Output query results in ijv sparse format '
-                             '(only for use with SCE)')
-    query.add_argument('--kNN',
-                        help='With --sparse, pick this number of nearest neighbours '
-                             'to include in the output sparse distance matrix',
-                        type=int,
-                        default=0)
-    query.add_argument('--threshold',
-                        help='With --sparse, pick distances under this value '
-                             'to include in the output sparse distance matrix',
-                        type=float,
-                        default=0)
-    query.add_argument('--accessory',
-                        help='With --sparse, use accessory distances instead '
-                             'of core',
-                        action='store_true',
-                        default=False)
-    query.add_argument('--no-correction',
-                        default=False,
-                        action='store_true',
-                        help='Disable correction for random matches at short k-mer lengths')
-    query.add_argument('--jaccard',
-                        default=False,
-                        action='store_true',
-                        help='Output adjusted Jaccard distances, not core and accessory distances')
-
-    optimisation = parser.add_argument_group('Optimisation options')
-    optimisation.add_argument('--cpus',
-                              type=int,
-                              default=1,
-                              help='Number of CPUs to use (--sketch and --query) '
-                                   '[default = 1]')
-    optimisation.add_argument('--use-gpu', default=False, action='store_true',
-                              help='Use GPU [default = False]')
-    optimisation.add_argument('--gpu-id',
-                              type=int,
-                              default=0,
-                              help='Device ID of GPU to use '
-                                   '[default = 0]')
-
-    other = parser.add_argument_group('Other')
-    other.add_argument('--version', action='version',
-                       version='%(prog)s '+__version__)
-
-    return parser.parse_args()
+    from docopt import docopt
+    arguments = docopt(__doc__, version="pp-sketchlib v"+__version__)
+
+    # .h5 is removed from the end of DB names due to sketchlib API
+    if arguments["<db1>"]:
+        arguments["db1"] = re.sub(r"\.h5$", "", arguments["<db1>"])
+    if arguments["<db2>"]:
+        arguments["db2"] = re.sub(r"\.h5$", "", arguments["<db2>"])
+
+    if arguments['--kmer']:
+        arguments['kmers'] = [int(arguments['--kmer'])]
+    else:
+        try:
+            (min_k, max_k, k_step) = [int(x) for x in arguments['-k'].split(",")]
+            if min_k >= max_k or min_k < 3 or max_k > 101 or k_step < 1:
+                raise RuntimeError("Invalid k-mer sizes")
+            arguments['kmers'] = np.arange(int(min_k), int(max_k) + 1, int(k_step))
+        except RuntimeError:
+            sys.stderr.write("Minimum kmer size must be smaller than maximum kmer size; " +
+                             "range must be between 3 and 101, step must be at least one\n")
+            sys.exit(1)
+
+    arguments['-s'] = int(round(float(arguments['-s'])/64))
+    arguments['--min-count'] = int(arguments['--min-count'])
+
+    if arguments['sparse']:
+        if arguments['--kNN']:
+            arguments['--kNN'] = int(arguments['--kNN'])
+            arguments['--threshold'] = 0
+        else:
+            arguments['--kNN'] = 0
+            arguments['--threshold'] = float(arguments['--threshold'])
+
+    arguments['--cpus'] = int(arguments['--cpus'])
+    arguments['--gpu'] = int(arguments['--gpu'])
+    if int(arguments['--gpu']) >= 0:
+        arguments['--use-gpu'] = True
+    else:
+        arguments['--use-gpu'] = False
 
+    return arguments
 
 def main():
     args = get_options()
 
-    if args.min_k >= args.max_k or args.min_k < 3 or args.max_k > 101 or args.k_step < 1:
-        sys.stderr.write("Minimum kmer size " + str(args.min_k) + " must be smaller than maximum kmer size " +
-                         str(args.max_k) + "; range must be between 3 and 101, step must be at least one\n")
-        sys.exit(1)
-    kmers = np.arange(args.min_k, args.max_k + 1, args.k_step)
-
     #
     # Create a database (sketch input)
     #
-    if args.sketch:
+    if args['sketch']:
         names = []
         sequences = []
 
-        with open(args.rfile, 'rU') as refFile:
-            for refLine in refFile:
-                refFields = refLine.rstrip().split("\t")
-                names.append(refFields[0])
-                sequences.append(list(refFields[1:]))
+        if args['-l']:
+            with open(args['-l'], 'rU') as refFile:
+                for refLine in refFile:
+                    refFields = refLine.rstrip().split("\t")
+                    names.append(refFields[0])
+                    sequences.append(list(refFields[1:]))
+        else:
+            for file in args['<files>']:
+                name = re.sub("/", "_", file)
+                names.append(name)
+                sequences.append([file])
 
 
         if len(set(names)) != len(names):
-            sys.stderr.write("Input contains duplicate names! All names must be unique\n")
+            sys.stderr.write("Input contains duplicate names. "
+                             "All names must be unique\n")
             sys.exit(1)
 
-        pp_sketchlib.constructDatabase(args.ref_db, names, sequences, kmers,
-                                       int(round(args.sketch_size/64)),
-                                       args.codon_phased,
-                                       not args.no_random,
-                                       args.strand,
-                                       args.min_count,
-                                       args.exact_counter,
-                                       args.cpus,
-                                       args.use_gpu,
-                                       args.gpu_id)
+        pp_sketchlib.constructDatabase(db_name=args['-o'],
+                                       samples=names,
+                                       files=sequences,
+                                       klist=args['kmers'],
+                                       sketch_size=args['-s'],
+                                       codon_phased=args['--codon-phased'],
+                                       calc_random=False,
+                                       use_rc=not args['--single-strand'],
+                                       min_count=args['--min-count'],
+                                       exact=args['--exact-counter'],
+                                       num_threads=args['--cpus'],
+                                       use_gpu=args['--use-gpu'],
+                                       device_id=args['--gpu'])
 
     #
     # Join two databases
     #
-    elif args.join:
-        join_name = args.output + ".h5"
-        db1_name = args.ref_db + ".h5"
-        db2_name = args.query_db + ".h5"
+    elif args['join']:
+        join_name = args['-o'] + ".h5"
+        db1_name = args['db1'] + ".h5"
+        db2_name = args['db2'] + ".h5"
 
         hdf1 = h5py.File(db1_name, 'r')
         hdf2 = h5py.File(db2_name, 'r')
@@ -255,6 +220,8 @@ def main():
                 sys.stderr.write("One database uses codon-phased seeds - cannot join "
                                  "with a standard seed database\n")
         except RuntimeError as e:
+            hdf1.close()
+            hdf2.close()
             sys.stderr.write("Unable to check sketch version\n")
 
         hdf_join = h5py.File(join_name + ".tmp", 'w') # add .tmp in case join_name exists
@@ -274,23 +241,29 @@ def main():
             sys.stderr.write("ERROR: " + str(e) + "\n")
             sys.stderr.write("Joining sketches failed\n")
             sys.exit(1)
+        finally:
+            hdf1.close()
+            hdf2.close()
+            hdf_join.close()
 
         # Clean up
-        hdf1.close()
-        hdf2.close()
-        hdf_join.close()
         os.rename(join_name + ".tmp", join_name)
 
     #
     # Query a database (calculate distances)
     #
-    elif args.query:
-        rList = getSampleNames(args.ref_db)
-        qList = getSampleNames(args.query_db)
+    elif args['query']:
+        rList = getSampleNames(args['db1'])
+
+        if 'db2' not in args:
+            args['db2'] = args['db1']
+            qList = rList
+        else:
+            qList = getSampleNames(args['db2'])
 
-        if args.subset != None:
+        if args['--subset'] != None:
             subset = []
-            with open(args.subset, 'r') as subset_file:
+            with open(args['--subset'], 'r') as subset_file:
                 for line in subset_file:
                     sample_name = line.rstrip().split("\t")[0]
                     subset.append(sample_name)
@@ -301,29 +274,71 @@ def main():
                 sys.exit(1)
 
         # Check inputs overlap
-        ref = h5py.File(args.ref_db + ".h5", 'r')
-        query = h5py.File(args.query_db + ".h5", 'r')
-        db_kmers = set(ref['sketches/' + rList[0]].attrs['kmers']).intersection(
-           query['sketches/' + qList[0]].attrs['kmers']
-        )
-        if args.read_k:
+        try:
+            ref = h5py.File(args['db1'] + ".h5", 'r')
+            query = h5py.File(args['db2'] + ".h5", 'r')
+            db_kmers = set(ref['sketches/' + rList[0]].attrs['kmers']).intersection(
+              query['sketches/' + qList[0]].attrs['kmers']
+            )
             query_kmers = sorted(db_kmers)
-        else:
-            query_kmers = sorted(set(kmers).intersection(db_kmers))
-            if (len(query_kmers) == 0):
-                sys.stderr.write("No requested k-mer lengths found in DB\n")
-                sys.exit(1)
-            elif (len(query_kmers) < len(query_kmers)):
-                sys.stderr.write("Some requested k-mer lengths not found in DB\n")
-        ref.close()
-        query.close()
-
-        if args.sparse:
-            sparseIdx = pp_sketchlib.queryDatabaseSparse(args.ref_db, args.query_db, rList, qList, query_kmers,
-                                                         not args.no_correction, args.threshold, args.kNN,
-                                                         not args.accessory, args.cpus, args.use_gpu, args.gpu_id)
-            if args.print:
-                if args.accessory:
+            if args['jaccard'] and len(args['kmers']) == 1:
+                if args['kmers'][0] not in query_kmers:
+                    raise RuntimeError(f"Input --kmer {args['kmers'][0]} not in "
+                                       f"(both) databases: {sorted(db_kmers)}\n")
+                else:
+                    query_kmers = args['kmers']
+            elif not query_kmers:
+                raise RuntimeError("No overlapping k-mers in db1 and db2")
+        finally:
+            ref.close()
+            query.close()
+
+        if args['sparse']:
+            # Set which distance to output (sparse only supports a single value)
+            dist_col = 0
+            if args['jaccard'] and len(query_kmers) != 1:
+                # Note default here is just to be sending a single-kmer length, so
+                # it does not need to be indexed into hence dist_col = 0 is correct
+                raise RuntimeError(f"For sparse jaccard a single k-mer must be "
+                                   f"selected {sorted(query_kmers)}")
+            elif args['--accessory']:
+                dist_col = 1
+
+            # Can use memory efficient version with kNN
+            if args['--threshold'] == 0:
+                sparseIdx = pp_sketchlib.querySelfSparse(ref_db_name=args['db1'],
+                                              rList=rList,
+                                              klist=query_kmers,
+                                              random_correct=args['--adj-random'],
+                                              jaccard=args['jaccard'],
+                                              kNN=args['--kNN'],
+                                              dist_cutoff=0,
+                                              dist_col=dist_col,
+                                              num_threads=args['--cpus'],
+                                              use_gpu=args['--use-gpu'],
+                                              device_id=args['--gpu'])
+            # Otherwise use general version (potentially large memory use as it
+            # actually calculates the dense distances, then sparsifies them)
+            else:
+                # This should be prohibited by docopt, but just in case
+                if args['jaccard']:
+                    raise ValueError("Cannot use sparse jaccard with non-self "
+                                     "or --threshold")
+                sparseIdx = pp_sketchlib.querySelfSparse(ref_db_name=args['db1'],
+                              rList=rList,
+                              klist=query_kmers,
+                              random_correct=args['--adj-random'],
+                              jaccard=args['jaccard'],
+                              kNN=0,
+                              dist_cutoff=args['--threshold'],
+                              dist_col=dist_col,
+                              num_threads=args['--cpus'],
+                              use_gpu=args['--use-gpu'],
+                              device_id=args['--gpu'])
+            if not args['-o']:
+                if args['jaccard']:
+                    distName = str(query_kmers[0])
+                elif args['--accessory']:
                     distName = 'Accessory'
                 else:
                     distName = 'Core'
@@ -335,17 +350,24 @@ def main():
 
             else:
                 coo_matrix = ijv_to_coo(sparseIdx, (len(rList), len(qList)), np.float32)
-                storePickle(rList, qList, rList == qList, coo_matrix, args.output)
+                storePickle(rList, qList, rList == qList, coo_matrix, args['-o'])
 
         else:
-            distMat = pp_sketchlib.queryDatabase(args.ref_db, args.query_db, rList, qList, query_kmers,
-                                                 not args.no_correction, args.jaccard, args.cpus, args.use_gpu,
-                                                 args.gpu_id)
+            distMat = pp_sketchlib.queryDatabase(ref_db_name=args['db1'],
+                                                 query_db_name=args['db2'],
+                                                 rList=rList,
+                                                 qList=qList,
+                                                 klist=query_kmers,
+                                                 random_correct=args['--adj-random'],
+                                                 jaccard=args['jaccard'],
+                                                 num_threads=args['--cpus'],
+                                                 use_gpu=args['--use-gpu'],
+                                                 device_id=args['--gpu'])
 
             # get names order
-            if args.print:
+            if not args['-o']:
                 names = iterDistRows(rList, qList, rList == qList)
-                if not args.jaccard:
+                if not args['jaccard']:
                     sys.stdout.write("\t".join(['Query', 'Reference', 'Core', 'Accessory']) + "\n")
                     for i, (ref, query) in enumerate(names):
                         sys.stdout.write("\t".join([query, ref, str(distMat[i,0]), str(distMat[i,1])]) + "\n")
@@ -354,18 +376,30 @@ def main():
                     for i, (ref, query) in enumerate(names):
                         sys.stdout.write("\t".join([query, ref] + [str(k) for k in distMat[i,]]) + "\n")
             else:
-                storePickle(rList, qList, rList == qList, distMat, args.output)
+                storePickle(rList, qList, rList == qList, distMat, args['-o'])
 
     #
     # Add random match chances to an older database
     #
-    elif args.add_random:
-        rList = getSampleNames(args.ref_db)
-        ref = h5py.File(args.ref_db + ".h5", 'r')
-        db_kmers = ref['sketches/' + rList[0]].attrs['kmers']
-        ref.close()
-
-        pp_sketchlib.addRandom(args.ref_db, rList, db_kmers, args.strand, args.cpus)
+    elif args['random']:
+        ref = h5py.File(args['db1'] + ".h5", 'r+')
+        if args['add']:
+            rList = getSampleNames(args['db1'])
+            db_kmers = ref['sketches/' + rList[0]].attrs['kmers']
+            ref.close()
+
+            pp_sketchlib.addRandom(db_name=args['db1'],
+                                   samples=rList,
+                                   klist=db_kmers,
+                                   use_rc=not args['--single-strand'],
+                                   num_threads=args['--cpus'])
+        elif args['remove']:
+            if 'random' in ref:
+                del ref['random']
+            else:
+                sys.stderr.write(args['db1'] + \
+                  " doesn't contain random match chances\n")
+            ref.close()
 
     sys.exit(0)
 
diff --git a/requirements.txt b/requirements.txt
index 8165e28d..398ec6ce 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -2,3 +2,4 @@ h5py
 numpy
 pybind11
 scipy
+docopt
diff --git a/setup.py b/setup.py
index b310a035..da41f68e 100755
--- a/setup.py
+++ b/setup.py
@@ -115,12 +115,12 @@ def build_extension(self, ext):
         'License :: OSI Approved :: Apache Software License',
         'Programming Language :: Python :: 3.8',
     ],
-    python_requires='>=3.6.0',
+    python_requires='>=3.8.0',
     keywords='bacteria genomics population-genetics k-mer',
     packages=['pp_sketch'],
     entry_points={
         "console_scripts": [
-            'poppunk_sketch = pp_sketch.__main__:main'
+            'sketchlib = pp_sketch.__main__:main'
             ]
     },
     setup_requires=['pybind11'],
diff --git a/pp_sketch-runner.py b/sketchlib-runner.py
similarity index 100%
rename from pp_sketch-runner.py
rename to sketchlib-runner.py
diff --git a/src/api.cpp b/src/api.cpp
index c8bee289..3219f044 100644
--- a/src/api.cpp
+++ b/src/api.cpp
@@ -16,12 +16,11 @@
 #include "gpu/gpu.hpp"
 #include "reference.hpp"
 #include "sketch/progress.hpp"
+#include "sketch/bitfuncs.hpp"
 
 using namespace Eigen;
 namespace py = pybind11;
 
-const int progressBitshift = 10;
-
 bool same_db_version(const std::string &db1_name, const std::string &db2_name) {
   // Open databases
   Database db1(db1_name + ".h5");
@@ -77,6 +76,7 @@ std::vector<Reference> create_sketches(
     KmerSeeds kmer_seeds = generate_seeds(kmer_lengths, codon_phased);
 
     ProgressMeter sketch_progress(names.size());
+    size_t done_count = 0;
     bool interrupt = false;
     std::vector<std::runtime_error> errors;
 #pragma omp parallel for schedule(dynamic, 5) num_threads(num_threads)
@@ -88,6 +88,8 @@ std::vector<Reference> create_sketches(
           SeqBuf seq_in(files[i], kmer_lengths.back());
           sketches[i] = Reference(names[i], seq_in, kmer_seeds, sketchsize64,
                                   codon_phased, use_rc, min_count, exact);
+#pragma omp atomic
+          ++done_count;
         } catch (const std::runtime_error &e) {
 #pragma omp critical
           {
@@ -98,7 +100,7 @@ std::vector<Reference> create_sketches(
       }
 
       if (omp_get_thread_num() == 0) {
-        sketch_progress.tick(num_threads);
+        sketch_progress.tick_count(done_count);
       }
     }
     sketch_progress.finalise();
@@ -167,28 +169,32 @@ NumpyMatrix query_db(std::vector<Reference> &ref_sketches,
     dist_cols = 2;
   }
 
+  size_t dist_rows;
+  // Set up progress meter
+  if (ref_sketches == query_sketches) {
+    // calculate dists
+    dist_rows = static_cast<size_t>(0.5 * (ref_sketches.size()) *
+                                           (ref_sketches.size() - 1));
+  } else {
+    dist_rows = ref_sketches.size() * query_sketches.size();
+  }
+  static const uint64_t n_progress_ticks = 1000;
+  uint64_t update_every = 1;
+  if (dist_rows > n_progress_ticks) {
+    update_every = dist_rows / n_progress_ticks;
+  }
+  ProgressMeter dist_progress(n_progress_ticks, true);
+  volatile int progress = 0;
+
   // These could be the same but out of order, which could be dealt with
   // using a sort, except the return order of the distances wouldn't be as
   // expected. self iff ref_names == query_names as input
   bool interrupt = false;
   if (ref_sketches == query_sketches) {
     // calculate dists
-    size_t dist_rows = static_cast<size_t>(0.5 * (ref_sketches.size()) *
-                                           (ref_sketches.size() - 1));
     distMat.resize(dist_rows, dist_cols);
-
     arma::mat kmer_mat = kmer2mat<std::vector<size_t>>(kmer_lengths);
 
-    // Set up progress meter
-    size_t progress_blocks = 1 << progressBitshift;
-    size_t update_every = dist_rows >> progressBitshift;
-    if (progress_blocks > dist_rows || update_every < 1) {
-      progress_blocks = dist_rows;
-      update_every = 1;
-    }
-    ProgressMeter dist_progress(progress_blocks, true);
-    int progress = 0;
-
     // Iterate upper triangle
 #pragma omp parallel for schedule(dynamic, 5) num_threads(num_threads) shared(progress)
     for (size_t i = 0; i < ref_sketches.size(); i++) {
@@ -209,18 +215,18 @@ NumpyMatrix query_db(std::vector<Reference> &ref_sketches,
                     ref_sketches[j], kmer_mat, random_chance);
           }
           if (pos % update_every == 0) {
-#pragma omp atomic
-              progress++;
-              dist_progress.tick(1);
+#pragma omp critical
+              {
+                progress += MAX(1, n_progress_ticks / dist_rows);
+                dist_progress.tick_count(progress);
+              }
           }
         }
       }
     }
-    dist_progress.finalise();
   } else {
     // If ref != query, make a thread queue, with each element one ref
     // calculate dists
-    size_t dist_rows = ref_sketches.size() * query_sketches.size();
     distMat.resize(dist_rows, dist_cols);
 
     // Prepare objects used in distance calculations
@@ -235,7 +241,6 @@ NumpyMatrix query_db(std::vector<Reference> &ref_sketches,
           random_chance.closest_cluster(query_sketches[q_idx]);
     }
 
-    ProgressMeter dist_progress(dist_rows, true);
 #pragma omp parallel for collapse(2) schedule(static) num_threads(num_threads)
     for (unsigned int q_idx = 0; q_idx < query_sketches.size(); q_idx++) {
       for (unsigned int r_idx = 0; r_idx < ref_sketches.size(); r_idx++) {
@@ -260,23 +265,144 @@ NumpyMatrix query_db(std::vector<Reference> &ref_sketches,
                 query_sketches[q_idx].core_acc_dist<std::vector<double>>(
                     ref_sketches[r_idx], kmer_mat, jaccard_random);
           }
-          if (omp_get_thread_num() == 0) {
-            dist_progress.tick(num_threads);
+          if ((q_idx * ref_sketches.size() + r_idx) % update_every == 0) {
+#pragma omp critical
+            {
+              progress += MAX(1, n_progress_ticks / dist_rows);
+              dist_progress.tick_count(progress);
+            }
           }
         }
       }
     }
-    dist_progress.finalise();
   }
 
   // Handle Ctrl-C from python
   if (interrupt) {
     throw py::error_already_set();
   }
+  dist_progress.finalise();
 
   return (distMat);
 }
 
+void check_sparse_inputs(const std::vector<Reference> &ref_sketches,
+                         const std::vector<size_t> &kmer_lengths,
+                         RandomMC &random_chance, const bool jaccard,
+                         const size_t dist_col) {
+  if (ref_sketches.size() < 1) {
+    throw std::runtime_error("Query with empty ref or query list!");
+  }
+  if (kmer_lengths[0] <
+      random_chance.min_supported_k(ref_sketches[0].seq_length())) {
+    throw std::runtime_error("Smallest k-mer has no signal above random "
+                             "chance; increase minimum k-mer length");
+  }
+  if ((jaccard && dist_col >= kmer_lengths.size()) || (dist_col != 0 && dist_col != 1)) {
+    throw std::runtime_error("dist_col out of range");
+  }
+
+  // Check all references are in the random object, add if not
+  if (random_chance.check_present(ref_sketches, true)) {
+    std::cerr
+        << "Some members of the reference database were not found "
+           "in its random match chances. Consider refreshing with addRandom"
+        << std::endl;
+  }
+}
+
+sparse_coo query_db_sparse(std::vector<Reference> &ref_sketches,
+                     const std::vector<size_t> &kmer_lengths,
+                     RandomMC &random_chance, const bool jaccard,
+                     const int kNN, const size_t dist_col,
+                     const size_t num_threads) {
+  check_sparse_inputs(ref_sketches, kmer_lengths, random_chance,
+                      jaccard, dist_col);
+
+  std::cerr << "Calculating distances using " << num_threads << " thread(s)"
+            << std::endl;
+
+  std::vector<float> dists(ref_sketches.size() * kNN);
+  std::vector<long> i_vec(ref_sketches.size() * kNN);
+  std::vector<long> j_vec(ref_sketches.size() * kNN);
+
+  bool interrupt = false;
+
+  // Set up progress meter
+  size_t dist_rows = static_cast<size_t>(ref_sketches.size() * ref_sketches.size());
+  static const uint64_t n_progress_ticks = 1000;
+  uint64_t update_every = 1;
+  if (dist_rows > n_progress_ticks) {
+    update_every = dist_rows / n_progress_ticks;
+  }
+  ProgressMeter dist_progress(n_progress_ticks, true);
+  volatile int progress = 0;
+
+  arma::mat kmer_mat = kmer2mat<std::vector<size_t>>(kmer_lengths);
+#pragma omp parallel for schedule(static) num_threads(num_threads) shared(progress)
+  for (size_t i = 0; i < ref_sketches.size(); i++) {
+    std::vector<float> row_dists(ref_sketches.size());
+    if (interrupt || PyErr_CheckSignals() != 0) {
+      interrupt = true;
+    } else {
+      for (size_t j = 0; j < ref_sketches.size(); j++) {
+        if (i != j) {
+          if (jaccard) {
+            // Need 1-J here to sort correctly
+            row_dists[j] = 1.0f - ref_sketches[i].jaccard_dist(
+                ref_sketches[j], kmer_lengths[dist_col], random_chance);
+          } else {
+            float core, acc;
+            std::tie(core, acc) =
+                ref_sketches[i].core_acc_dist<RandomMC>(
+                    ref_sketches[j], kmer_mat, random_chance);
+            if (dist_col == 0) {
+              row_dists[j] = core;
+            } else {
+              row_dists[j] = acc;
+            }
+          }
+        } else {
+          row_dists[j] = std::numeric_limits<float>::infinity();
+        }
+        if ((i * ref_sketches.size() + j) % update_every == 0) {
+#pragma omp critical
+          {
+            progress += MAX(1, n_progress_ticks / dist_rows);
+            dist_progress.tick_count(progress);
+          }
+        }
+        long offset = i * kNN;
+        std::vector<long> ordered_dists = sort_indexes(row_dists);
+        std::fill_n(i_vec.begin() + offset, kNN, i);
+        // std::copy_n(ordered_dists.begin(), kNN, j_vec.begin() + offset);
+
+        for (int k = 0; k < kNN; ++k) {
+          j_vec[offset + k] = ordered_dists[k];
+          dists[offset + k] = row_dists[ordered_dists[k]];
+        }
+
+      }
+    }
+  }
+  dist_progress.finalise();
+
+  // Handle Ctrl-C from python
+  if (interrupt) {
+    throw py::error_already_set();
+  }
+
+  // Revert to correct distance J = 1 - dist
+  if (jaccard) {
+    #pragma omp parallel for simd num_threads(num_threads)
+    for (size_t i = 0; i < dists.size(); ++i) {
+      dists[i] = 1.0f - dists[i];
+    }
+  }
+
+  return (std::make_tuple(i_vec, j_vec, dists));
+}
+
 // Load sketches from a HDF5 file
 // Returns empty vector on failure
 std::vector<Reference> load_sketches(const std::string &db_name,
diff --git a/src/api.hpp b/src/api.hpp
index 85741424..ac12cdf0 100644
--- a/src/api.hpp
+++ b/src/api.hpp
@@ -40,6 +40,12 @@ NumpyMatrix query_db(std::vector<Reference> &ref_sketches,
                      const bool jaccard,
                      const size_t num_threads);
 
+sparse_coo query_db_sparse(std::vector<Reference> &ref_sketches,
+                     const std::vector<size_t> &kmer_lengths,
+                     RandomMC &random_chance, const bool jaccard,
+                     const int kNN, const size_t dist_col,
+                     const size_t num_threads);
+
 #ifdef GPU_AVAILABLE
 // defined in gpu_api.cpp
 std::vector<Reference> create_sketches_cuda(const std::string &db_name,
@@ -58,6 +64,13 @@ NumpyMatrix query_db_cuda(std::vector<Reference> &ref_sketches,
                           RandomMC &random_chance,
                           const int device_id = 0,
                           const unsigned int num_cpu_threads = 1);
+
+sparse_coo query_db_sparse_cuda(std::vector<Reference> &ref_sketches,
+                          const std::vector<size_t> &kmer_lengths,
+                          RandomMC &random_match,
+                          const int kNN, const size_t dist_col,
+                          const int device_id,
+                          const unsigned int num_cpu_threads);
 #endif
 
 bool same_db_version(const std::string &db1_name,
diff --git a/src/dist/matrix.hpp b/src/dist/matrix.hpp
index cf98f8b7..e7935334 100644
--- a/src/dist/matrix.hpp
+++ b/src/dist/matrix.hpp
@@ -6,6 +6,8 @@
  */
 #pragma once
 
+#include <algorithm>
+#include <numeric>
 #include <vector>
 #include <cstdint>
 #include <cstddef>
@@ -13,11 +15,26 @@
 
 #include <Eigen/Dense>
 
+#include "matrix_types.hpp"
 #include "matrix_idx.hpp"
 
-typedef Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor> NumpyMatrix;
-typedef std::tuple<std::vector<long>, std::vector<long>, std::vector<float>> sparse_coo;
-typedef std::tuple<std::vector<long>, std::vector<long>, std::vector<long>> network_coo;
+// This type not used in any nvcc code
+using NumpyMatrix = Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>;
+
+//https://stackoverflow.com/a/12399290
+template <typename T>
+std::vector<long> sort_indexes(const T &v)
+{
+
+  // initialize original index locations
+  std::vector<long> idx(v.size());
+  std::iota(idx.begin(), idx.end(), 0);
+
+  std::stable_sort(idx.begin(), idx.end(),
+                   [&v](long i1, long i2) { return v[i1] < v[i2]; });
+
+  return idx;
+}
 
 NumpyMatrix long_to_square(const Eigen::VectorXf &rrDists,
                            const Eigen::VectorXf &qrDists,
diff --git a/src/dist/matrix_ops.cpp b/src/dist/matrix_ops.cpp
index 55b093dd..de48990e 100644
--- a/src/dist/matrix_ops.cpp
+++ b/src/dist/matrix_ops.cpp
@@ -4,59 +4,32 @@
  * Distance matrix transformations
  *
  */
-#include <vector>
-#include <numeric>
-#include <algorithm>
-#include <cstdint>
-#include <cstddef>
 #include <cmath>
-#include <string>
+#include <cstddef>
+#include <cstdint>
 #include <omp.h>
+#include <string>
+#include <vector>
 
 #include "api.hpp"
 
 const float epsilon = 1E-10;
 
 // Prototypes for cppying functions run in threads
-void square_block(const Eigen::VectorXf &longDists,
-                  NumpyMatrix &squareMatrix,
-                  const size_t n_samples,
-                  const size_t start,
-                  const size_t offset,
-                  const size_t max_elems);
+void square_block(const Eigen::VectorXf &longDists, NumpyMatrix &squareMatrix,
+                  const size_t n_samples, const size_t start,
+                  const size_t offset, const size_t max_elems);
 
 void rectangle_block(const Eigen::VectorXf &longDists,
-                     NumpyMatrix &squareMatrix,
-                     const size_t nrrSamples,
-                     const size_t nqqSamples,
-                     const size_t start,
+                     NumpyMatrix &squareMatrix, const size_t nrrSamples,
+                     const size_t nqqSamples, const size_t start,
                      const size_t max_elems);
 
-//https://stackoverflow.com/a/12399290
-template <typename T>
-std::vector<long> sort_indexes(const T &v)
-{
-
-  // initialize original index locations
-  std::vector<long> idx(v.size());
-  std::iota(idx.begin(), idx.end(), 0);
-
-  std::stable_sort(idx.begin(), idx.end(),
-                   [&v](long i1, long i2) { return v[i1] < v[i2]; });
-
-  return idx;
-}
-
-sparse_coo sparsify_dists(const NumpyMatrix &denseDists,
-                          const float distCutoff,
-                          const unsigned long int kNN)
-{
-  if (kNN > 0 && distCutoff > 0)
-  {
+sparse_coo sparsify_dists(const NumpyMatrix &denseDists, const float distCutoff,
+                          const unsigned long int kNN) {
+  if (kNN > 0 && distCutoff > 0) {
     throw std::runtime_error("Specify only one of kNN or distCutoff");
-  }
-  else if (kNN < 1 && distCutoff < 0)
-  {
+  } else if (kNN < 1 && distCutoff < 0) {
     throw std::runtime_error("kNN must be > 1 or distCutoff > 0");
   }
 
@@ -64,49 +37,36 @@ sparse_coo sparsify_dists(const NumpyMatrix &denseDists,
   std::vector<float> dists;
   std::vector<long> i_vec;
   std::vector<long> j_vec;
-  if (distCutoff > 0)
-  {
-    for (long i = 0; i < denseDists.rows(); i++)
-    {
-      for (long j = i + 1; j < denseDists.cols(); j++)
-      {
-        if (denseDists(i, j) < distCutoff)
-        {
+  if (distCutoff > 0) {
+    for (long i = 0; i < denseDists.rows(); i++) {
+      for (long j = i + 1; j < denseDists.cols(); j++) {
+        if (denseDists(i, j) < distCutoff) {
           dists.push_back(denseDists(i, j));
           i_vec.push_back(i);
           j_vec.push_back(j);
         }
       }
     }
-  }
-  else if (kNN >= 1)
-  {
+  } else if (kNN >= 1) {
     // Only add the k nearest (unique) neighbours
     // May be >k if repeats, often zeros
-    for (long i = 0; i < denseDists.rows(); i++)
-    {
+    for (long i = 0; i < denseDists.rows(); i++) {
       unsigned long unique_neighbors = 0;
       float prev_value = -1;
-      for (auto j : sort_indexes(denseDists.row(i)))
-      {
-        if (j == i)
-        {
+      for (auto j : sort_indexes(denseDists.row(i))) {
+        if (j == i) {
           continue; // Ignore diagonal which will always be one of the closest
         }
         bool new_val = abs(denseDists(i, j) - prev_value) < epsilon;
-        if (unique_neighbors < kNN || new_val)
-        {
+        if (unique_neighbors < kNN || new_val) {
           dists.push_back(denseDists(i, j));
           i_vec.push_back(i);
           j_vec.push_back(j);
-          if (!new_val)
-          {
+          if (!new_val) {
             unique_neighbors++;
             prev_value = denseDists(i, j);
           }
-        }
-        else
-        {
+        } else {
           break;
         }
       }
@@ -119,60 +79,52 @@ sparse_coo sparsify_dists(const NumpyMatrix &denseDists,
 NumpyMatrix long_to_square(const Eigen::VectorXf &rrDists,
                            const Eigen::VectorXf &qrDists,
                            const Eigen::VectorXf &qqDists,
-                           unsigned int num_threads)
-{
+                           unsigned int num_threads) {
   // Set up square matrix to move values into
   size_t nrrSamples = rows_to_samples(rrDists);
   size_t nqqSamples = 0;
-  if (qrDists.size() > 0)
-  {
+  if (qrDists.size() > 0) {
     nqqSamples = rows_to_samples(qqDists);
     // If one query, qqDists = [0] or None, so may come back as zero
-    if (nqqSamples < 1 && qrDists.size() > 0)
-    {
+    if (nqqSamples < 1 && qrDists.size() > 0) {
       nqqSamples = 1;
     }
-    if (qrDists.size() != nrrSamples * nqqSamples)
-    {
-      throw std::runtime_error("Shape of reference, query and ref vs query matrices mismatches");
+    if (qrDists.size() != nrrSamples * nqqSamples) {
+      throw std::runtime_error(
+          "Shape of reference, query and ref vs query matrices mismatches");
     }
   }
 
   // Initialise matrix and set diagonal to zero
   NumpyMatrix squareDists(nrrSamples + nqqSamples, nrrSamples + nqqSamples);
-  for (size_t diag_idx = 0; diag_idx < nrrSamples + nqqSamples; diag_idx++)
-  {
+  for (size_t diag_idx = 0; diag_idx < nrrSamples + nqqSamples; diag_idx++) {
     squareDists(diag_idx, diag_idx) = 0;
   }
 
 // Loop over threads for ref v ref square
 #pragma omp parallel for schedule(static) num_threads(num_threads)
-  for (long distIdx = 0; distIdx < rrDists.rows(); distIdx++)
-  {
+  for (long distIdx = 0; distIdx < rrDists.rows(); distIdx++) {
     unsigned long i = calc_row_idx(distIdx, nrrSamples);
     unsigned long j = calc_col_idx(distIdx, i, nrrSamples);
     squareDists(i, j) = rrDists[distIdx];
     squareDists(j, i) = rrDists[distIdx];
   }
 
-  if (qqDists.size() > 0)
-  {
+  if (qqDists.size() > 0) {
 #pragma omp parallel for schedule(static) num_threads(num_threads)
-    for (long distIdx = 0; distIdx < qqDists.rows(); distIdx++)
-    {
+    for (long distIdx = 0; distIdx < qqDists.rows(); distIdx++) {
       unsigned long i = calc_row_idx(distIdx, nqqSamples) + nrrSamples;
-      unsigned long j = calc_col_idx(distIdx, i - nrrSamples, nqqSamples) + nrrSamples;
+      unsigned long j =
+          calc_col_idx(distIdx, i - nrrSamples, nqqSamples) + nrrSamples;
       squareDists(i, j) = qqDists[distIdx];
       squareDists(j, i) = qqDists[distIdx];
     }
   }
 
   // Query vs ref rectangles
-  if (qrDists.size() > 0)
-  {
+  if (qrDists.size() > 0) {
 #pragma omp parallel for schedule(static) num_threads(num_threads)
-    for (long distIdx = 0; distIdx < qrDists.rows(); distIdx++)
-    {
+    for (long distIdx = 0; distIdx < qrDists.rows(); distIdx++) {
       unsigned long i = distIdx % nrrSamples;
       unsigned long j = distIdx / nrrSamples + nrrSamples;
       squareDists(i, j) = qrDists[distIdx];
@@ -184,10 +136,8 @@ NumpyMatrix long_to_square(const Eigen::VectorXf &rrDists,
 }
 
 Eigen::VectorXf square_to_long(const NumpyMatrix &squareDists,
-                               const unsigned int num_threads)
-{
-  if (squareDists.rows() != squareDists.cols())
-  {
+                               const unsigned int num_threads) {
+  if (squareDists.rows() != squareDists.cols()) {
     throw std::runtime_error("square_to_long input must be a square matrix");
   }
 
@@ -198,10 +148,8 @@ Eigen::VectorXf square_to_long(const NumpyMatrix &squareDists,
 // due to reverse direction
 // guided schedules inversely proportional to outer index
 #pragma omp parallel for schedule(guided, 1) num_threads(num_threads)
-  for (long i = n - 2; i >= 0; i--)
-  {
-    for (long j = i + 1; j < n; j++)
-    {
+  for (long i = n - 2; i >= 0; i--) {
+    for (long j = i + 1; j < n; j++) {
       longDists(square_to_condensed(i, j, n)) = squareDists(i, j);
     }
   }
diff --git a/src/dist/matrix_types.hpp b/src/dist/matrix_types.hpp
new file mode 100644
index 00000000..c7599b6f
--- /dev/null
+++ b/src/dist/matrix_types.hpp
@@ -0,0 +1,10 @@
+#pragma once
+
+// This file can be used by gcc and nvcc (i.e. no Eigen)
+// Compound types used by many functions in api, bindings etc
+
+#include <vector>
+#include <tuple>
+
+using sparse_coo = std::tuple<std::vector<long>, std::vector<long>, std::vector<float>>;
+using network_coo = std::tuple<std::vector<long>, std::vector<long>, std::vector<long>>;
\ No newline at end of file
diff --git a/src/gpu/align.hpp b/src/gpu/align.hpp
new file mode 100644
index 00000000..0a17ca2c
--- /dev/null
+++ b/src/gpu/align.hpp
@@ -0,0 +1,13 @@
+#pragma once
+
+// Align structs
+// https://stackoverflow.com/a/12779757
+#if defined(__CUDACC__) // NVCC
+#define ALIGN(n) __align__(n)
+#elif defined(__GNUC__) // GCC
+#define ALIGN(n) __attribute__((aligned(n)))
+#elif defined(_MSC_VER) // MSVC
+#define ALIGN(n) __declspec(align(n))
+#else
+#error "Please provide a definition for MY_ALIGN macro for your host compiler!"
+#endif
diff --git a/src/gpu/containers.cuh b/src/gpu/containers.cuh
new file mode 100644
index 00000000..4180b1fd
--- /dev/null
+++ b/src/gpu/containers.cuh
@@ -0,0 +1,176 @@
+#pragma once
+
+#include <vector>
+
+#include "cuda.cuh"
+
+template <typename T> class device_array {
+public:
+  // Default constructor
+  device_array() : data_(nullptr), size_(0) {}
+
+  // Constructor to allocate empty memory
+  device_array(const size_t size) : size_(size) {
+    CUDA_CALL(cudaMalloc((void **)&data_, size_ * sizeof(T)));
+    CUDA_CALL(cudaMemset(data_, 0, size_ * sizeof(T)));
+  }
+
+  // Constructor from vector
+  device_array(const std::vector<T> &data) : size_(data.size()) {
+    CUDA_CALL(cudaMalloc((void **)&data_, size_ * sizeof(T)));
+    CUDA_CALL(
+        cudaMemcpy(data_, data.data(), size_ * sizeof(T), cudaMemcpyDefault));
+  }
+
+  // Copy
+  device_array(const device_array &other) : size_(other.size_) {
+    CUDA_CALL(cudaMalloc((void **)&data_, size_ * sizeof(T)));
+    CUDA_CALL(
+        cudaMemcpy(data_, other.data_, size_ * sizeof(T), cudaMemcpyDefault));
+  }
+
+  // Copy assign
+  device_array &operator=(const device_array &other) {
+    if (this != &other) {
+      size_ = other.size_;
+      CUDA_CALL(cudaFree(data_));
+      CUDA_CALL(cudaMalloc((void **)&data_, size_ * sizeof(T)));
+      CUDA_CALL(
+          cudaMemcpy(data_, other.data_, size_ * sizeof(T), cudaMemcpyDefault));
+    }
+    return *this;
+  }
+
+  // Move
+  device_array(device_array &&other) : data_(nullptr), size_(0) {
+    data_ = other.data_;
+    size_ = other.size_;
+    other.data_ = nullptr;
+    other.size_ = 0;
+  }
+
+  // Move assign
+  device_array &operator=(device_array &&other) {
+    if (this != &other) {
+      CUDA_CALL(cudaFree(data_));
+      data_ = other.data_;
+      size_ = other.size_;
+      other.data_ = nullptr;
+      other.size_ = 0;
+    }
+    return *this;
+  }
+
+  ~device_array() { CUDA_CALL_NOTHROW(cudaFree(data_)); }
+
+  void swap(device_array& other) {
+    T* data_tmp = other.data_;
+    size_t size_tmp = other.size_;
+    other.data_ = data_;
+    other.size_ = size_;
+    data_ = data_tmp;
+    size_ = size_tmp;
+  }
+
+  void get_array(std::vector<T> &dst) const {
+    CUDA_CALL(cudaMemcpy(dst.data(), data_, dst.size() * sizeof(T),
+                         cudaMemcpyDefault));
+  }
+
+  void set_array(const std::vector<T> &src) {
+    size_ = src.size();
+    CUDA_CALL(
+        cudaMemcpy(data_, src.data(), size_ * sizeof(T), cudaMemcpyDefault));
+  }
+
+  void set_array(const T *src) {
+    CUDA_CALL(cudaMemcpy(data_, src, size_ * sizeof(T), cudaMemcpyDefault));
+  }
+
+  // Special functions for working with streams/graphs
+  // (prefer above functions unless using these)
+  void get_array_async(T *value, const size_t size, cudaStream_t stream) const {
+    if (size > size_) {
+      throw std::runtime_error("Size of host memory too big to copy");
+    }
+    CUDA_CALL(cudaMemcpyAsync((void *)value, data_, sizeof(T) * size,
+                              cudaMemcpyDefault, stream));
+  }
+
+  void set_array_async(const T *value, const size_t size, cudaStream_t stream) {
+    if (size > size_) {
+      throw std::runtime_error("Size of host memory too big to copy");
+    }
+    CUDA_CALL(cudaMemcpyAsync(data_, (void *)value, sizeof(T) * size,
+                              cudaMemcpyDefault, stream));
+  }
+
+  T *data() { return data_; }
+
+  size_t size() const { return size_; }
+
+private:
+  T *data_;
+  size_t size_;
+};
+
+// Specialisation of the above for void* memory needed by some cub functions
+// Construct once and use set_size() to modify
+template <> class device_array<void> {
+public:
+  // Default constructor
+  device_array() : data_(nullptr), size_(0) {}
+  // Constructor to allocate empty memory
+  device_array(const size_t size) : size_(size) {
+    if (size_ > 0) {
+      CUDA_CALL(cudaMalloc((void **)&data_, size_));
+    }
+  }
+
+  ~device_array() { CUDA_CALL_NOTHROW(cudaFree(data_)); }
+
+  void set_size(size_t size) {
+    if (size != size_) {
+      size_ = size;
+      CUDA_CALL(cudaFree(data_));
+      if (size_ > 0) {
+        CUDA_CALL(cudaMalloc((void **)&data_, size_));
+      } else {
+        data_ = nullptr;
+      }
+    }
+  }
+
+  void *data() { return data_; }
+
+  size_t size() const { return size_; }
+
+private:
+  device_array(const device_array<void> &) = delete;
+  device_array(device_array<void> &&) = delete;
+
+  void *data_;
+  size_t size_;
+};
+
+class cuda_stream {
+public:
+  cuda_stream() { CUDA_CALL(cudaStreamCreate(&stream_)); }
+
+  ~cuda_stream() {
+    if (stream_ != nullptr) {
+      CUDA_CALL_NOTHROW(cudaStreamDestroy(stream_));
+    }
+  }
+
+  cudaStream_t stream() { return stream_; }
+
+  void sync() { CUDA_CALL(cudaStreamSynchronize(stream_)); }
+
+private:
+  // Delete copy and move
+  cuda_stream(const cuda_stream &) = delete;
+  cuda_stream(cuda_stream &&) = delete;
+
+  cudaStream_t stream_;
+};
diff --git a/src/gpu/cuda.cuh b/src/gpu/cuda.cuh
index b9cb87af..57fe6da0 100644
--- a/src/gpu/cuda.cuh
+++ b/src/gpu/cuda.cuh
@@ -3,11 +3,16 @@
 #include <stdexcept>
 #include <stdio.h>
 
+#include <type_traits>
+static_assert(__CUDACC_VER_MAJOR__ >= 11, "CUDA >=11.0 required");
+
 #include <cuda.h>
 #include <cuda_runtime.h>
 #include <device_launch_parameters.h>
 
-const int progressBitshift = 10; // Update every 2^10 = 1024 dists
+#include "align.hpp"
+
+const int progress_blocks = 1000; // Update 1000
 
 static void HandleCUDAError(const char *file, int line,
                             cudaError_t status = cudaGetLastError()) {
@@ -28,32 +33,52 @@ static void HandleCUDAError(const char *file, int line,
 #define CUDA_CALL(err) (HandleCUDAError(__FILE__, __LINE__, err))
 #define CUDA_CALL_NOTHROW( err ) (err)
 
-// A bit lazy... should be a class and just use the destructor
-struct progress_atomics {
+struct ALIGN(16) progress_ptrs {
   volatile int *blocks_complete;
   bool *kill_kernel;
+};
 
-  void init() {
-    CUDA_CALL(cudaMallocManaged(&blocks_complete, sizeof(int)));
-    CUDA_CALL(cudaMallocManaged(&kill_kernel, sizeof(bool)));
-    *blocks_complete = 0;
-    *kill_kernel = false;
+class progress_atomics {
+public:
+  progress_atomics() {
+    CUDA_CALL(cudaMallocManaged((void**)&(managed_ptrs.blocks_complete), sizeof(int)));
+    CUDA_CALL(cudaMallocManaged((void**)&(managed_ptrs.kill_kernel), sizeof(bool)));
+    *(managed_ptrs.blocks_complete) = 0;
+    *(managed_ptrs.kill_kernel) = false;
   }
 
-  void free() {
-    CUDA_CALL(cudaFree((void *)blocks_complete));
-    CUDA_CALL(cudaFree(kill_kernel));
+  ~progress_atomics() {
+    CUDA_CALL(cudaFree((void *)managed_ptrs.blocks_complete));
+    CUDA_CALL(cudaFree(managed_ptrs.kill_kernel));
   }
+
+  void set_kill() {
+    *(managed_ptrs.kill_kernel) = true;
+  }
+
+  progress_ptrs get_ptrs() {
+    return managed_ptrs;
+  }
+
+  __host__
+  int complete() {
+    return *(managed_ptrs.blocks_complete);
+  }
+
+private:
+  progress_atomics(const progress_atomics &) = delete;
+  progress_atomics(progress_atomics &&) = delete;
+
+  progress_ptrs managed_ptrs;
 };
 
 // Use atomic add to update a counter, so progress works regardless of
 // dispatch order
 __device__ inline void update_progress(long long dist_idx, long long dist_n,
-                                       progress_atomics progress) {
+                                       progress_ptrs &progress) {
   // Progress indicator
-  // The >> progressBitshift is a divide by 1024 - update roughly every 0.1%
-  if (dist_idx % (dist_n >> progressBitshift) == 0) {
-    if (*(progress.kill_kernel) == true) {
+  if (dist_idx % (dist_n / progress_blocks) == 0) {
+    if (*(progress.kill_kernel)) {
       __trap();
     }
     atomicAdd((int *)progress.blocks_complete, 1);
diff --git a/src/gpu/device_memory.cuh b/src/gpu/device_memory.cuh
index 2795010d..563d1cfa 100644
--- a/src/gpu/device_memory.cuh
+++ b/src/gpu/device_memory.cuh
@@ -1,26 +1,15 @@
 #pragma once
 
+#include "align.hpp"
 #include "reference.hpp"
 
-// Align structs
-// https://stackoverflow.com/a/12779757
-#if defined(__CUDACC__) // NVCC
-#define ALIGN(n) __align__(n)
-#elif defined(__GNUC__) // GCC
-#define ALIGN(n) __attribute__((aligned(n)))
-#elif defined(_MSC_VER) // MSVC
-#define ALIGN(n) __declspec(align(n))
-#else
-#error "Please provide a definition for MY_ALIGN macro for your host compiler!"
-#endif
-
 struct ALIGN(8) RandomStrides {
   size_t kmer_stride;
   size_t cluster_inner_stride;
   size_t cluster_outer_stride;
 };
 
-typedef std::tuple<RandomStrides, std::vector<float>> FlatRandom;
+using FlatRandom = std::tuple<RandomStrides, std::vector<float>>;
 
 #ifdef GPU_AVAILABLE
 
diff --git a/src/gpu/dist.cu b/src/gpu/dist.cu
index eea5c21e..ff3fa66e 100644
--- a/src/gpu/dist.cu
+++ b/src/gpu/dist.cu
@@ -9,6 +9,7 @@
 // std
 #include <algorithm>
 #include <assert.h>
+#include <cfloat>
 #include <cmath>
 #include <cstdint>
 #include <iomanip>
@@ -23,14 +24,14 @@
 namespace py = pybind11;
 
 // memcpy_async
-#if __CUDACC_VER_MAJOR__ >= 11
+#include <cub/cub.cuh>
 #include <cuda/barrier>
 #include <cooperative_groups.h>
-#pragma diag_suppress static_var_with_dynamic_init
-#endif
+#pragma nv_diag_suppress static_var_with_dynamic_init
 
 // internal headers
 #include "cuda.cuh"
+#include "containers.cuh"
 #include "dist/matrix_idx.hpp"
 #include "gpu.hpp"
 #include "sketch/bitfuncs.hpp"
@@ -85,15 +86,15 @@ __device__ float jaccard_dist(const uint64_t *sketch1, const uint64_t *sketch2,
 // Simple linear regression, exact solution
 // Avoids use of dynamic memory allocation on device, or
 // linear algebra libraries
-__device__ void simple_linear_regression(float *core_dist,
-                                         float *accessory_dist,
+__device__ void simple_linear_regression(float dists[],
                                          const float xsum, const float ysum,
                                          const float xysum,
                                          const float xsquaresum,
-                                         const float ysquaresum, const int n) {
+                                         const float ysquaresum, const int n,
+                                         const int dist_col) {
   if (n < 2) {
-    *core_dist = 0;
-    *accessory_dist = 0;
+    dists[0] = 0.0f;
+    dists[1] = 0.0f;
   } else {
     // CUDA fast-math intrinsics on floats, which give comparable accuracy
     // Speed gain is fairly minimal, as most time spent on Jaccard distance
@@ -113,16 +114,16 @@ __device__ void simple_linear_regression(float *core_dist,
     // Store core/accessory in dists, truncating at zero
     // Memory should be initialised to zero so else block not strictly
     // necessary, but better safe than sorry!
-    if (beta < 0) {
-      *core_dist = 1 - __expf(beta);
+    if (beta < 0.0f) {
+      dists[0] = 1.0f - __expf(beta);
     } else {
-      *core_dist = 0;
+      dists[0] = 0.0f;
     }
 
-    if (alpha < 0) {
-      *accessory_dist = 1 - __expf(alpha);
+    if (alpha < 0.0f) {
+      dists[1] = 1.0f - __expf(alpha);
     } else {
-      *accessory_dist = 0;
+      dists[1] = 0.0f;
     }
   }
 }
@@ -136,6 +137,30 @@ __device__ void simple_linear_regression(float *core_dist,
 // Main kernel functions run on the device,
 // but callable from the host
 
+__global__ void set_idx(long* idx, size_t row_samples, size_t col_samples, size_t col_offset) {
+  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < row_samples * col_samples;
+    i += blockDim.x * gridDim.x) {
+    idx[i] = col_offset + i % col_samples;
+  }
+}
+
+__global__ void copy_top_k(float* sorted_dists, long* sorted_idx,
+  float* all_sorted_dists, long* all_sorted_idx, int segment_size, int n_out,
+    int kNN, int sketch_block_idx, int n_chunks, bool second_sort) {
+  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n_out;
+    i += blockDim.x * gridDim.x) {
+    const int offset_in = (i / kNN) * segment_size + i % kNN;
+    int offset_out = i;
+    // If copying from the sorted kNN * n_chunk list into the final sparse matrix
+    if (!second_sort) {
+      // If copying from the sorted n_chunk dense list into the kNN * n_chunk list
+      offset_out += ((i / kNN) * (kNN * (n_chunks - 1))) + (kNN * sketch_block_idx);
+    }
+    all_sorted_dists[offset_out] = sorted_dists[offset_in];
+    all_sorted_idx[offset_out] = sorted_idx[offset_in];
+  }
+}
+
 __global__ void calculate_dists(
     const bool self, const uint64_t *ref, const long ref_n,
     const uint64_t *query, const long query_n, const int *kmers,
@@ -143,7 +168,8 @@ __global__ void calculate_dists(
     const float *random_table, const uint16_t *ref_idx_lookup,
     const uint16_t *query_idx_lookup, const SketchStrides ref_strides,
     const SketchStrides query_strides, const RandomStrides random_strides,
-    progress_atomics progress, const bool use_shared) {
+    progress_ptrs progress, const bool use_shared,
+    const int dist_col, const bool max_diagonal) {
   // Calculate indices for query, ref and results
   int ref_idx, query_idx, dist_idx;
   if (self) {
@@ -181,11 +207,11 @@ __global__ void calculate_dists(
 
   // Calculate Jaccard distances over k-mer lengths
   int kmer_used = 0;
-  float xsum = 0;
-  float ysum = 0;
-  float xysum = 0;
-  float xsquaresum = 0;
-  float ysquaresum = 0;
+  float xsum = 0.0f;
+  float ysum = 0.0f;
+  float xysum = 0.0f;
+  float xsquaresum = 0.0f;
+  float ysquaresum = 0.0f;
   bool stop = false;
   for (int kmer_idx = 0; kmer_idx < kmer_n; kmer_idx++) {
     // Copy query sketch into __shared__ mem
@@ -199,27 +225,21 @@ __global__ void calculate_dists(
     const uint64_t *query_ptr;
     extern __shared__ uint64_t query_shared[];
     int query_bin_strides;
-#if __CUDACC_VER_MAJOR__ >= 11
     auto block = cooperative_groups::this_thread_block();
     __shared__ cuda::barrier<cuda::thread_scope::thread_scope_block> barrier;
     if (block.thread_rank() == 0) {
       init(&barrier, block.size()); // Friend function initializes barrier
     }
     block.sync();
-#endif
     if (use_shared) {
       size_t sketch_bins = query_strides.bbits * query_strides.sketchsize64;
       size_t sketch_stride = query_strides.bin_stride;
       if (threadIdx.x < warp_size) {
         for (int lidx = threadIdx.x; lidx < sketch_bins; lidx += warp_size) {
-#if __CUDACC_VER_MAJOR__ >= 11
           cuda::memcpy_async(query_shared + lidx,
                              query_start + (lidx * sketch_stride),
                              sizeof(uint64_t),
                              barrier);
-#else
-          query_shared[lidx] = query_start[lidx * sketch_stride];
-#endif
         }
       }
       query_ptr = query_shared;
@@ -228,11 +248,7 @@ __global__ void calculate_dists(
       query_ptr = query_start;
       query_bin_strides = query_strides.bin_stride;
     }
-#if __CUDACC_VER_MAJOR__ >= 11
     barrier.arrive_and_wait();
-#else
-    __syncthreads();
-#endif
 
     // Some threads at the end of the last block will have nothing to do
     // Need to have conditional here to avoid block on __syncthreads() above
@@ -280,8 +296,21 @@ __global__ void calculate_dists(
 
   if (ref_idx < ref_n) {
     // Run the regression, and store results in dists
-    simple_linear_regression(dists + dist_idx, dists + dist_n + dist_idx, xsum,
-                             ysum, xysum, xsquaresum, ysquaresum, kmer_used);
+    float fitted_dists[2];
+    // Set diagonal if you wish to ignore diagonals (set them to max)
+    if (max_diagonal && ref_idx == query_idx) {
+      fitted_dists[0] = INFINITY;
+      fitted_dists[1] = INFINITY;
+    } else {
+      simple_linear_regression(fitted_dists, xsum, ysum, xysum, xsquaresum,
+                               ysquaresum, kmer_used, dist_col);
+    }
+    if (dist_col < 0) {
+      dists[dist_idx] = fitted_dists[0];
+      dists[dist_idx + dist_n] = fitted_dists[1];
+    } else {
+      dists[dist_idx] = fitted_dists[dist_col];
+    }
 
     update_progress(dist_idx, dist_n, progress);
   }
@@ -304,7 +333,7 @@ std::tuple<size_t, size_t> getBlockSize(const size_t ref_samples,
   // We take the next multiple of 32 that is larger than the number of
   // reference sketches, up to a maximum of 512
   size_t blockSize =
-      std::min(512, 32 * static_cast<int>((ref_samples + 32 - 1) / 32));
+      std::min(256, 32 * static_cast<int>((ref_samples + 32 - 1) / 32));
   size_t blockCount = 0;
   if (self) {
     for (int i = 0; i < ref_samples; i++) {
@@ -319,18 +348,18 @@ std::tuple<size_t, size_t> getBlockSize(const size_t ref_samples,
 
 // Writes a progress meter using the device int which keeps
 // track of completed jobs
-void reportDistProgress(progress_atomics progress, long long dist_rows) {
-  long long progress_blocks = 1 << progressBitshift;
+void reportDistProgress(progress_atomics& progress, long long dist_rows) {
   int now_completed = 0;
   float kern_progress = 0;
   if (dist_rows > progress_blocks) {
     while (now_completed < progress_blocks - 1) {
       if (PyErr_CheckSignals() != 0) {
-        *(progress.kill_kernel) = true;
+        progress.set_kill();
         throw py::error_already_set();
       }
-      if (*(progress.blocks_complete) > now_completed) {
-        now_completed = *(progress.blocks_complete);
+      int complete = progress.complete();
+      if (complete > now_completed) {
+        now_completed = complete;
         kern_progress = now_completed / (float)progress_blocks;
         fprintf(stderr, "%cProgress (GPU): %.1lf%%", 13, kern_progress * 100);
       } else {
@@ -354,6 +383,24 @@ std::tuple<size_t, size_t, size_t> initialise_device(const int device_id) {
       std::make_tuple(mem_free, mem_total, static_cast<size_t>(shared_size)));
 }
 
+std::tuple<bool, size_t> check_shared_size(const SketchStrides& strides, const size_t shared_size) {
+  size_t sketch_size_bytes =
+    strides.sketchsize64 * strides.bbits * sizeof(uint64_t);
+  bool use_shared = true;
+  if (sketch_size_bytes > shared_size) {
+    std::cerr << "You are using a large sketch size, which may slow down "
+                "computation on this device"
+              << std::endl;
+    std::cerr << "Reduce sketch size to "
+              << std::floor(64 * shared_size /
+                            (strides.bbits * sizeof(uint64_t)))
+              << " or less for better performance" << std::endl;
+    sketch_size_bytes = 0;
+    use_shared = false;
+  }
+  return std::make_tuple(use_shared, sketch_size_bytes);
+}
+
 // Main function to run the distance calculations, reading/writing into
 // device_arrays Cache preferences: Upper dist memory access is hard to predict,
 // so try and cache as much as possible Query uses on-chip cache (__shared__) to
@@ -374,7 +421,6 @@ std::vector<float> dispatchDists(std::vector<Reference> &ref_sketches,
 
   // Progress meter
   progress_atomics progress;
-  progress.init();
 
   RandomStrides random_strides = std::get<0>(flat_random);
   long long dist_rows;
@@ -391,22 +437,14 @@ std::vector<float> dispatchDists(std::vector<Reference> &ref_sketches,
                              ref_random_idx, query_random_idx, kmer_lengths,
                              dist_rows, self, cpu_threads);
 
-  size_t sketch_size_bytes =
-      query_strides.sketchsize64 * query_strides.bbits * sizeof(uint64_t);
-  bool use_shared = true;
-  if (sketch_size_bytes > shared_size) {
-    std::cerr << "You are using a large sketch size, which may slow down "
-                 "computation on this device"
-              << std::endl;
-    std::cerr << "Reduce sketch size to "
-              << std::floor(64 * shared_size /
-                            (query_strides.bbits * sizeof(uint64_t)))
-              << " or less for better performance" << std::endl;
-    sketch_size_bytes = 0;
-    use_shared = false;
-  }
+  bool use_shared;
+  size_t shared_size_bytes;
+  std::tie(use_shared, shared_size_bytes) =
+    check_shared_size(query_strides, shared_size);
 
   size_t blockSize, blockCount;
+  bool max_diagonal = false;
+  int dist_col = -1;
   if (self) {
     std::tie(blockSize, blockCount) = getBlockSize(
         sketch_subsample.ref_size, sketch_subsample.ref_size, dist_rows, self);
@@ -414,13 +452,13 @@ std::vector<float> dispatchDists(std::vector<Reference> &ref_sketches,
     // Third argument is the size of __shared__ memory needed by a thread block
     // This is equal to the query sketch size in bytes (at a single k-mer
     // length)
-    calculate_dists<<<blockCount, blockSize, sketch_size_bytes>>>(
+    calculate_dists<<<blockCount, blockSize, shared_size_bytes>>>(
         self, device_arrays.ref_sketches(), sketch_subsample.ref_size,
         device_arrays.ref_sketches(), sketch_subsample.ref_size,
         device_arrays.kmers(), kmer_lengths.size(), device_arrays.dist_mat(),
         dist_rows, device_arrays.random_table(), device_arrays.ref_random(),
         device_arrays.ref_random(), ref_strides, ref_strides, random_strides,
-        progress, use_shared);
+        progress.get_ptrs(), use_shared, dist_col, max_diagonal);
   } else {
     std::tie(blockSize, blockCount) =
         getBlockSize(sketch_subsample.ref_size, sketch_subsample.query_size,
@@ -429,20 +467,19 @@ std::vector<float> dispatchDists(std::vector<Reference> &ref_sketches,
     // Third argument is the size of __shared__ memory needed by a thread block
     // This is equal to the query sketch size in bytes (at a single k-mer
     // length)
-    calculate_dists<<<blockCount, blockSize, sketch_size_bytes>>>(
+    calculate_dists<<<blockCount, blockSize, shared_size_bytes>>>(
         self, device_arrays.ref_sketches(), sketch_subsample.ref_size,
         device_arrays.query_sketches(), sketch_subsample.query_size,
         device_arrays.kmers(), kmer_lengths.size(), device_arrays.dist_mat(),
         dist_rows, device_arrays.random_table(), device_arrays.ref_random(),
         device_arrays.query_random(), ref_strides, query_strides,
-        random_strides, progress, use_shared);
+        random_strides, progress.get_ptrs(), use_shared, dist_col, max_diagonal);
   }
 
   // Check for error in kernel launch
   CUDA_CALL(cudaGetLastError());
   reportDistProgress(progress, dist_rows);
   fprintf(stderr, "%cProgress (GPU): 100.0%%\n", 13);
-  progress.free();
 
   // Copy results back to host
   CUDA_CALL(cudaDeviceSynchronize());
@@ -450,3 +487,305 @@ std::vector<float> dispatchDists(std::vector<Reference> &ref_sketches,
 
   return (dist_results);
 }
+
+
+// Function which sparsifies distances on the fly. Distances are calculated in
+// blocks, sorted and top k stored.
+// NB cuda graph not needed as API calls faster than ops here
+sparse_coo sparseDists(const dist_params params,
+  const std::vector<std::vector<uint64_t>> &ref_sketches,
+  const std::vector<SketchStrides> &ref_strides,
+  const FlatRandom &flat_random,
+  const std::vector<uint16_t> &ref_random_idx,
+  const std::vector<size_t> &kmer_lengths,
+  const int kNN,
+  const size_t dist_col,
+  const size_t samples_per_chunk,
+  const size_t num_big_chunks,
+  const int cpu_threads) {
+  /*
+   *
+   *   Device setup
+   *
+   */
+  CUDA_CALL(cudaDeviceSetCacheConfig(cudaFuncCachePreferL1));
+  CUDA_CALL(cudaDeviceSetSharedMemConfig(cudaSharedMemBankSizeEightByte));
+
+  // Progress meter (not used for printing in this function)
+  progress_atomics progress;
+
+  // const parameters for functions
+  const size_t n_chunks = ref_sketches.size();
+  const bool self = false;
+  const double total_blocks = static_cast<double>(n_chunks) * n_chunks;
+  const size_t idx_blockSize = 64;
+  const size_t copy_blockSize = 64;
+  const int begin_sort_bit = 0;
+  const int end_sort_bit = 8 * sizeof(float);
+
+  bool use_shared;
+  size_t shared_size_bytes;
+  std::tie(use_shared, shared_size_bytes) =
+    check_shared_size(ref_strides[0], params.shared_size);
+
+  /*
+   *
+   *   Memory allocations
+   *   All use the big_chunk size, for smaller chunks the end of the alloc
+   *   is not used
+   *
+   */
+
+  // Storage for results on host
+  std::vector<float> host_dists(params.n_samples * kNN);
+  std::vector<long> i_vec(params.n_samples * kNN);
+  std::vector<long> j_vec(params.n_samples * kNN);
+
+  //   sketch block 1
+  //   sketch block 2
+  //   sketch block 3
+  //   sketch block 4
+  device_array<uint64_t> block1(ref_sketches[0].size());
+  device_array<uint64_t> block2(block1.size());
+  device_array<uint64_t> block3(block1.size());
+  device_array<uint64_t> block4(block1.size());
+  //   dists (chunk size^2)
+  //   dists idx (dists.size())
+  device_array<float> dists((samples_per_chunk + 1) * (samples_per_chunk + 1));
+  device_array<long> dists_idx(dists.size());
+
+  // inner loop sorting
+  //   sorted dists (dists.size())
+  //   sorted dists idx (dists.size())
+  //   offsets
+  //   tmp space
+  device_array<float> sorted_dists(dists.size());
+  device_array<long> sorted_dists_idx(dists.size());
+  device_array<int> dist_partitions((samples_per_chunk + 1) + 1);
+  device_array<void> dist_sort_tmp;
+
+  // inner loop kNN pick (copy_top_k)
+  //   sorted dists
+  //   sorted dists idx
+  device_array<float> all_sorted_dists(kNN * n_chunks * (samples_per_chunk + 1));
+  device_array<long> all_sorted_dists_idx(all_sorted_dists.size());
+
+  // outer loop sorting
+  //   doubly sorted dists (staging for final results, copied to host)
+  //   doubly sorted dists idx (as above)
+  //   offsets
+  //   tmp space
+  device_array<float> doubly_sorted_dists(all_sorted_dists.size());
+  device_array<long> doubly_sorted_dists_idx(all_sorted_dists.size());
+  std::vector<int> host_partitions;
+  for (int partition_idx = 0; partition_idx < dist_partitions.size(); ++partition_idx) {
+    host_partitions.push_back(partition_idx * kNN * n_chunks);
+  }
+  device_array<int> second_sort_partitions(host_partitions);
+  host_partitions.clear();
+  device_array<void> outer_dist_sort_tmp;
+
+  // out loop kNN pick (copy_top_k)
+  //   final dists (staging for final results, copied to host)
+  //   final dists idx (as above)
+  device_array<float> final_dists(kNN * (samples_per_chunk + 1));
+  device_array<long> final_dists_idx(kNN * (samples_per_chunk + 1));
+
+  //   kmers
+  const std::vector<int> kmer_ints(kmer_lengths.begin(), kmer_lengths.end());
+  device_array<int> kmers(kmer_ints);
+
+  //   random
+  RandomStrides random_strides = std::get<0>(flat_random);
+  device_array<float> random_table(std::get<1>(flat_random));
+  device_array<uint16_t> random_idx(ref_random_idx);
+
+  /*
+   *
+   *   Memory setup
+   *   Set i_vec on host
+   *   Register sketches in host memory
+   *   Copy first sketch blocks to device
+   *
+   */
+
+  // Set i_vec (it just needs to be 0, 0, 0..., 1, 1, 1...)
+  #pragma omp parallel for num_threads(cpu_threads)
+  for (size_t sample_idx = 0; sample_idx < params.n_samples; ++sample_idx) {
+    std::fill_n(i_vec.begin() + sample_idx * kNN, kNN, sample_idx);
+  }
+
+  // Register sketches on host so they can be copied async
+  for (size_t chunk_idx = 0; chunk_idx < n_chunks; ++chunk_idx) {
+    CUDA_CALL(cudaHostRegister((void *)ref_sketches[chunk_idx].data(),
+                               ref_sketches[chunk_idx].size() * sizeof(uint64_t),
+                               cudaHostRegisterReadOnly));
+  }
+
+  // Four CUDA streams used in loops
+  cuda_stream dist_stream, idx_stream, mem_stream, sort_stream;
+
+  // Copy first set of refs in. Block 4 is used to store this permanently to stop
+  // two copies being needed when entering a new row
+  block1.set_array_async(ref_sketches[0].data(), ref_sketches[0].size(), mem_stream.stream());
+  block4.set_array_async(block1.data(), block1.size(), mem_stream.stream());
+
+  /*
+   *
+   *   Loop over chunks
+   *   Outer loop over chunks of rows
+   *   Inner loop over chunks of columns
+   *
+   */
+  fprintf(stderr, "%cProgress (GPU): %.1lf%%", 13, 0.0f);
+
+  // OUTER LOOP over n_chunks lots of refs
+  size_t row_offset = 0;
+  for (size_t row_chunk_idx = 0; row_chunk_idx < n_chunks; ++row_chunk_idx) {
+    size_t row_samples = samples_per_chunk + (row_chunk_idx < num_big_chunks ? 1 : 0);
+    size_t col_offset = 0;
+
+    //  INNER LOOP over n_chunks lots of queries
+    for (size_t col_chunk_idx = 0; col_chunk_idx < n_chunks; ++col_chunk_idx) {
+      size_t col_samples = samples_per_chunk + (col_chunk_idx < num_big_chunks ? 1 : 0);
+      size_t dist_rows = row_samples * col_samples;
+
+      // Check for interrupts
+      if (PyErr_CheckSignals() != 0) {
+        throw py::error_already_set();
+      }
+
+      // Only need to set new sort partitions if moving to a smaller chunk
+      if (host_partitions.size() != col_samples + 1) {
+        host_partitions.clear();
+        for (int partition_idx = 0; partition_idx < col_samples + 1; ++partition_idx) {
+          host_partitions.push_back(partition_idx * col_samples);
+        }
+        dist_partitions.set_array_async(host_partitions.data(), host_partitions.size(), sort_stream.stream());
+      }
+
+      //    (stream 1 async) Run dists on 1 vs 2
+      size_t blockSize, blockCount;
+      std::tie(blockSize, blockCount) =
+        getBlockSize(row_samples, col_samples, dist_rows, self);
+      uint64_t* query_ptr = col_chunk_idx == 0 ? block4.data() : block2.data();
+      bool max_diagonal = col_chunk_idx == row_chunk_idx;
+      mem_stream.sync();
+      // NB in calculate_dists ref idx changes fastest (so should be the column)
+      // so ref and query are 'backwards'
+      calculate_dists<<<blockCount, blockSize, shared_size_bytes, dist_stream.stream()>>>(
+        self,
+        query_ptr, col_samples,
+        block1.data(), row_samples,
+        kmers.data(), kmers.size(),
+        dists.data(), dist_rows,
+        random_table.data(), random_idx.data() + col_offset, random_idx.data() + row_offset,
+        ref_strides[col_chunk_idx], ref_strides[row_chunk_idx],
+        random_strides, progress.get_ptrs(), use_shared, dist_col, max_diagonal
+      );
+
+      //    (stream 2 async) Load next into block 3
+      //    swap ptrs for block 2 <-> 3
+      if (col_chunk_idx + 1 < n_chunks) {
+        block3.set_array_async(ref_sketches[col_chunk_idx + 1].data(), ref_sketches[col_chunk_idx + 1].size(), mem_stream.stream());
+        block2.swap(block3);
+      } else if (row_chunk_idx + 1 < n_chunks) {
+        block3.set_array_async(ref_sketches[row_chunk_idx + 1].data(), ref_sketches[row_chunk_idx + 1].size(), mem_stream.stream());
+        block1.swap(block3);
+      }
+
+      //    (stream 3 async) Set dist idx via kernel
+      const size_t idx_blockCount = (dist_rows + idx_blockSize - 1) / idx_blockSize;
+      set_idx<<<idx_blockCount, idx_blockSize, 0, idx_stream.stream()>>>(
+        dists_idx.data(), row_samples, col_samples, col_offset
+      );
+
+      //    (stream 4) cub::DeviceSegmentedRadixSort::SortPairs on dists, dists idx -> sorted dists, sorted dists idx
+      const int num_items = dist_rows;
+      const int num_segments = host_partitions.size() - 1;
+      int *d_offsets = dist_partitions.data();
+      float *d_keys_in = dists.data();
+      float *d_keys_out = sorted_dists.data();
+      long *d_values_in = dists_idx.data();
+      long *d_values_out = sorted_dists_idx.data();
+      // Determine temporary device storage requirements (first run only)
+      if (row_chunk_idx == 0 && col_chunk_idx == 0) {
+        size_t temp_storage_bytes = 0;
+        cub::DeviceSegmentedRadixSort::SortPairs(dist_sort_tmp.data(), temp_storage_bytes,
+            d_keys_in, d_keys_out, d_values_in, d_values_out,
+            num_items, num_segments, d_offsets, d_offsets + 1,
+            begin_sort_bit, end_sort_bit, sort_stream.stream());
+        dist_sort_tmp.set_size(temp_storage_bytes);
+      }
+      // Run sorting operation
+      dist_stream.sync();
+      idx_stream.sync();
+      size_t temp_storage_bytes = dist_sort_tmp.size();
+      cub::DeviceSegmentedRadixSort::SortPairs(dist_sort_tmp.data(), temp_storage_bytes,
+          d_keys_in, d_keys_out, d_values_in, d_values_out,
+          num_items, num_segments, d_offsets, d_offsets + 1,
+          begin_sort_bit, end_sort_bit, sort_stream.stream());
+
+      //    (stream 4) D->D copy of top kNN dists to start of dists
+      const bool second_sort = false;
+      const size_t dist_out_size = kNN * num_segments;
+      const size_t copy_blockCount = (dist_out_size + copy_blockSize - 1) / copy_blockSize;
+      copy_top_k<<<copy_blockCount, copy_blockSize, 0, sort_stream.stream()>>>(
+        sorted_dists.data(), sorted_dists_idx.data(), all_sorted_dists.data(), all_sorted_dists_idx.data(),
+        col_samples, dist_out_size, kNN, col_chunk_idx, n_chunks, second_sort
+      );
+
+      // Update progress
+      col_offset += col_samples;
+      const size_t blocks_done = row_chunk_idx * n_chunks + col_chunk_idx + 1;
+      fprintf(stderr, "%cProgress (GPU): %.1lf%%", 13, 100 * blocks_done / total_blocks);
+    }
+
+    const int num_items = all_sorted_dists.size();
+    const int num_segments = row_samples;
+    int *d_offsets = second_sort_partitions.data();
+    float *d_keys_in = all_sorted_dists.data();
+    float *d_keys_out = doubly_sorted_dists.data();
+    long *d_values_in = all_sorted_dists_idx.data();
+    long *d_values_out = doubly_sorted_dists_idx.data();
+    // Determine temporary device storage requirements
+    if (row_chunk_idx == 0) {
+      size_t temp_storage_bytes = 0;
+      cub::DeviceSegmentedRadixSort::SortPairs(outer_dist_sort_tmp.data(), temp_storage_bytes,
+          d_keys_in, d_keys_out, d_values_in, d_values_out,
+          num_items, num_segments, d_offsets, d_offsets + 1,
+          begin_sort_bit, end_sort_bit, sort_stream.stream());
+      outer_dist_sort_tmp.set_size(temp_storage_bytes);
+    }
+    // Run sorting operation
+    size_t temp_storage_bytes = outer_dist_sort_tmp.size();
+    cub::DeviceSegmentedRadixSort::SortPairs(outer_dist_sort_tmp.data(), temp_storage_bytes,
+        d_keys_in, d_keys_out, d_values_in, d_values_out,
+        num_items, num_segments, d_offsets, d_offsets + 1,
+        begin_sort_bit, end_sort_bit, sort_stream.stream());
+
+    // take top kNN
+    const bool second_sort = true;
+    const int block_offset = 0;
+    const size_t dist_out_size = kNN * num_segments;
+    const size_t copy_blockCount = (dist_out_size + copy_blockSize - 1) / copy_blockSize;
+    copy_top_k<<<copy_blockCount, copy_blockSize, 0, sort_stream.stream()>>>(
+      doubly_sorted_dists.data(), doubly_sorted_dists_idx.data(), final_dists.data(), final_dists_idx.data(),
+      kNN * n_chunks, dist_out_size, kNN, block_offset, n_chunks, second_sort
+    );
+    // Copy chunk of results back to host
+    final_dists.get_array_async(host_dists.data() + row_offset * kNN, dist_out_size, sort_stream.stream());
+    final_dists_idx.get_array_async(j_vec.data() + row_offset * kNN, dist_out_size, sort_stream.stream());
+
+    row_offset += row_samples;
+  }
+  fprintf(stderr, "%cProgress (GPU): 100.0%%\n", 13);
+  CUDA_CALL(cudaDeviceSynchronize());
+
+  // Unregister host memory
+  for (size_t chunk_idx = 0; chunk_idx < n_chunks; ++chunk_idx) {
+    CUDA_CALL(cudaHostUnregister((void *)ref_sketches[chunk_idx].data()));
+  }
+
+  return (std::make_tuple(i_vec, j_vec, host_dists));
+}
diff --git a/src/gpu/gpu.hpp b/src/gpu/gpu.hpp
index 1604cb40..4661222b 100644
--- a/src/gpu/gpu.hpp
+++ b/src/gpu/gpu.hpp
@@ -14,10 +14,24 @@
 #include "device_memory.cuh"
 #include "gpu_countmin.cuh"
 #include "device_reads.cuh"
+#include "dist/matrix_types.hpp"
 
 static const int warp_size = 32;
 
 #ifdef GPU_AVAILABLE
+
+// Small struct used in cuda_dists_init
+struct ALIGN(16) dist_params {
+  bool self;
+  SketchStrides ref_strides;
+  SketchStrides query_strides;
+  long long dist_rows;
+  long n_samples;
+  size_t shared_size;
+  size_t mem_free;
+  size_t mem_total;
+};
+
 // defined in dist.cu
 std::tuple<size_t, size_t, size_t> initialise_device(const int device_id);
 
@@ -46,6 +60,18 @@ std::vector<float> dispatchDists(std::vector<Reference> &ref_sketches,
                                  const bool self, const int cpu_threads,
                                  const size_t shared_size);
 
+sparse_coo sparseDists(const dist_params params,
+  const std::vector<std::vector<uint64_t>> &ref_sketches,
+  const std::vector<SketchStrides> &ref_strides,
+  const FlatRandom &flat_random,
+  const std::vector<uint16_t> &ref_random_idx,
+  const std::vector<size_t> &kmer_lengths,
+  const int kNN,
+  const size_t dist_col,
+  const size_t samples_per_chunk,
+  const size_t num_big_chunks,
+  const int cpu_threads);
+
 // in sketch.cu
 void copyNtHashTablesToDevice();
 
diff --git a/src/gpu/gpu_api.cpp b/src/gpu/gpu_api.cpp
index 4329d5d3..8458c8eb 100644
--- a/src/gpu/gpu_api.cpp
+++ b/src/gpu/gpu_api.cpp
@@ -303,65 +303,90 @@ flatten_by_samples(const std::vector<Reference> &sketches,
   return flat_ref;
 }
 
-// Main function callable via API
-// Checks inputs
-// Flattens sketches
-// Copies flattened sketches to device
-// Runs kernel function across distance elements
-// Copies and returns results
-NumpyMatrix query_db_cuda(std::vector<Reference> &ref_sketches,
-                          std::vector<Reference> &query_sketches,
-                          const std::vector<size_t> &kmer_lengths,
-                          RandomMC &random_match, const int device_id,
-                          const unsigned int num_cpu_threads) {
-  size_t mem_free, mem_total, shared_size;
-  std::tie(mem_free, mem_total, shared_size) = initialise_device(device_id);
+
+// Helper function which checks device memory size and whether chunking
+// of input is necessary
+dist_params cuda_dists_init(const std::vector<Reference> &ref_sketches,
+                     const std::vector<Reference> &query_sketches,
+                     const std::vector<size_t> &kmer_lengths,
+                     const int device_id) {
+  dist_params params;
+
+  std::tie(params.mem_free, params.mem_total, params.shared_size) =
+    initialise_device(device_id);
   std::cerr << "Calculating distances on GPU device " << device_id << std::endl;
 
   // Check sketches are compatible
-  bool self = false;
+  params.self = false;
   size_t bbits = ref_sketches[0].bbits();
   size_t sketchsize64 = ref_sketches[0].sketchsize64();
   checkSketchParamsMatch(ref_sketches, kmer_lengths, bbits, sketchsize64);
 
   // Set up sketch information and sizes
-  SketchStrides ref_strides;
-  ref_strides.bbits = bbits;
-  ref_strides.sketchsize64 = sketchsize64;
-  SketchStrides query_strides = ref_strides;
+  params.ref_strides.bbits = bbits;
+  params.ref_strides.sketchsize64 = sketchsize64;
+  params.query_strides = params.ref_strides;
 
-  long long dist_rows;
-  long n_samples = 0;
   if (ref_sketches == query_sketches) {
-    self = true;
-    dist_rows = static_cast<long long>(0.5 * (ref_sketches.size()) *
+    params.self = true;
+    params.dist_rows = static_cast<long long>(0.5 * (ref_sketches.size()) *
                                        (ref_sketches.size() - 1));
-    n_samples = ref_sketches.size();
+    params.n_samples = ref_sketches.size();
   } else {
     // Also check query sketches are compatible
     checkSketchParamsMatch(query_sketches, kmer_lengths, bbits, sketchsize64);
-    dist_rows = ref_sketches.size() * query_sketches.size();
-    n_samples = ref_sketches.size() + query_sketches.size();
+    params.dist_rows = ref_sketches.size() * query_sketches.size();
+    params.n_samples = ref_sketches.size() + query_sketches.size();
   }
+
+  return params;
+}
+
+// Main function callable via API
+// Checks inputs
+// Flattens sketches
+// Copies flattened sketches to device
+// Runs kernel function across distance elements
+// Copies and returns results
+NumpyMatrix query_db_cuda(std::vector<Reference> &ref_sketches,
+                          std::vector<Reference> &query_sketches,
+                          const std::vector<size_t> &kmer_lengths,
+                          RandomMC &random_match, const int device_id,
+                          const unsigned int num_cpu_threads) {
+  dist_params params =
+    cuda_dists_init(ref_sketches, query_sketches, kmer_lengths, device_id);
+
   double est_size =
-      (bbits * sketchsize64 * kmer_lengths.size() * n_samples *
+      (params.ref_strides.bbits * params.ref_strides.sketchsize64 * kmer_lengths.size() * params.n_samples *
            sizeof(uint64_t) + // Size of sketches
        kmer_lengths.size() * std::pow(random_match.n_clusters(), 2) *
            sizeof(float) + // Size of random matches
-       n_samples * sizeof(uint16_t) +
-       dist_rows * 2 * sizeof(float)); // Size of distance matrix
+       params.n_samples * sizeof(uint16_t) +
+       params.dist_rows * 2 * sizeof(float)); // Size of distance matrix
   std::cerr << "Estimated device memory required: " << std::fixed
             << std::setprecision(0) << est_size / (1048576) << "Mb"
             << std::endl;
   std::cerr << "Total device memory: " << std::fixed << std::setprecision(0)
-            << mem_total / (1048576) << "Mb" << std::endl;
+            << params.mem_total / (1048576) << "Mb" << std::endl;
   std::cerr << "Free device memory: " << std::fixed << std::setprecision(0)
-            << mem_free / (1048576) << "Mb" << std::endl;
+            << params.mem_free / (1048576) << "Mb" << std::endl;
 
-  if (est_size > mem_free * (1 - mem_epsilon) && !self) {
-    throw std::runtime_error(
-        "Using greater than device memory is unsupported for query mode. "
-        "Split your input into smaller chunks");
+  unsigned int chunks = 1;
+  if (est_size > params.mem_free * (1 - mem_epsilon)) {
+    if (params.self) {
+      // To prevent memory being exceeded, total distance matrix is split up into
+      // chunks which do fit in memory. The most is needed in the 'corners' where
+      // two separate lots of sketches are loaded, hence the factor of two below
+
+      // These are iterated over in the same order as a square distance matrix.
+      // The i = j chunks are 'self', i < j can be skipped
+      // as they contain only lower triangle values, i > j work as query vs ref
+      chunks = floor((est_size * 2) / (params.mem_free * (1 - mem_epsilon))) + 1;
+    } else {
+      throw std::runtime_error(
+          "Using greater than device memory is unsupported for query mode. "
+          "Split your input into smaller chunks");
+    }
   }
 
   // Turn the random matches into an array (same for any ref, query or subsample
@@ -380,27 +405,15 @@ NumpyMatrix query_db_cuda(std::vector<Reference> &ref_sketches,
 
   // Ready to run dists on device
   SketchSlice sketch_subsample;
-  unsigned int chunks = 1;
-  std::vector<float> dist_results(dist_rows * 2);
+  std::vector<float> dist_results(params.dist_rows * 2);
   NumpyMatrix coreSquare, accessorySquare;
-  if (self) {
-    // To prevent memory being exceeded, total distance matrix is split up into
-    // chunks which do fit in memory. The most is needed in the 'corners' where
-    // two separate lots of sketches are loaded, hence the factor of two below
-
-    // These are iterated over in the same order as a square distance matrix.
-    // The i = j chunks are 'self', i < j can be skipped
-    // as they contain only lower triangle values, i > j work as query vs ref
-    if (est_size > mem_free * (1 - mem_epsilon)) {
-      chunks = floor((est_size * 2) / (mem_free * (1 - mem_epsilon))) + 1;
-    }
-    size_t calc_per_chunk = n_samples / chunks;
-    unsigned int num_big_chunks = n_samples % chunks;
-
+  if (params.self) {
+    size_t calc_per_chunk = params.n_samples / chunks;
+    unsigned int num_big_chunks = params.n_samples % chunks;
     // Only allocate these square matrices if they are needed
     if (chunks > 1) {
-      coreSquare.resize(n_samples, n_samples);
-      accessorySquare.resize(n_samples, n_samples);
+      coreSquare.resize(params.n_samples, params.n_samples);
+      accessorySquare.resize(params.n_samples, params.n_samples);
     }
     unsigned int total_chunks = (chunks * (chunks + 1)) >> 1;
     unsigned int chunk_count = 0;
@@ -425,9 +438,9 @@ NumpyMatrix query_db_cuda(std::vector<Reference> &ref_sketches,
         }
 
         dist_results = dispatchDists(
-            ref_sketches, ref_sketches, ref_strides, query_strides, flat_random,
+            ref_sketches, ref_sketches, params.ref_strides, params.query_strides, flat_random,
             ref_random_idx, ref_random_idx, sketch_subsample, kmer_lengths,
-            chunk_i == chunk_j, num_cpu_threads, shared_size);
+            chunk_i == chunk_j, num_cpu_threads, params.shared_size);
 
         // Read intermediate dists out
         if (chunks > 1) {
@@ -455,13 +468,13 @@ NumpyMatrix query_db_cuda(std::vector<Reference> &ref_sketches,
     sketch_subsample.query_offset = 0;
 
     dist_results = dispatchDists(
-        ref_sketches, query_sketches, ref_strides, query_strides, flat_random,
+        ref_sketches, query_sketches, params.ref_strides, params.query_strides, flat_random,
         ref_random_idx, query_random_idx, sketch_subsample, kmer_lengths, false,
-        num_cpu_threads, shared_size);
+        num_cpu_threads, params.shared_size);
   }
 
   NumpyMatrix dists_ret_matrix;
-  if (self && chunks > 1) {
+  if (params.self && chunks > 1) {
     // Chunked computation yields square matrix, which needs to be converted
     // back to long form
     Eigen::VectorXf core_dists = square_to_long(coreSquare, num_cpu_threads);
@@ -477,4 +490,41 @@ NumpyMatrix query_db_cuda(std::vector<Reference> &ref_sketches,
   }
 
   return dists_ret_matrix;
-}
\ No newline at end of file
+}
+
+sparse_coo query_db_sparse_cuda(std::vector<Reference> &ref_sketches,
+                          const std::vector<size_t> &kmer_lengths,
+                          RandomMC &random_match,
+                          const int kNN, const size_t dist_col,
+                          const int device_id,
+                          const unsigned int num_cpu_threads) {
+  dist_params params =
+    cuda_dists_init(ref_sketches, ref_sketches, kmer_lengths, device_id);
+  FlatRandom flat_random =
+      random_match.flattened_random(kmer_lengths, ref_sketches[0].seq_length());
+  std::vector<uint16_t> ref_random_idx =
+      random_match.lookup_array(ref_sketches);
+
+  // Flatten all of the sketches in blocks
+  //   (as stride is bins) - return vector<vector<uint64_t>>
+  // Use a fixed chunk size of up to 3000 samples per chunk
+  const size_t chunks = 1 + params.n_samples / 3000;
+  const size_t samples_per_chunk = params.n_samples / chunks;
+  const unsigned int num_big_chunks = params.n_samples % chunks;
+  std::vector<std::vector<uint64_t>> sample_blocks;
+  std::vector<SketchStrides> sample_block_strides(chunks, params.ref_strides);
+  size_t start_idx = 0;
+  for (size_t chunk_idx = 0; chunk_idx < chunks; ++chunk_idx) {
+    size_t end_idx = start_idx + samples_per_chunk +
+      (chunk_idx < num_big_chunks ? 1 : 0);
+    sample_blocks.push_back(
+      flatten_by_samples(ref_sketches, kmer_lengths,
+              sample_block_strides[chunk_idx], start_idx, end_idx, num_cpu_threads));
+    start_idx = end_idx;
+  }
+
+  return sparseDists(params, sample_blocks, sample_block_strides,
+                     flat_random, ref_random_idx, kmer_lengths,
+                     kNN, dist_col, samples_per_chunk, num_big_chunks,
+                     num_cpu_threads);
+}
diff --git a/src/gpu/sketch.cu b/src/gpu/sketch.cu
index ece14eb9..d14475cf 100644
--- a/src/gpu/sketch.cu
+++ b/src/gpu/sketch.cu
@@ -14,7 +14,7 @@ namespace py = pybind11;
 // memcpy_async
 #include <cooperative_groups.h>
 #include <cooperative_groups/memcpy_async.h>
-#pragma diag_suppress static_var_with_dynamic_init
+#pragma nv_diag_suppress static_var_with_dynamic_init
 
 #include "cuda.cuh"
 #include "gpu.hpp"
diff --git a/src/sketchlib_bindings.cpp b/src/sketchlib_bindings.cpp
index 70b5ecc4..d1edd162 100644
--- a/src/sketchlib_bindings.cpp
+++ b/src/sketchlib_bindings.cpp
@@ -5,44 +5,38 @@
  */
 
 // pybind11 headers
+#include <pybind11/eigen.h>
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
-#include <pybind11/eigen.h>
 
 namespace py = pybind11;
 
 #include <highfive/H5Exception.hpp>
 
-#include "version.h"
 #include "api.hpp"
+#include "version.h"
 
 NumpyMatrix longToSquare(const Eigen::Ref<Eigen::VectorXf> &distVec,
-                         const unsigned int num_threads)
-{
+                         const unsigned int num_threads) {
   Eigen::VectorXf dummy_query_ref;
   Eigen::VectorXf dummy_query_query;
-  NumpyMatrix converted = long_to_square(distVec,
-                                         dummy_query_ref,
-                                         dummy_query_query,
-                                         num_threads);
+  NumpyMatrix converted =
+      long_to_square(distVec, dummy_query_ref, dummy_query_query, num_threads);
   return (converted);
 }
 
-NumpyMatrix longToSquareMulti(const Eigen::Ref<Eigen::VectorXf> &distVec,
-                              const Eigen::Ref<Eigen::VectorXf> &query_ref_distVec,
-                              const Eigen::Ref<Eigen::VectorXf> &query_query_distVec,
-                              const unsigned int num_threads)
-{
-  NumpyMatrix converted = long_to_square(distVec,
-                                         query_ref_distVec,
-                                         query_query_distVec,
-                                         num_threads);
+NumpyMatrix
+longToSquareMulti(const Eigen::Ref<Eigen::VectorXf> &distVec,
+                  const Eigen::Ref<Eigen::VectorXf> &query_ref_distVec,
+                  const Eigen::Ref<Eigen::VectorXf> &query_query_distVec,
+                  const unsigned int num_threads) {
+  NumpyMatrix converted = long_to_square(distVec, query_ref_distVec,
+                                         query_query_distVec, num_threads);
   return (converted);
 }
 
 Eigen::VectorXf squareToLong(const Eigen::Ref<NumpyMatrix> &distMat,
-                             const unsigned int num_threads)
-{
+                             const unsigned int num_threads) {
   Eigen::VectorXf converted = square_to_long(distMat, num_threads);
   return (converted);
 }
@@ -54,73 +48,38 @@ void constructDatabase(const std::string &db_name,
                        std::vector<size_t> kmer_lengths,
                        const size_t sketch_size,
                        const bool codon_phased = false,
-                       const bool calc_random = true,
-                       const bool use_rc = true,
-                       size_t min_count = 0,
-                       const bool exact = false,
-                       const size_t num_threads = 1,
-                       const bool use_gpu = false,
-                       const int device_id = 0)
-{
+                       const bool calc_random = true, const bool use_rc = true,
+                       size_t min_count = 0, const bool exact = false,
+                       const size_t num_threads = 1, const bool use_gpu = false,
+                       const int device_id = 0) {
   std::vector<Reference> ref_sketches;
 #ifdef GPU_AVAILABLE
-  if (use_gpu)
-  {
-    if (codon_phased)
-    {
+  if (use_gpu) {
+    if (codon_phased) {
       throw std::runtime_error(
           "Codon phased seeds not yet implemented for GPU sketching");
     }
-    ref_sketches = create_sketches_cuda(db_name,
-                                        sample_names,
-                                        file_names,
-                                        kmer_lengths,
-                                        sketch_size,
-                                        use_rc,
-                                        min_count,
-                                        num_threads,
-                                        device_id);
-  }
-  else
-  {
-    ref_sketches = create_sketches(db_name,
-                                   sample_names,
-                                   file_names,
-                                   kmer_lengths,
-                                   sketch_size,
-                                   codon_phased,
-                                   use_rc,
-                                   min_count,
-                                   exact,
-                                   num_threads);
+    ref_sketches = create_sketches_cuda(db_name, sample_names, file_names,
+                                        kmer_lengths, sketch_size, use_rc,
+                                        min_count, num_threads, device_id);
+  } else {
+    ref_sketches = create_sketches(db_name, sample_names, file_names,
+                                   kmer_lengths, sketch_size, codon_phased,
+                                   use_rc, min_count, exact, num_threads);
   }
 #else
-  ref_sketches = create_sketches(db_name,
-                                 sample_names,
-                                 file_names,
-                                 kmer_lengths,
-                                 sketch_size,
-                                 codon_phased,
-                                 use_rc,
-                                 min_count,
-                                 exact,
-                                 num_threads);
+  ref_sketches = create_sketches(db_name, sample_names, file_names,
+                                 kmer_lengths, sketch_size, codon_phased,
+                                 use_rc, min_count, exact, num_threads);
 #endif
-  if (calc_random)
-  {
-    if (ref_sketches.size() >= default_n_clusters)
-    {
-      RandomMC random = calculate_random(ref_sketches,
-                                         db_name,
-                                         default_n_clusters,
-                                         default_n_MC,
-                                         codon_phased,
-                                         use_rc,
-                                         num_threads);
-    }
-    else
-    {
-      std::cerr << "Too few input genomes to calculate random match chances" << std::endl;
+  if (calc_random) {
+    if (ref_sketches.size() >= default_n_clusters) {
+      RandomMC random =
+          calculate_random(ref_sketches, db_name, default_n_clusters,
+                           default_n_MC, codon_phased, use_rc, num_threads);
+    } else {
+      std::cerr << "Too few input genomes to calculate random match chances"
+                << std::endl;
     }
   }
 }
@@ -133,15 +92,12 @@ NumpyMatrix queryDatabase(const std::string &ref_db_name,
                           const bool random_correct = true,
                           const bool jaccard = false,
                           const size_t num_threads = 1,
-                          const bool use_gpu = false,
-                          const int device_id = 0)
-{
-  if (jaccard && use_gpu)
-  {
-    throw std::runtime_error("Extracting Jaccard distances not supported on GPU");
+                          const bool use_gpu = false, const int device_id = 0) {
+  if (use_gpu && (jaccard || kmer_lengths.size() < 2)) {
+    throw std::runtime_error(
+        "Extracting Jaccard distances not supported on GPU");
   }
-  if (!same_db_version(ref_db_name, query_db_name))
-  {
+  if (!same_db_version(ref_db_name, query_db_name)) {
     std::cerr << "WARNING: versions of input databases sketches are different,"
                  " results may not be compatible"
               << std::endl;
@@ -150,226 +106,183 @@ NumpyMatrix queryDatabase(const std::string &ref_db_name,
   std::vector<Reference> ref_sketches =
       load_sketches(ref_db_name, ref_names, kmer_lengths, false);
   std::vector<Reference> query_sketches;
-  if (ref_db_name == query_db_name && ref_names == query_names)
-  {
+  if (ref_db_name == query_db_name && ref_names == query_names) {
     query_sketches = ref_sketches;
-  }
-  else
-  {
+  } else {
     query_sketches =
         load_sketches(query_db_name, query_names, kmer_lengths, false);
   }
 
   RandomMC random;
-  if (random_correct)
-  {
+  if (random_correct) {
     random = get_random(ref_db_name, ref_sketches[0].rc());
-  }
-  else
-  {
+  } else {
     random = RandomMC();
   }
 
   NumpyMatrix dists;
 #ifdef GPU_AVAILABLE
-  if (use_gpu)
-  {
-    dists = query_db_cuda(ref_sketches,
-                          query_sketches,
-                          kmer_lengths,
-                          random,
-                          device_id,
-                          num_threads);
-  }
-  else
-  {
-    dists = query_db(ref_sketches,
-                     query_sketches,
-                     kmer_lengths,
-                     random,
-                     jaccard,
-                     num_threads);
+  if (use_gpu) {
+    dists = query_db_cuda(ref_sketches, query_sketches, kmer_lengths, random,
+                          device_id, num_threads);
+  } else {
+    dists = query_db(ref_sketches, query_sketches, kmer_lengths, random,
+                     jaccard, num_threads);
   }
 #else
-  dists = query_db(ref_sketches,
-                   query_sketches,
-                   kmer_lengths,
-                   random,
-                   jaccard,
+  dists = query_db(ref_sketches, query_sketches, kmer_lengths, random, jaccard,
                    num_threads);
 #endif
   return (dists);
 }
 
-sparse_coo sparseQuery(const std::string &ref_db_name,
-                       const std::string &query_db_name,
-                       const std::vector<std::string> &ref_names,
-                       const std::vector<std::string> &query_names,
-                       std::vector<size_t> kmer_lengths,
-                       const bool random_correct = true,
-                       const float dist_cutoff = 0,
-                       const unsigned long int kNN = 0,
-                       const bool core = true,
-                       const size_t num_threads = 1,
-                       const bool use_gpu = false,
-                       const int device_id = 0)
-{
-  NumpyMatrix dists = queryDatabase(ref_db_name, query_db_name,
-                                    ref_names, query_names, kmer_lengths,
-                                    random_correct, false, num_threads,
-                                    use_gpu, device_id);
-
-  unsigned int dist_col = 0;
-  if (!core)
-  {
-    dist_col = 1;
+sparse_coo
+querySelfSparse(const std::string &ref_db_name, const std::vector<std::string> &ref_names,
+            std::vector<size_t> kmer_lengths, const bool random_correct = true,
+            const bool jaccard = false, const unsigned long int kNN = 0,
+            const float dist_cutoff = 0.0f, const size_t dist_col = 0,
+            const size_t num_threads = 1, const bool use_gpu = false,
+            const int device_id = 0) {
+  if (use_gpu && (jaccard || kmer_lengths.size() < 2 || dist_col > 1)) {
+    throw std::runtime_error(
+        "Extracting Jaccard distances not supported on GPU");
   }
-  Eigen::VectorXf dummy_query_ref;
-  Eigen::VectorXf dummy_query_query;
-  NumpyMatrix long_form = long_to_square(dists.col(dist_col),
-                                         dummy_query_ref,
-                                         dummy_query_query,
-                                         num_threads);
-  sparse_coo sparse_return = sparsify_dists(long_form,
-                                            dist_cutoff,
-                                            kNN);
-
-  return (sparse_return);
+  std::vector<Reference> ref_sketches =
+      load_sketches(ref_db_name, ref_names, kmer_lengths, false);
+
+  RandomMC random;
+  if (random_correct) {
+    random = get_random(ref_db_name, ref_sketches[0].rc());
+  } else {
+    random = RandomMC();
+  }
+
+  sparse_coo sparse_dists;
+  if (dist_cutoff > 0) {
+    NumpyMatrix dists = queryDatabase(ref_db_name, ref_db_name, ref_names,
+                                    ref_names, kmer_lengths, random_correct,
+                                    false, num_threads, use_gpu, device_id);
+
+    Eigen::VectorXf dummy_query_ref;
+    Eigen::VectorXf dummy_query_query;
+    NumpyMatrix long_form = long_to_square(dists.col(dist_col), dummy_query_ref,
+                                          dummy_query_query, num_threads);
+    sparse_dists = sparsify_dists(long_form, dist_cutoff, kNN);
+  } else {
+#ifdef GPU_AVAILABLE
+    if (use_gpu) {
+      sparse_dists = query_db_sparse_cuda(ref_sketches, kmer_lengths, random,
+        kNN, dist_col, device_id, num_threads);
+    } else {
+      sparse_dists = query_db_sparse(ref_sketches, kmer_lengths, random,
+        jaccard, kNN, dist_col, num_threads);
+    }
+#else
+    sparse_dists = query_db_sparse(ref_sketches, kmer_lengths, random,
+      jaccard, kNN, dist_col, num_threads);
+#endif
+  }
+
+  return sparse_dists;
 }
 
 void addRandomToDb(const std::string &db_name,
                    const std::vector<std::string> &sample_names,
                    const std::vector<size_t> kmer_lengths,
-                   const bool use_rc = true,
-                   const size_t num_threads = 1)
-{
+                   const bool use_rc = true, const size_t num_threads = 1) {
   std::string db_version;
   bool codon_phased;
   std::tie(db_version, codon_phased) = get_db_attr(db_name);
 
   std::vector<Reference> ref_sketches =
       load_sketches(db_name, sample_names, kmer_lengths, false);
-  RandomMC random = calculate_random(ref_sketches,
-                                     db_name,
-                                     default_n_clusters,
-                                     default_n_MC,
-                                     codon_phased,
-                                     use_rc,
-                                     num_threads);
+  RandomMC random =
+      calculate_random(ref_sketches, db_name, default_n_clusters, default_n_MC,
+                       codon_phased, use_rc, num_threads);
 }
 
-double jaccardDist(const std::string &db_name,
-                   const std::string &sample1,
-                   const std::string &sample2,
-                   const size_t kmer_size,
-                   bool random_correct)
-{
-  auto sketch_vec = load_sketches(db_name,
-                                  {sample1, sample2},
-                                  {kmer_size},
-                                  false);
+double jaccardDist(const std::string &db_name, const std::string &sample1,
+                   const std::string &sample2, const size_t kmer_size,
+                   bool random_correct) {
+  auto sketch_vec =
+      load_sketches(db_name, {sample1, sample2}, {kmer_size}, false);
   RandomMC random;
-  if (random_correct)
-  {
+  if (random_correct) {
     random = get_random(db_name, sketch_vec[0].rc());
-  }
-  else
-  {
+  } else {
     random = RandomMC();
   }
-  return (sketch_vec.at(0).jaccard_dist(sketch_vec.at(1),
-                                        kmer_size,
-                                        random));
+  return (sketch_vec.at(0).jaccard_dist(sketch_vec.at(1), kmer_size, random));
 }
 
 // Wrapper which makes a ref to the python/numpy array
 sparse_coo sparsifyDists(const Eigen::Ref<NumpyMatrix> &denseDists,
-                         const float distCutoff,
-                         const unsigned long int kNN)
-{
-  sparse_coo coo_idx = sparsify_dists(denseDists,
-                                      distCutoff,
-                                      kNN);
+                         const float distCutoff, const unsigned long int kNN) {
+  sparse_coo coo_idx = sparsify_dists(denseDists, distCutoff, kNN);
   return coo_idx;
 }
 
-PYBIND11_MODULE(pp_sketchlib, m)
-{
+PYBIND11_MODULE(pp_sketchlib, m) {
   m.doc() = "Sketch implementation for PopPUNK";
 
   // Exported functions
   m.def("constructDatabase", &constructDatabase, "Create and save sketches",
-        py::arg("db_name"),
-        py::arg("samples"),
-        py::arg("files"),
-        py::arg("klist"),
-        py::arg("sketch_size"),
-        py::arg("codon_phased") = false,
-        py::arg("calc_random") = true,
-        py::arg("use_rc") = true,
-        py::arg("min_count") = 0,
-        py::arg("exact") = false,
-        py::arg("num_threads") = 1,
-        py::arg("use_gpu") = false,
-        py::arg("device_id") = 0);
+        py::arg("db_name"), py::arg("samples"), py::arg("files"),
+        py::arg("klist"), py::arg("sketch_size"),
+        py::arg("codon_phased") = false, py::arg("calc_random") = true,
+        py::arg("use_rc") = true, py::arg("min_count") = 0,
+        py::arg("exact") = false, py::arg("num_threads") = 1,
+        py::arg("use_gpu") = false, py::arg("device_id") = 0);
 
-  m.def("queryDatabase", &queryDatabase, py::return_value_policy::reference_internal, "Find distances between sketches; return all distances",
-        py::arg("ref_db_name"),
-        py::arg("query_db_name"),
-        py::arg("rList"),
-        py::arg("qList"),
-        py::arg("klist"),
-        py::arg("random_correct") = true,
-        py::arg("jaccard") = false,
-        py::arg("num_threads") = 1,
-        py::arg("use_gpu") = false,
-        py::arg("device_id") = 0);
+  m.def("queryDatabase", &queryDatabase,
+        py::return_value_policy::reference_internal,
+        "Find distances between sketches; return all distances",
+        py::arg("ref_db_name"), py::arg("query_db_name"), py::arg("rList"),
+        py::arg("qList"), py::arg("klist"), py::arg("random_correct") = true,
+        py::arg("jaccard") = false, py::arg("num_threads") = 1,
+        py::arg("use_gpu") = false, py::arg("device_id") = 0);
 
-  m.def("queryDatabaseSparse", &sparseQuery, py::return_value_policy::reference_internal, "Find distances between sketches; return a sparse matrix",
-        py::arg("ref_db_name"),
-        py::arg("query_db_name"),
-        py::arg("rList"),
-        py::arg("qList"),
-        py::arg("klist"),
-        py::arg("random_correct") = true,
-        py::arg("dist_cutoff") = 0,
-        py::arg("kNN") = 0,
-        py::arg("core") = true,
-        py::arg("num_threads") = 1,
-        py::arg("use_gpu") = false,
+  m.def("querySelfSparse", &querySelfSparse,
+        py::return_value_policy::reference_internal,
+        "Find self-self distances between sketches; return a sparse matrix",
+        py::arg("ref_db_name"), py::arg("rList"), py::arg("klist"), py::arg("random_correct") = true,
+        py::arg("jaccard") = false,
+        py::arg("kNN") = 0, py::arg("dist_cutoff") = 0.0f,
+        py::arg("dist_col") = 0,
+        py::arg("num_threads") = 1, py::arg("use_gpu") = false,
         py::arg("device_id") = 0);
 
-  m.def("addRandom", &addRandomToDb, "Add random match chances into older databases",
-        py::arg("db_name"),
-        py::arg("samples"),
-        py::arg("klist"),
-        py::arg("use_rc") = true,
+  m.def("addRandom", &addRandomToDb,
+        "Add random match chances into older databases", py::arg("db_name"),
+        py::arg("samples"), py::arg("klist"), py::arg("use_rc") = true,
         py::arg("num_threads") = 1);
 
   m.def("jaccardDist", &jaccardDist, "Calculate a raw Jaccard distance",
-        py::arg("db_name"),
-        py::arg("ref_name"),
-        py::arg("query_name"),
-        py::arg("kmer_length"),
-        py::arg("random_correct") = false);
-
-  m.def("squareToLong", &squareToLong, py::return_value_policy::reference_internal, "Convert dense square matrices to long form",
-        py::arg("distMat").noconvert(),
-        py::arg("num_threads") = 1);
+        py::arg("db_name"), py::arg("ref_name"), py::arg("query_name"),
+        py::arg("kmer_length"), py::arg("random_correct") = false);
 
-  m.def("longToSquare", &longToSquare, py::return_value_policy::reference_internal, "Convert dense long form distance matrices to square form",
-        py::arg("distVec").noconvert(),
-        py::arg("num_threads") = 1);
+  m.def("squareToLong", &squareToLong,
+        py::return_value_policy::reference_internal,
+        "Convert dense square matrices to long form",
+        py::arg("distMat").noconvert(), py::arg("num_threads") = 1);
 
-  m.def("longToSquareMulti", &longToSquareMulti, py::return_value_policy::reference_internal, "Convert dense long form distance matrices to square form (in three blocks)",
+  m.def("longToSquare", &longToSquare,
+        py::return_value_policy::reference_internal,
+        "Convert dense long form distance matrices to square form",
+        py::arg("distVec").noconvert(), py::arg("num_threads") = 1);
+
+  m.def("longToSquareMulti", &longToSquareMulti,
+        py::return_value_policy::reference_internal,
+        "Convert dense long form distance matrices to square form (in three "
+        "blocks)",
         py::arg("distVec").noconvert(),
         py::arg("query_ref_distVec").noconvert(),
-        py::arg("query_query_distVec").noconvert(),
-        py::arg("num_threads") = 1);
+        py::arg("query_query_distVec").noconvert(), py::arg("num_threads") = 1);
 
-  m.def("sparsifyDists", &sparsifyDists, py::return_value_policy::reference_internal, "Transform all distances into a sparse matrix",
-        py::arg("distMat").noconvert(),
-        py::arg("distCutoff") = 0,
+  m.def("sparsifyDists", &sparsifyDists,
+        py::return_value_policy::reference_internal,
+        "Transform all distances into a sparse matrix",
+        py::arg("distMat").noconvert(), py::arg("distCutoff") = 0,
         py::arg("kNN") = 0);
 
   m.attr("version") = VERSION_INFO;


Test patch:
diff --git a/docker/test b/docker/test
index 17b93b6f..8a4655cf 100755
--- a/docker/test
+++ b/docker/test
@@ -8,4 +8,4 @@ HERE=$(dirname $0)
 [ ! -z $(docker images -q $TAG_SHA) ] || docker pull $TAG_SHA
 
 ## Just check that we can bring up the container and run something
-docker run -it -w /src --rm $TAG_SHA poppunk_sketch --version
+docker run -it -w /src --rm $TAG_SHA sketchlib --version
diff --git a/src/test/gpu_dist_test.cpp b/src/test/gpu_dist_test.cpp
deleted file mode 100644
index 90ede0d6..00000000
--- a/src/test/gpu_dist_test.cpp
+++ /dev/null
@@ -1,50 +0,0 @@
-
-#include <fstream>
-#include <iostream>
-
-#include "reference.hpp"
-#include "database/database.hpp"
-#include "random/random_match.hpp"
-#include "api.hpp"
-#include "sketch/countmin.hpp"
-#include "sketch/seqio.hpp"
-
-int main(int argc, char* argv[])
-{
-    // Runs a test of functionality
-    std::ifstream rfile(argv[1]);
-    std::string name, file;
-    std::vector<std::string> names;
-    std::vector<std::vector<std::string>> files;
-    while (rfile >> name >> file)
-    {
-        names.push_back(name);
-        std::vector<std::string> file_list = {file};
-        files.push_back(file_list);
-    }
-
-    HighFive::File h5_db = open_h5("listeria.h5");
-    Database listeria_db(h5_db);
-    std::vector<Reference> listeria_sketches;
-    for (auto name_it = names.cbegin(); name_it != names.cend(); name_it++)
-    {
-        listeria_sketches.push_back(listeria_db.load_sketch(*name_it));
-    }
-    RandomMC random_retrived = listeria_db.load_random(true);
-
-    std::vector<size_t> kmer_lengths {15, 17, 19, 21, 23, 25, 27, 29};
-
-#ifdef GPU_AVAILABLE
-    NumpyMatrix listeria_dists =
-      query_db_cuda(
-        listeria_sketches,
-	    listeria_sketches,
-	    kmer_lengths,
-	    random_retrived,
-	    0,
-	    4);
-    std::cout << listeria_dists << std::endl;
-#endif
-    return 0;
-}
-
diff --git a/src/test/main.cpp b/src/test/main.cpp
deleted file mode 100644
index 76f5a0a7..00000000
--- a/src/test/main.cpp
+++ /dev/null
@@ -1,124 +0,0 @@
-
-#include <fstream>
-#include <iostream>
-
-#include "reference.hpp"
-#include "database/database.hpp"
-#include "api.hpp"
-
-int main (int argc, char* argv[])
-{
-    // Runs a test of functionality
-
-    std::vector<size_t> kmer_lengths {15, 17, 19, 21, 23, 25, 27, 29};
-    KmerSeeds kmer_seeds = generate_seeds(kmer_lengths, false);
-    SeqBuf ref_seq({argv[2]}, kmer_lengths.back());
-    Reference ref(argv[1], ref_seq, kmer_seeds, 156, false, true, 0, false);
-    // Reference ref_copy(argv[1], argv[2], kmer_lengths);
-    SeqBuf query_seq({argv[4]}, kmer_lengths.back());
-    Reference query(argv[3], query_seq, kmer_seeds, 156, false, true, 0, false);
-
-    RandomMC random_match(true);
-
-    std::cout << ref.jaccard_dist(ref, 15, random_match) << std::endl;      // Should be 1
-    std::cout << ref.jaccard_dist(query, 15, random_match) << std::endl;
-    std::cout << ref.jaccard_dist(query, 29, random_match) << std::endl;
-    std::cout << query.jaccard_dist(ref, 29, random_match) << std::endl;
-
-    auto core_acc = ref.core_acc_dist<RandomMC>(query, random_match);
-    std::cout << std::get<0>(core_acc) << "\t" << std::get<1>(core_acc) << std::endl;
-
-    Database sketch_db("sketch.h5");
-    sketch_db.add_sketch(ref);
-    sketch_db.add_sketch(query);
-
-    Reference ref_read = sketch_db.load_sketch(argv[1]);
-    Reference query_read = sketch_db.load_sketch(argv[3]);
-    std::cout << ref_read.jaccard_dist(query_read, 15, random_match) << std::endl;
-
-    std::vector<Reference> ref_sketches = create_sketches("full",
-                               {argv[1], argv[3]},
-                               {{argv[2]}, {argv[4]}},
-                               kmer_lengths,
-                               156,
-                               false,
-                               true,
-                               0,
-                               false,
-                               2);
-    RandomMC basic_adjust(true);
-    NumpyMatrix dists = query_db(ref_sketches,
-                              ref_sketches,
-                              kmer_lengths,
-                              basic_adjust,
-                              false,
-                              2);
-
-    std::cout << dists << std::endl;
-
-    std::ifstream rfile(argv[5]);
-    std::string name, file;
-    std::vector<std::string> names;
-    std::vector<std::vector<std::string>> files;
-    while (rfile >> name >> file)
-    {
-        names.push_back(name);
-        std::vector<std::string> file_list = {file};
-        files.push_back(file_list);
-    }
-
-    create_sketches("listeria",
-                    names,
-                    files,
-                    kmer_lengths,
-                    156,
-                    false,
-                    true,
-                    0,
-                    false,
-                    4);
-
-    HighFive::File h5_db = open_h5("listeria.h5");
-    Database listeria_db(h5_db);
-    std::vector<Reference> listeria_sketches;
-    for (auto name_it = names.cbegin(); name_it != names.cend(); name_it++)
-    {
-        listeria_sketches.push_back(listeria_db.load_sketch(*name_it));
-    }
-
-    // Save random matches to db
-    RandomMC random(listeria_sketches, 3, 5, false, true, 4);
-    try {
-        listeria_db.save_random(random);
-    } catch (const std::exception& e) {
-        std::cerr << "Not writing random matches" << std::endl;
-        std::cout << e.what() << std::endl;
-    }
-    RandomMC random_retrived = listeria_db.load_random(true);
-    if (random_retrived != random) {
-        throw std::runtime_error("Saving and loading random from DB failed!");
-    }
-
-    NumpyMatrix listeria_dists = query_db(listeria_sketches,
-                            listeria_sketches,
-                            kmer_lengths,
-                            random,
-                            false,
-                            4);
-    std::cout << listeria_dists << std::endl;
-
-    // Check ref v query works
-    std::vector<Reference> query_sketches(2);
-    std::copy(listeria_sketches.begin(), listeria_sketches.begin() + 2,
-              query_sketches.begin());
-    NumpyMatrix query_dists = query_db(listeria_sketches,
-                            query_sketches,
-                            kmer_lengths,
-                            random,
-                            false,
-                            4);
-    std::cout << query_dists << std::endl;
-
-    return 0;
-}
-
diff --git a/src/test/matrix_test.cpp b/src/test/matrix_test.cpp
deleted file mode 100644
index a618f9bb..00000000
--- a/src/test/matrix_test.cpp
+++ /dev/null
@@ -1,38 +0,0 @@
-
-#include <fstream>
-#include <iostream>
-
-#include "api.hpp"
-
-int main (int argc, char* argv[])
-{
-    // Runs a test of functionality
-    NumpyMatrix square(6,6);
-    square << 0, 1, 2, 3, 10, 20,
-              1, 0, 4, 5, 10, 20,
-              2, 4, 0, 6, 10, 20,
-              3, 5, 6, 0, 10, 20,
-              10, 10, 10, 10, 0, 8,
-              20, 20, 20, 20, 8, 0;
-    std::vector<long> i_vec, j_vec;
-    std::vector<float> dists;
-    sparse_coo sparse_dists_cutoff = sparsify_dists(square, 7, 0, 1);
-    std::tie(i_vec, j_vec, dists) = sparse_dists_cutoff;
-    std::cout << "Dists < 7" << std::endl;
-    for (unsigned int it = 0; it < i_vec.size(); it++) {
-        std::cout << i_vec[it] << "\t" << j_vec[it] << "\t" << dists[it] << std::endl;
-    }
-
-    sparse_coo sparse_dists_3nn = sparsify_dists(square, 0, 3, 1);
-    std::tie(i_vec, j_vec, dists) = sparse_dists_3nn;
-    std::cout << "3-NN" << std::endl;
-    for (unsigned int it = 0; it < i_vec.size(); it++) {
-        std::cout << i_vec[it] << "\t" << j_vec[it] << "\t" << dists[it] << std::endl;
-    }
-
-
-
-
-    return 0;
-}
-
diff --git a/src/test/read_test.cpp b/src/test/read_test.cpp
deleted file mode 100644
index c28eeb67..00000000
--- a/src/test/read_test.cpp
+++ /dev/null
@@ -1,45 +0,0 @@
-
-#include <fstream>
-#include <iostream>
-
-#include "reference.hpp"
-#include "database/database.hpp"
-#include "random/random_match.hpp"
-#include "api.hpp"
-#include "sketch/countmin.hpp"
-#include "sketch/seqio.hpp"
-
-int main (int argc, char* argv[])
-{
-    // Runs a test of functionality
-
-
-    std::vector<size_t> kmer_lengths {15, 29};
-
-    SeqBuf ref_seq({argv[2], argv[3]}, kmer_lengths.back());
-    Reference ref(argv[1], ref_seq, kmer_lengths, 15, true, 0, false);
-    SeqBuf query_seq({argv[5], argv[6]}, kmer_lengths.back());
-    Reference query(argv[4], query_seq, kmer_lengths, 156, true, 20, false);
-
-    RandomMC random(true);
-    std::cout << ref.jaccard_dist(query, 15, random) << std::endl;
-    std::cout << ref.jaccard_dist(query, 29, random) << std::endl;
-
-    auto core_acc = ref.core_acc_dist<RandomMC>(query, random);
-    std::cout << std::get<0>(core_acc) << "\t" << std::get<1>(core_acc) << std::endl;
-
-#ifdef GPU_AVAILABLE
-    printf("GPU\n");
-    create_sketches_cuda("gpu_db",
-                   {argv[1], argv[4]},
-                   {{argv[2], argv[3]}, {argv[5], argv[6]}},
-                   kmer_lengths,
-                   156,
-                   true,
-                   20,
-				   4,
-                   0);
-#endif
-    return 0;
-}
-
diff --git a/test/clean_test.py b/test/clean_test.py
index 12e5e4d1..3086b9db 100755
--- a/test/clean_test.py
+++ b/test/clean_test.py
@@ -9,7 +9,7 @@
 
 def deleteDir(dirname):
     if os.path.isdir(dirname):
-        shutil.rmtree(dirname)    
+        shutil.rmtree(dirname)
 
 sys.stderr.write("Cleaning up tests\n")
 refs = []
@@ -26,6 +26,10 @@ def deleteDir(dirname):
     "db1.h5",
     "db2.h5",
     "joined.h5",
+    "ppsketch.npy",
+    "ppsketch.pkl",
+    "sparse_query.npz",
+    "sparse_query.pkl",
     "tmp.cpu.query.dists.txt",
     "tmp.gpu.query.dists.txt",
     "tmp.cpu.self.dists.txt",
diff --git a/test/ppsketch_ref.npy b/test/ppsketch_ref.npy
index f62c2618..0cb430f5 100644
Binary files a/test/ppsketch_ref.npy and b/test/ppsketch_ref.npy differ
diff --git a/test/ppsketch_ref.pkl b/test/ppsketch_ref.pkl
index 34bd2314..7cdc9ce4 100644
Binary files a/test/ppsketch_ref.pkl and b/test/ppsketch_ref.pkl differ
diff --git a/test/ppsketch_ref_jaccard.npy b/test/ppsketch_ref_jaccard.npy
index cec1e906..0d0739f8 100644
Binary files a/test/ppsketch_ref_jaccard.npy and b/test/ppsketch_ref_jaccard.npy differ
diff --git a/test/ppsketch_ref_phased.npy b/test/ppsketch_ref_phased.npy
index 7d15f23a..c7bd1885 100644
Binary files a/test/ppsketch_ref_phased.npy and b/test/ppsketch_ref_phased.npy differ
diff --git a/test/ppsketch_ref_phased.pkl b/test/ppsketch_ref_phased.pkl
index 34bd2314..7cdc9ce4 100644
Binary files a/test/ppsketch_ref_phased.pkl and b/test/ppsketch_ref_phased.pkl differ
diff --git a/test/ppsketch_ref_phased_jaccard.npy b/test/ppsketch_ref_phased_jaccard.npy
index a5bc0bf2..204871ad 100644
Binary files a/test/ppsketch_ref_phased_jaccard.npy and b/test/ppsketch_ref_phased_jaccard.npy differ
diff --git a/test/ppsketch_ref_phased_raw_jaccard.npy b/test/ppsketch_ref_phased_raw_jaccard.npy
index c9f00c25..18320a1f 100644
Binary files a/test/ppsketch_ref_phased_raw_jaccard.npy and b/test/ppsketch_ref_phased_raw_jaccard.npy differ
diff --git a/test/ppsketch_ref_raw_jaccard.npy b/test/ppsketch_ref_raw_jaccard.npy
index 5c824545..01f7d31e 100644
Binary files a/test/ppsketch_ref_raw_jaccard.npy and b/test/ppsketch_ref_raw_jaccard.npy differ
diff --git a/test/ppsketch_ref_raw_jaccard.pkl b/test/ppsketch_ref_raw_jaccard.pkl
index 34bd2314..7cdc9ce4 100644
Binary files a/test/ppsketch_ref_raw_jaccard.pkl and b/test/ppsketch_ref_raw_jaccard.pkl differ
diff --git a/test/run_test.py b/test/run_test.py
index add8063e..6f418cd0 100755
--- a/test/run_test.py
+++ b/test/run_test.py
@@ -12,11 +12,6 @@
 
 parser = argparse.ArgumentParser(description='Test pp-sketchlib',
                                  prog='run_test')
-parser.add_argument('--no-cpp',
-                    action='store_true',
-                    default=False,
-                    help='Do not run C++ tests')
-args = parser.parse_args()
 
 if not os.path.isfile("12754_4#89.contigs_velvet.fa"):
     sys.stderr.write("Extracting example dataset\n")
@@ -31,36 +26,42 @@
 
 # create sketches
 sys.stderr.write("Sketch smoke test\n")
-subprocess.run("poppunk_sketch --sketch --rfile references.txt --ref-db test_db --sketch-size 10000 --min-k 15 --k-step 4 --cpus 2", shell=True, check=True)
-subprocess.run("poppunk_sketch --sketch --rfile references.txt --ref-db test_db_phased --codon-phased --sketch-size 10000 --min-k 15 --k-step 4 --cpus 2", shell=True, check=True)
+subprocess.run("python ../sketchlib-runner.py sketch -l references.txt -o test_db -s 10000 -k 15,29,4 --cpus 2", shell=True, check=True)
+subprocess.run("python ../sketchlib-runner.py sketch -l references.txt -o test_db_phased --codon-phased --cpus 2", shell=True, check=True)
+subprocess.run("python ../sketchlib-runner.py sketch 12673_8#24.contigs_velvet.fa 12673_8#34.contigs_velvet.fa -o test_db_small -s 1000 --kmer 14", shell=True, check=True)
+subprocess.run("python ../sketchlib-runner.py add random test_db --cpus 2", shell=True, check=True)
 # calculate distances
 sys.stderr.write("Distance integration test\n")
-subprocess.run("poppunk_sketch --query --ref-db test_db --query-db test_db --min-k 15 --k-step 4 --cpus 2", shell=True, check=True) # checks if can be run
+subprocess.run("python ../sketchlib-runner.py query dist test_db --cpus 2", shell=True, check=True) # checks if can be run
+subprocess.run("python ../sketchlib-runner.py query dist test_db -o ppsketch --cpus 2", shell=True, check=True) # checks if can be run
+subprocess.run("python ../sketchlib-runner.py query jaccard test_db_small --cpus 2", shell=True, check=True) # checks if can be run
 subprocess.run("python test-dists.py --ref-db test_db --results ppsketch_ref", shell=True, check=True) # checks results match
-subprocess.run("poppunk_sketch --query --ref-db test_db_phased --query-db test_db_phased --min-k 15 --k-step 4 --cpus 2", shell=True, check=True) # checks if can be run
+subprocess.run("python ../sketchlib-runner.py query dist test_db_phased --cpus 2", shell=True, check=True) # checks if can be run
 subprocess.run("python test-dists.py --ref-db test_db_phased --results ppsketch_ref_phased", shell=True, check=True) # checks results match
 
+sys.stderr.write("Sparse distance smoke test\n")
+subprocess.run("python ../sketchlib-runner.py query sparse test_db --kNN 2", shell=True, check=True) # checks if can be run
+subprocess.run("python ../sketchlib-runner.py query sparse test_db -o sparse_query --kNN 2", shell=True, check=True) # checks if can be run
+subprocess.run("python ../sketchlib-runner.py query sparse test_db --threshold 0.01", shell=True, check=True) # checks if can be run
+subprocess.run("python ../sketchlib-runner.py query sparse jaccard test_db --kNN 2 --kmer 19", shell=True, check=True) # checks if can be run
+
 sys.stderr.write("Ref v query distance smoke test\n")
-subprocess.run("poppunk_sketch --sketch --rfile rlist.txt --ref-db r_db --sketch-size 10000 --min-k 15 --k-step 4 --cpus 2", shell=True, check=True)
-subprocess.run("poppunk_sketch --sketch --rfile qlist.txt --ref-db q_db --sketch-size 10000 --min-k 15 --k-step 4 --cpus 2", shell=True, check=True)
-subprocess.run("poppunk_sketch --query --ref-db r_db --query-db q_db --read-k", shell=True, check=True) # checks if can be run
+subprocess.run("python ../sketchlib-runner.py sketch -l rlist.txt -o r_db --cpus 2", shell=True, check=True)
+subprocess.run("python ../sketchlib-runner.py sketch -l qlist.txt -o q_db --cpus 2", shell=True, check=True)
+subprocess.run("python ../sketchlib-runner.py query dist r_db q_db.h5", shell=True, check=True) # checks if can be run
+subprocess.run("python ../sketchlib-runner.py query jaccard r_db q_db", shell=True, check=True) # checks if can be run
 
-sys.stderr.write("Sparse distance smoke test\n")
-subprocess.run("poppunk_sketch --query --ref-db test_db --query-db test_db --read-k --sparse --kNN 2", shell=True, check=True) # checks if can be run
-subprocess.run("poppunk_sketch --query --ref-db test_db --query-db test_db --read-k --sparse --threshold 0.01", shell=True, check=True) # checks if can be run
 # Joining
 sys.stderr.write("Join smoke test\n")
-subprocess.run("poppunk_sketch --sketch --rfile db1_refs.txt --ref-db db1 --sketch-size 10000 --min-k 15 --k-step 4 --cpus 2", shell=True, check=True)
-subprocess.run("poppunk_sketch --sketch --rfile db2_refs.txt --ref-db db2 --sketch-size 10000 --min-k 15 --k-step 4 --cpus 2", shell=True, check=True)
-subprocess.run("poppunk_sketch --join --ref-db db1 --query-db db2 --output joined", shell=True, check=True)
+subprocess.run("python ../sketchlib-runner.py sketch -l db1_refs.txt -o db1 --cpus 2", shell=True, check=True)
+subprocess.run("python ../sketchlib-runner.py sketch -l db2_refs.txt -o db2 --cpus 2", shell=True, check=True)
+subprocess.run("python ../sketchlib-runner.py join db1.h5 db2.h5 -o joined", shell=True, check=True)
+# Random
+sys.stderr.write("Random test\n")
+subprocess.run("python ../sketchlib-runner.py remove random test_db --cpus 2", shell=True, check=True)
 # Matrix
 sys.stderr.write("Matrix integration test\n")
 subprocess.run("python test-matrix.py", shell=True, check=True)
 
-# C++ test
-if not args.no_cpp:
-    sys.stderr.write("Testing functions via C++\n")
-    subprocess.run("sketch_test 12673_8#43 12673_8#43.contigs_velvet.fa 19183_4#69 19183_4#69.contigs_velvet.fa references.txt", shell=True, check=True)
-
 sys.stderr.write("Tests completed\n")
 
diff --git a/test/test-dists.py b/test/test-dists.py
index 64cc2912..2843a804 100644
--- a/test/test-dists.py
+++ b/test/test-dists.py
@@ -1,10 +1,22 @@
 import sys
+#sys.path.insert(0, '../build/lib.macosx-10.9-x86_64-3.10')
 import pp_sketchlib
 import numpy as np
 import h5py
 import pickle
 import argparse
 
+def dist_in_tolerance(d1, d2, dist_name, sample_name, tols):
+    if d1 != 0 and d2 != 0:
+            diff = d1 - d2
+            diff_fraction = 2*(diff)/(d1 + d2)
+            if (abs(diff) > tols['warn_diff'] and abs(diff_fraction) > tols['warn_diff_frac']):
+                sys.stderr.write(dist_name + " mismatches for " + sample_name + "\n")
+                sys.stderr.write("expected: " + str(d1) + "; calculated: " + str(d2) + "\n")
+            if (abs(diff) > tols['err_diff'] and abs(diff_fraction) > tols['err_diff_frac']):
+                sys.stderr.write("Difference outside tolerance")
+                sys.exit(1)
+
 def iterDistRows(refSeqs, querySeqs, self=True):
     if self:
         if refSeqs != querySeqs:
@@ -36,6 +48,11 @@ def iterDistRows(refSeqs, querySeqs, self=True):
 
 args = parser.parse_args()
 
+tols = {"warn_diff": args.warn_diff,
+        "warn_diff_frac": args.warn_diff_frac,
+        "err_diff": args.error_diff,
+        "err_diff_frac": args.error_diff_frac}
+
 # Generate distances
 rList = []
 ref = h5py.File(args.ref_db + ".h5", 'r')
@@ -45,13 +62,25 @@ def iterDistRows(refSeqs, querySeqs, self=True):
 db_kmers = ref['sketches/' + rList[0]].attrs['kmers']
 ref.close()
 
-distMat = pp_sketchlib.queryDatabase(args.ref_db, args.ref_db, rList, rList, db_kmers)
-jaccard_dists = pp_sketchlib.queryDatabase(args.ref_db, args.ref_db, rList, rList,
-                                           db_kmers, jaccard = True)
-jaccard_dists_raw = pp_sketchlib.queryDatabase(args.ref_db, args.ref_db, rList, rList,
-                                           db_kmers, jaccard = True,
-                                           random_correct = False)
-distMat = np.hstack((distMat, jaccard_dists, jaccard_dists_raw))
+distMat = pp_sketchlib.queryDatabase(ref_db_name=args.ref_db,
+                                     query_db_name=args.ref_db,
+                                     rList=rList,
+                                     qList=rList,
+                                     klist=db_kmers)
+jaccard_dists = pp_sketchlib.queryDatabase(ref_db_name=args.ref_db,
+                                           query_db_name=args.ref_db,
+                                           rList=rList,
+                                           qList=rList,
+                                           klist=db_kmers,
+                                           jaccard = True)
+jaccard_dists_raw = pp_sketchlib.queryDatabase(ref_db_name=args.ref_db,
+                                               query_db_name=args.ref_db,
+                                               rList=rList,
+                                               qList=rList,
+                                               klist=db_kmers,
+                                               jaccard = True,
+                                               random_correct = False)
+distMat_all = np.hstack((distMat, jaccard_dists, jaccard_dists_raw))
 
 # Read old distances
 with open(args.results + ".pkl", 'rb') as pickle_file:
@@ -59,7 +88,7 @@ def iterDistRows(refSeqs, querySeqs, self=True):
 oldDistMat = np.load(args.results + ".npy")
 oldJaccardDistMat = np.load(args.results + "_jaccard.npy")
 oldRawJaccardDistMat = np.load(args.results + "_raw_jaccard.npy")
-oldDistMat = np.hstack((oldDistMat, oldJaccardDistMat, oldRawJaccardDistMat))
+oldDistMat_all = np.hstack((oldDistMat, oldJaccardDistMat, oldRawJaccardDistMat))
 
 # Check both match
 if (rList_old != rList):
@@ -71,14 +100,41 @@ def iterDistRows(refSeqs, querySeqs, self=True):
 names = iterDistRows(rList, rList, True)
 for i, (ref, query) in enumerate(names):
     for j, (dist) in enumerate(['core', 'accessory'] + [str(x) for x in db_kmers]):
-        if oldDistMat[i, j] != 0 and distMat[i, j] != 0:
-            diff = distMat[i, j] - oldDistMat[i, j]
-            diff_fraction = 2*(diff)/(oldDistMat[i, j] + distMat[i, j])
-            if (abs(diff) > args.warn_diff and abs(diff_fraction) > args.warn_diff_frac):
-                sys.stderr.write(dist + " mismatches for " + ref + "," + query + "\n")
-                sys.stderr.write("expected: " + str(oldDistMat[i, j]) + "; calculated: " + str(distMat[i, j]) + "\n")
-            if (abs(diff) > args.error_diff and abs(diff_fraction) > args.error_diff_frac):
-                sys.stderr.write("Difference outside tolerance")
-                sys.exit(1)
+        dist_in_tolerance(oldDistMat_all[i, j], distMat_all[i, j], dist, ref + "," + query, tols)
+
+# Test sparse queries
+square_core_mat = pp_sketchlib.longToSquare(distVec=distMat[:, [0]])
+
+kNN=3
+sparseDistMat = pp_sketchlib.querySelfSparse(ref_db_name=args.ref_db,
+                                             rList=rList,
+                                             klist=db_kmers,
+                                             kNN=kNN)
+sparse_knn = pp_sketchlib.sparsifyDists(distMat=square_core_mat, distCutoff=0, kNN=kNN)
+if (sparseDistMat[0] != sparse_knn[0] or
+    sparseDistMat[1] != sparse_knn[1]):
+    sys.stderr.write("Sparse distances (kNN) mismatching\n")
+    print(sparseDistMat)
+    print(sparse_knn)
+    sys.exit(1)
+
+for idx, (d1, d2) in enumerate(zip(sparseDistMat[2], sparseDistMat[2])):
+    dist_in_tolerance(d1, d2, "sparse distances (kNN)", str(idx), tols)
+
+cutoff=0.01
+sparseDistMat = pp_sketchlib.querySelfSparse(ref_db_name=args.ref_db,
+                                             rList=rList,
+                                             klist=db_kmers,
+                                             dist_cutoff=cutoff)
+sparse_threshold = pp_sketchlib.sparsifyDists(distMat=square_core_mat, distCutoff=cutoff, kNN=0)
+if (sparseDistMat[0] != sparse_threshold[0] or
+    sparseDistMat[1] != sparse_threshold[1]):
+    sys.stderr.write("Sparse distances (cutoff) mismatching\n")
+    print(sparseDistMat)
+    print(sparse_threshold)
+    sys.exit(1)
+
+for idx, (d1, d2) in enumerate(zip(sparseDistMat[2], sparseDistMat[2])):
+    dist_in_tolerance(d1, d2, "sparse distances (cutoff)", str(idx), tols)
 
 sys.exit(0)
diff --git a/test/test-matrix.py b/test/test-matrix.py
index ad3364c3..3cf227fe 100644
--- a/test/test-matrix.py
+++ b/test/test-matrix.py
@@ -3,14 +3,14 @@
 from math import sqrt
 
 # testing without install
-#sys.path.insert(0, '../build/lib.macosx-10.9-x86_64-3.8')
+#sys.path.insert(0, '../build/lib.macosx-10.9-x86_64-3.10')
 import pp_sketchlib
 try:
     from pp_sketch.matrix import sparsify
 except ImportError as e:
     from scipy.sparse import coo_matrix
     def sparsify(distMat, cutoff, kNN, threads):
-        sparse_coordinates = pp_sketchlib.sparsifyDists(distMat,
+        sparse_coordinates = pp_sketchlib.sparsifyDists(distMat=distMat,
                                                         distCutoff=cutoff,
                                                         kNN=kNN)
         sparse_scipy = coo_matrix((sparse_coordinates[2],
@@ -36,8 +36,10 @@ def check_res(res, expected):
 qr_mat = np.array([10, 11, 12, 13, 14, 15, 16, 17], dtype=np.float32)
 # NB r_idx is inner/fastest: q0r0, q0r1, q0r2, q0r3, q1r0, q1r1, q1r2, q1r3
 
-square1 = pp_sketchlib.longToSquare(rr_mat, 2)
-square2 = pp_sketchlib.longToSquareMulti(rr_mat, qr_mat, qq_mat)
+square1 = pp_sketchlib.longToSquare(distVec=rr_mat, num_threads=2)
+square2 = pp_sketchlib.longToSquareMulti(distVec=rr_mat,
+                                         query_ref_distVec=qr_mat,
+                                         query_query_distVec=qq_mat)
 
 square1_res = np.array([[0, 1, 2, 3],
                         [1, 0, 4, 5],
@@ -55,7 +57,7 @@ def check_res(res, expected):
 check_res(square1_res, square1)
 check_res(square2_res, square2)
 
-check_res(pp_sketchlib.squareToLong(square1_res, 2), rr_mat)
+check_res(pp_sketchlib.squareToLong(distMat=square1_res, num_threads=2), rr_mat)
 
 # sparsification
 sparse1 = sparsify(square2_res, cutoff=5, kNN=0, threads=2)