Best way to apply a two-dimensional function in numpy

user_7053813 0
Viewed 58

I often find myself applying two dimensional functions in numpy. I have done it several different ways, however, none seem to be very elegant. Is there a "correct" way to do the following in numpy?

import numpy as np

X = np.linspace(0, 1, 100)
Y = np.linspace(0, 1, 100)

mg = np.meshgrid(X, Y)
zipped = np.c_[mg[0].ravel(), mg[1].ravel()]

def harmonic_avg(row):
    return 2 * row[0] * row[1] / (row[0] + row[1])

result = np.apply_along_axis(harmonic_avg, 1, zipped)

result.reshape(100,100)
numpy python
1 Answers

In this case the most elegant way would imho be using the meshgrids directly

result = 2*mg[0]*mg[1]/(mg[0]+mg[1])

Or as a function

def harmonic_avg(mg):
    return 2*mg[0]*mg[1]/(mg[0]+mg[1])
result = harmonic_avg(mg)
user_55246370
Related