Can vectorization be applied to a function Z= f(X,Y) defined by two or more functions

user_13611363 0
Viewed 59

I am using np.meshgrid() to plot an scalar function: Z= f(X,Y).

I have no problems when f(X,Y) is defined by a single equation. However, when f(X,Y) is split in two or more equations, such as:

X, Y = np.meshgrid(x, y, sparse = True, indexing = 'ij')
Z = f1(X,Y), y1 < y <= y2
Z = f2(X,Y), y > y2

I have no idea how to extract the value of 'Y' from the meshgrid array in order to check it against y1 and y2 in the conditional statement. I have tried several ways, like the one below, to no avail:

row1, col1 = np.where(Y == y1)
row2, col2 = np.where(Y == y2)
if( Y > Y[row1][0] and Y <= Y[row2][0] ):
    Z = f1(X, Y)
elif ( Y > Y[row2][0] ):
    Z = f2(X, Y)

When I use arrays for the calculations of Z, the execution is really slow, at least 10 to 20 times slower, depending on the grid size, when compared with vectorization. Is it possible to apply vectorization in this case? how? Thanks!

numpy python
2 Answers

If the split is for a given constant x or y value, as in your example, you could compute the function for both regions (y < y1 and y >= y1) and then stack the result:

import matplotlib.pyplot as plt
import numpy as np

x = np.arange(-5, 5, 0.1)
y1 = np.arange(-5, 2, 0.1)
y2 = np.arange(2, 5, 0.1)

xx, yy1 = np.meshgrid(x, y1, sparse=True)
xx, yy2 = np.meshgrid(x, y2, sparse=True)

z1 = np.sin(xx**2 + yy1**2) / (xx**2 + yy1**2)
z2 = xx + yy2

z = np.vstack([z1,z2])
y = np.hstack([y1,y2])
h = plt.contourf(x,y,z)

enter image description here

This doen't work, however, if the regions are divided not by x or y alone but by say x+y.

user_39443220

Use np.where in general. 

z = np.where(Y <= y2, f1(X, Y), f2(X, Y))

If you have a null case (as it seems with Y < y1) you may need to initialize with np.zeros_like

Z = np.zeros_like(Y)  # same shape and dtype as `Y`
np.where(y1 < Y <= y2, f1(X, Y), out = Z)
np.where(y< y2, f2(X,Y), out = Z)

Be aware that if one of your functions raises an error outside its range, you will get errors since np.where calculates the function over the entire X, Y range. In that case you'll want to try @Stef 's method

user_44277770
Related