#!/usr/bin/env python
# -*- coding: utf-8 -*-
''' Ejercicio 16, guía 1 de la materia Física 2 (Cátedra Depine)
    ------------------------------------------------------------
    
    Basado en un ejemplo del libro "Waves. Berkeley Physics Course - Vol 3" de
    F. S. Crawford. El mismo puede verse en la fig. 3.12, pág. 140.
    Encontramos la solución estacionaria del problema, sin otras aproximaciones.
    Trabajamos algebraicamente proponiendo soluciones en modos normales.

    Alberto Camjayi, 07/09/2015 '''

#   Parámetros del problema
g = 10                 # aceleración de la gravedad
m = 1                  # masa de los péndulos
l_largos = 1           # largo de los primeros péndulos
l_cortos = 0.5         # idem de los del final de la cadena
k = 5                  # constante de los resortes
Gama = 0.1             # constante de amortiguamiento
N_largos = 30          # número de péndulos largos
N_cortos = 200         # número de péndulos cortos
F0_sobre_m = 4         # fuerza sobre la primera masa
Omega = 3              # frecuencia angular de forzado
# ----------------------------------------------------------------
# Importamos algunas herramientas necesarias
from scipy import *
from scipy.linalg import eigh
from scipy.sparse import spdiags, hstack, vstack, bmat, dia_matrix
from pylab import figure, grid, plot, axis, xlabel, ylabel, title, legend, show

def get_M(N1,l1,N2,l2):             # Aquí construímos la matriz
    Nm = N1+N2
    diag0 = zeros(Nm); diag1 = zeros(Nm)
    diag0[:N1] = g/l1+2*k/m          # Elementos de la diagonal
    diag0[N1:] = g/l2+2*k/m          # Elementos de la diagonal
    diag1[1:] = -k/m                 # Elementos extradiagonales

    Mi = spdiags([diag0,diag1],[0,1],Nm,Nm)
    return Mi

def get_sol_estacionaria(w0,Fm,Omega):
    Aes = zeros(w0.size)
    Bes = zeros(w0.size)

    Aes =         Omega*Gama*Fm/((w0**2-Omega**2)**2 + Omega**2*Gama**2)
    Bes = (w0**2 - Omega**2)*Fm/((w0**2-Omega**2)**2 + Omega**2*Gama**2)

    return Aes, Bes


M = get_M(N_largos,l_largos,N_cortos,l_cortos).toarray()
w0_cuadrado, Dmatrix = eigh(M,lower=False)

F = zeros(N_largos+N_cortos); F[0] = F0_sobre_m

F_en_modos = dot(transpose(Dmatrix),F)


for i in range(7):
    Omega = Omega + i*3 

    A_en_modos, B_en_modos = get_sol_estacionaria(w0_cuadrado,F_en_modos,Omega)

    A = dot(Dmatrix,A_en_modos)

    B = dot(Dmatrix,B_en_modos)

    figure()
    plot(A,'ok-',label='Amplitud absorbente')
    plot(B,'sr-',label=u'Amplitud elástica')
    plot((N_largos,N_largos),(1.2*min(A.min(),B.min()),1.2*max(A.max(),B.max())),'k-')
    grid()
    axis([0,40,1.2*min(A.min(),B.min()),1.2*max(A.max(),B.max())])
    xlabel('Masa')
    ylabel('Amplitud')
    legend(loc='lower right')
    w0_min = str(round(w0_cuadrado.min(),2)) 
    w0_max = str(round(w0_cuadrado.max(),2)) 
    title(r'$\Omega = $'+str(Omega)+', '+r'$\omega_0^{min} =$'+w0_min+', '+r'$\omega_0^{max} =$'+w0_max)

show()