{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0ff7abb1-c698-480b-a894-bdbb5d151fe6",
   "metadata": {},
   "source": [
    "## Assignment 5\n",
    "In this assignment, we look at spectral estimation and study the effects of acausal filters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "156c8712-11be-4721-9d39-ff2c47a66bcb",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from setupFigure import SetupFigure\n",
    "from dftFast import dft_fast_coeff, dft_fast_synthesis\n",
    "from seismicTrace import SeismicTrace"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f98951f5-96dc-4ef5-b365-85bffa21f700",
   "metadata": {},
   "source": [
    "### Task 1: Functions for Hann window and filter transfer functions\n",
    "Write a function that calculates samples of a Hann window for given sampling interval and window length."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c36eb15e-26e7-4e6e-bef9-8d566dc07411",
   "metadata": {},
   "outputs": [],
   "source": [
    "def hann_window(dt, tlen):\n",
    "    \"\"\"\n",
    "    Calculate samples of the Hann window for given DT and TLEN.\n",
    "    Number of samples is assumed to be int(tlen/dt)\n",
    "    :param dt: sampling interval\n",
    "    :param tlen: length of window\n",
    "    \"\"\"\n",
    "    ns = int(tlen/dt)\n",
    "    return 2*np.sin(np.pi*dt*np.arange(0, ns)/tlen)**2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60dffc7b-17a1-40aa-8a6f-918504e06af2",
   "metadata": {},
   "source": [
    "Write a function that calculates samples of the boxcar lowpass filter transfer function for positive frequencies only."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be77830b-6b2f-4371-9711-e83a1598c363",
   "metadata": {},
   "outputs": [],
   "source": [
    "def boxcar_lowpass_filter(fc, df, nf):\n",
    "    \"\"\"\n",
    "    Calculates samples of the boxcar lowpass filter transfer function (positive frequencies only)\n",
    "    :param fc: corner frequency in Hz\n",
    "    :param df: frequency stepping\n",
    "    :param nf: number of frequencies\n",
    "    \"\"\"\n",
    "    return np.where(df*np.arange(0, nf) > fc, 0, 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33eb7b1d-4abd-42a2-99ec-cdd1765fbfde",
   "metadata": {},
   "source": [
    "Write a function that calculates samples of the boxcar highpass filter transfer function for positive frequencies only."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29c56269-025a-43e9-8748-c0514a72f70d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def boxcar_highpass_filter(fc, df, nf):\n",
    "    \"\"\"\n",
    "    Calculates samples of the boxcar highpass filter transfer function (positive frequencies only)\n",
    "    :param fc: corner frequency in Hz\n",
    "    :param df: frequency stepping\n",
    "    :param nf: number of frequencies\n",
    "    \"\"\"\n",
    "    return np.where(df*np.arange(0, nf) < fc, 0, 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "937b6b14-4d8f-40ed-bc71-86888d6e726c",
   "metadata": {},
   "source": [
    "Write a function that calculates samples of the Hann filter transfer function for positive frequencies only."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0101ff10-d14d-4b96-83f8-b26f96e5ec8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def hann_lowpass_filter(fc, df, nf):\n",
    "    \"\"\"\n",
    "    Calculates samples of the Hann filter transfer function (positive frequencies only)\n",
    "    :param fc: corner frequency in Hz\n",
    "    :param df: frequency stepping\n",
    "    :param nf: number of frequencies\n",
    "    \"\"\"\n",
    "    f = df*np.arange(0, nf)\n",
    "    return np.where(f < fc, np.cos(np.pi*f/fc)**2, 0)    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d641182d-5526-45a9-a69a-0a6965fd942b",
   "metadata": {},
   "source": [
    "### Task 2: Spectral estimation and windowing\n",
    "Here, we first setup a time series that consists of a superposition of different exponentially damped sine functions simulating a series of free oscillations:\n",
    "\\begin{align}\n",
    "s(t) = \\sum_{k=0}^N A_k\\sin(2\\pi f_k t)\\,\\exp(-\\pi f_k t/Q_k) \\,.\n",
    "\\end{align}\n",
    "The aim is to find the frequencies by Fourier analysis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eba91f73-1651-4fd7-ad90-70980b1c5d4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "feig = np.array([0.3, 0.48, 0.64, 0.68, 0.82, 0.85, 0.93, 0.94])*1.e-3         # eigenfreqeuncies in Hz\n",
    "amp =  np.array([1.0, 0.30, 2.00, 0.61, 4.00, 0.53, 3.00, 0.24])               # amplitudes\n",
    "phi =  np.array([ 11,   46,  271,  123,  337,  173,   65,  221])*np.pi/180.    # phases\n",
    "q =    np.array([600,  200, 2000, 1000,  600,  100,  800,  900])               # quality factors"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e455172a-50ad-4aa4-a1b3-d29d50a52b49",
   "metadata": {},
   "source": [
    "Write code that impements the above equation using the provided frequencies, amplitudes, phases and Q values. Choose dt=100 s and allow a varaible time series length of tens of hours. Plot the resulting time series."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c808e6c4-8e39-4e7c-a689-860230ff775a",
   "metadata": {},
   "outputs": [],
   "source": [
    "nhour = 100\n",
    "dt = 100.0\n",
    "tlen = 3600*nhour\n",
    "ns = int(tlen/dt)\n",
    "print(\"Number of samples: \", ns)\n",
    "t = dt*np.arange(0, ns)\n",
    "seis = np.zeros(ns)\n",
    "for j, f in enumerate(feig):\n",
    "    seis = seis + amp[j]*np.sin(2*np.pi*f*t+phi[j])*np.exp(-np.pi*f*t/q[j])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d0dbe36-34e8-486d-bc63-f5be4aff2d0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig1, ax1 = SetupFigure(12, 5, \"Time [h|\", \"s(t)\", \"Ground motion\", 14)\n",
    "ax1.plot(t/3600, seis, color='blue', ls='-')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ae65e6e-0f73-42fc-b251-878f0ef7f36c",
   "metadata": {},
   "source": [
    "Produce a second time series by multiplying the original one with a Hann window and also plot it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b62c1a43-f57e-4397-a8f9-7ec813db1e9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "win = hann_window(dt, tlen)\n",
    "seisw = seis*win"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41193b3d-110b-42a4-b2ee-9ee0dbd759af",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig2, ax2 = SetupFigure(12, 5, \"Time [h|\", \"swin(t)\", \"Windowed ground motion\", 14)\n",
    "ax2.plot(t/3600, seisw, color='blue', ls='-')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26feab1b-41a6-464f-b578-95814286b915",
   "metadata": {},
   "source": [
    "Perform a fast Fourier analysis, calculate the associated frequencies and plot the resulting amplitude spectra of both the original and the windowed time series. Compare the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "164cd4c2-72ea-44e8-a236-95475118aad1",
   "metadata": {},
   "outputs": [],
   "source": [
    "spec = dft_fast_coeff(seis)\n",
    "f = 1./tlen*np.arange(0, spec.size)\n",
    "print(\"Number of frequencies: \", spec.size, \" fmax = \", f[-1]*1000, \"mHz\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c88a619-8b9c-44ed-bc5f-f6ecce938af3",
   "metadata": {},
   "outputs": [],
   "source": [
    "specw = dft_fast_coeff(seisw)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73b8fae2-e712-47d7-99a5-ebe5375084cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig3, ax3 = SetupFigure(12, 5, \"Frequency [mHz]\", \"S(f)\", \"Amplitude spectrum of ground motion\", 14)\n",
    "ax3.set_xlim(0,1.2)\n",
    "ax3.plot(f*1000, np.absolute(spec), color='blue', ls='-', label='boxcar')\n",
    "ax3.plot(f*1000, np.absolute(specw), color='red', ls='-', label='hann')\n",
    "ax3.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "689b3af4-1146-4041-9f4e-baed098d2a23",
   "metadata": {},
   "source": [
    "### Task 2: Read a seismogram from file and compute spectrum\n",
    "Here, we first read a recorded seismogram form a file, do a Fourier analysis and later apply diverse acausal filters. We make use of the class \"SeismicTrace\" that defines a data structure for a seismic trace and also some useful functions related to traces. You may want to look at the provided code in \"seismicTrace.py\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "912ae1f7-8cbc-4446-9ef0-b446b57d789d",
   "metadata": {},
   "outputs": [],
   "source": [
    "seisfile = \"CH.PANIX.HHZ\"\n",
    "with open(seisfile, 'r') as fd:\n",
    "    seistrace = SeismicTrace.readFromAscii(fd)\n",
    "t = seistrace.tanf+seistrace.dt*np.arange(0, seistrace.nsamp)\n",
    "seistrace.printInfo()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eaee9262-f577-49a1-acee-9b9c27f996be",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig4, ax4 = SetupFigure(12, 5, \"Time [s]\", \"s(t)\", \"Unfiltered seismic data\", 14)\n",
    "ax4.plot(t, seistrace.vals, color='blue', ls='-')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "934a4fa1-7f8d-4977-bf86-bfe094abc5e5",
   "metadata": {},
   "source": [
    "Perform a fast Fourier analysis, calculate the associated frequencies and plot the resulting amplitude spectrum."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4a0c8c9-b51a-4116-9988-73d75b2f6704",
   "metadata": {},
   "outputs": [],
   "source": [
    "spec = dft_fast_coeff(seistrace.vals)\n",
    "f = 1./seistrace.tlen*np.arange(0, spec.size)\n",
    "print(\"Number of frequencies: \", spec.size, \" fmax = \", f[-1], \"Hz\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82fb96d5-39ac-4c6a-a360-3544167ec26b",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig5, ax5 = SetupFigure(12, 5, \"Frequency [Hz]\", \"S(f)\", \"Amplitude spectrum of seismogram\", 14)\n",
    "ax5.set_xlim(0, 1)\n",
    "ax5.plot(f, np.absolute(spec), color='blue', ls='-', label='boxcar')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d36d0c3-6faf-4e0b-b934-8822100b2b2f",
   "metadata": {},
   "source": [
    "### Task 3: Acausal low pass filtering"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36b37ff1-3080-4651-a344-9cb4ede9109d",
   "metadata": {},
   "source": [
    "Apply the boxcar lowpass filter to the data and plot the result."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1e59a76-d251-4d66-b6ba-7c29a0626b3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "fc = 0.1\n",
    "df = 1./seistrace.tlen\n",
    "specfil = spec*boxcar_lowpass_filter(fc, df, spec.size)\n",
    "seisfil = dft_fast_synthesis(specfil)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "532ad374-2728-435d-8dca-afd7fcf9c475",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig6, ax6 = SetupFigure(12, 5, \"Time [s]\", \"sf(t)\", \"Boxcar lowpass filtered seismic data, FC = {:5.2f} Hz\".format(fc), 14)\n",
    "#ax6.set_xlim(64000, 64500)\n",
    "ax6.plot(t, seisfil, color='blue', ls='-')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1097e6d-c459-489d-aecd-75911419a319",
   "metadata": {},
   "source": [
    "Apply the Hann low pass filter to the data and plot the result."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be2e1d6d-0c4e-4bfa-8612-ce626393bdff",
   "metadata": {},
   "outputs": [],
   "source": [
    "specfil2 = spec*hann_lowpass_filter(fc, df, spec.size)\n",
    "seisfil2 = dft_fast_synthesis(specfil2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5d4bf72-9bc2-4679-9178-4f7bc650e4ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig7, ax7 = SetupFigure(12, 5, \"Time [s]\", \"sf(t)\", \"Hann lowpass filtered seismic data, FC = {:5.2f}\".format(fc), 14)\n",
    "#ax7.set_xlim(64000, 64500)\n",
    "ax7.plot(t, seisfil2, color='blue', ls='-')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3bbec473-2bf7-41d0-828d-9413dff147e8",
   "metadata": {},
   "source": [
    "Apply the boxcar highpass filter to the data and plot the result."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9ec91f3-cc1c-48e5-ac53-882002fc74f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "specfil3 = spec*boxcar_highpass_filter(fc, df, spec.size)\n",
    "seisfil3 = dft_fast_synthesis(specfil3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a24950e-8438-4e3b-903e-75f2f4408105",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig8, ax8 = SetupFigure(12, 5, \"Time [s]\", \"sf(t)\", \"Boxcar highpass filtered seismic data, FC = {:5.2f}\".format(fc), 14)\n",
    "#ax7.set_xlim(64000, 64500)\n",
    "ax8.plot(t, seisfil3, color='blue', ls='-')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92e0ebc5-9f7d-4302-8ad1-9aed6c7f2d6a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}