Efficient numpy array manipulation to convert an identity matrix to a permutation matrix Task

I want to be able to generate the permutation matrix that splits a 1D array of consecutive numbers (i.e. even, odd, even, odd, even, odd, ...) into a 1D array where the first half are the evens, and the second half are the odds. So (even1, odd1, even2, odd2, even3, odd3) goes to (even1, even2, even3, odd1, odd2, odd3).

For example, with N=6, the permutation matrix would be:

M = array([1, 0, 0, 0, 0, 0],
          [0, 0, 1, 0, 0, 0],
          [0, 0, 0, 0, 1, 0],
          [0, 1, 0, 0, 0, 0],
          [0, 0, 0, 1, 0, 0],
          [0, 0, 0, 0, 0, 1])


You can check that multiplying this with M * array([0, 1, 2, 3, 4, 5]) = array([0, 2, 4, 1, 3, 5]).

My approach in pseudocode

(Full code below.) This is the mathematically correct way to generate this:

I = NxN identity matrix
for i in [0:N-1]:
    if i &lt; N/2:
        shift the 1 in row i by 2*i to the right
    if i &gt;= N/2:
        shift the 1 in row i by 2*(i - N/2)+1 to the right


You can see how that works to generate M above.

Code (Python)

I implement the above pseudocode by using numpy array manipulation (this code is copy-and-pasteable):

import numpy as np

def permutation_matrix(N):
    N_half = int(N/2) #This is done in order to not repeatedly do int(N/2) on each array slice
    I = np.identity(N) 
    I_even, I_odd = I[:N_half], I[N_half:] #Split the identity matrix into the top and bottom half, since they have different shifting formulas

    #Loop through the row indices
    for i in range(N_half):
        # Apply method to the first half
        i_even = 2 * i #Set up the new (shifted) index for the 1 in the row
        zeros_even = np.zeros(N) #Create a zeros array (will become the new row)
        zeros_even[i_even] = 1. #Put the 1 in the new location
        I_even[i] = zeros_even #Replace the row in the array with our new, shifted, row

        # Apply method to the second half
        i_odd = (2 * (i - N_half)) + 1
        zeros_odd = np.zeros(N)
        zeros_odd[i_odd] = 1.
        I_odd[i] = zeros_odd

    M = np.concatenate((I_even, I_odd), axis=0) 

    return M

N = 8
M = permutation_matrix(N)
print(M)

Output:
array([[1., 0., 0., 0., 0., 0., 0., 0.],
       [0., 0., 1., 0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 1., 0., 0., 0.],
       [0., 0., 0., 0., 0., 0., 1., 0.],
       [0., 1., 0., 0., 0., 0., 0., 0.],
       [0., 0., 0., 1., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0., 1., 0., 0.],
       [0., 0., 0., 0., 0., 0., 0., 1.]])


My issues

I have a feeling that there are more efficient ways to do this. To summarise what I am doing to each matrix:


Looping through the rows
At each row, identify where the 1 needs to be moved to, call it idx
Create a separate zeros array, and insert a 1 into index idx
Replace the row we are evaluating with our modified zeros array


Is it necessary to split the array in two?

Is there an Pythonic way to implement two different functions on two halves of the same array without splitting them up?

Is there an approach where I can shift the 1s without needing to create a separate zeros array in memory?

Do I even need to loop through the rows?

Are there more efficient libraries than numpy for this?
Are there more efficient libraries than numpy for this?


Since permutation matrices are rather sparse, the scipy.sparse library is helpful.
Using its coo_matrix method we can build a matrix which contains given values at given indices. 

From here it's just a matter of building the right lists of indices.

from itertools import chain
from scipy.sparse import coo_matrix

def permutation_matrix(n):
    # row and column indices - first even, then odd numbers in the latter
    I, J = range(n), list(chain(range(0, n, 2), range(1, n, 2)))
    # the following also works, if you are so inclined. 
    # J = [m*2 + d for d, m in map(lambda k: divmod(k, n//2 + n % 2), range(n))]

    return coo_matrix(([1]*n, (I, J)))


If needed, we can use the .A property (short for .toarray()) to build a full matrix from this: e.g.
permutation_matrix(10).A.



Why bother with sparse matrices?

Multiplication with sparse matrices will be much faster: e.g. matrix-vector products can be computed in O(n) time instead of O(n^2).
Similarly, the memory requirements for storing these matrices in sparse format is O(n) instead of O(n^2).




  Is there an approach where I can shift the 1s without needing to
  create a separate zeros array in memory?


Sure, numpy.roll does circular shifting:

numpy.roll([1,0,0], 4) == numpy.array([0, 1, 0])



  Is there an Pythonic way to implement two different functions on two
  halves of the same array without splitting them up?


Not sure, but you can always cook something up.

from collections import deque
from itertools import chain

def apply_to_parts(part_selector, funs):

    def wrapper(vals):
        results = [deque() for _ in funs]

        for ix, val in enumerate(vals):
            part_ix = part_selector(ix, val)
            results[part_ix].append(funs[part_ix](val))

        return list(chain(*results))

    return wrapper

# multiply elements at even indices by 2
# divide   elements at odd  indices by 4    
# return a list of elements in part 1 followed by elements in part 2
example = apply_to_parts(
    lambda ix, val: ix % 2,
    [lambda x:2*x, lambda y: y/4]
)

# should return [2, 6, 0.5, 1.0]
example([1,2,3,4])