import numpy as np
from sympy import *
from scipy.linalg import eigh
from scipy.special import jn_zeros as Jn_zeros, jn as Jn
import matplotlib.pyplot as plt

%matplotlib notebook

init_session()

k, L = symbols('k L', positive=True)

def u_fun(x, m):
    """ Trial function. """
    c = symbols('c0:%i' % m)
    w = (1 - z)**2 * sum(c[k]*x**k for k in range (0, m))
    return w, c

m = 10
w, coef = u_fun(z, m)
display(w)

T_inte = w**2
U_inte = z*diff(w, z)**2

T = integrate(T_inte, (z, 0, 1))
U = integrate(U_inte, (z, 0, 1))

K = Matrix(m, m, lambda ii, jj: diff(U, coef[ii], coef[jj]))
M = Matrix(m, m, lambda ii, jj: diff(T, coef[ii], coef[jj]))

Kn = np.array(K).astype(np.float64)
Mn = np.array(M).astype(np.float64)

"""
Explanation: Vibrations modes of a hanging chain using the Ritz method
The equation of the vibration of a hanging chain
$$\mu\frac{\partial^2}{\partial t^2} u(z, t) = \frac{\partial}{\partial z}\left(gz\frac{\partial}{\partial z}u(z,t)\right)$$
If we assume a solution of the form $u(z, t) = e^{i\omega t} w(z)$, we can rewrite the
differential equation as
$$-k^2 w(z) = \frac{\partial}{\partial z}\left(z\frac{\partial}{\partial z}w(z)\right)$$
with $k^2 = \mu \omega^2/g$, and $z(0)=0$, $\partial w(L)/\partial z = 0$.
This can be written in variational form as
$$J = \int_0^L z \left[\frac{\partial w}{\partial z}\right]^2 dz + k^2\int_0^L w^2 dz$$
End of explanation
"""


vals, vecs = eigh(Kn, Mn, eigvals=(0, m-1))
np.sqrt(vals)

lam = Jn_zeros(0, m)/2
lam

z_vec = np.linspace(0, 1, 100)
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(8, 4))
for num in range(5):
    u_num = lambdify((z), w.subs({coef[kk]: vecs[kk, num] for kk in range(m)}), "numpy")  
    ax1.plot(z_vec, u_num(z_vec)/u_num(0))
    ax2.plot(z_vec, Jn(0, 2*lam[num]*np.sqrt(z_vec)), label=r"$m=%i$"%num)

ax1.set_title("Approximated solution")
ax2.set_title("Exact solution")
plt.legend(loc="best", framealpha=0);

"""
Explanation: We can compare the eigenvalues obtained using the Ritz method and the analytic ones,
that are given by
$$k_n = \frac{\alpha}{2\sqrt{L}}\, ,$$
where $\alpha_n$ is the $n$th zero of the Bessel function of the first kind and zero order
$\mathrm{J}_0$.
End of explanation
"""


