time_series_analysis.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Análise Básica de Séries Temporais\n",
    "\n",
    "O estudo de fenômenos que variam no tempo é uma prática bastante utilizada em Astronomia. Este Jupyter Notebook traz alguns exemplos modestos de como utilizar algumas ferramentas baseadas em Python para encontrar o período presente em séries temporais.\n",
    "\n",
    "[1. O conceito de dobra em fase para busca de períodos](#dobra_fase)\n",
    "\n",
    "[2. O periodograma de Lomb-Scargle](#lombscargle)\n",
    "\n",
    "[3. PDM - Phase Dispersion Minimization](#pdm)\n",
    "\n",
    "[4. Ajustando um modelo utilizando o módulo LMFIT](#lmfit)\n",
    "\n",
    "[5. Visualizando a Transformada de Fourier](#dft_visualization)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Bibliotecas utilizadas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'PyAstronomy'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-1-92f4d73e1ffc>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     12\u001b[0m \u001b[0;31m# PyAstronomy\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     13\u001b[0m \u001b[0;31m# biblioteca com algumas opções para séries temporais\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 14\u001b[0;31m \u001b[0;32mfrom\u001b[0m \u001b[0mPyAstronomy\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpyTiming\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mpyPDM\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     15\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     16\u001b[0m \u001b[0;31m# Scipy\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'PyAstronomy'"
     ]
    }
   ],
   "source": [
    "# manipulação de arrays\n",
    "import numpy as np\n",
    "\n",
    "# biblioteca gráfica\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Astropy\n",
    "# biblioteca de rotinas astronômicas\n",
    "from astropy.stats import LombScargle\n",
    "from astropy.io import fits\n",
    "\n",
    "# PyAstronomy\n",
    "# biblioteca com algumas opções para séries temporais\n",
    "from PyAstronomy.pyTiming import pyPDM\n",
    "\n",
    "# Scipy\n",
    "# biblioteca de rotinas matemáticas e estatísticas\n",
    "from scipy.signal import savgol_filter\n",
    "\n",
    "# LMFIT\n",
    "# biblioteca de rotinas de minimização\n",
    "from lmfit import minimize, Minimizer, Parameters, report_fit\n",
    "\n",
    "# Bokeh\n",
    "# biblioteca de visualização interativa\n",
    "from bokeh.layouts import column\n",
    "from bokeh.models import CustomJS, ColumnDataSource, Slider\n",
    "from bokeh.plotting import figure, output_notebook, show"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Algumas funções utilizadas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# funcao que le arquivo de texto com tres colunas numéricas (float)\n",
    "# (sugestão: use numpy.genfromtxt se necessitar ler floats e strings)\n",
    "def read_txt(file):\n",
    "    time, flux, err = np.loadtxt(file, usecols=(0,1,2), \\\n",
    "                                 unpack=True, comments=\"#\")\n",
    "    return time, flux, err\n",
    "\n",
    "# calcula o periodograma de lombscargle\n",
    "# retorna melhor frequência, melhor potência, \n",
    "# probabilidade de falso alarme,\n",
    "# periodograma e modelo do ajuste\n",
    "def scargle(x, y, err, frequencia, plot=False):\n",
    "    from astropy.stats import LombScargle\n",
    "    # objeto contendo as propriedades do calculo\n",
    "    ls = LombScargle(x, y, err)\n",
    "\n",
    "    # calcula o o periodograma\n",
    "    power = ls.power(frequencia)\n",
    "\n",
    "    # localiza a o índice do maior pico\n",
    "    idx = np.argmax(power)\n",
    "\n",
    "    # melhor frequência e pico correspondente\n",
    "    best_freq = frequencia[idx]\n",
    "    best_power = power[idx]\n",
    "    modelo = ls.model(x,best_freq)\n",
    "\n",
    "    # calcula a probabilidade de falso alarme\n",
    "    fap = ls.false_alarm_probability(best_power, method='baluev', \\\n",
    "                                     minimum_frequency=frequencia.min(), \\\n",
    "                                     maximum_frequency=frequencia.max())\n",
    "\n",
    "    print('Frequência: {0:.8f} (Período: {1:.8f})\\nFAP: {2:6.2f}%'.format(best_freq,1/best_freq,fap*100))\n",
    "    if plot:\n",
    "        # plota resultados\n",
    "        phase = (x*best_freq) % 1\n",
    "        fig, ax = plt.subplots(1,2, figsize=(14,6))\n",
    "        ax[0].errorbar(phase, y, err, fmt='.', color='blue', label='Medidas')\n",
    "        srt = np.argsort(phase) # obtem o indice da fase ordenada\n",
    "        ax[0].plot(phase[srt], modelo[srt], '-', color='red', label='Modelo')\n",
    "        ax[0].set_xlabel(\"Fase\")\n",
    "        ax[0].set_ylabel(\"Fluxo\")\n",
    "        ax[0].legend()\n",
    "\n",
    "        ax[1].plot(frequencia, power, color='darkgreen')\n",
    "        ax[1].set_xlabel(\"Frequência\")\n",
    "        ax[1].set_ylabel(r'$P(\\omega)$')\n",
    "        plt.show()\n",
    "    return best_freq, best_power, fap, power, modelo\n",
    "\n",
    "\n",
    "# Ler arquivo CoRoT do campo sismo\n",
    "def read_an2_corot(file):\n",
    "    hdul = fits.open(file)\n",
    "    data = hdul[1].data\n",
    "    status = data['RAWSTATUS']  # status 0 == valid measurements\n",
    "    time  = data['DATETT']\n",
    "    flux  = data['RAWFLUX']\n",
    "    fluxerror = data['RAWFLUXDEV']\n",
    "    mask = (status == 0)\n",
    "\n",
    "    time = np.ascontiguousarray(time[mask], dtype=np.float64)\n",
    "    flux = np.ascontiguousarray(flux[mask], dtype=np.float64)\n",
    "    fluxerror = np.ascontiguousarray(fluxerror[mask], dtype=np.float64)\n",
    "                                \n",
    "    return time, flux, fluxerror\n",
    "\n",
    "\n",
    "# Ler arquivo CoRoT do campo exo\n",
    "def read_en2_corot(file):\n",
    "    hdul = fits.open(file)\n",
    "    data = hdul[1].data\n",
    "    status = data['STATUS']  # status 0 == valid measurements\n",
    "    time  = data['DATEBARTT']\n",
    "    flux  = data['WHITEFLUX']\n",
    "    mask = (status == 0)\n",
    "\n",
    "    time = np.ascontiguousarray(time[mask], dtype=np.float64)\n",
    "    flux = np.ascontiguousarray(flux[mask], dtype=np.float64)\n",
    "    return time, flux\n",
    "\n",
    "\n",
    "# retorna modelo e residuos (respectivamente)\n",
    "def box_model(pars, x):\n",
    "    v = pars.valuesdict()\n",
    "    yfit = np.zeros(len(x)) + v['baseflux']\n",
    "    indice = (x > v['t1']) & (x < v['t2'])\n",
    "    yfit[indice] = v['depth']    \n",
    "    return yfit\n",
    "\n",
    "def boxfit(pars, x, y):\n",
    "    yfit = box_model(pars,x)\n",
    "    return (y - yfit)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. O conceito de dobra em fase para busca de períodos\n",
    "<a id=\"dobra_fase\"></a>\n",
    "Quando buscamos períodos em séries temporais não sabemos de antemão, exceto quando temos informações a priori do sistema estudado, em qual região ou intervalo de períodos ele será mais provável de ser encontrado. Dessa forma, todo o espaço de parâmetros precisa ser buscado. Uma maneira conveniente de se fazer isso é dobrar a série temporal em uma fase de período teste e utlizando ajuste de funções (por exemplo, Lomb-Scargle, BLS...) ou estudo de de dispersão (por exemplo, PDM) pode-se estimar a qualidade do ajuste com respeito a algum tipo de referência.\n",
    "\n",
    "Neste bloco é mostrado visualmente o conceito. A seguir iremos utiliza estas técnicas para detectar o período nesta série de exemplo, uma curva de luz obtida com medidas do satélite CoRoT."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# gera saída do bokeh no notebook\n",
    "output_notebook()\n",
    "\n",
    "arquivo = 'CoRoT_id_0102609164.txt'\n",
    "x, y, err = read_txt(arquivo)\n",
    "\n",
    "c = int(len(x)/3)\n",
    "x, y, err = x[0:c], y[0:c], err[0:c]\n",
    "\n",
    "ph = np.zeros(len(x))\n",
    "ph_dummy = np.zeros(len(x))\n",
    "\n",
    "# gera objeto que contém a interatividade\n",
    "source = ColumnDataSource(data=dict(x=x, y=y, ph=ph))\n",
    "\n",
    "# gera os gráficos\n",
    "p0 = figure(plot_width=900, plot_height=300, x_range=(x.min(),x.max()))\n",
    "p0.circle('x', 'y', source=source, size=2, alpha=0.5, color='red')\n",
    "p0.xaxis.axis_label = \"Data Juliana (dias)\"\n",
    "p0.yaxis.axis_label = \"Fluxo Arbitrário\"\n",
    "\n",
    "p1 = figure(plot_width=900, plot_height=300, x_range=(0,1))\n",
    "p1.circle('ph', 'y', source=source, size=8, alpha=0.3)\n",
    "p1.xaxis.axis_label = \"Fase\"\n",
    "p1.yaxis.axis_label = \"Fluxo Arbitrário\"\n",
    "\n",
    "\n",
    "# função que dá origem à interação (em javascript para funcionar em JupNtbk)\n",
    "callback = CustomJS(args=dict(source=source), code=\"\"\"\n",
    "        var data = source.data;\n",
    "        var f = cb_obj.value\n",
    "        var x = data['x']\n",
    "        var y = data['y']\n",
    "        var ph = data['ph']\n",
    "        for (var i = 0; i < x.length; i++){\n",
    "            ph[i] = (x[i]/f) % 1\n",
    "        }\n",
    "        source.change.emit();\n",
    "    \"\"\")\n",
    "# cria botão de interação e respectivos limites, valor inicial e passo\n",
    "slider = Slider(start=0.5, end=1.5, value=0, step=0.0001, title=\"Período teste\")\n",
    "slider.js_on_change('value', callback)\n",
    "\n",
    "# mostra os resultados\n",
    "layout = column(slider, p0, p1)\n",
    "show(layout)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. O periodograma de Lomb-Scargle\n",
    "<a id=\"lombscargle\"></a>\n",
    "\n",
    "O periodograma de Lomb-Scargle é uma técnica que testa frequências (ou períodos) a partir de ajustes de mínimos quadrados. A \"potência\" do espectro é calculada a partir da relação:\n",
    "$$\n",
    "P_{x}(\\omega)=\\frac{1}{2}\\left(\\frac{\\left[\\sum_{j} X_{j} \\cos \\omega\\left(t_{j}-\\tau\\right)\\right]^{2}}{\\sum_{j} \\cos ^{2} \\omega\\left(t_{j}-\\tau\\right)}+\\frac{\\left[\\sum_{j} X_{j} \\sin \\omega\\left(t_{j}-\\tau\\right)\\right]^{2}}{\\sum_{j} \\sin ^{2} \\omega\\left(t_{j}-\\tau\\right)}\\right)\n",
    "$$\n",
    "onde, $\\tau$ é dado por $\\tan 2 \\omega \\tau=\\frac{\\sum_{j} \\sin 2 \\omega t_{j}}{\\sum_{j} \\cos 2 \\omega t_{j}}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "arquivo = 'GJ876_rv.txt'\n",
    "\n",
    "x, y, err = read_txt(arquivo)\n",
    "# remove o valor médio\n",
    "# y -= np.mean(y)\n",
    "\n",
    "# cria array de frequencias a ser estudadas\n",
    "frequencia = np.arange(0.001,12,0.0001)\n",
    "\n",
    "# objeto contendo as propriedades do calculo\n",
    "ls = LombScargle(x, y, err)\n",
    "\n",
    "# calcula o o periodograma\n",
    "power = ls.power(frequencia)\n",
    "\n",
    "# localiza a o índice do maior pico\n",
    "idx = np.argmax(power)\n",
    "\n",
    "# melhor frequência e pico correspondente\n",
    "best_freq = frequencia[idx]\n",
    "best_power = power[idx]\n",
    "modelo = ls.model(x,best_freq)\n",
    "\n",
    "# calcula a probabilidade de falso alarme\n",
    "fap = ls.false_alarm_probability(best_power, method='baluev', \\\n",
    "                                 minimum_frequency=frequencia.min(), \\\n",
    "                                 maximum_frequency=frequencia.max())\n",
    "\n",
    "print('Frequência: {0:.8f} (Período: {1:.8f})\\nFAP: {2:6.2f}%'.format(best_freq,1/best_freq,fap*100))\n",
    "\n",
    "# plota resultados\n",
    "phase = (x*best_freq) % 1\n",
    "fig, ax = plt.subplots(1,2, figsize=(14,6))\n",
    "ax[0].errorbar(phase, y, err, fmt='.', color='blue', label='Medidas')\n",
    "srt = np.argsort(phase) # obtem o indice da fase ordenada\n",
    "ax[0].plot(phase[srt], modelo[srt], '-', color='red', label='Modelo')\n",
    "ax[0].set_xlabel(\"Fase\")\n",
    "ax[0].set_ylabel(\"Fluxo\")\n",
    "ax[0].legend()\n",
    "\n",
    "ax[1].plot(frequencia, power, color='darkgreen')\n",
    "ax[1].set_xlabel(\"Frequência\")\n",
    "ax[1].set_ylabel(r'$P(\\omega)$')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Refinando frequências e analisando resíduos"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Z\n",
    "frequencia = np.arange(0.001,0.5,0.000001)\n",
    "best_freq, best_power, fap, power, modelo = scargle(x, y, err, frequencia, plot=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Removendo o pico principal\n",
    "frequencia = np.arange(0.01,0.05,0.000001)\n",
    "y2 = y - modelo\n",
    "best_freq2, best_power2, fap2, power2, modelo2 = scargle(x, y2, err, frequencia, plot=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### fazendo um laço\n",
    "### para remover frequências iterativamente\n",
    "nf = 5   # número de frequências\n",
    "frequencias = np.arange(0.001,12,0.0001)\n",
    "residuos = np.zeros(len(x))\n",
    "yfit = np.zeros(len(x))\n",
    "y2 = y\n",
    "best_freq, best_power, fap = np.zeros(int(nf)), np.zeros(int(nf)), np.zeros(int(nf))\n",
    "\n",
    "for i in range(nf):\n",
    "    print('=== Análise da {0}ª frequência ==='.format(i+1))\n",
    "    best_freq[i], best_power[i], fap[i], power, modelo = scargle(x, y2, err, frequencia, plot=False)\n",
    "    yfit += modelo\n",
    "    residuos = y2-modelo\n",
    "    wrms = np.sqrt(np.sum(err*(y2 - modelo)**2))\n",
    "    print('WRMS: {0:.4f}\\n\\n\\n'.format(wrms))\n",
    "    y2 = residuos\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. PDM - Phase Dispersion Minimization\n",
    "<a id=\"#pdm\"></a>\n",
    "\n",
    "O PDM é uma técnica que testa qual a dispersão total que uma determinada curva de fase possui quando é dobrada em um determinado período teste. \n",
    "\n",
    "É relativamente simples verificar que quando encontramos o período correto de uma determinada variabilidade periódica haverá o ‘encaixe’ adequado ao longo da fase do período. Isso significa que se dividirmos este espaço de fase em pequenas partes, que normalmente chamamos de **bins**, e na ausência de fatores sistemáticos, se medirmos o desvio padrão dentro de cada bin, ele será o menor possível quando o período testado for o período correto da variabilidade. Podemos então produzir um diagrama da soma dos resíduos de todos os bins como uma função do período testado e então verificar qual o período teste que possui a menor dispersão.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "arquivo = 'CoRoT_id_0102609164.txt'\n",
    "x, y, err = read_txt(arquivo)\n",
    "n_bins = 10\n",
    "\n",
    "P = pyPDM.PyPDM(x, y)\n",
    "scan = pyPDM.Scanner(minVal=0.01, maxVal=12, dVal=0.01, mode=\"frequency\")\n",
    "\n",
    "# calcula o pdm\n",
    "frequencia, theta = P.pdmEquiBinCover(n_bins, 5, scan)\n",
    "\n",
    "# localiza a o índice do maior pico\n",
    "idx = np.argmin(theta)\n",
    "\n",
    "# melhor frequência e pico correspondente\n",
    "best_freq = frequencia[idx]\n",
    "best_power = power[idx]\n",
    "\n",
    "# calcular modelo utilizando o filtro de Savitzky-Golay\n",
    "phase = x % (1/best_freq)\n",
    "idx_tmp = np.argsort(phase) # ordena em fase\n",
    "# Modelo baseado no filtro de Savitzky-Golay\n",
    "result = savgol_filter(y[idx_tmp], 35, 4, mode='nearest')\n",
    "modelo = result[np.argsort(idx_tmp)]\n",
    "wrms = np.sqrt(np.sum(err*(y - modelo)**2))    \n",
    "print('Frequência: {0:.8f} (Período: {1:.8f})'.format(best_freq,1/best_freq))\n",
    "print('WRMS: {0:.4f}'.format(wrms))\n",
    "\n",
    "# plota resultados\n",
    "fig, ax = plt.subplots(1,2, figsize=(14,6))\n",
    "ax[0].errorbar(phase, y, err, fmt='.', color='blue', label='Medidas')\n",
    "srt = np.argsort(phase) # obtem o indice da fase ordenada\n",
    "ax[0].plot(phase[srt], modelo[srt], '-', color='red', label='Modelo', zorder=20)\n",
    "ax[0].set_xlabel(\"Fase\")\n",
    "ax[0].set_ylabel(\"Fluxo\")\n",
    "ax[0].legend()\n",
    "\n",
    "ax[1].plot(frequencia, theta, color='darkgreen')\n",
    "ax[1].set_xlabel(\"Frequência\")\n",
    "ax[1].set_ylabel(r'$P(\\Theta)$')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Ajustando um modelo utilizando a biblioteca LMFIT\n",
    "<a id=\"lmfit\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "arquivo = 'EN2_STAR_CHR_0102890318_20070206T133547_20070402T070302.fits'\n",
    "time, flux = read_en2_corot(arquivo)\n",
    "\n",
    "periodo = 1.5089682\n",
    "# Remover trends com um filtro Savitzky-Golay\n",
    "modelo = savgol_filter(flux, 9501, 2, mode='nearest')\n",
    "\n",
    "# Plota dados\n",
    "plt.title(arquivo)\n",
    "plt.plot(time, flux, '.')\n",
    "plt.plot(time, modelo, '.')\n",
    "plt.xlabel('Data Juliana Modificada')\n",
    "plt.ylabel('Fluxo CoRoT')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Normaliza medidas\n",
    "flux_norm = flux/modelo\n",
    "\n",
    "phase = time % periodo\n",
    "idx = np.argsort(phase)\n",
    "x2 = phase[idx]\n",
    "y2 = flux_norm[idx]\n",
    "\n",
    "\n",
    "# cria dicionário de parâmetros a ser ajustados\n",
    "p = Parameters()\n",
    "p.add_many(('baseflux', 1, True, 0.9, 1.1), \\\n",
    "           ('t1', 1, True, 0, periodo), \\\n",
    "           ('t2', 1.2, True, 0, periodo), \\\n",
    "           ('depth', 0.9, True, 0.8, 1.0))\n",
    "\n",
    "init =  Minimizer(boxfit, p, nan_policy='omit', fcn_args=(x2,y2))\n",
    "result = init.minimize(method='differential_evolution')\n",
    "print('Relatório do ajuste:')\n",
    "print(report_fit(result))\n",
    "\n",
    "yfit = box_model(result.params,x2)\n",
    "plt.plot( x2, y2, '.')\n",
    "plt.plot( x2, yfit, 'r-')\n",
    "plt.xlabel('Fase do Período')\n",
    "plt.ylabel('Fluxo Normalizado')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Visualizando a Transformada de Fourier (EPO application)\n",
    "<a id=\"dft_visualization\"></a>\n",
    "(Baseado em: [But what is the Fourier Transform? A visual introduction.](https://youtu.be/spUNpyF58BY))\n",
    "\n",
    "**Transformada de Fourier** (equação de análise):\n",
    "\n",
    "$$ X(\\omega)=\\int_{-\\infty}^{+\\infty} x(t) e^{-j \\omega t} d t $$\n",
    "\n",
    "onde, a partir da fórmula de Euler, $e^{-j \\omega t} = \\cos{\\omega t} -j\\sin{\\omega t}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from bokeh.layouts import column, row\n",
    "from bokeh.models import CustomJS, ColumnDataSource, Slider\n",
    "from bokeh.plotting import figure, output_notebook, show\n",
    "\n",
    "output_notebook()\n",
    "\n",
    "# definir parametros do sinal\n",
    "ruido = False\n",
    "subtrair_media = True\n",
    "\n",
    "amp = [1.0, 2.0, 3.6]\n",
    "f = [2.0, 5.0, 3.5]\n",
    "phase = [0.1, 0.27, 0.7]\n",
    "h = 10.0    # valor médio adicionado ao sinal\n",
    "\n",
    "sampling = 0.01    # amostragem do sinal em seg/ciclo\n",
    "\n",
    "fini = 0.0   # frequência inicial a ser buscada\n",
    "fend = 30    # frequência final a ser buscada\n",
    "fstep = 0.01 # passo em frequência\n",
    "\n",
    "nyqst = 1/(2*sampling)\n",
    "print('Limite de Nyqst: {0:5.1f} ciclos/segundo'.format(nyqst))\n",
    "# construir variáveis\n",
    "nf = len(f)\n",
    "x = np.arange(0,5,sampling)  # variavel independente x\n",
    "ys = np.zeros((nf,len(x)))   # array para conter sinais individuais\n",
    "y = np.zeros(len(x))         # array contendo sinais somados\n",
    "for i in range(nf):\n",
    "    ys[i,:] = amp[i]*np.sin(2*np.pi*(f[i]*x + phase[i]))\n",
    "    y += ys[i,:]\n",
    "y += h                       # adiciona valor médio do sinal\n",
    "\n",
    "\n",
    "# definir funcoes usadas dentro do js (devem ser declaradas antes do callback)\n",
    "im, re = np.zeros(len(x)),np.zeros(len(x))\n",
    "mgi =  np.linspace(0,1,len(x))\n",
    "mgr =  np.linspace(0,1,len(x))\n",
    "magf = np.linspace(0,1,len(x))\n",
    "fmag = np.zeros(len(x))\n",
    "yfit = np.zeros(len(x))\n",
    "\n",
    "freq = np.arange(fini,fend,fstep)\n",
    "mag = np.zeros(len(freq))\n",
    "theta = np.zeros(len(freq))\n",
    "\n",
    "noise = np.random.normal(0,2,len(x))  # ruido gaussiano com média 0 e desvio padrao 2\n",
    "if ruido:\n",
    "    y += noise\n",
    "\n",
    "if subtrair_media:\n",
    "    y -= np.mean(y)\n",
    "\n",
    "    \n",
    "\n",
    "    \n",
    "# calculo do objeto interativo\n",
    "source = ColumnDataSource(data=dict(x=x, y=y, im=im, re=re, mgi=mgi, \\\n",
    "                                    mgr=mgr, magf=magf, fmag=fmag, yfit=yfit))\n",
    "# plota sinais separados\n",
    "s0 = figure(plot_width=600, plot_height=165, x_range=(x.min(),x.max()), title='Sinal Arbitrário')\n",
    "\n",
    "ran = np.zeros(len(x))\n",
    "for i in range(nf):\n",
    "    ran += ys[i,:].max()-ys[i,:].min()\n",
    "    s0.line(x, ys[i,:]+ran+3)\n",
    "s0.axis.visible=False\n",
    "# plota sinais somados\n",
    "s3 = figure(plot_width=600, plot_height=165, x_range=(x.min(),x.max()))\n",
    "s3.line(x, y)\n",
    "s3.line('x', 'yfit', source=source, line_alpha=0.8, color='red')\n",
    "s3.yaxis.axis_label = \"Intensidade\"\n",
    "s3.xaxis.axis_label = \"Tempo (segundos)\"\n",
    "\n",
    "\n",
    "rg = 1.25*y.max()\n",
    "\n",
    "# calcula dft\n",
    "for j,i in enumerate(freq):\n",
    "    r1 = y*np.cos(-2*np.pi*x*i)/len(x)\n",
    "    im1 = y*np.sin(-2*np.pi*x*i)/len(x)  \n",
    "    r1 = np.sum(r1)\n",
    "    im1 = np.sum(im1)\n",
    "    mag[j] = 2*np.sqrt(r1**2 + im1**2)\n",
    "\n",
    "# plota plano\n",
    "s1 = figure(plot_width=300, plot_height=270,x_range=(-rg,rg), y_range=(-rg,rg))\n",
    "s1.line(x=[0,0], y=[-rg,rg], color='black')\n",
    "s1.line(x=[-rg,rg], y=[0,0], color='black')\n",
    "s1.line('re', 'im', source=source, color='blue', line_alpha=0.2)\n",
    "s1.circle('re', 'im', source=source, color='blue', alpha=0.2)\n",
    "s1.circle('mgr','mgi', source=source, color='red', size=5)\n",
    "s1.xaxis.axis_label = \"Real\"\n",
    "s1.yaxis.axis_label = \"Imaginário\"\n",
    "\n",
    "\n",
    "nf = np.array([nyqst,nyqst])\n",
    "np = np.array([0,1.2*mag.max()])\n",
    "s2 = figure(plot_width=900, plot_height=250,x_range=(fini,fend), y_range=(0,1.05*mag.max()))\n",
    "s2.line(freq, mag, color='green', line_width=1.5)\n",
    "s2.circle('fmag','magf', source=source, color='red', size=5)\n",
    "s2.line(x=nf, y=np, color='blue', legend='Limite de Nyqst')\n",
    "s2.legend.location = \"top_right\"\n",
    "s2.yaxis.axis_label = \"Amplitude\"\n",
    "s2.xaxis.axis_label = \"Frequência (ciclos/segundo)\"\n",
    "\n",
    "\n",
    "callback = CustomJS(args=dict(source=source), code=\"\"\"\n",
    "        const arrSum = arr => arr.reduce((a,b) => a + b, 0);\n",
    "        var data = source.data;\n",
    "        var f = cb_obj.value;\n",
    "        var x = data['x'];\n",
    "        var y = data['y'];\n",
    "        var im = data['im'];\n",
    "        var re = data['re'];\n",
    "        var mgi = data['mgi'];\n",
    "        var mgr = data['mgr'];\n",
    "        var magf = data['magf']\n",
    "        var fmag = data['fmag']\n",
    "        var yfit = data['yfit']\n",
    "        for (var i = 0; i < x.length; i++){\n",
    "            re[i] = y[i]*Math.cos(-2*Math.PI*x[i]*f)\n",
    "            im[i] = y[i]*Math.sin(-2*Math.PI*x[i]*f)\n",
    "        }\n",
    "        let ret = arrSum(re)/x.length;\n",
    "        let imt = arrSum(im)/x.length;\n",
    "        let mg = Math.sqrt(ret*ret + imt*imt);\n",
    "\n",
    "        var tfix = 0;\n",
    "        if(imt>0 && ret<0){\n",
    "            tfix = Math.PI;\n",
    "        } else if(imt<0 && ret<0){\n",
    "            tfix = 3*Math.PI/2;\n",
    "        } else if(imt<0 && ret>0){\n",
    "            tfix = 2*Math.PI;\n",
    "        }\n",
    "        let theta = Math.atan(imt/ret) + tfix\n",
    "\n",
    "\n",
    "        for (var i = 0; i < x.length; i++){\n",
    "            mgi[i] = imt;\n",
    "            mgr[i] = ret;\n",
    "        //    magf[i] = mg*2*(i/x.length)\n",
    "            magf[i] = mg*2\n",
    "            fmag[i] = f;\n",
    "        }\n",
    "\n",
    "        for (var i = 0; i < x.length; i++){\n",
    "            yfit[i] = 2*mg*Math.cos(-2*Math.PI*((f*x[i]) + theta/(2*Math.PI)))\n",
    "        }\n",
    "        source.change.emit();\n",
    "    \"\"\")\n",
    "\n",
    "slider = Slider(start=fini, end=fend, value=0, step=fstep, title=\"Ciclos por segundo\")\n",
    "slider.js_on_change('value', callback)\n",
    "\n",
    "layout = column(row(column(s0,s3), column(slider, s1)), s2, width=900)\n",
    "show(layout)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}