CODEWITHSUNDEEP
pythonpython-3.xnumpy
TragedyStruck
TragedyStruck

Reputation: 596

Using two np.linspace, how to fill 2D array with complex values?

I'm trying to fill a 2D array with complex(x,y), where x and y are from two two arrays:

xstep = np.linspace(xmin, xmax, Nx)
ystep = np.linspace(ymin, ymax, Ny)

However I can't figure out how to "spread" these values out on a 2D array.

So far my attempts are not really working out. I was hoping for something along the lines of:

result = np.array(xstep + (1j * ystep))

Maybe something from fromfunction, meshgrid or full, but I can't quite make it work.

As an example, say I do this:

xstep = np.linspace(0, 1, 2)  # array([0., 1.])
ystep = np.linspace(0, 1, 3)  # array([0. , 0.5, 1. ])

I'm trying to construct an answer:

array([
[0+0j, 0+0.5j, 0+1j],
[1+0j, 1+0.5j, 1+1j]
])

Note that I am not married to the linspace, so any quicker method would also do, it is just my natural starting point for creating this array, being new to Numpy.

Upvotes: 2

Views: 2702

Answers (3)

unutbu
unutbu

Reputation: 880887

In [4]: xstep = np.linspace(0, 1, 2)

In [5]: ystep = np.linspace(0, 1, 3)

In [6]: xstep[:, None] + 1j*ystep
Out[6]: 
array([[0.+0.j , 0.+0.5j, 0.+1.j ],
       [1.+0.j , 1.+0.5j, 1.+1.j ]])

xstep[:, None] is equivalent to xstep[:, np.newaxis] and its purpose is to add a new axis to xstep on the right. Thus, xstep[:, None] is a 2D array of shape (2, 1).

In [19]: xstep[:, None].shape
Out[19]: (2, 1)

xstep[:, None] + 1j*ystep is thus the sum of a 2D array of shape (2, 1) and a 1D array of shape (3,).

NumPy broadcasting resolves this apparent shape conflict by automatically adding new axes (of length 1) on the left. So, by NumPy broadcasting rules, 1j*ystep is promoted to an array of shape (1, 3). (Notice that xstep[:, None] is required to explicitly add new axes on the right, but broadcasting will automatically add axes on the left. This is why 1j*ystep[None, :] was unnecessary though valid.)

Broadcasting further promotes both arrays to the common shape (2, 3) (but in a memory-efficient way, without copying the data). The values along the axes of length 1 are broadcasted repeatedly:

In [15]: X, Y = np.broadcast_arrays(xstep[:, None], 1j*ystep)

In [16]: X
Out[16]: 
array([[0., 0., 0.],
       [1., 1., 1.]])

In [17]: Y
Out[17]: 
array([[0.+0.j , 0.+0.5j, 0.+1.j ],
       [0.+0.j , 0.+0.5j, 0.+1.j ]])

Upvotes: 4

Paul Panzer
Paul Panzer

Reputation: 53109

You can use np.ogrid with imaginary "step" to obtain linspace semantics:

y, x = np.ogrid[0:1:2j, 0:1:3j]                                                                               
y + 1j*x
# array([[0.+0.j , 0.+0.5j, 0.+1.j ],                                                                                 
#        [1.+0.j , 1.+0.5j, 1.+1.j ]])

Here the ogrid line means make an open 2D grid. axis 0: 0 to 1, 2 steps, axis 1: 0 to 1, 3 steps. The type of the slice "step" acts as a switch, if it is imaginary (in fact anything of complex type) its absolute value is taken and the expression is treated like a linspace. Otherwise range semantics apply.

The return values

y, x
# (array([[0.],                                                                                                       
#         [1.]]), array([[0. , 0.5, 1. ]]))

are "broadcast ready", so in the example we can simply add them and obtain a full 2D grid.

If we allow ourselves an imaginary "stop" parameter in the second slice (which only works with linspace semantics, so depending on your style you may prefer to avoid it) this can be condensed to one line:

sum(np.ogrid[0:1:2j, 0:1j:3j])
# array([[0.+0.j , 0.+0.5j, 0.+1.j ],
#        [1.+0.j , 1.+0.5j, 1.+1.j ]])

A similar but potentially more performant method would be preallocation and then broadcasting:

out = np.empty((y.size, x.size), complex)
out.real[...], out.imag[...] = y, x
out
# array([[0.+0.j , 0.+0.5j, 0.+1.j ],
#        [1.+0.j , 1.+0.5j, 1.+1.j ]])

And another one using outer sum:

np.add.outer(np.linspace(0,1,2), np.linspace(0,1j,3))
# array([[0.+0.j , 0.+0.5j, 0.+1.j ],
#        [1.+0.j , 1.+0.5j, 1.+1.j ]])

Upvotes: 3

Space Impact
Space Impact

Reputation: 13255

Use reshape(-1,1) for xstep as:

xstep = np.linspace(0, 1, 2)  # array([0., 1.])
ystep = np.linspace(0, 1, 3)  # array([0. , 0.5, 1. ])

result = np.array(xstep.reshape(-1,1) + (1j * ystep))

result

array([[0.+0.j , 0.+0.5j, 0.+1.j ],
       [1.+0.j , 1.+0.5j, 1.+1.j ]])

Upvotes: 1

Related Questions