Slice a 3d numpy array using a 1d lookup between indices - python

Slice a 3d numpy array using a 1d lookup between indices
import numpy as np
a = np.arange(12).reshape(2, 3, 2)
b = np.array([2, 0])
b maps i to j where i and j are the first 2 indexes of a, so ​a[i,j,k]
Desired result after applying b to a is:
[[4 5]
​ [6 7]]
Naive solution:
c = np.empty(shape=(2, 2), dtype=int)
for i in range(2):
​j = b[i]
​c[i, :] = a[i, j, :]
Question: Is there a way to do this using a numpy or scipy routine or routines or fancy indexing?
Application: Reinforcement Learning finite MDPs where b is a deterministic policy vector pi(a|s), a is the state transition probabilities p(s'|s,a) and c is the state transition matrix for that policy vector p(s'|s). The arrays will be large and this operation will be repeated a large number of times so needs to be scaleable and fast.
What I have tried:
Compiling using numba but line profiler suggests my code is slower compared to a similarly sized numpy routine. Also numpy is more widely understood and used.
Maintaining pi(a|s) as a sparse matrix (all zero except one 1 per row) b_as_a_matrix and then using einsum but this involves storing and updating the matrix and creates more work (an extra loop over j and sum operation).
c = np.einsum('ij,ijk->ik', b_as_a_matrix, a)

Numpy arrays can be indexed using other arrays as indices. See also: NumPy selecting specific column index per row by using a list of indexes.
With that in mind, we can vectorize your loop to simply use b for indexing:
>>> import numpy as np
>>> a = np.arange(12).reshape(2, 3, 2)
>>> b = np.array([2, 0])
>>> i = np.arange(len(b))
>>> i
array([0, 1])
>>> a[i, b, :]
array([[4, 5],
[6, 7]])

Related

Pad numpy array so each row is offset by one from the previous row

