Develop Reference

Fastest way to find nearest neighbours in NumPy array

What is the fastest way to perform operations on adjacent elements of an mxn array within distance $l$ (where m, n are large). If this was an image, it would equate to an operation on the surrounding pixels. To make things clearer, I've created a new array with the neighbours of the corresponding source.
Given some array like
x = [[1,2,3],
[4,5,6],
[7,8,9]]
if I were to take the [0,0] element, and want the surrounding elements at $l$=1, I'd need the [0,1] and [1,0] elements (namley 2 and 4). The desired output would look something like this
y = [[[2,4], [1,3,5], [2,6]],
[[1,5,7], [4,6,2,8], [3,9,5]],
[[4,8], [7,5,9], [8,6]]]
I've tried playing around with kdTree from scipy.spatial, and am aware of https://stackoverflow.com/a/45742628/20451990, but as far as I can tell this is actually finding the nearest data points, whereas I want to find the nearest array elements. I guess it could be naively done by iterating through, but that is very slow...
The end goal here is to generate combinations of nearby array elements which I will be taking the product of. For the example above this could be
[[1*2, 1*4], [2*1, 2*3, 2*5], [3*2, 3*6]],...]

Key takeaways
With numba, it is possible to get roughly 690x times faster algorithms than with naïve python code with for-loops and list appends.
With numba, functions have signature; you tell explicitly what is the datatype.
Avoid memory (re-)allocations. Try to allocate memory for any arrays in advance. Reuse the data containers whenever possible (See: cell_result in the numbafied process_cell())
Numba is not super handy with classes (at least, OOP style code), stuff which is dynamically typed, containers with mixed types or containers changing in size. Prefer simple functions and typed structures with defined size. See also: Supported Python features
Numba likes for-loops, and they're fast!
Prewords
You asked for a fastest way to calculate this. I had no baseline, so I created first a pure python for-loop solution as a baseline. Then, I used numba to make the code run fast. It most probably is not the fastest implementation but at least it is way faster than the naïve pure python for-loop approach.
So, if you are not familiar with numba this is a good way to learn about it a bit :)
Used test data
I use two pieces of test data. First, the simple array given in the question. I call this myarr, and it is used for easy comparison of the output:
import numpy as np
myarr = np.array(
[
[1, 2, 3],
[4, 5, 6],
[7, 8, 9],
],
dtype=np.float32,
)
The second dataset is for benchmarking. You mentioned that the arrays will be of size 30 x 30 and the distance I will be less than 4.
arr_large = np.arange(1, 30 * 30 + 1, 1, dtype=np.float32).reshape(30, 30)
In other words, the arr_large is a 30 x 30 2d-array:
>>> arr_large
array([[ 1., 2., 3., 4., 5., 6., 7., 8., 9., 10., 11.,
12., 13., 14., 15., 16., 17., 18., 19., 20., 21., 22.,
23., 24., 25., 26., 27., 28., 29., 30.],
...
[871., 872., 873., 874., 875., 876., 877., 878., 879., 880., 881.,
882., 883., 884., 885., 886., 887., 888., 889., 890., 891., 892.,
893., 894., 895., 896., 897., 898., 899., 900.]], dtype=float32)
I specified the dtype because specifying datatype is needed at the optimization step. For the pure python solution this is of course not necessary at all.
Baseline solution: Pure python with for-loops
I implemented the baseline soution with a python class and for-loops. The output from it looks like this (source for NeighbourProcessor below):
Example output with 3 x 3 input array (I=1)
n = NeighbourProcessor()
output = n.process(myarr, max_distance=1)
The output is then
>>> output
{(0, 0): [2, 4],
(0, 1): [2, 6, 10],
(0, 2): [6, 18],
(1, 0): [4, 20, 28],
(1, 1): [10, 20, 30, 40],
(1, 2): [18, 30, 54],
(2, 0): [28, 56],
(2, 1): [40, 56, 72],
(2, 2): [54, 72]}
which is same as
{(0, 0): [1 * 2, 1 * 4],
(0, 1): [2 * 1, 2 * 3, 2 * 5],
(0, 2): [3 * 2, 3 * 6],
(1, 0): [4 * 1, 4 * 5, 4 * 7],
(1, 1): [5 * 2, 5 * 4, 5 * 6, 5 * 8],
(1, 2): [6 * 3, 6 * 5, 6 * 9],
(2, 0): [7 * 4, 7 * 8],
(2, 1): [8 * 5, 8 * 7, 8 * 9],
(2, 2): [9 * 6, 9 * 8]}
This is basically what was asked in the question; the target ouput was
[[1*2, 1*4], [2*1, 2*3, 2*5], [3*2, 3*6]],...]
Here I used a dictionary with (row, column) as the key because that way you can more easily find the output for each cell.
Baseline performance
For the largest input of 30 x 30, and largest distance (I=4), the calculation takes about 0.188 seconds on my laptop:
>>> %timeit n.process(arr_large, max_distance=4)
188 ms ± 19.2 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
Code for NeighbourProcessor
import math
import numpy as np
class NeighbourProcessor:
def __init__(self):
self.arr = None
def process(self, arr, max_distance=1):
self.arr = arr
output = dict()
rows, columns = self.arr.shape
for current_row in range(rows):
for current_col in range(columns):
cell_result = self.process_cell(current_row, current_col, max_distance)
output[(current_row, current_col)] = cell_result
return output
def row_col_is_within_array(self, row, col):
if row < 0 or col < 0:
return False
if row > self.arr.shape[0] - 1 or col > self.arr.shape[1] - 1:
return False
return True
def distance(self, row, col, current_row, current_col):
distance_squared = (current_row - row) ** 2 + (current_col - col) ** 2
return np.sqrt(distance_squared)
def are_neighbours(self, row, col, current_row, current_col, max_distance):
if row == current_row and col == current_col:
return False
if not self.row_col_is_within_array(row, col):
return False
return self.distance(row, col, current_row, current_col) <= max_distance
def neighbours(self, current_row, current_col, max_distance):
start_row = math.floor(current_row - max_distance)
start_col = math.floor(current_col - max_distance)
end_row = math.ceil(current_row + max_distance)
end_col = math.ceil(current_col + max_distance)
for row in range(start_row, end_row + 1):
for col in range(start_col, end_col + 1):
if self.are_neighbours(
row, col, current_row, current_col, max_distance
):
yield row, col
def process_cell(self, current_row, current_col, max_distance):
cell_output = []
current_cell_value = self.arr[current_row][current_col]
for row, col in self.neighbours(current_row, current_col, max_distance):
neighbour_cell_value = self.arr[row][col]
cell_output.append(current_cell_value * neighbour_cell_value)
return cell_output
Short explanation
So what the NeighbourProcessor.process does is goes through the rows and columns of the input array, starting from (0,0), which is left top corner, and processing from left to right, top to bottom until the bottom right corner, which is (n_rows, n_columns), each time marking the cell as current cell; (current_row, current_column).
For each current cell, process it in process_cell. That will form an iterator with neighbours() which iterates all the neighbours at within maximum distance of I from the current cell. You can check how the logic goes in are_neighbours
Faster solution: Using numba and memory pre-allocation
Now I will make a functions-only version with numba, and try to make the processing as fast as possible. There is possibility also to use classes in numba, but they are still bit more experimental and complex, and this problem can be solved with functions only. The readability of the code suffers a bit, but that's the price we sometimes pay for speed optimization.
I'll start with the process function. Now it will have to create a a three dimensional array instead of a dict. The reason we want to create the array ahead of time because we memory allocation is a costly process and we want to do that exactly once. So, instead of having this as output for myarr:
# output[(row,column)]
#
output[(0,0)] # [2,4]
output[(0,1)] # [2, 6, 10]
#..etc
I want constant-sized output:
# output[row][column]
#
output[0][0] # [2, 4, nan, nan]
output[0][1] # [2, 6, 10, nan]
#..etc
Notice that after all the "pairs", the output is np.nan (not a number). Any postprocessing script must then just simply ignore the extra nans.
Solving for the required size for the pre-allocated array
How I know the size of the third dimension, i.e. the number of neighbours for given max. distance I? Well, I don't. It seems this is quite a complicated problem. See, for example this, this or the Gauss circle problem in Wikipedia. Nevertheless, I can quite easily calculate an upper bound for the number of neighbours. In the following I assume that neighbour is a neighbour if and only if the distance of the middle point of the cells is less or equal to I. If you create sketches with pen and paper, you will notice that when you increase the number of neighbours, the maximum number of neighbours grows as:
I = 1 -> max_number_neighbours = 4
I = 2 -> max_number_neighbours = 9
I = 3 -> max_number_neighbours = 28
Here is an example sketch with 10 x 10 2d-array and distance I=3, when current cell is (4,5), the number of neighbours must be less or equal to 28:
This pattern is represented as a function of max distance (I): (2*I-1)**2 + 4 -1, or
n_third_dimension = max_number_neighbours = (2*I-1)**2 + 3
Refactoring the code to work with numba
We start with creating the function signature of the entry point. In this case, we create a function process with the function signature:
#numba.jit("f4[:,:,:](f4[:,:], f4)")
def process(arr, max_distance):
...
See the docs for the other available types. The f4[:,:] just means that the input is 2d-array of float32 and f4[:,:,:](....) means that the function output is 3d-array of float32. Next, we create the output with the formula we invented above. Here is one part of the magic: memory pre-allocation with np.empty:
n_third_dimension = (2 * math.ceil(max_distance) - 1) ** 2 + 3
output = np.empty((*arr.shape, n_third_dimension), dtype=np.float32)
cell_result = np.empty(n_third_dimension, dtype=np.float32)
Numbafied code
I will not walk though the rest of the code hand-in-hand, but you can see below that it is a bit modified version of the pure python for-loop baseline.
import math
import numba
import numpy as np
#numba.njit("f4(i4,i4,i4,i4)")
def distance(row, col, current_row, current_col):
distance_squared = (current_row - row) ** 2 + (current_col - col) ** 2
return np.sqrt(distance_squared)
#numba.njit("boolean(i4,i4, i4,i4)")
def row_col_is_within_array(
row,
col,
arr_rows,
arr_cols,
):
if row < 0 or col < 0:
return False
if row > arr_rows - 1 or col > arr_cols - 1:
return False
return True
#numba.njit("boolean(i4,i4,i4,i4,f4,i4,i4)")
def are_neighbours(
neighbour_row,
neighbour_col,
current_row,
current_col,
max_distance,
arr_rows,
arr_cols,
):
if neighbour_row == current_row and neighbour_col == current_col:
return False
if not row_col_is_within_array(
neighbour_row,
neighbour_col,
arr_rows,
arr_cols,
):
return False
return (
distance(neighbour_row, neighbour_col, current_row, current_col) <= max_distance
)
#numba.njit("f4[:](f4[:,:], f4[:], i4,i4,i4,f4)")
def process_cell(
arr, cell_result, current_row, current_col, n_third_dimension, max_distance
):
for i in range(n_third_dimension):
cell_result[i] = np.nan
current_cell_value = arr[current_row][current_col]
# Potential cell neighbour area
start_row = math.floor(current_row - max_distance)
start_col = math.floor(current_col - max_distance)
end_row = math.ceil(current_row + max_distance)
end_col = math.ceil(current_col + max_distance)
arr_rows, arr_cols = arr.shape
cell_pointer = 0
for neighbour_row in range(start_row, end_row + 1):
for neighbour_col in range(start_col, end_col + 1):
if are_neighbours(
neighbour_row,
neighbour_col,
current_row,
current_col,
max_distance,
arr_rows,
arr_cols,
):
neighbour_cell_value = arr[neighbour_row][neighbour_col]
cell_result[cell_pointer] = current_cell_value * neighbour_cell_value
cell_pointer += 1
return cell_result
#numba.njit("f4[:,:,:](f4[:,:], f4)")
def process(arr, max_distance):
n_third_dimension = (2 * math.ceil(max_distance) - 1) ** 2 + 3
output = np.empty((*arr.shape, n_third_dimension), dtype=np.float32)
cell_result = np.empty(n_third_dimension, dtype=np.float32)
rows, columns = arr.shape
for current_row in range(rows):
for current_col in range(columns):
cell_result = process_cell(
arr,
cell_result,
current_row,
current_col,
n_third_dimension,
max_distance,
)
output[current_row][current_col][:] = cell_result
return output
Example output
>>> output = process(myarr, max_distance=1.0)
>>> output
array([[[ 2., 4., nan, nan],
[ 2., 6., 10., nan],
[ 6., 18., nan, nan]],
[[ 4., 20., 28., nan],
[10., 20., 30., 40.],
[18., 30., 54., nan]],
[[28., 56., nan, nan],
[40., 56., 72., nan],
>>> output[0]
array([[ 2., 4., nan, nan],
[ 2., 6., 10., nan],
[ 6., 18., nan, nan]], dtype=float32)
>>> output[0][1]
array([ 2., 6., 10., nan], dtype=float32)
# Above is the same as target: [2 * 1, 2 * 3, 2 * 5]
Speed of the numbafied code and closing words
The baseline approach rxecution time was 188 ms. Now, it is 271 µs. That is only 0.00144 times of what the original code took! (99.85% reduction in execution time. Some would say 693x faster.).
>>> %timeit process(arr_large, max_distance=4.0)
271 µs ± 2.88 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
Note that you might want to calculate the distance differently, or add there weighting, or some more complex logic, aggregation functions, etc. This could be still further optimized a bit by creating better estimate for the maximum number of neighbors, for example. Have fun with numba, and I hope you learned something! :)
Bonus tip: There is also ahead of time compilation in numba which you can use to make also the first function call fast!

Scipy function that can do np.diff() with compressed sparse column matrix

I want to compute discrete difference of identity matrix.
The code below use numpy and scipy.
import numpy as np
from scipy.sparse import identity
from scipy.sparse import csc_matrix
x = identity(4).toarray()
y = csc_matrix(np.diff(x, n=2))
print(y)
I would like to improve performance or memory usage.
Since identity matrix produce many zeros, it would reduce memory usage to perform calculation in compressed sparse column(csc) format. However, np.diff() does not accept csc format, so converting between csc and normal format using csc_matrix would slow it down a bit.
Normal format
x = identity(4).toarray()
print(x)
[[1. 0. 0. 0.]
[0. 1. 0. 0.]
[0. 0. 1. 0.]
[0. 0. 0. 1.]]
csc format
x = identity(4)
print(x)
(0, 0) 1.0
(1, 1) 1.0
(2, 2) 1.0
(3, 3) 1.0
Thanks

Here is my hacky solution to get the sparse matrix as you want.
L - the length of the original identity matrix,
n - the parameter of np.diff.
In your question they are:
L = 4
n = 2
My code produces the same y as your code, but without the conversions between csc and normal formats.
Your code:
from scipy.sparse import identity, csc_matrix
x = identity(L).toarray()
y = csc_matrix(np.diff(x, n=n))
My code:
from scipy.linalg import pascal
def get_data(n, L):
nums = pascal(n + 1, kind='lower')[-1].astype(float)
minuses_from = n % 2 + 1
nums[minuses_from : : 2] *= -1
return np.tile(nums, L - n)
data = get_data(n, L)
row_ind = (np.arange(n + 1) + np.arange(L - n).reshape(-1, 1)).flatten()
col_ind = np.repeat(np.arange(L - n), n + 1)
y = csc_matrix((data, (row_ind, col_ind)), shape=(L, L - n))
I have noticed that after applying np.diff to the identity matrix n times, the values of the columns are the binomial coefficients with their signs alternating. This is my variable data.
Then I am just constructing the csc_matrix.

Unfortunately, it does not seem that SciPy provides any tools for this kind of sparse matrix manipulation. Regardless, by cleverly manipulating the indices and data of the entries one can emulate np.diff(x,n) in a straightforward fashion.
Given 2D NumPy array (matrix) of dimension MxN, np.diff() multiplies each column (of column index y) with -1 and adds the next column to it (column index y+1). Difference of order k is just the iterative application of k differences of order 1. A difference of order 0 is just the returns the input matrix.
The method below makes use of this, iterateively eliminating duplicate entries by addition through sum_duplicates(), reducing the number of columns by one, and filtering non-valid indices.
def csc_diff(x, n):
'''Emulates np.diff(x,n) for a sparse matrix by iteratively taking difference of order 1'''
assert isinstance(x, csc_matrix) or (isinstance(x, np.ndarray) & len(x.shape) == 2), "Input matrix must be a 2D np.ndarray or csc_matrix."
assert isinstance(n, int) & n >= 0, "Integer n must be larger or equal to 0."
if n >= x.shape[1]:
return csc_matrix(([], ([], [])), shape=(x.shape[0], 0))
if isinstance(x, np.ndarray):
x = csc_matrix(x)
# set-up of data/indices via column-wise difference
if(n > 0):
for k in range(1,n+1):
# extract data/indices of non-zero entries of (current) sparse matrix
M, N = x.shape
idx, idy = x.nonzero()
dat = x.data
# difference: this row (y) * (-1) + next row (y+1)
idx = np.concatenate((idx, idx))
idy = np.concatenate((idy, idy-1))
dat = np.concatenate(((-1)*dat, dat))
# filter valid indices
validInd = (0<=idy) & (idy<N-1)
# x_diff: csc_matrix emulating np.diff(x,1)'s output'
x_diff = csc_matrix((dat[validInd], (idx[validInd], idy[validInd])), shape=(M, N-1))
x_diff.sum_duplicates()
x = x_diff
return x
Moreover, the method outputs an empty csc_matrix of dimension Mx0 when the difference order is larger or equal to the number of columns of the input matrix. This is why the output is identical, see
csc_diff(x, 2).toarray()
> array([[ 1., 0.],
[-2., 1.],
[ 1., -2.],
[ 0., 1.]])
which is identical to
np.diff(x.toarray(), 2)
> array([[ 1., 0.],
[-2., 1.],
[ 1., -2.],
[ 0., 1.]])
This identity holds for other difference orders, too
(csc_diff(x, 0).toarray() == np.diff(x.toarray(), 0)).all()
>True
(csc_diff(x, 3).toarray() == np.diff(x.toarray(), 3)).all()
>True
(csc_diff(x, 13).toarray() == np.diff(x.toarray(), 13)).all()
>True

CVXPY constraints formulation

I am trying to solve the Isoperimetric problem (7.14) from Additional Exercises for Convex Optimization by Stephen Boyd using CVXPY. The problem formulation is:
The code of constraints is given below:
constraints = [ y[1] == 0,
y[N] == 0,
y[F] == yfixed[F],
cp.abs(( y[i+2] - 2 * y[i+1] + y[i]) / h**2) <= C for i in np.arange(1,199) ] #not using the first constraint here
The constraints have for loops, and when I tried to formulate the problem according to the CVXPY documentation, I got the following error
Invalid syntax
cp.abs(( y[i+2] - 2 * y[i+1] + y[i]) / h**2) <= C for i in np.arange(1,199) ]
^
How to use the loops in CVXPY constraints?

You need to express the constraints in terms of matrix-vector equalities and inequalities which follow the DCP protocol for cvxpy.
To elaborate, I can see three kinds of constraints in this problem: (I am going to assume 0-based indexing for programming convenience for the rest of the answer.)
Take y as the N+1 dimensional optimization variable.
Fixed point equality constraints: These basically set some indices of the y vector to a set of given values. Note that the zero boundary conditions y[0] == 0 and y[N] == 0 also fall under this.
Perimeter constraint: This is to be computed using successive differences. And finally we set something like the sum of the square roots of 1 plus the squares of the differences to be less than L. This part is probably the most involved one to write following the cvxpy protocols.
Curvature constraints: This also involves a calculation similar to successive differences like above but this is much easier to write as you'll see this is simply a matrix-vector multiplication type constraint like the first one.
Now let's write the constraints in code.
Necessary imports:
import numpy as np
import cvxpy as cp
import matplotlib.pyplot as plt
from scipy.linalg import circulant
1. Equality constraints:
These basically pick some indices from y and set those to given values. This can be implemented as follows:
def equality_constraints(N, F, vals):
'''
Sets some indices (F) in the y vector to given values. Also sets
y[0] == 0 and y[N] == 0.
Returns E (a matrix) and e (a vector) such that E # y == e.
'''
E = np.zeros((F.shape[0]+2, N+1)) # '+2' for selecting y[0] and y[N] also
e = np.zeros(F.shape[0]+2)
E[0, 0] = 1
E[-1, -1] = 1
if F.shape[0]:
E[1:-1:1][np.arange(F.shape[0]), F] = 1
e[1:-1:1] = vals
return E, e
E is a binary matrix of shape (F.shape[0] + 2) x (N+1) and it has exactly one column set to 1 in each row, giving an index for the (N+1) dimensional vector y and e contains the value for that index of y.
2. Perimeter constraint:
For this, we need successive differences of the form y[i+1] - y[i] for i = 0, 1, 2, . . . , N-1. Note that we can similarly construct a vector having this N successive differences as its elements. We can perform the square root and other operations on this vector easily using vectorized computations. Here we are constructing an N x (N+1) matrix M, which when multiplied by y will give the N differences.
def length_matrix(N):
'''
Returns L with [-1, 1, 0, . . . , 0] as first row and sliding it
to the right to get the following rows.
'''
val = np.array([-1, 1])
offsets = np.array([0, 1])
col0 = np.zeros(N+1)
col0[offsets] = val
M = circulant(col0).T[:-(len(val) - 1)]
return M
The matrix M will be a circulant matrix. I simply transposed it and removed the last row to get the desired matrix. You can see this post to know how to create one such matrix. M looks like this:
array([[-1., 1., 0., ..., 0., 0., 0.],
[ 0., -1., 1., ..., 0., 0., 0.],
[ 0., 0., -1., ..., 0., 0., 0.],
...,
[ 0., 0., 0., ..., 1., 0., 0.],
[ 0., 0., 0., ..., -1., 1., 0.],
[ 0., 0., 0., ..., 0., -1., 1.]])
3. Curvature constraints:
Exactly same matrix calculation like the last one. Just repeat and slide [1, -2, 1] along the rows!
def curvature_constraints(N, C, h):
'''
Returns D and C_vec to be used as D # y <= C and D # y >= -C
as curvature constraints.
'''
val = np.array([1, -2, 1])
offsets = np.array([0, 1, 2])
col0 = np.zeros(N+1)
col0[offsets] = val
D = circulant(col0).T[:-(len(val) - 1)]
D /= h**2
C_vec = np.ones(D.shape[0]) * C
return D, C_vec
I am dividing by h**2 in the matrix itself.
Example:
I have taken this example from the site of the book itself. The data is also available here.
L = 1.5
a = 1
C = 15
N = 200
h = a/N
F = np.array([20,40,140,180]) # fixed points
vals = np.array([0.1, 0.15, 0.15, 0.2])
# Declare an array for plotting purposes
yfixed = np.zeros(N+1)
yfixed[F] = vals
x = np.linspace(0, a, N+1)
Problem Formulation and Solution:
I am leaving it for you to understand how I have assembled the matrices in formulating the constraints, especially the perimeter one. This is not difficult, but might require you some practice depending on how comfortable you are with vectorization. The DCP page is a very good place to start.
y = cp.Variable(N+1)
E, e = equality_constraints(N, F, vals)
M = length_matrix(N)
D, d = curvature_constraints(N, C, h)
constraints = [
E # y == e,
h * cp.sum(cp.norm(cp.vstack([(M # y)/h, np.ones(N)]), p = 2, axis = 0)) <= L,
D # y <= d,
D # y >= -d
]
objective_function = h * cp.sum(y)
objective = cp.Maximize(objective_function)
problem = cp.Problem(objective, constraints)
problem.solve()
plt.plot(0, 0, 'ko')
plt.plot(a, 0, 'ko')
for i in F:
plt.plot(x[i], yfixed[i], 'bo')
plt.plot(x, y.value) # y.value gives the value of the cp Variable
plt.savefig('curve.png')
I got the answer as 0.1594237500556726 for the above example and the curve looks like this:
I have checked this solution with few other contrived test cases to verify correctness. However, there might be other more efficient solutions formulating this problem differently or there might even be some unexpected or embarrassing errors here! Feel free to let me know in case there is some error or you find anything difficult to understand in the answer.

Try to split in two:
constraints = [ y[1] == 0,
y[N] == 0,
y[F] == yfixed[F] ] +
[ cp.abs(( y[i+2] - 2 * y[i+1] + y[i]) / h**2) <= C for i in np.arange(1,199) ]

Column/row slicing a torch sparse tensor

I have a pytorch sparse tensor that I need sliced row/column wise using this slice [idx][:,idx] where idx is a list of indexes, using the mentioned slice yields my desired result on an ordinary float tensor. Is it possible applying the same slicing on a sparse tensor? Example here:
#constructing sparse matrix
i = np.array([[0,1,2,2],[0,1,2,1]])
v = np.ones(4)
i = torch.from_numpy(i.astype("int64"))
v = torch.from_numpy(v.astype("float32"))
test1 = torch.sparse.FloatTensor(i, v)
#constructing float tensor
test2 = np.array([[1,0,0],[0,1,0],[0,1,1]])
test2 = autograd.Variable(torch.cuda.FloatTensor(test2), requires_grad=False)
#slicing
idx = [1,2]
print(test2[idx][:,idx])
output:
Variable containing:
1 0
1 1
[torch.cuda.FloatTensor of size 2x2 (GPU 0)]
I am holding a 250.000 x 250.000 adjacency matrix, where I need to slice n rows and n columns, using the random idx, by simply sampling n random idx's. Since the dataset is so large it is not realistic to convert to a more convenient datatype.
can I achieve the same slicing result on test1? Is it even possible? If not, are there any work-arounds?
Right now I am running my model with the following "hack" of a solution:
idx = sorted(random.sample(range(0, np.shape(test1)[0]), 9000))
test1 = test1AsCsr[idx][:,idx].todense().astype("int32")
test1 = autograd.Variable(torch.cuda.FloatTensor(test1), requires_grad=False)
Where test1AsCsr is my test1 converted to a numpy CSR matrix. This solution works, it is however very slow, and makes my GPU utilization very low, since it needs to read/write from CPU memory, constantly.
Edit: Its fine with a non-sparse tensor as result

Well it's been a couple of years since there was activity on this question, but better late than never.
This is the function I use for slicing sparse tensors. (Helper functions are below)
def slice_torch_sparse_coo_tensor(t, slices):
"""
params:
-------
t: tensor to slice
slices: slice for each dimension
returns:
--------
t[slices[0], slices[1], ..., slices[n]]
"""
t = t.coalesce()
assert len(args) == len(t.size())
for i in range(len(args)):
if type(args[i]) is not torch.Tensor:
args[i] = torch.tensor(args[i], dtype= torch.long)
indices = t.indices()
values = t.values()
for dim, slice in enumerate(args):
invert = False
if t.size(0) * 0.6 < len(slice):
invert = True
all_nodes = torch.arange(t.size(0))
unique, counts = torch.cat([all_nodes, slice]).unique(return_counts=True)
slice = unique[counts==1]
if slice.size(0) > 400:
mask = ainb_wrapper(indices[dim], slice)
else:
mask = ainb(indices[dim], slice)
if invert:
mask = ~mask
indices = indices[:, mask]
values = values[mask]
return torch.sparse_coo_tensor(indices, values, t.size()).coalesce()
Usage (took 2.4s on my machine):
indices = torch.randint(low= 0, high= 200000, size= (2, 1000000))
values = torch.rand(size=(1000000,))
t = torch.sparse_coo_tensor(indices, values, size=(200000, 200000))
idx = torch.arange(1000)
slice_coo(t, [idx, idx])
out:
tensor(indices=tensor([[ 13, 62, 66, 78, 134, 226, 233, 266, 299, 344, 349,
349, 369, 396, 421, 531, 614, 619, 658, 687, 769, 792,
810, 840, 926, 979],
[255, 479, 305, 687, 672, 867, 444, 559, 772, 96, 788,
980, 423, 699, 911, 156, 267, 721, 381, 781, 97, 271,
840, 292, 487, 185]]),
values=tensor([0.4260, 0.4816, 0.8001, 0.8815, 0.3971, 0.4914, 0.7068,
0.2329, 0.4038, 0.1757, 0.7758, 0.3210, 0.2593, 0.8290,
0.1320, 0.4322, 0.7529, 0.8341, 0.8128, 0.4457, 0.4100,
0.1618, 0.4097, 0.3088, 0.6942, 0.5620]),
size=(200000, 200000), nnz=26, layout=torch.sparse_coo)
Timings for slice_torch_sparse_coo_tensor:
%timeit slice_torch_sparse_coo_tensor(t, torch.randperm(200000)[:500], torch.arange(200000))
output:
1.08 s ± 447 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
for the built-in torch.index_select (implemented here):
%timeit t.index_select(0, torch.arange(100))
output:
56.7 s ± 4.87 s per loop (mean ± std. dev. of 7 runs, 1 loop each)
These are the helper functions I use for this purpose, the function "ainb" finds the elements of a that are in b. I found this function in the internet a while ago but I can't find the post to link it.
import torch
def ainb(a,b):
"""gets mask for elements of a in b"""
size = (b.size(0), a.size(0))
if size[0] == 0: # Prevents error in torch.Tensor.max(dim=0)
return torch.tensor([False]*a.size(0), dtype= torch.bool)
a = a.expand((size[0], size[1]))
b = b.expand((size[1], size[0])).T
mask = a.eq(b).max(dim= 0).values
return mask
def ainb_wrapper(a, b, splits = .72):
inds = int(len(a)**splits)
tmp = [ainb(a[i*inds:(i+1)*inds], b) for i in list(range(inds))]
return torch.cat(tmp)
Since the function scales quadratically with the amount of elements I added a wrapper that splits the input into chunks and then concatenates the output. It's more efficient using only CPU, but I am not sure whether this holds when using a GPU, would appreciate it if someone could test it :)
It's my first time posting, so feedback on the quality of the post is also appreciated.

Possible answer for 2-dimentional sparse indices
Find an answer below, playing with several pytorch methods (torch.eq(), torch.unique(), torch.sort(), etc.) in order to output a compact, sliced tensor of shape (len(idx), len(idx)).
I tested several edge cases (unordered idx, v with 0s, i with multiple same index pairs, etc.), though I may have forgot some. Performance should also be checked.
import torch
import numpy as np
def in1D(x, labels):
"""
Sub-optimal equivalent to numpy.in1D().
Hopefully this feature will be properly covered soon
c.f. https://github.com/pytorch/pytorch/issues/3025
Snippet by Aron Barreira Bordin
Args:
x (Tensor): Tensor to search values in
labels (Tensor/list): 1D array of values to search for
Returns:
Tensor: Boolean tensor y of same shape as x, with y[ind] = True if x[ind] in labels
Example:
>>> in1D(torch.FloatTensor([1, 2, 0, 3]), [2, 3])
FloatTensor([False, True, False, True])
"""
mapping = torch.zeros(x.size()).byte()
for label in labels:
mapping = mapping | x.eq(label)
return mapping
def compact1D(x):
"""
"Compact" values 1D uint tensor, so that all values are in [0, max(unique(x))].
Args:
x (Tensor): uint Tensor
Returns:
Tensor: uint Tensor of same shape as x
Example:
>>> densify1D(torch.ByteTensor([5, 8, 7, 3, 8, 42]))
ByteTensor([1, 3, 2, 0, 3, 4])
"""
x_sorted, x_sorted_ind = torch.sort(x, descending=True)
x_sorted_unique, x_sorted_unique_ind = torch.unique(x_sorted, return_inverse=True)
x[x_sorted_ind] = x_sorted_unique_ind
return x
# Input sparse tensor:
i = torch.from_numpy(np.array([[0,1,4,3,2,1],[0,1,3,1,4,1]]).astype("int64"))
v = torch.from_numpy(np.arange(1, 7).astype("float32"))
test1 = torch.sparse.FloatTensor(i, v)
print(test1.to_dense())
# tensor([[ 1., 0., 0., 0., 0.],
# [ 0., 8., 0., 0., 0.],
# [ 0., 0., 0., 0., 5.],
# [ 0., 4., 0., 0., 0.],
# [ 0., 0., 0., 3., 0.]])
# note: test1[1, 1] = v[i[1,:]] + v[i[6,:]] = 2 + 6 = 8
# since both i[1,:] and i[6,:] are [1,1]
# Input slicing indices:
idx = [4,1,3]
# Getting the elements in `i` which correspond to `idx`:
v_idx = in1D(i, idx).byte()
v_idx = v_idx.sum(dim=0).squeeze() == i.size(0) # or `v_idx.all(dim=1)` for pytorch 0.5+
v_idx = v_idx.nonzero().squeeze()
# Slicing `v` and `i` accordingly:
v_sliced = v[v_idx]
i_sliced = i.index_select(dim=1, index=v_idx)
# Building sparse result tensor:
i_sliced[0] = compact1D(i_sliced[0])
i_sliced[1] = compact1D(i_sliced[1])
# To make sure to have a square dense representation:
size_sliced = torch.Size([len(idx), len(idx)])
res = torch.sparse.FloatTensor(i_sliced, v_sliced, size_sliced)
print(res)
# torch.sparse.FloatTensor of size (3,3) with indices:
# tensor([[ 0, 2, 1, 0],
# [ 0, 1, 0, 0]])
# and values:
# tensor([ 2., 3., 4., 6.])
print(res.to_dense())
# tensor([[ 8., 0., 0.],
# [ 4., 0., 0.],
# [ 0., 3., 0.]])
Previous answer for 1-dimentional sparse indices
Here is a (probably sub-optimal and not covering all edge cases) solution, following the intuitions shared in a related open issue (hopefully this feature will be properly covered soon):
# Constructing a sparse tensor a bit more complicated for the sake of demo:
i = torch.LongTensor([[0, 1, 5, 2]])
v = torch.FloatTensor([[1, 3, 0], [5, 7, 0], [9, 9, 9], [1,2,3]])
test1 = torch.sparse.FloatTensor(i, v)
# note: if you directly have sparse `test1`, you can get `i` and `v`:
# i, v = test1._indices(), test1._values()
# Getting the slicing indices:
idx = [1,2]
# Preparing to slice `v` according to `idx`.
# For that, we gather the list of indices `v_idx` such that i[v_idx[k]] == idx[k]:
i_squeeze = i.squeeze()
v_idx = [(i_squeeze == j).nonzero() for j in idx] # <- doesn't seem optimal...
v_idx = torch.cat(v_idx, dim=1)
# Slicing `v` accordingly:
v_sliced = v[v_idx.squeeze()][:,idx]
# Now defining your resulting sparse tensor.
# I'm not sure what kind of indexing you want, so here are 2 possibilities:
# 1) "Dense" indixing:
test1x = torch.sparse.FloatTensor(torch.arange(v_idx.size(1)).long().unsqueeze(0), v_sliced)
print(test1x)
# torch.sparse.FloatTensor of size (3,2) with indices:
#
# 0 1
# [torch.LongTensor of size (1,2)]
# and values:
#
# 7 0
# 2 3
# [torch.FloatTensor of size (2,2)]
# 2) "Sparse" indixing using the original `idx`:
test1x = torch.sparse.FloatTensor(autograd.Variable(torch.LongTensor(idx)).unsqueeze(0), v_sliced)
# note: this indexing would fail if elements of `idx` were not in `i`.
print(test1x)
# torch.sparse.FloatTensor of size (3,2) with indices:
#
# 1 2
# [torch.LongTensor of size (1,2)]
# and values:
#
# 7 0
# 2 3
# [torch.FloatTensor of size (2,2)]

How to create a rectangular matrix in Python?

m,n=input()
a=[0]*n
for i in range(0,m):
a[i]=[0]*m
for i in range(0,m):
for j in range(0,n):
a[i][j]=input()
print a
Consider the above piece of code written in Python 2.7.4 to accept a 2-D array and then print it. This code functions well but it should accept any 2-D array means for example the values of m and could be 3,2 respectively let's say but it only accepts a square matrix. We cannot create a rectangular matrix because it gives the error: index out of range if the values of m and n are not equal. Is there any way to create a rectangular matrix just like we can do in C/C++ easily?

Numpy is a great module for fast linear algebra operations. You can create a rectangular array with Numpy which is essentially a matrix. (Numpy also has matrix functions too but they are a little more tedious to work with).
As an example, create a 3x4 array as follows
import numpy as np
input = np.zeros((3, 4)) #This creates a 3x4 array. It is good practice to initialize your array with 0's
input[0][3] = 5 #Fill out your array. The 0,0 index is the top left corner
In [42]:input
Out[42]:
array([[ 0., 0., 0., 5.],
[ 0., 0., 0., 0.],
[ 0., 0., 0., 0.]])
Proceed to fill out the rest of your rectangular matrix as per normal.

You can do something like this:
rows = input()
cols = input()
arr = [[None] * cols for _ in range(rows)]
for i in range(rows):
for j in range(cols):
print(arr)
arr[i][j] = input()
print(arr)
You should see this similar question: How to initialize a two-dimensional array in Python?
As far as getting your exact code to work, you just have m and n switched on lines 2 & 4:
m = input() # rows
n = input() # cols
a = [0] * m # Initialize rows
for i in range(0, m): # For each row (must be n rows)
a[i] = [0] * n # Initialize a column
for i in range(0, m): # For each row
for j in range(0, n): # For each column
a[i][j] = input() # Input a value
print a