# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Planet
#     language: python
#     name: planet
# ---

# %env CUDA_VISIBLE_DEVICES=0

import numpy as np
import faiss

particle_count = 10**5
particles = np.random.uniform(size=(particle_count, 3)).astype(dtype='float32')

# GPU
gpu_resources = faiss.StandardGpuResources()
gpu_device_id = 0

# %%timeit
index_gpu = faiss.IndexFlatL2(particles.shape[1])
index_gpu = faiss.index_cpu_to_gpu(gpu_resources, gpu_device_id, index_gpu)
index_gpu.add(particles)
distances, indices = index_gpu.search(particles, 50)

# %%timeit
# CPU
index_cpu = faiss.IndexHNSWFlat(particles.shape[1], 32)
# Parameters require empirical tuning
# index_flat.efConstruction = 40
index_cpu.efSearch = 100
index_cpu.add(particles)
distances, indices = index_cpu.search(particles, 50)

# +
# (indices[:, 0] == np.arange(len(points))).mean()

StarcoderdataJupyterScriptsDedupFiltered