initialize numpy array or arrays with particular shape

user_353337 0
Viewed 308

I have a Python function that, given a scalar or vector a, does some sort of symmetry expansion, e.g.,

import numpy


def expand(a):
    zero = numpy.zeros_like(a)
    return numpy.array([
        [zero, a, a],
        [a, zero, a],
        [a, a, zero],
        #
        [zero, -a, -a],
        [-a, zero, -a],
        [-a, -a, zero],
    ])

In this example, the resulting array will have shape (6, 3, a.shape). Eventually I'll need shape (6, a.shape, 3) or (3, a.shape, 6) and contiguity, so I'll numpy.moveaxis() around. Even (3, 6, a.shape) would be an improvement. I can get this with

import numpy


def expand(a):
    zero = numpy.zeros_like(a)
    return numpy.array([
        [zero, a, a, zero, -a, -a],
        [a, zero, a, -a, zero, -a],
        [a, a, zero, -a, -a, zero]
    ])

but this isn't as readable as the first version, especially if the transformation is more complex. (There always 3 columns.)

Is there a way to initialize out to have the right shape straight away? (Note that reshape()ing won't do it, the data would be in wrong order.)

numpy python
3 Answers

Is there a way to initialize out to have the right shape straight away?

Yes, there actually is: Simply allocate with the shape you want and then do your initialization on a transposed view as needed. As the view is---well---a view all changes will propagate to the original array while maintaing its memory layout.

import numpy
    
def expand(a):
    zero = numpy.zeros_like(a)
    out = numpy.empty((6,*a.shape,3),a.dtype)
    out.transpose(0,2,1)[...] = [
        [zero, a, a],
        [a, zero, a],
        [a, a, zero],
        #
        [zero, -a, -a],
        [-a, zero, -a],
        [-a, -a, zero],
    ]
    return out

Example:

>>> out = expand(numpy.arange(1,3))
>>> out
array([[[ 0,  1,  1],
        [ 0,  2,  2]],

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

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

       [[ 0, -1, -1],
        [ 0, -2, -2]],

       [[-1,  0, -1],
        [-2,  0, -2]],

       [[-1, -1,  0],
        [-2, -2,  0]]])
>>> out.flags
  C_CONTIGUOUS : True
  F_CONTIGUOUS : False
  OWNDATA : True
  WRITEABLE : True
  ALIGNED : True
  WRITEBACKIFCOPY : False
  UPDATEIFCOPY : False
user_72073920

Let's try (3,6) + a.shape:

def expand(a):
    sub = 1-np.eye(3, dtype=int)
    kernel = np.hstack([sub,-sub])
    return np.multiply.outer(kernel, a)

expand(a).shape
# (3, 6, 2, 5)

Now if you want (3,a.shape, 6), you can do a transpose:

def expand2(a):
    sub = 1-np.eye(3, dtype=int)
    kernel = np.hstack([sub,-sub])
    return np.multiply.outer(kernel, a).transpose(0,2,3,1)

expand2(a).shape
# (3, 2, 5, 6)
user_42384080

HEre is a straight forward numpy method

>>> a = np.arange(1,3)
>>> 
>>> out = np.empty((6,*a.shape,3),a.dtype)
>>> aux = out.reshape(2,3,-1,3)
>>> aux[0] = a[:,None]
>>> aux[1] = -a[:,None]
>>> np.einsum("ijkj->ijk",aux)[...] = 0
>>> 
>>> out
array([[[ 0,  1,  1],
        [ 0,  2,  2]],

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

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

       [[ 0, -1, -1],
        [ 0, -2, -2]],

       [[-1,  0, -1],
        [-2,  0, -2]],

       [[-1, -1,  0],
        [-2, -2,  0]]])
>>> 
>>> out.flags
  C_CONTIGUOUS : True
  F_CONTIGUOUS : False
  OWNDATA : True
  WRITEABLE : True
  ALIGNED : True
  WRITEBACKIFCOPY : False
  UPDATEIFCOPY : False
user_72073920
Related