I am trying to figure out how to pad an array using the pattern shown below:
0000
0 000
00 00
000 0
0000
For example:
[[1,2,3],
[4,5,6],
[7,8,9]]
Would become:
[[1,2,3,0,0],
[0,4,5,6,0],
[0,0,7,8,9]]
I figured out how to do it using manual looping but I feel like there is probably a much faster way that uses numpy's array manipulation functions, I just can't figure out how to do it.
AFAIK, there is no function that does directly that in Numpy. You can use a loop iterating over the rows, but this solution is inefficient unless the 2D array is huge.
Numba solution
One solution to do this very efficiently is to use Numba and trivial loops:
import numba as nb
#nb.njit
def row_shifts_numba(arr):
n, m = arr.shape
out = np.zeros((n, n+m-1), dtype=arr.dtype)
for i in range(n):
for j in range(m):
out[i, i+j] = arr[i, j]
return out
data = np.array([[1,2,3],
[4,5,6],
[7,8,9]])
row_shifts_numba(data)
Note that the first execution is slower due to the jut-in-time compilation. If you do not want to pay this compilation time during the first execution, then you can specify the type of the array in the signature (eg. #nb.njit('int32[:,::1](int32[:,::1])') where int32 is the input/output array type and ::1 means the axis is contiguous).
Alternative pure-Numpy solution
Another solution consists in using a Numpy reshape trick so to generate the output 2D array. The idea is to create a bigger 2D array and then reshape it so to produce the shifts:
def row_shifts_numpy(arr):
n, m = arr.shape
out = np.zeros((n, n+m), dtype=arr.dtype)
out[:n,:m] = arr
return out.reshape(-1)[:-n].reshape(n, n+m-1)
This should be a bit slower than Numba due to the (implicit) creation of temporary arrays, but it is fully vectorized and only use Numpy.
After creating a zero array with the desired shape, we can create an index array in which the values are shifted row by row. Then filling the zero array by the original array values:
n, m = arr.shape
result = np.zeros((n, n+m-1), dtype=np.int64)
first_col_ind = np.array(np.arange(m))
ind = first_col_ind[:, None] + np.arange(n)
# [[0 1 2]
# [1 2 3]
# [2 3 4]]
result[np.arange(n), ind] = arr.T
# result:
# [[1 2 3 0 0]
# [0 4 5 6 0]
# [0 0 7 8 9]]

Is there a way to apply a numpy function that takes two 1d arrays as arguments on each row of two 2d arrays together?

I am trying to run something like:
np.bincount(array1, weights = array2, minlength=7)
where both array1 and array2 are 2d n numpy arrays of shape (m,n). My desired goal is that np.bincount() is run n times with each row of array1 and array2
I have tried using np.apply_along_axis() but as far as I can tell this only allows for the function to be run on each row of array1 without using each row of array2 as arguments for np.bincount. I was hoping to find a way to do this cleanly with a numpy function rather than iteration as this is a performance critical function but so far can't find another way.
For Example, given these arrays:
array1 = [[1,2,3],[4,5,6]]
array2 = [[7,8,9],[10,11,12]]
I would want to compute:
[np.bincounts([1,2,3], weights = [7,8,9],minlength=7), np.bincounts([4,5,6], weights = [10,11,12], minlength=7)]
A simple solution is simply to use comprehension lists:
result = [np.bincount(v, weights=w) for v,w in zip(array1, array2)]
Because the resulting arrays can have a different size (and actually do have a different size in your example), the result cannot be a Numpy array but a regular list. Most Numpy function are not able to work on a list of variable-sized arrays or even produce them.
If you have a lot of row in the arrays, you can mitigate the cost of the CPython interpreter loops using the Numba's JIT (or eventually Cython in this case). Note that the input arrays must be converted in Numpy arrays before calling the Numba function for sake of performance. If you know that all the arrays are of the same size, you can write a more efficient implementation using Numba (by preallocating the resulting array and doing the bincount yourself).
Update
With fixed-size arrays, here is a fast implementation in Numba:
import numpy as np
import numba as nb
array1 = np.array([[1,2,3],[4,5,6]], dtype=np.int32)
array2 = np.array([[7,8,9],[10,11,12]], dtype=np.int32)
#nb.njit('i4[:,::1](i4[:,::1],i4[:,::1])')
def compute(array1, array2):
assert array1.shape == array2.shape
n, m = array1.shape
res = np.zeros((n, 7), dtype=np.int32)
for i in range(n):
for j in range(m):
v = array1[i, j]
assert v>=0 and v<7 # Can be removed if the input is safe
res[i, v] += array2[i, j]
return res
result = compute(array1, array2)
# result is
# array([[ 0, 7, 8, 9, 0, 0, 0],
# [ 0, 0, 0, 0, 10, 11, 12]])

Numpy: smart matrix multiplication to sparse result matrix

In python with numpy, say I have two matrices:
S, a sparse x*x matrix
M, a dense x*y matrix
Now I want to do np.dot(M, M.T) which will return a dense x*x matrix S_.
However, I only care about the cells that are nonzero in S, which means that it would not make a difference for my application if I did
S_ = S*S_
Obviously, that would be a waste of operations as I would like to leave out the irrelevant cells given in S alltogether. Remember that in matrix multiplication
S_[i,j] = np.sum(M[i,:]*M[:,j])
So I want to do this operation only for i,j such that S[i,j]=True.
Is this supported somehow by numpy implementations that run in C so that I do not need to implement it with python loops?
EDIT 1 [solved]: I still have this problem, actually M is now also sparse.
Now, given rows and cols of S, I implemented it like this:
data = np.array([ M[rows[i],:].dot(M[cols[i],:]).data[0] for i in xrange(len(rows)) ])
S_ = csr( (data, (rows,cols)) )
... but it is still slow. Any new ideas?
EDIT 2: jdehesa has given a great solution, but I would like to save more memory.
The solution was to do the following:
data = M[rows,:].multiply(M[cols,:]).sum(axis=1)
and then build a new sparse matrix from rows, cols and data.
However, when running the above line, scipy builds a (contiguous) numpy array with as many elements as nnz of the first submatrix plus nnz of the second submatrix, which can lead to MemoryError in my case.
In order to save more memory, I would like to multiply iteratively each row with its respective 'partner' column, then sum over and discard the result vector. Using simple python to implement this, basically I am back to the extremely slow version.
Is there a fast way of solving this problem?
Here is how you can do it with NumPy/SciPy, both for dense and sparse M matrices:
import numpy as np
import scipy.sparse as sp
# Coordinates where S is True
S = np.array([[0, 1],
[3, 6],
[3, 4],
[9, 1],
[4, 7]])
# Dense M matrix
# Random big matrix
M = np.random.random(size=(1000, 2000))
# Take relevant rows and compute values
values = np.sum(M[S[:, 0]] * M[S[:, 1]], axis=1)
# Make result matrix from values
result = np.zeros((len(M), len(M)), dtype=values.dtype)
result[S[:, 0], S[:, 1]] = values
# Sparse M matrix
# Construct sparse M as COO matrix or any other way
M = sp.coo_matrix(([10, 20, 30, 40, 50], # Data
([0, 1, 3, 4, 6], # Rows
[4, 4, 5, 5, 8])), # Columns
shape=(1000, 2000))
# Convert to CSR for fast row slicing
M_csr = M.tocsr()
# Take relevant rows and compute values
values = M_csr[S[:, 0]].multiply(M_csr[S[:, 1]]).sum(axis=1)
values = np.squeeze(np.asarray(values))
# Construct COO sparse matrix from values
result = sp.coo_matrix((values, (S[:, 0], S[:, 1])), shape=(M.shape[0], M.shape[0]))

Scale rows of 3D-tensor

I have an n-by-3-by-3 numpy array A and an n-by-3 numpy array B. I'd now like to multiply every row of every one of the n 3-by-3 matrices with the corresponding scalar in B, i.e.,
import numpy as np
A = np.random.rand(10, 3, 3)
B = np.random.rand(10, 3)
for a, b in zip(A, B):
a = (a.T * b).T
print(a)
Can this be done without the loop as well?
You can use NumPy broadcasting to let the elementwise multiplication happen in a vectorized manner after extending B to 3D after adding a singleton dimension at the end with np.newaxis or its alias/shorthand None. Thus, the implementation would be A*B[:,:,None] or simply A*B[...,None].

Vectorize eigenvalue calculation in Numpy

I would like a numpy-sh way of vectorizing the calculation of eigenvalues, such that I can feed it a matrix of matrices and it would return a matrix of the respective eigenvalues.
For example, in the code below, B is the block 6x6 matrix composed of 4 copies of the 3x3 matrix A.
C is what I would like to see as output, i.e. an array of dimension (2,2,3) (because A has 3 eigenvalues).
This is of course a very simplified example, in the general case the matrices A can have any size (although they are still square), and the matrix B is not necessarily formed of copies of A, but different A1, A2, etc (all of same size but containing different elements).
import numpy as np
A = np.array([[0, 1, 0],
[0, 2, 0],
[0, 0, 3]])
B = np.bmat([[A, A], [A,A]])
C = np.array([[np.linalg.eigvals(B[0:3,0:3]),np.linalg.eigvals(B[0:3,3:6])],
[np.linalg.eigvals(B[3:6,0:3]),np.linalg.eigvals(B[3:6,3:6])]])
Edit: if you're using a version of numpy >= 1.8.0, then np.linalg.eigvals operates over the last two dimensions of whatever array you hand it, so if you reshape your input to an (n_subarrays, nrows, ncols) array you'll only have to call eigvals once:
import numpy as np
A = np.array([[0, 1, 0],
[0, 2, 0],
[0, 0, 3]])
# the input needs to be an array, since matrices can only be 2D.
B = np.repeat(A[np.newaxis,...], 4, 0)
# for arbitrary input arrays you could do something like:
# B = np.vstack(a[np.newaxis,...] for a in input_arrays)
# but for this to work it will be necessary for each element in
# 'input_arrays' to have the same shape
# eigvals will operate over the last two dimensions of the array and return
# a (4, 3) array of eigenvalues
C = np.linalg.eigvals(B)
# reshape this output so that it matches your original example
C.shape = (2, 2, 3)
If your input arrays don't all have the same dimensions, e.g. input_arrays[0].shape == (2, 2), input_arrays[1].shape == (3, 3) etc. then you could only vectorize this calculation across subsets with matching dimensions.
If you're using an older version of numpy then unfortunately I don't think there's any way to vectorize the calculation of the eigenvalues over multiple input arrays - you'll just have to loop over your inputs in Python instead.
You could just do something like this
C = np.array([[np.linalg.eigvals(B[i:i+3, j:j+3])
for i in xrange(0, B.shape[0], 3)]
for j in xrange(0, B.shape[1], 3)])
Perhaps a nicer approach is to use the block_view function from https://stackoverflow.com/a/5078155/1352250:
B_blocks = block_view(B)
C = np.array([[np.linalg.eigvals(m) for m in v] for v in B_blocks])
Update
As ali_m points out, this method is a form of syntactic sugar that will not reduce the overhead incurred from calling eigvals a large number of times. While this overhead should be small if each matrix it is applied to is large-ish, for the 6x6 matrices that the OP is interested in, it is not trivial (see the comments below; according to ali_m, there might be a factor of three difference between the version I give above, and the version he posted that uses Numpy >= 1.8.0).

Categories