{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Research example\n",
    "This notebook preserves a 2019 exploratory method by Andrew Dyck.\n",
    "Read README.md before running it. Full scientific validation is still needed.\n",
    "Outputs below have been cleared and do not represent a new execution.\n"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "from pathlib import Path\n",
    "Path(\"generated\").mkdir(exist_ok=True)\n"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# FRB Pipeline Simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib notebook\n",
    "import numpy as np\n",
    "from matplotlib.pyplot import *\n",
    "import gc\n",
    "from scipy import signal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def axis_labels(x, y, z):\n",
    "    '''\n",
    "    Creates the labels for the x, y axes and title as well as triggers the legend\n",
    "    \n",
    "    Note: Inputs must be strings\n",
    "    \n",
    "    \n",
    "    INPUTS:\n",
    "    \n",
    "    x: (string) The label for the x-axis\n",
    "    y: (string) The y-axis label\n",
    "    z: (string) The title for the plot\n",
    "    \n",
    "    \n",
    "    OUTPUTS:\n",
    "    \n",
    "    xlab   : The x-axis label\n",
    "    ylab   : The y-axis label\n",
    "    titles : The title\n",
    "    legends: The legend of the plot\n",
    "    \n",
    "    '''\n",
    "    \n",
    "    xlab = xlabel(x)\n",
    "    ylab = ylabel(y)\n",
    "    titles = title(z)\n",
    "    legends = legend\n",
    "    return xlab, ylab, titles, legends\n",
    "\n",
    "\n",
    "def peakfinder_pulsar_DM(x, threshold):\n",
    "    '''\n",
    "    Finds the peaks in a dedispersed pulsar that has been run through a matched filter with a Gaussian.\n",
    "    \n",
    "    Note: The array or list used for x must have two dimensions. The first specifying which DM we are looking at, and \n",
    "          the second representing the number of data points.\n",
    "          \n",
    "          The threshold must be appropriate for every DM. This means that it should have the same length as the \n",
    "          number of DMs\n",
    "    \n",
    "    \n",
    "    INPUTS:\n",
    "    \n",
    "    x         : (array or list) The power of the pulse for every DM that has been looked at.\n",
    "    threshold : (float) The y value at which peaks can be designated as peaks.\n",
    "    \n",
    "    \n",
    "    OUTPUTS:\n",
    "    \n",
    "    peaksy: (list) The values ,for every DM, that can be called peaks over the threshold. \n",
    "    \n",
    "    \n",
    "    '''\n",
    "    \n",
    "    peaksx = [] #This is the values that the peaks are centred around\n",
    "    peaksy = []\n",
    "    for j in range(np.shape(x)[0]):\n",
    "        peakssy= []\n",
    "        DM_intensity = x[j]\n",
    "        for i in range(len(DM_intensity)-2):\n",
    "            if DM_intensity[i] > threshold[j] and DM_intensity[i-2]<DM_intensity[i] and DM_intensity[i-1] <DM_intensity[i] and DM_intensity[i+1] < DM_intensity[i] and DM_intensity[i+2] < DM_intensity[i]:\n",
    "                peakssy.append(DM_intensity[i])\n",
    "        peaksy.append(peakssy)\n",
    "    return peaksy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following steps will simulate the anologue process that occurs before the data is downsampled. Any pulsar will work for this, however, the pulsar being simulated will be made using the paramters below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Define the parameters for the noise that will be mixed with the pulse\n",
    "Cf = 1400e6 #This is the centre frequency in MHz, and where we will be viewing from\n",
    "Bw = 10e6 #This is the bandwidth that we are observing with our telescope\n",
    "sample_rate = 12e9 # Frequency running simulation at.\n",
    "period = .005 #seconds "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With a period of 5ms, we will use a time frame with this length to simulate one pulse. The bins in this simulation is determined by the sampling rate, which can be changed above. The current value is at such that 12GB ram computerss will just be able to simulate. The higher value the sampling rate is given, the more accurate the simulation becomes, but at the expense of what the computer is able to process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Define the timescale in which we will be observing\n",
    "t = np.linspace(-.0025,.0025, int(period*sample_rate)) #Decreasing the period for this will decrease the samples obtained"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To create the pulsar, first noise is simulated using a normal distribution of values. This distribution is made such that the standard deviation is 1, and is centred around 0. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Create the Gaussian noise profile for the pulsar\n",
    "length = len(t)     #This defines the number of timesamples we have, ie the number of samples collected\n",
    "mu = 0              #This indicates an offset of 0, meaning the signal is perfetly centred about 0\n",
    "sigma = 1           #This idicates an std of 1\n",
    "gnoise = np.random.normal(mu, sigma, length) #This is the GAussian noise for the signal\n",
    "pw = .0004 #seconds\n",
    "print(length)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This noise is then multiplied to a gaussian that has a standard deviation equal to that of the desired pulse width.The gaussian was made using the equation \\begin{equation} e^{\\frac{-1}{2}(\\frac{t}{pw})^2} \\end{equation}."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gnoise = np.exp(.5*-(t/pw)**2)*gnoise #Define the pulse of the signal. The gaussian profile takes an input of the "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Observe now that the resulting signal looks like a Gausian distribution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Now plot the signal, and observe the pulse shape. To see the pulse without destroying the computer, \n",
    "#plot every 10000th point\n",
    "figure()\n",
    "plot(np.linspace(0,.005, int(period*sample_rate))[0::10000],gnoise[0::10000], ls='none', marker='.')\n",
    "axis_labels('time(s)', 'value(unitless)', 'Gaussian noise pulse')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, to create the pulsar, we take the Fourier transform of the Gaussian, so that we may eliminate all but one frequency of the Gaussian"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Now take the fourier transform of the pulse to get it in frequency space\n",
    "fourier_pulse = np.fft.fft(gnoise)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The x-axis for the resulting graph is caluclated here. This tells what the frequency step for the broadband noise of the Gaussian is. This is done by taking the fftfreq of the length of the Gaussian, and having it step by the sample rate in units of MHz."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#The fourier transformed data should look like broadband noise, so define the axes properly before plotting\n",
    "broad_freq = np.fft.fftfreq(length, 1/(sample_rate/1e6)) #This gives the signal in MHz, (specifically from the 1e6)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The plot of the Fourier transform of the Gaussian wave packet is a broadband noise plot as expected."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "figure()\n",
    "plot(broad_freq[::1000], np.abs(fourier_pulse[::1000])) #Square the abs value to get better view\n",
    "axis_labels(\"frequency(MHz)\", 'amplitude(unitless)', 'Fourier transformed  gaussian pulse')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, all frequencies outside of our bandwidth are turned to zero. This means that the only values with any values are between 1395 and 1405 MHz. This is true for both the negative frequencies and positive ones between this range."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Now we modify the noise, so that there are only +- of the same frequency in it of it. This is the pulsar!\n",
    "#Everything outside the two reference frequencies is turned to zero\n",
    "highpass_freq = (Cf-Bw/2.0) #This sets the centre frequency to 0, with a bandwidth deifned above\n",
    "highpass_index1 = int(highpass_freq/(sample_rate/2)*length/2) #This finds the point in the array with the value\n",
    "highpass_index2 = length - highpass_index1 #Same but for different point\n",
    "fourier_pulse[:highpass_index1] = 0 #This is turning everything outside the two points to zero\n",
    "fourier_pulse[highpass_index2:] = 0 #Same as above\n",
    "\n",
    "#Everything outside of the two frequencies is turned to 0\n",
    "nyquest_index = int(length/2) #This value is determined by the nyquist sampling theorem(helps figure out where the chinkc are)\n",
    "#The following is the same as above, but now for outside the frequency values\n",
    "lowpass_freq = Cf+Bw/2.0 \n",
    "lowpass_index1 = int(lowpass_freq/(sample_rate/2)*length/2)\n",
    "lowpass_index2 = length - lowpass_index1\n",
    "fourier_pulse[lowpass_index1:nyquest_index] = 0\n",
    "fourier_pulse[nyquest_index:lowpass_index2] = 0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below is the plot of the new frequencies, As stated before, the only frequencies with nonzero values are those within the bandwidth. This now defines our pulsar, and the frequencies that we will be observing at"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "figure()\n",
    "plot(broad_freq[::1000], np.abs(fourier_pulse[::1000]))\n",
    "axis_labels(\"frequency(MHz)\", 'amplitude(unitless)', 'Newly created pulsar')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This reduced frequency pulsar is now turned back into a timestream and observed to still have a Gaussian shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#Now turn the pulsar back into time space from frequency space\n",
    "blimited_pulse = np.fft.ifft(fourier_pulse)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#If this pulse is plotted again, it sill looks like the same as before\n",
    "figure()\n",
    "plot(np.linspace(0,.005, int(period*sample_rate))[0::10000], blimited_pulse[::10000])\n",
    "axis_labels('time(s)', 'value(unitless)', 'Pulsar wave packet')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To simulate what the computer recieves after the telescope detects the signal, we need to create a mixing signal which will be used to downsample the data. This mixing signal is created using the complex Euler equation:\n",
    "$e^{-2\\pi j *\\frac{Cf}{samp}*len}$, where Cf is defined as above, samp is the sample rate, and len is the array with size length, with step size of 1. It is also assumed that the pulse currently is centred around the same point, x=0, which eliminates the \\omega that is in the original equation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Now need to create the mixing signal:\n",
    "mixing_signal = np.exp(-2.00j*np.pi* Cf/sample_rate *np.arange(length)) #With this we create the mixing signal using the euler equation for the signal\n",
    "\n",
    "#We assume the pulse is always centred around the same point, so the only part of the signal equaiton needed is \n",
    "#omega*t. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The pulse is then multiplied to the mixing signal. The fourier transform of this new signal is then taken, and a lowpass filter is applied to it. This is to create an IQ data set for the bandwidth. The inverse fourier transform of this signal is then taken, and it is downsampled by taking only samples that have spcings of the $\\frac{sample\\_ rate}{bandwidth}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "#Mix and sample down the pusar signal so that we can have an I Q data set for the bandwidth\n",
    "mixedd_down_s = blimited_pulse*mixing_signal\n",
    "N_cutoff = int(Bw/2/sample_rate*length) #This is from the lowpass fileter definition\n",
    "fmixed = np.fft.fft(mixedd_down_s) #This creates the fourier transform of the mixed signal\n",
    "fmixed[N_cutoff:-N_cutoff] = 0.0 # this makes it so that the imaginary magnitudes are comparable to the reals. This is\n",
    "                                    #also our filter\n",
    "filtered_mixed_downs_s = np.fft.ifft(fmixed) #This brings the fourier space pulse back into a packet\n",
    "#Downsample so that is...\n",
    "downsampled_filtered_mixed_down_s = filtered_mixed_downs_s[::int(sample_rate/(Bw))]\n",
    "#complex sampled at 10MHz instead of 12GHz\n",
    "#This gives us 50000 indicies instead of 60000000"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To open up more RAM in the computer, variables that will no longer be used are forgotten"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#%reset_selective -f 't'\n",
    "%reset_selective -f gnoise \n",
    "%reset_selective -f fourier_pulse\n",
    "%reset_selective -f broad_freq\n",
    "%reset_selective -f blimited_pulse\n",
    "%reset_selective -f mixing_signal\n",
    "%reset_selective -f mixedd_down_s\n",
    "%reset_selective -f N_cutoff\n",
    "%reset_selective -f fmixed\n",
    "%reset_selective -f filtered_mixed_downs_s\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The pulse is now downsampled, and realistically, this is where the simulation can begin. The previous lines of code were used to demonstrate how the anologue to digital converter chain works. Now, the new downsampled signal will be plotted below. Again, this pulsar signal still looks like a wave packet. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "#The simulation can be run from this point\n",
    "\n",
    "\n",
    "# Use gaussian noise, filter with same time width.\n",
    "#Now plot the signal, and observe that the pulse shape is the same, but now using smaller increments between samples\n",
    "figure()\n",
    "plot(np.linspace(0,.005, len(downsampled_filtered_mixed_down_s))[0:], downsampled_filtered_mixed_down_s[:].real)\n",
    "plot(np.linspace(0,.005, len(downsampled_filtered_mixed_down_s))[0:], downsampled_filtered_mixed_down_s[:].imag)\n",
    "axis_labels('Time(s)', 'Amplitude', 'Downsampled Pulsar')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Right now, there is a single pulse, so this is more of an FRB than a pulsar. In order to make this a repeating pulse, use numpy to repeat it. Below, this was done 30 times to get a nice sample of pulses."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#We repeat the pulse to get about 0.15 seconds of data\n",
    "repetition = 30 #Gives total number of desired pulses\n",
    "time_length = .005*repetition\n",
    "downsampledd_filtered_mixed_down_s_repeated = np.tile(downsampled_filtered_mixed_down_s, 30)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now finally, the pulsar is dispersed, which is what the telescope would show the pulsar as. The number of time samples of the new repeated pulse is found, and this is used to find the frequency steps for the pulse. A Dispersion measure (DM) is arbitrarilly chosen. Then the fourier transform of the delay time, given by the equation\n",
    "\\begin{equation}4.15e6 *DM*(\\frac{1}{\\nu_1^2}-\\frac{1}{\\nu_2^2}) \\end{equation} is taken. This equation assumes a frequency given in MHz, however here, it is given in Hz. Therefore, the equation is multiplied by 1e6^2 which, when multiplied to the original 4.15e3 constant, will give the appropriate constant below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#This is where the pulsar is dispersed\n",
    "#We do a convolution in fourier space, which must be long wrt the dispersion delay time\n",
    "dis_length = len(downsampledd_filtered_mixed_down_s_repeated) #This gives us the number of time samples of the \n",
    "                                                             #downsampled repeated pulse\n",
    "freq_step = np.linspace(-Bw/2, Bw/2, dis_length) #This is the new frequency step for the pulse\n",
    "DM = 600 #This is the dispersion measure and it is completely arbitrary\n",
    "\n",
    "#Create the signal to be convolved with the data. This signal is what disperses the pulsar\n",
    "H = np.exp(2j*np.pi*4.15e15*DM*freq_step**2/((Cf+freq_step)*Cf**2)) #Where the part that looks similar to the delay \n",
    "#time is the 'time' portion of the equation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As a precaution, the delay time is checked to ensure that it is less than the dispersion time of the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Check how long the dispersion is\n",
    "delay_time = 4.15e15 * (1/(Cf-Bw/2)**2 - 1/(Cf+Bw/2)**2) * DM\n",
    "\n",
    "#The delay time must be significantly less than the dispersion of the data\n",
    "print(delay_time) #This time must be less thah H"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The pulse signal and the dispersion function are convolved together, and the outcome is the dispersed pulse train. However when taking the fft and ifft, a overall gain is added to the signal, unless the norm of the fft and ifft is set to 'ortho' which will ensure no gain is applied to the signal. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dispersed_downsampled_filtered_mixed_sig = np.fft.ifft(np.fft.fft(downsampledd_filtered_mixed_down_s_repeated, norm='ortho')*np.fft.fftshift(H), norm='ortho')\n",
    "#The pulse signal and the dispersion of the pulse were convolved to give a dispersed pulse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The pulse has successfully been dispersed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "#Plotting the convolved signal...\n",
    "figure()\n",
    "plot(np.linspace(0,.15, dis_length)[:], dispersed_downsampled_filtered_mixed_sig[:].real)\n",
    "plot(np.linspace(0,.15, dis_length)[:], dispersed_downsampled_filtered_mixed_sig[:].imag)\n",
    "axis_labels('time(s)', 'Amplitude', 'Dispersed downsampled mixed pulsar (DM=20) ')\n",
    "#savefig('generated/despirsed_pulse_DM=20.png')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To simulate a real pulse, noise is added to the pulse train. This is white noise generated by choosing a random selection of values within a amplitude range of the pulse,and adding imaginary and real noise to the pulses. The noise_and_signal is the original signal with 1500000 points with en equal number of points appended to the beginning of the array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Noise is added to the pulse to simulate the RFI environment\n",
    "noise_interval = 1 #This tells us how many samples from the noise to concatinate\n",
    "noise_varience = .05\n",
    "\n",
    "noise1 = np.random.normal(0,noise_varience,len(dispersed_downsampled_filtered_mixed_sig),)\n",
    "noise2 =  np.random.normal(0,noise_varience,len(dispersed_downsampled_filtered_mixed_sig),)*1j\n",
    "dispersed_downsampled_filtered_mixed_sig = dispersed_downsampled_filtered_mixed_sig+noise1+noise2\n",
    "all_noise = noise1+noise2\n",
    "\n",
    "noise_and_signal = np.concatenate([all_noise[::noise_interval],dispersed_downsampled_filtered_mixed_sig])#The noise and signal have been concatenated as if \n",
    "                                              #They were observed right after another \n",
    "shapes = all_noise[::noise_interval].shape[0] #Tells us where the end of just noise is and where the signal begins\n",
    "pure_signal = noise_and_signal[shapes:] #This gives the signal with noise\n",
    "pure_noise = noise_and_signal[:shapes]  #This gives pure noise"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The new pulse with noise is shown below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "figure()\n",
    "plot(pure_signal.real)\n",
    "plot(pure_signal.imag)\n",
    "axis_labels('time(s)', 'Amplitude', 'Dispersed downsampled mixed pulsar (DM=20) ')\n",
    "#savefig('generated/despirsed_pulse_with_noise_DM=20.png')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "From this point, the pulse can be saved to a file and used as a simulated pulse inside GNUradio. The signal is divided by the maximum value of the real and imaginary parts and multiplied by the length of an int16 unit. This is then passed to the rest of the pipeline. For saving the file, an array with twice the length of the data is created and the real values are saved as the even indicies while the imaginary values are saved as the odd indicies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "pure_signal.dtype\n",
    "rmax = pure_signal.real.max()\n",
    "print(rmax)\n",
    "immax = pure_signal.imag.max()\n",
    "print(immax)\n",
    "int_max = 2**15-2\n",
    "print(int_max)\n",
    "dmax = max(rmax,immax)\n",
    "print(dmax)\n",
    "dispersed_downsampled_filtered_mixed_signal = pure_signal*int_max/dmax\n",
    "dispersed_downsampled_filtered_mixed_down_s_real = dispersed_downsampled_filtered_mixed_signal.real.astype(np.int16)\n",
    "dispersed_downsampled_filtered_mixed_down_s_imag = dispersed_downsampled_filtered_mixed_signal.imag.astype(np.int16)\n",
    "out = np.zeros(2*len(dispersed_downsampled_filtered_mixed_down_s_real), dtype=np.int16)\n",
    "out[::2] = dispersed_downsampled_filtered_mixed_down_s_real\n",
    "out[1::2] = dispersed_downsampled_filtered_mixed_down_s_imag\n",
    "out.tofile('generated/pulse_sim_10mhz_int16_5ms_period_60dm_1400MHz_center_150ms_long.bin')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now this is how this is how the pulse looks when it enters into the computer. The task is to now dedisperse it. We want an integration time that is dependent on the bandwidth. Therefore the inverse of the bandwidth is taken as the time that comes from the bandwidth. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#We want a 10Msps, since this is the bandwidth\n",
    "times = 1/10e6\n",
    "#Make the pure noise have the same apmlitude as the signal\n",
    "\n",
    "\n",
    "\n",
    "dispersed_downsampled_filtered_mixed_signal_noise = pure_noise*int_max/dmax\n",
    "dispersed_downsampled_filtered_mixed_down_s_real_noise = dispersed_downsampled_filtered_mixed_signal_noise.real.astype(np.int16)\n",
    "dispersed_downsampled_filtered_mixed_down_s_imag_noise = dispersed_downsampled_filtered_mixed_signal_noise.imag.astype(np.int16)\n",
    "out_noise = np.zeros(2*len(dispersed_downsampled_filtered_mixed_down_s_real_noise), dtype=np.int16)\n",
    "out_noise[::2] = dispersed_downsampled_filtered_mixed_down_s_real_noise\n",
    "out_noise[1::2] = dispersed_downsampled_filtered_mixed_down_s_imag_noise\n",
    "out_noise.tofile('generated/pulse_sim_10mhz_int16_5ms_period_60dm_1400MHz_center_150ms_long_noise.bin')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The number of desired frequency channels is set. This is an arbitrary choice. The integration size determins the signal to noise ratio of the data. An integration number of .2ms is chosen, since it is much smaller than the pulse width. Therefore, the resulting number of integrations done on the pulsar is 5. Finally, the number of timesteps for the pulse is calculated by the ratio of samples in the signal to the product of the number of frequency channels and integration size. This will dictate how many timesteps will be seen in the data.\n",
    "\n",
    "With these parameters, the FFT of the data per timestep is done, and arranged into an array as frequency by time.\n",
    "\n",
    "\n",
    "The same process is done with the noise, but using the length of the noise instead of the data for finding the number of timesteps. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.std(out_noise)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "#To check that this number of frequencies, integrations, and sampling rate work\n",
    "#integrate 5 samples\n",
    "integration_time = 2e-4 #seconds, this is the 1/.2ms integration\n",
    "step = 400 #This comes from the number of frequency channels desired. This is an educated, but arbitrary choice\n",
    "integration_size = int(integration_time/(times)/step) #This comes from the 0.2ms integration that we desire. It is .2 bc .1 was too small\n",
    "n_time = int(len(dispersed_downsampled_filtered_mixed_signal)/step/integration_size) #This finds the number of\n",
    "                                                            #timesteps\n",
    "rec_spect = np.zeros(((n_time), step), dtype=complex) #This the is the recieved 3d plot of frequency vs time\n",
    "for i in range(n_time): #This is the number of 'rows'\n",
    "    for j in range(integration_size):\n",
    "        spect = np.fft.fft(dispersed_downsampled_filtered_mixed_signal[(integration_size*i+j)*step : (integration_size*i+j+1)*step])\n",
    "        #This is an easy way to get the row size that we desire\n",
    "        rec_spect[i] += spect*spect.conjugate() #The first timestep of the frequencies needs to be multiplied by its conjugate\n",
    "        #The spetra are added together because we are 'folding' the amplitudes together. If we don't, then there appears \n",
    "        #a jump in the data.\n",
    "\n",
    "#Define the constants for the noise data set\n",
    "n_noise_time = int(len(dispersed_downsampled_filtered_mixed_signal_noise)/step/integration_size) #Number of timesteps for the noise\n",
    "        \n",
    "bb_noise = np.zeros((n_noise_time,step), dtype = complex) #This is the 3d array of the noise        \n",
    "for i in range(n_noise_time): #This is the number of 'rows'\n",
    "    for j in range(integration_size):\n",
    "        spect_n = np.fft.fft(dispersed_downsampled_filtered_mixed_signal_noise[(integration_size*i+j)*step : (integration_size*i+j+1)*step])\n",
    "        #This is an easy way to get the row size that we desire\n",
    "        bb_noise[i] += spect_n*spect_n.conjugate() #The first timestep of the frequencies needs to be multiplied by its conjugate\n",
    "        #The spetra are added together because we are 'folding' the amplitudes together. If we don't, then there appears \n",
    "        #a jump in the data.    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The estimated signal to noise of the pulse can now be calculated. This will be compared to the SNR that will be obtained at the end. The amplitude however has to be estimated by hand, since it will always change."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "estimated_noise = 1/(noise_varience/.069)**2*5\n",
    "print(estimated_noise)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For bench testing purposes, the pulse that has now gone through an FFT and been integrated will be compared to the integration and FFT blocks in GNUradio. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rmax_FFT = rec_spect.real.max()\n",
    "print(rmax)\n",
    "immax_FFT = rec_spect.imag.max()\n",
    "print(immax)\n",
    "int_max = 2**15-2\n",
    "print(int_max)\n",
    "dmax_FFT = max(rmax_FFT,immax_FFT)\n",
    "recieved_spect = (np.fft.fftshift(rec_spect))\n",
    "rec_spect_real = recieved_spect.real.astype(np.float32)\n",
    "out_spec = np.zeros((len(rec_spect_real),step), dtype=np.float32)\n",
    "for i in range(step):\n",
    "    out_spec[:,i] = rec_spect_real[:,i]\n",
    "out_spec.tofile('generated/pulse_sim_10mhz_int16_5ms_period_60dm_1400MHz_center_150ms_long_integrated_FFT.bin')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "out_spec"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to get this into a more conventional array, the transpose of the FFTed and integrated data is taken, so that it is arranged frequency by time. Then the y axis of the plot is created, where first the number of frequency channels is found, then the centre frequency of that array is set at 1400MHz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_frequencies = np.fft.fftfreq(step, 1/(10e6/1e6)) #This is the number of frequency channles.The first parameter \n",
    "                                #should be the number of chanels desired, and the second is the sample rate in MHz.\n",
    "                                #This will lable the the x axis\n",
    "pulse_freqs = np.fft.fftshift(n_frequencies)+Cf/1e6 #This will place the centre frequency of the plot at 1400MHz\n",
    "dispersed = np.transpose(np.fft.fftshift((rec_spect)))#The array of the data must be flipped to match \n",
    "                                    #the units in the plot,and it must be transposed so that time is on the x axis.\n",
    "bb_flipped = np.flip(np.transpose(np.fft.fftshift(bb_noise)),0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dispersed"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The 3d graph of the dispersed pulsar is then plotted"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "#This is the graph of the pulsar, where we can see that it is dispersed\n",
    "figure()\n",
    "imshow(abs(dispersed), extent=[np.amin(np.linspace(0,.15, dis_length)),\\\n",
    "                            np.amax(np.linspace(0,.15, dis_length)),\\\n",
    "                            np.amin(pulse_freqs),\\\n",
    "                            np.amax(pulse_freqs)],\\\n",
    "                              aspect='auto')\n",
    "axis_labels('time(s)','Frequency (MHz)',  'Dispersed Pulsar (DM=20)')\n",
    "colorbar(label = 'power')\n",
    "#savefig(\"generated/Dispersed_pulsar_DM=20_noisy.png\")\n",
    "#Congrats!, we now have a dispersed pulsar\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This the the 3d plot of the noise behind the plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "figure()\n",
    "imshow(abs(bb_flipped), extent=[np.amin(np.linspace(0,.15, dis_length)),\\\n",
    "                            np.amax(np.linspace(0,.15, dis_length)),\\\n",
    "                            np.amin(pulse_freqs),\\\n",
    "                            np.amax(pulse_freqs)],\\\n",
    "                              aspect='auto')\n",
    "axis_labels('time(s)','Frequency (MHz)',  'Dispersed Pulsar (DM=20)')\n",
    "colorbar(label = 'power')\n",
    "#savefig(\"generated/Noise_DM=20_noisy.png\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now the pulsar is attempted to be dedispersed. Assume that the DM, nor the period of the pulsar is known. An array of possible DM's is created, along with a timestep array of when the pulses arrive. Then using the time delay equation used to disperse the pulsar, the time changes for every frequency are calculated. This is repeated for each trial dispersion measure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DMs = np.linspace(400,800,401) #Assume that we do not know the DM for the pulsar. We must then Guess a range of DM\n",
    "\n",
    "time_change = np.zeros((len(DMs), step)) #This is the time after the first frequency which the second frequncy arrives\n",
    "for i in range(len(DMs)):\n",
    "    change = 4.148e3*DMs[i]*( 1/(np.min(pulse_freqs))**2 - 1/( pulse_freqs )**2) #This is the definition for the time that the second freq arrives\n",
    "    time_change[i,:] = change[:]/(time_length/n_time)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using the time delays, the pulsar is now dedispersed for every guessed DM. This is done by using numpy.roll. The threshold for the roll is set by finding the time that each element in the data corresponds to, and moving elements so that they line up with their corresponding elements(come back to give better explanation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dedispersion(DMss, frequency_channels, dispersed_pulse, time_change, n_time, integration_size): \n",
    "    \"\"\"The function takes a set of DMs and dedisperses the given dispersed pulse for all the DMs given.\"\"\"\n",
    "    dedispers = np.zeros((len(DMss),dispersed_pulse.shape[0], dispersed_pulse.shape[1]))  #The dedispersed pulse will have the same shape as the dispersed one\n",
    "    for i in range(len(DMss)):\n",
    "        dedis = dispersed_pulse*0\n",
    "        for j in range(frequency_channels):\n",
    "            dedis[j,:] = np.roll(dispersed_pulse[j,:],int(round(time_change[i,j]) ))#np.roll shifts \n",
    "                #the data over after the last position and reintroduces it to the first position\n",
    "        dedispers[i,:,:] = dedis\n",
    "    return dedispers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using the above definition, the pulse is dedispersed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dedisperse = dedispersion(DMs, step, dispersed, time_change, n_time, integration_size)  #The dedispersed pulse will have the same shape as the dispersed one\n",
    "dedisperse_bb = dedispersion(DMs,step,abs(bb_flipped), time_change, n_time, integration_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.std(dedisperse_bb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One DM value for the pulse is then plotted below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "figure()\n",
    "imshow(abs(dedisperse[200]), extent=[np.amin(np.linspace(0,.15, dis_length)),\\\n",
    "                            np.amax(np.linspace(0,.15, dis_length)),\\\n",
    "                            np.amin(pulse_freqs),\\\n",
    "                            np.amax(pulse_freqs)],\\\n",
    "                              aspect='auto')\n",
    "axis_labels('time(s)','Frequency (MHz)',  'Dedispersed Pulsar (DM=20)')\n",
    "colorbar(label = 'power') #This shows the effec the dedispersion had on the 200th DM\n",
    "#savefig(\"generated/Dedispersed_pulsar_DM=20_noisy.png\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To create a DM vs time plot, the array is summed along the frequency axis for every DM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DM_measure = np.sum(dedisperse, axis = 1)\n",
    "bb_measure = np.sum(dedisperse_bb, axis = 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This new array is then plotted below. Note that the graph appears to have a wave in the region with the highest amplitudes. This is caused by the moire effect which can be shown as the cause with the graph following the one below(see pic for no noise pulsar)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "figure()\n",
    "imshow((abs(DM_measure)), extent=[np.amin(np.linspace(0,.15, dis_length)),\\\n",
    "                            np.amax(np.linspace(0,.15, dis_length)),\\\n",
    "                            np.amax(DMs),np.amin(DMs)\n",
    "                            ],\\\n",
    "                              aspect='auto')\n",
    "axis_labels('time(s)','Dispersion Measure',  'Dispersion Measure vs Time of Pulsar (DM=20)')\n",
    "colorbar(label = 'power') #This is the plot of DM vs time\n",
    "#savefig(\"generated/dispersion_measure_vs_time_DM=20_noisy.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rmax_FFT = rec_spect.real.max()\n",
    "print(rmax)\n",
    "immax_FFT = rec_spect.imag.max()\n",
    "print(immax)\n",
    "int_max = 2**15-2\n",
    "print(int_max)\n",
    "dmax_FFT = max(rmax_FFT,immax_FFT)\n",
    "DMs_real = DM_measure.real.astype(np.float32)\n",
    "out_DMs = np.zeros((len(DMs_real[0])*5), dtype=np.float32)\n",
    "for i in range(200,205):\n",
    "    print(i)\n",
    "    out_DMs[(750*(i-200)):(750*(i-200+1))] = DMs_real[i,:]\n",
    "out_DMs.tofile('generated/pulse_sim_10mhz_int16_5ms_period_60dm_1400MHz_center_150ms_long_dispersed.bin')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "figure()\n",
    "plot(out_DMs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To detect whether or not a pulsar or FRB was detected, a matched filter will be implemented. Since we are dealing with an ideal case of white noise, then the filter that will be applied will be a pulse of the same width as the signal. This width is calculated by taking the original sigma of simulated pulse and dividing it by the final timestep after all modification of the pulse is done. This occurrs after integration. Then the gaussian pulse that will be used as the filter is created. The length is made such that packets are not significantly overlapping."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Now we will begin to try to impliment the matched filter\n",
    "#Define the standard deviation of the noise\n",
    "\n",
    "#Define the pulse shape by the standard deviation of one pulse\n",
    "packet_length = int(len(downsampled_filtered_mixed_down_s)/step/(integration_size))\n",
    "standards = pw/(period/packet_length)\n",
    "\n",
    "gau = np.exp(.5*-( np.arange(-int(packet_length/2),int(packet_length/2)) /standards)**2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.std(bb_measure[200])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The cross correlation between the data and the filter is the implemented below. The correlation for signal and filter is specified to give the same shaped array as the data given to it. The noise is also run through the matched filter, but the correlation is specified to give onyl real values. This ensures that the standard deviation will not be increased by a factor of ten from if just the real valued correlations were taken. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Now perform the cross correlation between the signal and the simulated\n",
    "correlates = []\n",
    "corr_bb = []\n",
    "for i in range(len(DM_measure)):\n",
    "    correlates.append(signal.correlate(DM_measure[i], gau, mode='same'))\n",
    "    corr_bb.append(signal.correlate(bb_measure[i],gau,mode='valid'))\n",
    "correlates = np.asarray(correlates)\n",
    "corr_bb = np.asarray(corr_bb)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "corr_real = correlates.real.astype(np.float32)\n",
    "out_corr = np.zeros((len(correlates[0])), dtype=np.float32)\n",
    "\n",
    "out_corr[:] = corr_real[200,:]\n",
    "out_corr.tofile('generated/pulse_sim_10mhz_int16_5ms_period_60dm_1400MHz_center_150ms_long_corr.bin')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The pulsar will be considered detected if the signal to noise ratio, given by\n",
    "\\begin{equation} \\frac{\\langle y \\rangle}{\\sigma_y} \\end{equation}\n",
    "is over a threshold. The average value of y is the average value of the correlated signal. This is found by using a peak finding function for each DM and then take the average of those peaks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#We can find the peaks\n",
    "mean_heights = np.zeros(len(correlates))\n",
    "for i in range(len(correlates)):\n",
    "    mean_heights[i] = np.mean(correlates[i])\n",
    "peak_heights = peakfinder_pulsar_DM(correlates, mean_heights)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "figure()\n",
    "plot(out_corr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "corr_std"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now the vaience of the non pulse signal is found. This is done by pointing the telescope somewhere else and observing the noise for a brief period. The data collected is assumed to not have a pulsar in it. This data, which by this point has been run through the matched filter, is used to now find the mean and standard deviation along the time axis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "maximums = np.zeros(len(peak_heights))\n",
    "for i in range(len(peak_heights)):\n",
    "    maximums[i] = np.mean(peak_heights[i])\n",
    "    \n",
    "#Find the mean and standard deviation\n",
    "corr_mean = np.mean(corr_bb, axis=1)\n",
    "corr_std = np.std(corr_bb, axis=1)\n",
    "\n",
    "SNR = (maximums-corr_mean)/(corr_std)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The plot of SNR to DM is plotted"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "figure()\n",
    "plot(DMs,SNR)\n",
    "axis_labels('DM', 'Signal to Noise Ratio', 'Signal to Noise ratio per DM = 20')\n",
    "#savefig(\"generated/SNR_vs_DM_DM=20_noisy.png\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally we will check if a pulsar or FRB has been found"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Checking for a pulsar:\n",
    "if len(np.where(SNR>10)[0]) >5:\n",
    "    Pulsar = True\n",
    "else:\n",
    "    Pulsar = False\n",
    "\n",
    "print(Pulsar)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "800000/400/5/4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "750*.2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}