We Keep Coding

Create a 3D (partially) diagonal array from a 2D array

I'd like to ask how can I efficiently generate a numpy 3D array from a 2D array with each row filling the diagonal part of the new array?
For example, the input 2D array is
array([[1, 2],
[3, 4],
[5, 6],
[7, 8]])
and I want the output to be
array([[[1, 0],
[0, 2]],
[[3, 0],
[0, 4]],
[[5, 0],
[0, 6]],
[[7, 0],
[0, 8]]])
Typically, the size of the first dimensional is very large. Thanks in advance.

Assuming a the input and using indexing with unravel_index:
x, y = np.unravel_index(np.arange(a.size), a.shape)
out = np.zeros(a.shape+(a.shape[-1],), dtype=a.dtype)
out[x, y, y] = a.flat
Output:
array([[[1, 0],
[0, 2]],
[[3, 0],
[0, 4]],
[[5, 0],
[0, 6]],
[[7, 0],
[0, 8]]])
timings:

arr = np.array([[1, 2], [3, 4], [5, 6], [7, 8]])
res = np.apply_along_axis(np.diag, 1, arr)

How to draw lines connecting only points with a distance equal to 1 in Python

For example, there is an array of coordinates of points:
coord = np.array([[0, 0], [0, 2], [0, 4], [1, 1], [1, 2], [1, 3], [1, 5], [2, 0], [2, 1], [2, 3], [2, 4], [3, 0], [3, 2], [3, 3], [3, 5], [4, 1], [4, 2], [4, 4]])
Plot:
plt.scatter(coord[:,0], coord[:,1])
And I want to draw lines between each pair of points with unit distance. Can you help me please to derive an algorithm to do that?
P.S. There is what I want to get:
(Yes, picture should include also diagonals of length sqrt(2).)
Thanks!

You can use scipy.spatial.distance.pdist to compute the pairwise distances and filter. As the output of pdist is in a condensed form you need triu_indices to convert:
from scipy.spatial.distance import pdist
idx = np.vstack(np.triu_indices(len(coord), k=1)).T
d = pdist(coord)
coord[idx[d<=np.sqrt(2)]]
Output (pairs of points):
array([[[0, 0],
[1, 1]],
[[0, 2],
[1, 1]],
[[0, 2],
[1, 2]],
[[0, 2],
[1, 3]],
[[0, 4],
[1, 3]],
[[0, 4],
[1, 5]],
...
[[4, 1],
[4, 2]]])

Row-wise replacement of numpy array with values of another numpy array

I have 0s and 1s store in a 3-dimensional numpy array:
g = np.array([[[0, 1], [0, 1], [1, 0]], [[0, 0], [1, 0], [1, 1]]])
# array([
# [[0, 1], [0, 1], [1, 0]],
# [[0, 0], [1, 0], [1, 1]]])
and I'd like to replace these values by those in another array using a row-wise replacement strategy. For example, replacing the vales of g by x:
x = np.array([[2, 3], [4, 5]])
array([[2, 3],
[4, 5]])
to obtain:
array([
[[2, 3], [2, 3], [3, 2]],
[[4, 4], [5, 4], [5, 5]]])
The idea here would be to have the first row of g replaced by the first elements of x (0 becomes 2 and 1 becomes 3) and the same for the other row (the first dimension - number of "rows" - will always be the same for g and x)
I can't seem to be able to use np.where because there's a ValueError: operands could not be broadcast together with shapes (2,3,2) (2,2) (2,2).

IIUC,
np.stack([x[i, g[i]] for i in range(x.shape[0])])
Output:
array([[[2, 3],
[2, 3],
[3, 2]],
[[4, 4],
[5, 4],
[5, 5]]])

Vectorized approach with np.take_along_axis to index into the last axis of x with g using axis=-1 -
In [20]: np.take_along_axis(x[:,None],g,axis=-1)
Out[20]:
array([[[2, 3],
[2, 3],
[3, 2]],
[[4, 4],
[5, 4],
[5, 5]]])
Or with manual integer-based indexing -
In [27]: x[np.arange(len(g))[:,None,None],g]
Out[27]:
array([[[2, 3],
[2, 3],
[3, 2]],
[[4, 4],
[5, 4],
[5, 5]]])

One solution, is to simply use comprehension directly here:
>>> np.array([[x[i][c] for c in r] for i, r in enumerate(g)])
array([[[2, 3],
[2, 3],
[3, 2]],
[[4, 4],
[5, 4],
[5, 5]]])

From what I understand, g is an array of indexes (indexes being 0 or 1) and x is the array to who's values you use.
Something like this should work (tested quickly)
import numpy as np
def swap_indexes(index_array, array):
out_array = []
for i, row in enumerate(index_array):
out_array.append([array[i,indexes] for indexes in row])
return np.array(out_array)
index_array = np.array([[[0, 1], [0, 1], [1, 0]], [[0, 0], [1, 0], [1, 1]]])
x = np.array([[2, 3], [4, 5]])
print(swap_indexes(index_array, x))
[EDIT: fixed typo that created duplicates]

