ASD
ASD

Reputation: 19

Is there a way to construct 3D array from 2D array in python using numpy?

I've got a matrix that values are positive integers and want to find the bit representation unpacked into a tensor using numpy

e.g

[[1 , 2],
 [3 , 4]]

to

[[0 , 0],
 [0 , 1]], 

[[0 , 1],
 [1 , 0]],

[[1 , 0],
 [1 , 0]]

Upvotes: 0

Views: 146

Answers (1)

Michael Szczesny
Michael Szczesny

Reputation: 5036

This solution is limited to 2D arrays but works with elements larger than np.uint8 and calculates the necessary bit depth.

import numpy as np

M = np.array([[1,2],
              [3,4]])
bits = int(np.log2(M.max())) + 1

(np.where(
     M.reshape(-1,1) & 2**np.arange(bits)[::-1], 1, 0)
   .reshape(*M.shape, -1)
   .transpose(2,0,1))

Output

array([[[0, 0],
        [0, 1]],

       [[0, 1],
        [1, 0]],

       [[1, 0],
        [1, 0]]])

How this works

Construct a range with powers of 2

2**np.arange(bits)[::-1]

Broadcast this range with logical_and over the input elements

(M.reshape(-1,1) & 2**np.arange(bits)[::-1])

Output

array([[0, 0, 1],
       [0, 2, 0],
       [0, 2, 1],
       [4, 0, 0]])

Convert to 1,0 bool array with np.where

array([[0, 0, 1],   # 1 in binary
       [0, 1, 0],   # 2 in binary
       [0, 1, 1],   # 3 in binary
       [1, 0, 0]])  # 4 in binary

Shape to desired output.

Upvotes: 1

Related Questions