Ohm
Ohm

Reputation: 2442

Matplotlib challenge - smart way to plot only in the existence range

I am trying to plot the following plot using matplotlib, for the function:

y1 = (x1/2.) + np.sqrt(np.power(x1/2.,2) + a)
y2 = (x2/2.) - np.sqrt(np.power(x2/2.,2) + a)

where a < 0, so I took for this example a = -1 I tried to restrict the calculation for the range of values between -10 and +10, omitting the values in which the root argument is bellow zero.

The plot should look like this - plot_function

import numpy as np
import matplotlib.pyplot as plt

plt.figure(figsize=(12,3))

# Example 1
plt.subplot(1,3,1)
plt.ylim([-10,10])
plt.xlim([-10,10])
a = -1
x1 = np.linspace(np.sqrt(-2.*a),10,100)
x2 = np.linspace(-10,-1*np.sqrt(-2.*a),100)
x12 = x1 + x2
x1 = x12
x2 = x12
print x12

y1 = (x1/2.) + np.sqrt(np.power(x1/2.,2) + a)
y2 = (x2/2.) - np.sqrt(np.power(x2/2.,2) + a)
# put the spines to the middle
ax = plt.gca()
ax.spines['left'].set_position('center')
ax.spines['right'].set_color('none')
ax.spines['bottom'].set_position('center')
ax.spines['top'].set_color('none')
ax.spines['left'].set_smart_bounds(True)
ax.spines['bottom'].set_smart_bounds(True)
ax.xaxis.set_ticks_position('bottom')
ax.yaxis.set_ticks_position('left')
#some sort of split condition:
solid1 = y1 > x1/2.
solid2 = y2 > x2/2.
#plot each line in two seperate plot calls
plt.plot(x1[solid1], y1[solid1], 'r', linewidth=4)
plt.plot(x1[np.logical_not(solid1)], y1[np.logical_not(solid1)], 'r--', linewidth=4)
plt.plot(x2[solid2], y2[solid2], 'g', linewidth=4)
plt.plot(x2[np.logical_not(solid2)], y2[np.logical_not(solid2)], 'g--', linewidth=4)

plt.show()

Instead I got this beast beast

Someone has a good idea?

Upvotes: 2

Views: 270

Answers (1)

elyase
elyase

Reputation: 40973

Don't worry about negative values, numpy will only issue a Warning:

import numpy as np
import matplotlib.pyplot as plt

a = -1
x = np.linspace(-19, 19, 10000)
y1 = (x/2.) + np.sqrt(np.power(x/2., 2) + a)
y2 = (x/2.) - np.sqrt(np.power(x/2., 2) + a)
plt.plot(x, y1, x, y2, c='k')

# customization
plt.ylim(-12,12)
ax = plt.gca()
ax.spines['left'].set_position('center')
ax.spines['right'].set_color('none')
ax.spines['bottom'].set_position('center')
ax.spines['top'].set_color('none')
ax.xaxis.set_ticks_position('bottom')
ax.yaxis.set_ticks_position('left')
ax.grid(True)

plt.show()

Result:

enter image description here

If you want to avoid the numpy warning just do:

with np.errstate(invalid='ignore'):
    y1 = (x/2.) + np.sqrt(np.power(x/2.,2) + a)
    y2 = (x/2.) - np.sqrt(np.power(x/2.,2) + a)

Upvotes: 3

Related Questions