How to fill multidimensional array using an equation in python

user_13029118 0
Viewed 699

I am new to python and I would like to fill a Numpy multidimensional array using an equation. In Fortran I can use the index of the array to fill it up, is this possible to do in python? Say I have an equation a=i*j where i and j are the row and column position respectively. So if I have an n by n array then the array would be filled with the results from the equation, so first value would be 1 since a=1*1 and so on.

numpy arrays python
2 Answers

You could use numpy broadcasting to get this result:

i = np.arange(3)[:, np.newaxis]  # (3, 1)
# array([0, 1, 2])
j = np.arange(4)[np.newaxis, :]  # (1, 4)
# array([0, 1, 2, 3])
arr = i * j                      # (3 ,4)
# array([[0, 0, 0, 0],
#        [0, 1, 2, 3],
#        [0, 2, 4, 6]])

You can perform most calculations you could want with the index arrays i [n x 1] and j [1 x m] and the result will be always be [n x m]

arr = np.sin(i)**2 + np.cos(i)**2 + (j-i)
# array([[ 1.,  2.,  3.,  4.],
#        [ 0.,  1.,  2.,  3.],
#        [-1.,  0.,  1.,  2.]])

You could also use np.meshgrid() to explicitly repeat the indices along the other dimension to get full 2d arrays for i and j:

i, j = np.meshgrid(np.arange(3), np.arange(4), indexing='ij')
# i
# array([[0, 1, 2, 3],
#        [0, 1, 2, 3],
#        [0, 1, 2, 3]])
# j
# array([[0, 0, 0, 0],
#        [1, 1, 1, 1],
#        [2, 2, 2, 2]])

arr = i * j
# array([[0, 0, 0, 0],
#        [0, 1, 2, 3],
#        [0, 2, 4, 6]])

This is a good visualisation of what happens during broadcasting automatically.

Note the indexing argument ij of np.meshgrid() for matrix indexing; from the docs:

Giving the string ‘ij’ returns a meshgrid with matrix indexing, while ‘xy’ returns a meshgrid with Cartesian indexing. In the 2-D case with inputs of length M and N, the outputs are of shape (N, M) for ‘xy’ indexing and (M, N) for ‘ij’ indexing.

user_75708170

To do a=i*j, or something similar, it is best to first create an array containing i values and an array containing j values, and than simply do a=i*j:

For example, if you want a 5 x 10 matrix:

import numpy as np

i, j = np.meshgrid(np.arange(5), np.arange(10))
a = i * j

The result is:

array([[ 0,  0,  0,  0,  0],
       [ 0,  1,  2,  3,  4],
       [ 0,  2,  4,  6,  8],
       [ 0,  3,  6,  9, 12],
       [ 0,  4,  8, 12, 16],
       [ 0,  5, 10, 15, 20],
       [ 0,  6, 12, 18, 24],
       [ 0,  7, 14, 21, 28],
       [ 0,  8, 16, 24, 32],
       [ 0,  9, 18, 27, 36]])
user_3892890
Related