import sympy as sym
import numpy as np
import matplotlib.pyplot as plt
from scipy.signal import find_peaks
from IPython.display import display,Math

print('Determine which values of "a" give the function a critical value at x=1 or x=2')
display(Math('f(x,a) = x^2e^{-ax^2}'))
display(Math('a \\: \\epsilon \\: \{0,.25,.5,...,2\}'))

Determine which values of "a" give the function a critical value at x=1 or x=2

$\displaystyle f(x,a) = x^2e^{-ax^2}$

$\displaystyle a \: \epsilon \: \{0,.25,.5,...,2\}$

a,x = sym.symbols('a,x')

baseexpr = x**2 * sym.exp(-a*x**2)

arange = np.arange(0,2.25,.25)
print(arange)
xrange = np.linspace(-3,3,100)

# setup the figure
fig,ax = plt.subplots(1,2) # (row, col)
print(fig)
print(ax)

for ai in arange:
    fx = baseexpr.subs(a,ai)
    dfx = sym.diff(fx)
#     print(dfx)
    critpnts = sym.solve(dfx)
#     print(critpnts)
    
    ax[0].plot(xrange,sym.lambdify(x,fx)(xrange))
    ax[1].plot(xrange,sym.lambdify(x,dfx)(xrange))
    
#     print(critpnts)
    if 1 in critpnts:
        display(Math('\\Rightarrow %s \\text{ has a critical point at } x=1' %sym.latex(fx)))
    elif 2 in critpnts:
        display(Math('\\Rightarrow %s \\text{ has a critical point at } x=2' %sym.latex(fx)))
    else:
        display(Math('\\quad %s \\text{ has NO critical point at x=1 or x=2}' %sym.latex(fx)))
        

    
ax[0].set_ylim([0,2])
ax[0].set_title('Function')
ax[0].plot([1,1],[0,2],'--',color='gray')
ax[0].plot([2,2],[0,2],'--',color='gray')

ax[1].set_ylim([-2,2])
ax[1].set_title('Derivative')
ax[1].plot([1,1],[-2,2],'--',color='gray')
ax[1].plot([2,2],[-2,2],'--',color='gray')

plt.show()

[0.   0.25 0.5  0.75 1.   1.25 1.5  1.75 2.  ]
Figure(432x288)
[<matplotlib.axes._subplots.AxesSubplot object at 0x1112ab4e0>
 <matplotlib.axes._subplots.AxesSubplot object at 0x1112f8ac8>]

$\displaystyle \quad x^{2} \text{ has NO critical point at x=1 or x=2}$

$\displaystyle \Rightarrow x^{2} e^{- 0.25 x^{2}} \text{ has a critical point at } x=2$

$\displaystyle \quad x^{2} e^{- 0.5 x^{2}} \text{ has NO critical point at x=1 or x=2}$

$\displaystyle \quad x^{2} e^{- 0.75 x^{2}} \text{ has NO critical point at x=1 or x=2}$

$\displaystyle \Rightarrow x^{2} e^{- 1.0 x^{2}} \text{ has a critical point at } x=1$

$\displaystyle \quad x^{2} e^{- 1.25 x^{2}} \text{ has NO critical point at x=1 or x=2}$

$\displaystyle \quad x^{2} e^{- 1.5 x^{2}} \text{ has NO critical point at x=1 or x=2}$

$\displaystyle \quad x^{2} e^{- 1.75 x^{2}} \text{ has NO critical point at x=1 or x=2}$

$\displaystyle \quad x^{2} e^{- 2.0 x^{2}} \text{ has NO critical point at x=1 or x=2}$