python help - conditional math operation on arrays - python-2.7

Question

I'm trying to learn python - but new to OOP. I'd like to make the functions fA, fB, fC operate on the whole r-theta space at once, instead of one point at a time. My problem is the conditional (r<=1). The code below is quite ugly (!) but it works.
How can I make this more pythonesque? Thanks!
(in this simplified example, note the math for (r>1) diverges when r goes to zero)
from math import pi, sin, cos, exp
import numpy as np
import matplotlib.pyplot as plt
def fA(rr,th,a,b,c):
if (rr<=1):
fx = a * sin(th)
fy = b * rr * cos(th)
fz = c * rr
else:
fx = (a / rr) * sin(th)
fy = (b / rr) * cos(th)
fz = (c / rr)
return(fx,fy,fz)
def fB(rr,th,a,b,c):
if (rr<=1):
fx = b * sin(2.*th)
fy = a * rr * cos(2.*th)
fz = c * rr
else:
fx = (b / rr) * sin(2.*th)
fy = (a / rr) * cos(2.*th)
fz = c
return(fx,fy,fz)
def fC(rr,th,a,b,c):
if (rr<=1):
fx = exp(rr - 1.) * cos(th)
fy = exp(rr - 1.) * sin(th)
fz = c
else:
fx = exp(1. - rr) * cos(th)
fy = exp(1. - rr) * sin(th)
fz = c / rr
return(fx,fy,fz)
nx = 101
ny = 101
dx = 4. / (nx-1)
dy = 4. / (ny-1)
X = np.zeros((ny,nx))
Y = np.zeros((ny,nx))
for ix in range(nx):
for iy in range(ny):
X[iy,ix] = dx*(ix - (nx-1)/2)
Y[iy,ix] = dy*(iy - (ny-1)/2)
r = np.sqrt(X**2. + Y**2.)
theta = np.arctan2(Y,X)
Ax = np.zeros((ny,nx))
Ay = np.zeros((ny,nx))
Az = np.zeros((ny,nx))
Bx = np.zeros((ny,nx))
By = np.zeros((ny,nx))
Bz = np.zeros((ny,nx))
Cx = np.zeros((ny,nx))
Cy = np.zeros((ny,nx))
Cz = np.zeros((ny,nx))
for ix in range (nx):
for iy in range(ny):
Ax[iy,ix], Ay[iy,ix], Az[iy,ix] = fA(r[iy,ix], theta[iy,ix], 1.0, 1.0, 1.5)
Bx[iy,ix], By[iy,ix], Bz[iy,ix] = fB(r[iy,ix], theta[iy,ix], 1.5, 0.8, 1.0)
Cx[iy,ix], Cy[iy,ix], Cz[iy,ix] = fC(r[iy,ix], theta[iy,ix], 0.9, 1.1, 1.2)
plt.figure()
plt.subplot(3,3,1)
plt.imshow(Ax)
plt.colorbar()
plt.title('Ax')
plt.subplot(3,3,2)
plt.imshow(Ay)
plt.colorbar()
plt.title('Ay')
plt.subplot(3,3,3)
plt.imshow(Az)
plt.colorbar()
plt.title('Az')
plt.subplot(3,3,4)
plt.imshow(Bx)
plt.colorbar()
plt.title('Bx')
plt.subplot(3,3,5)
plt.imshow(By)
plt.colorbar()
plt.title('By')
plt.subplot(3,3,6)
plt.imshow(Bz)
plt.colorbar()
plt.title('Bz')
plt.subplot(3,3,7)
plt.imshow(Cx)
plt.colorbar()
plt.title('Cx')
plt.subplot(3,3,8)
plt.imshow(Cy)
plt.colorbar()
plt.title('Cy')
plt.subplot(3,3,9)
plt.imshow(Cz)
plt.colorbar()
plt.title('Cz')
plt.show()

Answer 1

