how to plot 2d condition function

Question

I want to plot a 2d condition function, and the code as below:

from numpy import exp,arange
from pylab import meshgrid,cm,imshow,show

def z_func(x,y):
    zParity = x % 2
    nParity = y % 2
    if zParity == 0 and nParity == 0:
        return x+y+1
    elif zParity!=0 and nParity!=0:
        return x+y-1
    else:
        return x+y

x = arange(1,100,1)
y = arange(1,100,1)
X,Y = meshgrid(x, y) # grid of point
Z = z_func(X, Y) # evaluation of the function on the grid
im = imshow(Z,cmap=cm.RdBu) # drawing the function
show()

I get an error: ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all() How can I change the code to work?

Answer 1

Try this:

from numpy import exp,arange
from pylab import meshgrid,cm,imshow,show
import numpy as np

def z_func(x,y):
    zParity = x % 2
    nParity = y % 2
    if np.all((zParity == 0, nParity == 0)):
        return x+y+1
    elif np.all((zParity!=0, nParity!=0)):
        return x+y-1
    else:
        return x+y

x = arange(1,100,1)
y = arange(1,100,1)
X,Y = meshgrid(x, y) # grid of point
Z = z_func(X, Y) # evaluation of the function on the grid
im = imshow(Z,cmap=cm.RdBu) # drawing the function
show()

This is what I get after running the code (I hope I have understood what you want to get, otherwise I will edit my answer accordingly):