CODEWITHSUNDEEP
pythonarraysnumpymatrix
PumpkinPie
PumpkinPie

Reputation: 141

An array of matrices, built using values from an array in Python

Here is my code. What I want it to return is an array of matrices

[[1,1],[1,1]], [[2,4],[8,16]], [[3,9],[27,81]]

I know I can probably do it using for loop and looping through my vector k, but I was wondering if there is a simple way that I am missing. Thanks!

from numpy import *
import numpy as np

k=np.arange(1,4,1)
print k

def exam(p):
   return np.array([[p,p**2],[p**3,p**4]])

print exam(k)

The output:

[1 2 3]

[[[ 1  2  3]

  [ 1  4  9]]

 [[ 1  8 27]

  [ 1 16 81]]]

Upvotes: 1

Views: 46

Answers (1)

shx2
shx2

Reputation: 64368

The key is to play with the shapes and broadcasting.

b = np.arange(1,4) # the base
e = np.arange(1,5) # the exponent

b[:,np.newaxis] ** e
=> 
array([[ 1,  1,  1,  1],
       [ 2,  4,  8, 16],
       [ 3,  9, 27, 81]])

(b[:,None] ** e).reshape(-1,2,2) 
=>
array([[[ 1,  1],
        [ 1,  1]],

       [[ 2,  4],
        [ 8, 16]],

       [[ 3,  9],
        [27, 81]]])

If you must have the output as a list of matrices, do:

m = (b[:,None] ** e).reshape(-1,2,2)
[ np.mat(a) for a in m ]
=>
[matrix([[1, 1],
        [1, 1]]),
 matrix([[ 2,  4],
        [ 8, 16]]),
 matrix([[ 3,  9],
        [27, 81]])]

Upvotes: 2

Related Questions