Like #mdurant said np.where and np.meshgrid will be useful. Here, let me reorganize your code and provide another way to avoid Python looping by using numpy advanced indexing:
import sys
from math import pi, sin, cos, exp
import numpy as np
import matplotlib.pyplot as plt
def _generate_coordinate(nx, ny):
"""
Generate coordinate data points in a function to prevent namespace
pollution.
"""
dx = 4. / (nx-1)
dy = 4. / (ny-1)
X = np.zeros((ny,nx))
Y = np.zeros((ny,nx))
for ix in range(nx):
for iy in range(ny):
X[iy,ix] = dx*(ix - (nx-1)/2)
Y[iy,ix] = dy*(iy - (ny-1)/2)
return np.sqrt(X**2 + Y**2), np.arctan2(Y,X)
nx = ny = 101
r, theta = _generate_coordinate(101, 101)
def calculate_numpy():
# Helper methods for vector-based calculator.
def fA(rr, th, a, b, c):
# Calculate every value because mulplication doesn't give NaN.
arrx = a * np.sin(th)
arry = b * rr * np.cos(th)
arrz = c * rr
# Override value with a certain condition.
slct = rr > 1
rr = rr[slct]
th = th[slct]
arrx[slct] = a / rr * np.sin(th)
arry[slct] = b / rr * np.cos(th)
arrz[slct] = c / rr
return arrx, arry, arrz
def fB(rr, th, a, b, c):
# Calculate every value because mulplication doesn't give NaN.
arrx = b * np.sin(2.*th)
arry = a * rr * np.cos(2.*th)
arrz = c * rr
# Override value with a certain condition.
slct = rr > 1
rr = rr[slct]
th = th[slct]
arrx[slct] = b / rr * np.sin(2.*th)
arry[slct] = a / rr * np.cos(2.*th)
arrz[slct] = c
return arrx, arry, arrz
def fC(rr,th,a,b,c):
# Calculate every value because mulplication doesn't give NaN.
arrx = np.exp(rr-1) * np.cos(th)
arry = np.exp(rr-1) * np.sin(th)
arrz = np.empty_like(rr)
arrz.fill(c)
# Override value with a certain condition.
slct = rr > 1
rr = rr[slct]
th = th[slct]
arrx[slct] = np.exp(1.-rr) * np.cos(th)
arry[slct] = np.exp(1.-rr) * np.sin(th)
arrz[slct] = c / rr
return arrx, arry, arrz
# Carry out calculation.
Ax, Ay, Az = fA(r, theta, 1.0, 1.0, 1.5)
Bx, By, Bz = fB(r, theta, 1.5, 0.8, 1.0)
Cx, Cy, Cz = fC(r, theta, 0.9, 1.1, 1.2)
return Ax, Ay, Az, Bx, By, Bz, Cx, Cy, Cz
def calculate_loop():
# Helper methods for loop calculator.
def fA(rr,th,a,b,c):
if (rr<=1):
fx = a * sin(th)
fy = b * rr * cos(th)
fz = c * rr
else:
fx = (a / rr) * sin(th)
fy = (b / rr) * cos(th)
fz = (c / rr)
return(fx,fy,fz)
def fB(rr,th,a,b,c):
if (rr<=1):
fx = b * sin(2.*th)
fy = a * rr * cos(2.*th)
fz = c * rr
else:
fx = (b / rr) * sin(2.*th)
fy = (a / rr) * cos(2.*th)
fz = c
return(fx,fy,fz)
def fC(rr,th,a,b,c):
if (rr<=1):
fx = exp(rr - 1.) * cos(th)
fy = exp(rr - 1.) * sin(th)
fz = c
else:
fx = exp(1. - rr) * cos(th)
fy = exp(1. - rr) * sin(th)
fz = c / rr
return(fx,fy,fz)
# Create buffer arrays for loops.
Ax = np.zeros((ny,nx))
Ay = np.zeros((ny,nx))
Az = np.zeros((ny,nx))
Bx = np.zeros((ny,nx))
By = np.zeros((ny,nx))
Bz = np.zeros((ny,nx))
Cx = np.zeros((ny,nx))
Cy = np.zeros((ny,nx))
Cz = np.zeros((ny,nx))
# Carry out calculation with Python loops. This is slow.
for ix in range (nx):
for iy in range(ny):
Ax[iy,ix], Ay[iy,ix], Az[iy,ix] = fA(r[iy,ix], theta[iy,ix], 1.0, 1.0, 1.5)
Bx[iy,ix], By[iy,ix], Bz[iy,ix] = fB(r[iy,ix], theta[iy,ix], 1.5, 0.8, 1.0)
Cx[iy,ix], Cy[iy,ix], Cz[iy,ix] = fC(r[iy,ix], theta[iy,ix], 0.9, 1.1, 1.2)
return Ax, Ay, Az, Bx, By, Bz, Cx, Cy, Cz
def main():
calculate = globals()["calculate_" + sys.argv[1]]
Ax, Ay, Az, Bx, By, Bz, Cx, Cy, Cz = calculate()
plt.figure()
plt.subplot(3,3,1)
plt.imshow(Ax)
plt.colorbar()
plt.title('Ax')
plt.subplot(3,3,2)
plt.imshow(Ay)
plt.colorbar()
plt.title('Ay')
plt.subplot(3,3,3)
plt.imshow(Az)
plt.colorbar()
plt.title('Az')
plt.subplot(3,3,4)
plt.imshow(Bx)
plt.colorbar()
plt.title('Bx')
plt.subplot(3,3,5)
plt.imshow(By)
plt.colorbar()
plt.title('By')
plt.subplot(3,3,6)
plt.imshow(Bz)
plt.colorbar()
plt.title('Bz')
plt.subplot(3,3,7)
plt.imshow(Cx)
plt.colorbar()
plt.title('Cx')
plt.subplot(3,3,8)
plt.imshow(Cy)
plt.colorbar()
plt.title('Cy')
plt.subplot(3,3,9)
plt.imshow(Cz)
plt.colorbar()
plt.title('Cz')
plt.show()
if __name__ == '__main__':
main()
The function calculate_numpy() is where I demonstrate the advanced indexing. If you want to totally avoid duplicated calculation, you will need to create buffers like what you did in calculate_loop(). But I would say in terms of runtime, the duplication calculation is OK.
Let's say the program is saved in a file draw.py. We have both numpy ndarray and loop versions of the code, and can use timeit to benchmark them:
$ python -m timeit -s "import draw" "draw.calculate_loop()"
10 loops, best of 3: 95.2 msec per loop
$ python -m timeit -s "import draw" "draw.calculate_numpy()"
100 loops, best of 3: 2.11 msec per loop
As you can see, the numpy version is 45 times faster than your loop version. Most of the cases it's good enough.

Answer 2

Choosing one of your relations at random, you can use numpy.where instead of the if... semantics:
fx = where( rr<=1, exp(rr - 1.) * cos(th), exp(1. - rr) * cos(th))
This effectively does if/else for a set of arrays rather than one at a time individual numbers. You would then be able to do Ax, Ay, Az = fA(...) without the looping.
X and Y you can make using meshgrid or mgrid.
To not evaluate all elements, you could use slice notation
fx = empty_like(rr)
fx[rr<=1] = exp(rr[rr<=1] - 1.) * cos(th[rr<=1])
fx[rr>1] = exp(1. - rr[rr>1]) * cos(th[rr>1]))