How to split array by indices where the splitted sub-arrays include the split point

I have a 2D array containing values and a 1D array with index values where I would like to split the 2D matrix, where the splitted sub-arrays include the 'split-point'.
I know I can use the numpy.split function to split by indices and I know I can use stride_tricks to split an array for creating consecutive overlapping subset-views.
But it seems the stride_ticks only applies if we want to split an array into equal sized sub-arrays.
Minimal example, I can do the following:
>>> import numpy as np
>>> array = np.random.randint(0,10, (10,2))
>>> indices = np.array([2,3,8])
>>> array
array([[8, 1],
[1, 0],
[2, 0],
[8, 8],
[1, 6],
[7, 8],
[4, 4],
[9, 4],
[6, 7],
[6, 4]])
>>> split_array = np.split(array, indices, axis=0)
>>> split_array
[array([[8, 1],
[1, 0]]),
array([[2, 0]]),
array([[8, 8],
[1, 6],
[7, 8],
[4, 4],
[9, 4]]),
array([[6, 7],
[6, 4]])]
But I'm merely looking for an option within the split function where I could define include_split_point=True, which would give me a result as such:
[array([[8, 1],
[1, 0],
[2, 0]]),
array([[2, 0],
[8, 8]]),
array([[8, 8],
[1, 6],
[7, 8],
[4, 4],
[9, 4],
[6, 7]]),
array([[6, 7],
[6, 4]])]

Create a new array with the index elements repeated
new_indices = np.zeros(array.shape[0], dtype = int)
new_indices[indices] = 1
new_indices += 1
new_array = np.repeat(array, new_indices, axis = 0)
Update indices to account for the changed array
indices = indices + np.arange(1, len(indices)+1)
Split using the indices as usual
np.split(new_array, indices, axis = 0)
output:
[array([[8, 1],
[1, 0],
[2, 0]]),
array([[2, 0],
[8, 8]]),
array([[8, 8],
[1, 6],
[7, 8],
[4, 4],
[9, 4],
[6, 7]]),
array([[6, 7],
[6, 4]])]

numpy array operation method

>>> c= array([[[1, 2],
[3, 4]],
[[2, 1],
[4, 3]],
[[3, 2],
[1, 4]]])
>>> x
array([[0, 1, 2],
[3, 4, 5]])
return me a matrix such that each column is the product of each matrix in c multiply the each corresponding column of x in regular matrix multiplication. I'm trying to figure out a way to vectorized it or at least not using for loop to solve it.
array([[6, 6, 16]
12, 16, 22]])
to extends this operation further let's say that I have an array of matrices,say
>>> c
array([[[1, 2],
[3, 4]],
[[2, 1],
[4, 3]],
[[3, 2],
[1, 4]]])
>>> x
array([[[1, 2, 3],
[1, 2, 3]],
[[1, 0, 2],
[1, 0, 2]],
[[2, 3, 1],
[0, 1, 0]]])
def fun(c,x):
for i in range(len(x)):
np.einsum('ijk,ki->ji',c,x[i])
##something
So basically, I want to have each matrix in x multiply with all of c. return a structure similar to c without introducing this for loop
The reason I'm doing this because I've encounter a problem to solve a problem ,trying to vectorized
Xc (the operation follows the normal matrix column vector multiplication), c is 3D array; like the c from above-- a column vector that each element is a matrix (in numpy its the form in the above). X is the matrix with each elements is a 1D array. The output of the Xc should be 1D array.

You can use np.einsum -
np.einsum('ijk,ki->ji',c,x)
Sample run -
In [155]: c
Out[155]:
array([[[1, 2],
[3, 4]],
[[2, 1],
[4, 3]],
[[3, 2],
[1, 4]]])
In [156]: x
Out[156]:
array([[0, 1, 2],
[3, 4, 5]])
In [157]: np.einsum('ijk,ki->ji',c,x)
Out[157]:
array([[ 6, 6, 16],
[12, 16, 22]])
For the 3D case of x, simply append the new dimension at the start of the string notation for x and correspondingly at the output string notation too, like so -
np.einsum('ijk,lki->lji',c,x)
Sample run -
In [151]: c
Out[151]:
array([[[1, 2],
[3, 4]],
[[2, 1],
[4, 3]],
[[3, 2],
[1, 4]]])
In [152]: x
Out[152]:
array([[[1, 2, 3],
[1, 2, 3]],
[[1, 0, 2],
[1, 0, 2]],
[[2, 3, 1],
[0, 1, 0]]])
In [153]: np.einsum('ijk,lki->lji',c,x)
Out[153]:
array([[[ 3, 6, 15],
[ 7, 14, 15]],
[[ 3, 0, 10],
[ 7, 0, 10]],
[[ 2, 7, 3],
[ 6, 15, 1]]])