Speed up batch generation for EstimatorNeighborHood

[ivirshup]

Following up on discussion with @davidsebfischer

The current code for generating batches of data for EstimatorNeighborHood could be much faster. The current code looks like:

def previous(g, indices, features, max_indices): h_neighbors = [] a_neighborhood = np.zeros((len(indices), max_indices), "float32") for i, j in enumerate(indices): a_j = np.asarray(g[j, :].todense()).flatten() idx_neighbors = np.where(a_j > 0.)[0] h_neighbors_j = features[idx_neighbors, :] # if self.h0_in: # h_neighbors_j = self.h_0[key][idx_neighbors, :] # else: # h_neighbors_j = self.h_1[key][idx_neighbors, :][:, self.idx_neighbor_features] h_neighbors_j = np.expand_dims(h_neighbors_j, axis=0) # Pad neighborhoods: diff = max_indices - h_neighbors_j.shape[1] zeros = np.zeros((1, diff, h_neighbors_j.shape[2]), dtype="float32") h_neighbors_j = np.concatenate([h_neighbors_j, zeros], axis=1) h_neighbors.append(h_neighbors_j) a_neighborhood[i, :len(idx_neighbors)] = a_j[idx_neighbors] h_neighbors = np.concatenate(h_neighbors, axis=0) return a_neighborhood, h_neighbors

I'd suggest it instead looks something like:

import numpy as np from numba import njit @njit def pad_neighbors( data, indptr, max_indices: int ): out = np.zeros((len(indptr) - 1, max_indices), dtype=data.dtype) for i, vals in enumerate(np.split(data, indptr[1:-1])): out[i, :len(vals)] = vals return out @njit def pad_tensors( indices, indptr, features, max_indices: int, ): out = np.zeros( (len(indptr) - 1, max_indices, features.shape[1]), dtype=features.dtype ) for i, inds in enumerate(np.split(indices, indptr[1:-1])): out[i, :len(inds), :] = features[inds, :] return out def curr(g, indices, features, max_indices): g_sub = g[indices] a = pad_neighbors(g_sub.data, g_sub.indptr, max_indices) h = pad_tensors(g_sub.indices, g_sub.indptr, features, max_indices) return a, h

A quick demo:

import scanpy as sc adata = sc.datasets.pbmc3k_processed().raw.to_adata() g = adata.obsp["distances"] # k=10 test_inds = np.random.choice(adata.n_obs, 1000, replace=True) X = adata.obsm["X_pca"]

The values are the same:

a, h = curr(g, test_inds, X, 15) prev_a, prev_h = previous(g, test_inds, X, 15) assert np.array_equal(a.astype(np.float32), prev_a) assert np.array_equal(h, prev_h)

But the changed code is much faster:

In [115]: %timeit previous(g, test_inds, X, 15) 114 ms ± 2.72 ms per loop (mean ± std. dev. of 7 runs, 10 loops each) In [116]: %timeit curr(g, test_inds, X, 15) 1.5 ms ± 46.7 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)