I have a tensor of shape (m*n, m*n) and I want to extract a tensor of size (n, m*n) containing the m blocks of size n*n that are on the diagonal. For example:
>>> a
tensor([[1, 2, 0, 0],
[3, 4, 0, 0],
[0, 0, 5, 6],
[0, 0, 7, 8]])
I want to have a function extract(a, m, n) that will output:
>>> extract(a, 2, 2)
tensor([[1, 2, 5, 6],
[3, 4, 7, 8]])
I've thought of using some kind of slicing, because the blocks can be expressed by:
>>> for i in range(m):
... print(a[i*m: i*m + n, i*m: i*m + n])
tensor([[1, 2],
[3, 4]])
tensor([[5, 6],
[7, 8]])
You can take advantage of reshape and slicing:
import torch
import numpy as np
def extract(a, m, n):
s=(range(m), np.s_[:], range(m), np.s_[:]) # the slices of the blocks
a.reshape(m, n, m, n)[s] # reshaping according to blocks and slicing
return a.reshape(m*n, n).T # reshape to desired output format
Example:
a = torch.arange(36).reshape(6,6)
a
tensor([[ 0, 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, 31, 32, 33, 34, 35]])
extract(a, 3, 2)
tensor([[ 0, 6, 14, 20, 28, 34],
[ 1, 7, 15, 21, 29, 35]])
extract(a, 2, 3)
tensor([[ 0, 6, 12, 21, 27, 33],
[ 1, 7, 13, 22, 28, 34],
[ 2, 8, 14, 23, 29, 35]])
You can achieve this for a block diagonal matrix (of equally sized square blocks of width n) with torch.nonzero():
>>> n = 2
>>> a[a.nonzero(as_tuple=True)].view(n,n,-1)
tensor([[[1, 2],
[3, 4]],
[[5, 6],
[7, 8]]])
Related
How do I transpose the first 2 dimensions of a 3D array 'matrix?
matrix = np.random.rand(2,3,4)
In the third dimensions I want to swap 'rows' with 'columns', preferably without a loop.
You can use the .transpose() function.
matrix = matrix.transpose(1, 0, 2)
means swap the first and the second axis.
You can use swapaxes:
matrix2 = matrix.swapaxes(0,1)
example:
# input
array([[[ 0, 1, 2, 3],
[ 4, 5, 6, 7],
[ 8, 9, 10, 11]],
[[12, 13, 14, 15],
[16, 17, 18, 19],
[20, 21, 22, 23]]])
# output
array([[[ 0, 1, 2, 3],
[12, 13, 14, 15]],
[[ 4, 5, 6, 7],
[16, 17, 18, 19]],
[[ 8, 9, 10, 11],
[20, 21, 22, 23]]])
I have an array
xx = np.arange(24).reshape(2, 12)
array([[ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11],
[12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23]])
and I would like to reshape it, to obtain
array([[[ 0, 1, 2, 3],
[12, 13, 14, 15]],
[[ 4, 5, 6, 7],
[16, 17, 18, 19]],
[[ 8, 9, 10, 11],
[20, 21, 22, 23]]])
I can achieve it via
xx.T.reshape(3, 4, 2).transpose(0, 2, 1)
But it has to be transposed twice, which seems unnecessary to me. So could somebody confirm that this is the only way of doing it or provide more readable solution otherwise?
Thanks!
It is possible to do a single transpose:
data = np.arange(24).reshape(2, 12)
data = data.reshape(2, 3, 4).transpose(1, 0, 2)
Edit:
I checked this using itertools.permutations and itertools.product:
import itertools
import numpy as np
data = np.arange(24).reshape(2, 12)
desired_data = np.array([[[ 0, 1, 2, 3],
[12, 13, 14, 15]],
[[ 4, 5, 6, 7],
[16, 17, 18, 19]],
[[ 8, 9, 10, 11],
[20, 21, 22, 23]]])
shapes = [2, 3, 4]
transpose_dims = [0, 1, 2]
shape_permutations = itertools.permutations(shapes)
transpose_permutations = itertools.permutations(transpose_dims)
for shape, transpose in itertools.product(
list(shape_permutations),
list(transpose_permutations),
):
new_data = data.reshape(*shape).transpose(*transpose)
try:
np.allclose(new_data, desired_data)
except ValueError as e:
pass
else:
break
print(f"{shape=}, {transpose=}")
shape=(2, 3, 4), transpose=(1, 0, 2)
I would do it this way: first, generate two arrays (shown separated for the sake of decomposition):
xx.reshape(2, -1, 4)
# Output:
# array([[[ 0, 1, 2, 3],
# [ 4, 5, 6, 7],
# [ 8, 9, 10, 11]],
#
# [[12, 13, 14, 15],
# [16, 17, 18, 19],
# [20, 21, 22, 23]]])
From here, I would then stack along the second dimension in order to combine them like you want:
np.stack(xx.reshape(2, -1, 4), axis=1)
# Output:
# array([[[ 0, 1, 2, 3],
# [12, 13, 14, 15]],
#
# [[ 4, 5, 6, 7],
# [16, 17, 18, 19]],
#
# [[ 8, 9, 10, 11],
# [20, 21, 22, 23]]])
You'd avoid the transposition. Hopefully it's more readable, but in the end, that's highly subjective, right? '^^
To add on top of #Paul's answer, there is some speedup from removing one of the transpose. The time gain is of ~15%:
I have two matrices, A and B, with shapes (n, m, k) and (n, m) respectively. n is the batch size, m is the amount of data in a batch, and k is the feature size.
Each element of B is an index less than m (specifically B = torch.randint(high=m, shape=(n,m))).
I want to implement [A[i][B[i]] for i in range(n)] in a smarter way.
Is there a better way in pytorch to implement this without doing for loop?
You can use
a[torch.arange(n)[:, None], b]
An example:
>>> n, m, k = 3, 2, 5
>>> a = torch.arange(30).view(n, m, k)
>>> b = torch.randint(high=m, size=(n,m))
# first indexer (of shape (n, 1))
>>> torch.arange(n)[:, None]
tensor([[0],
[1],
[2]])
# second indexer
>>> b
tensor([[1, 0],
[0, 1],
[1, 1]])
The indexers have the shape (3, 1) and (3, 2) respectively so they'll be broadcasted to (3, 2) to effectively have
tensor([[0, 0],
[1, 1],
[2, 2]])
and
tensor([[1, 0],
[0, 1],
[1, 1]])
which says: for the first row, take 1st (k,) array and put the result and take 0th (k,) array and put the result. This fills in a (m, k) array in the output which is repeated n times for each row,
to get
>>> a[torch.arange(n)[:, None], b]
tensor([[[ 5, 6, 7, 8, 9],
[ 0, 1, 2, 3, 4]],
[[10, 11, 12, 13, 14],
[15, 16, 17, 18, 19]],
[[25, 26, 27, 28, 29],
[25, 26, 27, 28, 29]]])
comparing with list comprehension:
>>> [a[i][b[i]] for i in range(n)]
[tensor([[5, 6, 7, 8, 9],
[0, 1, 2, 3, 4]]),
tensor([[10, 11, 12, 13, 14],
[15, 16, 17, 18, 19]]),
tensor([[25, 26, 27, 28, 29],
[25, 26, 27, 28, 29]])]
I have a huge (N*20) matrix where every 5 rows is a valid sample, ie. every (5*20) matrix. I'm trying to reshape it into a (N/5,1,20,5) matrix where the dimension 20 is kept unchanged. I could do it in tensroflow using keep_dim, but how can I achieve this in numpy?
Thanks in advance.
Reshape and then swap the axes around:
arr1 = arr.reshape(N/5,5,1,20)
arr2 = arr1.transpose(0,2,3,1)
for example
In [476]: arr = np.arange(24).reshape(6,4)
In [477]: arr
Out[477]:
array([[ 0, 1, 2, 3],
[ 4, 5, 6, 7],
[ 8, 9, 10, 11],
[12, 13, 14, 15],
[16, 17, 18, 19],
[20, 21, 22, 23]])
In [478]: arr1 = arr.reshape(2,3,1,4)
In [479]: arr2 = arr1.transpose(0,2,3,1)
In [480]: arr2.shape
Out[480]: (2, 1, 4, 3)
In [482]: arr2
Out[482]:
array([[[[ 0, 4, 8],
[ 1, 5, 9],
[ 2, 6, 10],
[ 3, 7, 11]]],
[[[12, 16, 20],
[13, 17, 21],
[14, 18, 22],
[15, 19, 23]]]])
I have a numpy array of size nxm. I want the number of columns to be limited to k and rest of the columns to be extended in new rows. Following is the scenario -
Initial array: nxm
Final array: pxk
where p = (m/k)*n
Eg. n = 2, m = 6, k = 2
Initial array:
[[1, 2, 3, 4, 5, 6,],
[7, 8, 9, 10, 11, 12]]
Final array:
[[1, 2],
[7, 8],
[3, 4],
[9, 10],
[5, 6],
[11, 12]]
I tried using reshape but not getting the desired result.
Here's one way to do it
q=array([[1, 2, 3, 4, 5, 6,],
[7, 8, 9, 10, 11, 12]])
r=q.T.reshape(-1,2,2)
s=r.swapaxes(1,2)
t=s.reshape(-1,2)
as a one liner,
q.T.reshape(-1,2,2).swapaxes(1,2).reshape(-1,2)
array([[ 1, 2],
[ 7, 8],
[ 3, 4],
[ 9, 10],
[ 5, 6],
[11, 12]])
EDIT: for the general case, use
q=arange(1,1+n*m).reshape(n,m) #example input
r=q.T.reshape(-1,k,n)
s=r.swapaxes(1,2)
t=s.reshape(-1,k)
one liner is:
q.T.reshape(-1,k,n).swapaxes(1,2).reshape(-1,k)
example for n=3,m=12,k=4
q=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, 31, 32, 33, 34, 35, 36]])
result is
array([[ 1, 2, 3, 4],
[13, 14, 15, 16],
[25, 26, 27, 28],
[ 5, 6, 7, 8],
[17, 18, 19, 20],
[29, 30, 31, 32],
[ 9, 10, 11, 12],
[21, 22, 23, 24],
[33, 34, 35, 36]])
Using numpy.vstack and numpy.hsplit:
a = np.array([[1, 2, 3, 4, 5, 6,],
[7, 8, 9, 10, 11, 12]])
n, m, k = 2, 6, 2
np.vstack(np.hsplit(a, m/k))
result array:
array([[ 1, 2],
[ 7, 8],
[ 3, 4],
[ 9, 10],
[ 5, 6],
[11, 12]])
UPDATE As flebool commented, above code is very slow, because hsplit returns a python list, and then vstack reconstructs the final array from a list of arrays.
Here's alternative solution that is much faster.
a.reshape(-1, m/k, k).transpose(1, 0, 2).reshape(-1, k)
or
a.reshape(-1, m/k, k).swapaxes(0, 1).reshape(-1, k)