When the number of sphere is large, it is generally a good idea to use a data structure like a quadtree, k-D tree and especially ball tree. This is what Scipy uses internally and this result in an algorithm running in O(M log(N)) time where N is the number of sphere and M is the number of points.
The thing is implementing such data structure efficiently on GPU is pretty insane (because they are not SIMD friendly). In your case, the number of sphere is small enough so that using a naive algorithm running in O(N M) should be relatively fast.
The solution of @DanielF is a good start but is is not very efficient as is creates several temporary arrays and use a memory access pattern that is not SIMD friendly (which is critical on GPU so to get reasonable performances). The temporary arrays cause slow allocations and memory accesses while GPU shine for heavy computational operations (typically on large arrays).
To solve this problem, you can write your own kernel. The thing is the input array are not efficiently stored. They should be transposed so the computation can make more SIMD-friendly memory accesses and transposing arrays is relatively slow on GPU (as well as creating new arrays). Please consider using a more suitable input data layout. Additionally, computing the square root is not required since you can compare the result to the squared radius. Here is an example computing a "collision" matrix (where rows are spheres and columns are points):
import cupy as cp
points = cp.random.rand(300_000, 3)
spheres = cp.random.rand(200, 3)
radii = cp.random.rand(200)
checkInSphere = cp.RawKernel(r'''
extern "C" __global__
void checkInSphere(const double* points, int pointCount, const double* spheres, const double* radii, int sphereCount, bool* found) {
int tid = blockDim.x * blockIdx.x + threadIdx.x;
if(tid < pointCount)
{
const double px = points[tid];
const double py = points[tid+pointCount];
const double pz = points[tid+pointCount*2];
for(int i=0 ; i<sphereCount ; ++i)
{
const double r = radii[i];
const double sx = spheres[i];
const double sy = spheres[i+sphereCount];
const double sz = spheres[i+sphereCount*2];
const double dx = px - sx;
const double dy = py - sy;
const double dz = pz - sz;
const double sqDist = dx * dx + dy * dy + dz * dz;
found[i*pointCount+tid] = sqDist <= r * r;
}
}
}
''', 'checkInSphere')
threadPerBlock = 64
blockCount = (points.size + (threadPerBlock - 1)) // threadPerBlock
found = cp.empty((spheres.shape[0], points.shape[0]), dtype=cp.bool_)
checkInSphere((blockCount,), (threadPerBlock,), (points.T.copy(), points.shape[0], spheres.T.copy(), radii, spheres.shape[0], found))
This code is about 4 times faster in double precision than the one of @DanielF (ie. the last version without the square root and the check). The thing is my GPU like most GPUs is VERY slow to operate on double-precision (DP) floating-point numbers. In fact, mainstream client-side GPUs are not designed for such workload and are generally slower than CPUs for that. However, (expensive) server-side GPUs can efficiently operate on double-precision numbers. Thus, please consider working with simple-precision floating-point (SP) numbers if you are using a basic mainstream GPU. If you really want to use DP, then this computation is certainly faster on a CPU using tools like Numba. Note that using SP operations are generally also faster on server-side GPUs as well as CPUs. On my machine, the SP version of this code takes about 0.7 ms and is about 15 times faster than the one of @DanielF (updated version).
For better performance, please consider changing the input layout, pre-allocate all the buffers, pre-compute the squares of radii, and make use of shared memory so to load the sphere locations only once. Loop unrolling and register tiling should also help to improve performance even further though it will make the code barely readable. One could also store only the index of the closest sphere for each point, or even only a boolean to say if the point is in any sphere so to reduce the size of the output (memory accesses are expensive).
