initialize numpy array or arrays with particular shape

Question

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

Paul Panzer · Accepted Answer · 2020-07-19 20:15:06Z

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

Quang Hoang · Accepted Answer · 2020-07-17 19:41:03Z

1

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)

answered Jul 17, 2020 at 19:41

Quang Hoang

151k11 gold badges64 silver badges86 bronze badges

1 Comment

Nico Schlömer Over a year ago

I works indeed, but I'm not a big fan: It does numerical computation where it shouldn't be necessary, is particular to the special case in the question, and is arguably not as readable as the first version in the question.

Paul Panzer · Accepted Answer · 2020-07-17 20:36:02Z

0

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

edited Jul 17, 2020 at 20:36

answered Jul 17, 2020 at 20:28

Paul Panzer

53.3k3 gold badges59 silver badges103 bronze badges

Collectives™ on Stack Overflow

initialize numpy array or arrays with particular shape

3 Answers 3

Comments

1 Comment

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

3 Answers 3

Comments

1 Comment

Comments

Your Answer

Sign up or log in

Post as a guest

Related