Search code examples
pythonmatplotlibcontourfigurescatter

Matplotlib contour plot with conditions on the curve


Suppose I want to plot the contour of z=0 where z=(19*y^3-6*x*y-1). I can use the following code:

x = np.linspace(-2,2,1000)
y = np.linspace(-1,1,1000)
X,Y = np.meshgrid(x,y)
z = (19)*Y**3-6*X*Y-1
plt.figure()
plt.contour(z,0)
plt.show()

Now I want to display the portion of the curve where 3*19*y^2-6*x<0 in solid line, and the portion where 3*19*y^2-6*x>0 in dashed line. Basically I'm doing some sort of stability analysis of a system, and I want to show different regions of the z=0 curve differently depending on whether dz/dy is positive or negative.

What I can think of is to locate the coordinates of the two parts myself, and to use scatter plot to show the two parts of the curve using different colour (or line style). I also know how to do this easily in Mathematica. I just wonder whether there is a more elegant solution in matplotlib to do this job.


Solution

  • Maybe the following approach is interesting, although not perfect?

    A variable z2 is created with for the 3*19*y^2-6*x>0 condition. z2 is erased everywhere except where it is close to z. Then it is colored with a red-blue colormap, red for the positive part, blue for the negative, and white around 0.

    The background is set to black and the contour color to white in order to have enough contrast.

    Note that both the contour plot and the imshow need the extent parameter to be set in order to get informative axes.

    import matplotlib.pyplot as plt
    import numpy as np
    
    x = np.linspace(-2, 2, 1000)
    y = np.linspace(-1, 1, 1000)
    X, Y = np.meshgrid(x, y)
    z = (19) * Y ** 3 - 6 * X * Y - 1
    z2 = 3 * 19 * Y ** 2 - 6 * X
    z2 = np.where(np.abs(z) < 0.2, z2, np.NaN)
    plt.gca().set_facecolor('black')
    plt.imshow(z2, cmap='coolwarm', vmin=-1, vmax=+1, alpha=1, extent=[-2, 2, -1, 1], origin='lower')
    plt.contour(z, 0, extent=[-2, 2, -1, 1], zorder=3, colors='white')
    plt.show()
    

    resulting plot