Conditional function change in matplotlib's contour plot

Question

I would like to generate a contour plot, using a 2 variable function, Z(X,Y). However, I would like to impose a condition that changes the function when X is less/greater than a certain value. Such a change would allow me to use just a singular plt.contour line (i.e I do not want to create two separately defined functions, that results in using two separate plotting command lines).
I continue to run into (what I believe to be) truth/logic errors. My guess is that some aspect of the numpy meshgrid does not comply with that conditional "switch" of the function. Attached below is the short code to display the concept, along with the full Traceback error. If anything is unclear, or if what I have provided is insufficient to explain my problem, please feel free to comment below.

import numpy as np
import matplotlib.pyplot as plt


X = np.linspace(0,50,100)
Y = np.linspace(0,50,100)
X, Y = np.meshgrid(X,Y)


def z(x,y):

    if x < 20:
        return np.sin(x) + np.cos(y)
    
    else:
        return np.tan(x * y)


Z = z(X,Y)


plt.contourf(X, Y, Z)
plt.xlabel('x')
plt.ylabel('y')
plt.colorbar()

ValueError                                Traceback (most recent call last)
<ipython-input-29-7e200be093e6> in <module>
     16 
     17 
---> 18 Z = z(X,Y)
     19 
     20 plt.figure(figsize=(8,6))

<ipython-input-29-7e200be093e6> in z(x, y)
      9 
     10 def z(x,y):
---> 11     if x < 20:
     12         return np.sin(x) + np.cos(y)
     13 

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()```

score 2 · Answer 1 · answered Sep 20 '20 at 19:52

2

Everything is fine but u have to change the comparison on np.any(x<20). It means that if any of the element of x is greater than 20. You can also use np.all, if you want to do this just if every element of the array fulfill the condition

import numpy as np
import matplotlib.pyplot as plt


X = np.linspace(0,50,100)
Y = np.linspace(0,50,100)
X, Y = np.meshgrid(X,Y)


def z(x,y):

    if np.any(x < 20):
        return np.sin(x) + np.cos(y)
    
    else:
        return np.tan(x * y)


Z = z(X,Y)


plt.contourf(X, Y, Z)
plt.xlabel('x')
plt.ylabel('y')
plt.colorbar()

answered Sep 20 '20 at 19:52

Maciek Woźniak

346
2
10

Thank you! Although, I am having a bit of difficulty when thinking about whether I should use np.all(x<20) or np.any(x<20) . When using both of these changes, the error goes away which is great! But they both (of course) produce two different plots. Does my situation require one or the other, to your knowledge? – ryguy Sep 20 '20 at 20:19
consider this example, if any(np.any) element of the array is > 0 -> true, if all (np.all) elements of the array are >0 -> in the case below, false np.any(np.asarray([[1,0,4],[-3,4,-2]])>0) # true np.all(np.asarray([[1,0,4],[-3,4,-2]])>0) # false – Maciek Woźniak Sep 20 '20 at 21:32

score 2 · Answer 2 · answered Sep 20 '20 at 19:55

Truth testing in numpy array is element wise,

import numpy as np
X =  np.array([1,2,3,4])
print(X<3)

output: [ True True False False]

You can use .all() or .any() depending on whether you want all or any of the elements to satisfy the criterion

print((X<3).all())

output : False

score 2 · Accepted Answer · answered Sep 20 '20 at 20:10

You can simply use numpy.where().
You have to pass:

a condition as the first parameter
values from which to choose if that condition is true as the second parameter
values from which to choose if that condition is false as the third parameter

In this way, your z() function becomes:

def z(x, y):
    return np.where(x < 20,
                    np.sin(x) + np.cos(y),
                    np.tan(x*y))

Generated plot:

Conditional function change in matplotlib's contour plot

3 Answers3