CODEWITHSUNDEEP
pythonarrays
Joanne
Joanne

Reputation: 503

Calculate a function when 2D array as variables

I want to do a calculation through a function, which is related to kx and ky. kx and ky are two 2D arrays. How can I use python to do this?

Thank you!!!

enter image description here

Upvotes: 1

Views: 812

Answers (1)

cadolphs
cadolphs

Reputation: 9617

The de-facto standard package to do array math in python is numpy. So install the numpy package and then you'd do

import numpy as np  ## Pretty much everyone imports numpy as np, so might as well.


def my_function(kx, ky):
    return kx / (kx**2 + ky**2)

kx = np.array([[1, 2, 3], [1, 2, 3], [1, 2, 3]])
ky = np.array([[1, 1, 1], [2, 2, 2], [3, 3, 3]]) # Well we could also just have said ky = kx.T for the transpose, I guess.

output = my_function(kx, ky)

Note that math operations on numpy arrays are typically applied element-wise, so kx**2 just squares each element of kx individually, which is what you want here. For matrix-multiplication type squaring you'd do something like kx.dot(kx).

I assume your arrays kx and ky are just random examples; otherwise, there's more efficient ways to generate them than writing them out like this.

Of course nothing is wrong with the other posted solutions that use plain python objects and for loops, but for larger arrays you can expect numpy to be much more efficient.

Upvotes: 1

Related Questions