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

Question

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 < N/2:
        shift the 1 in row i by 2*i to the right
    if i >= 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:

1. Looping through the rows

2. At each row, identify where the 1 needs to be moved to, call it idx

3. Create a separate zeros array, and insert a 1 into index idx

4. 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?

Answer 1

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])

Answer 2

The permutation matrix always has the same form if you look at it the right way. Taking the indices of the elements as their identity you basically have the following "vector of vectors":

[0, n//2+1, 1, n//2+2, ..., n//2, n]

once you realize that it becomes a matter of "interweaving the two halves of the identity matrix". This Stack Overflow question gives an interesting suggestion on how to do that.

This should work just fine for your purposes vastly simplifying the array accesses by virtue of using slicing to a bit more of its potential:

def permutation_matrix(N):
    I = np.identity(N)
    P = np.empty((N,N))
    mid = N//2 if N % 2 == 0 else N//2+1
    P[0::2] = I[:mid]
    P[1::2] = I[mid:]
    return P

with this rather satisfying result:

>>> numbers
array([0, 1, 2, 3, 4, 5])
>>> numbers.dot(permutation_matrix(6))
array([0., 2., 4., 1., 3., 5.])

introducing the more appropriate mid that uses flooring division even allows handling an uneven amount of numbers:

>>> numbers = np.array([0,1,2,3,4])
>>> numbers.dot(permutation_matrix(5))
array([0., 2., 4., 1., 3.])

Answer 3

One way is to the sequence 0,1,2,3,4,5 or ...(N-1) and using these as row coords in a sparse-matrix (CSR) representation:

from scipy.sparse import csr_matrix

N = 6

csr_matrix(([1]*6, ([0,3,1,4,2,5], [0,1,2,3,4,5] ))).toarray()

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]], dtype=int64)

and for general N:

csr_matrix(([1]*N, ([0,3,1,4,2,5], list(range(N)) ))).toarray()

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]], dtype=int64)

and a roundrobin iterator to generate the low/hi values:

from itertools import chain, zip_longest, cycle

# If you know N is even, you can get away with this...
N = 6
[x for its in zip(range(N//2), range(N//2, N)) for x in its]
# [0, 3, 1, 4, 2, 5]

# But in the general case, N could be odd, and you need to handle one of the iterators being exhausted first and yielding None...
N = 7
[x for its in zip_longest(range(N//2), range(N//2, N)) for x in its if x is not None]
# [0, 3, 1, 4, 2, 5, 6]

(It turned out writing that roundrobin iterator was a world of pain. Could be less grief to use bitwise arithmetic, or imperative code like the other answers.)