LightCurveCreation.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from sunpy.io import fits\n",
    "import sunpy.map\n",
    "import os\n",
    "import datetime\n",
    "import time\n",
    "from sunpy.lightcurve import LightCurve\n",
    "from scipy.optimize import curve_fit\n",
    "from sunpy.physics.solar_rotation import mapcube_solar_derotate\n",
    "import scipy\n",
    "import batman\n",
    "from scipy.fftpack import fft, fftfreq, fftshift, ifft\n",
    "from scipy.interpolate import interp1d\n",
    "\n",
    "#Fixes a bug in sunpy.\n",
    "#This is fixed by https://github.com/sunpy/sunpy/pull/1853\n",
    "sunpy.time.time.TIME_FORMAT_LIST.append(\"%Y.%m.%d_%H:%M:%S_TAI\")\n",
    "\n",
    "#Select the correct matplotlib backend\n",
    "#Options: inline|qt|osx|gtx\n",
    "%matplotlib osx "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "personal_data_directory = '/home/gtaylor/'\n",
    "aia_data_directory = '/data/SDO/AIA/level1/'\n",
    "hmi_data_directory = '/data/SDO/HMI/intensitygram/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#Stuff for multithreading\n",
    "#From http://stackoverflow.com/a/35517484\n",
    "\n",
    "def jobs_manager():\n",
    "    from IPython.lib.backgroundjobs import BackgroundJobManager\n",
    "    from IPython.core.magic import register_line_magic\n",
    "    from IPython import get_ipython\n",
    "\n",
    "    jobs = BackgroundJobManager()\n",
    "\n",
    "    @register_line_magic\n",
    "    def job(line):\n",
    "        ip = get_ipython()\n",
    "        jobs.new(line, ip.user_global_ns)\n",
    "\n",
    "    return jobs\n",
    "\n",
    "def kill_thread(thread):\n",
    "    import ctypes\n",
    "\n",
    "    id = thread.ident\n",
    "    code = ctypes.pythonapi.PyThreadState_SetAsyncExc(\n",
    "        ctypes.c_long(id),\n",
    "        ctypes.py_object(SystemError)\n",
    "    )\n",
    "    if code == 0:\n",
    "        raise ValueError('invalid thread id')\n",
    "    elif code != 1:\n",
    "        ctypes.pythonapi.PyThreadState_SetAsyncExc(\n",
    "            ctypes.c_long(id),\n",
    "            ctypes.c_long(0)\n",
    "        )\n",
    "        raise SystemError('PyThreadState_SetAsyncExc failed')\n",
    "\n",
    "jobs = jobs_manager()\n",
    "\n",
    "#Use %job [stuff to do]\n",
    "#jobs.status() to check on jobs\n",
    "#jobs.traceback(num) for stack trace\n",
    "#kill_thread(jobs.all[num]) to kill\n",
    "#Kill all running threads with\n",
    "#for thread in jobs.running:\n",
    "#    kill_thread(thread)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#The angle to rotate the picture of the sun so that Venus goes horizontal through the image\n",
    "transit_rotate = {\"venus\":7.4,\"mercury\":2.0}\n",
    "\n",
    "no_rotate = 0\n",
    "#Splits are a dictionary of two tuples of two tuples each that are each in the format (w,[x,y,z,...]). \n",
    "#That is, it breaks the image into w slices and then uses slices x,y,z,...\n",
    "#The first tuple is for AIA images, the second tuple is for HMI images\n",
    "#The section of the sun that contains the transit.\n",
    "#(64,[13,14])\n",
    "inside_split = {\"venus\":((32,[8]), (16,[3])),\"mercury\":((256,[159,160,161]), (256,[166,167,168]))}\n",
    "#The section of the sun that doesn't contain the transit\n",
    "outside_split = {\"venus\":(((32,range(8) + range(9,32)), (16,range(3) + range(4,16)))),\"mercury\":((256,range(159) + range(162,256)), ((256,range(166) + range(169,256))))}\n",
    "no_split = {\"venus\":(((1,[0]), (1,[0]))),\"mercury\":(((1,[0]), (1,[0])))}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "hours = {\"venus\":[\"21\",\"22\",\"23\",\"00\",\"01\",\"02\",\"03\",\"04\",\"05\"],\"mercury\":[\"10\",\"11\",\"12\",\"13\",\"14\",\"15\",\"16\",\"17\",\"18\",\"19\",\"20\"]}\n",
    "\n",
    "def checkIfInHours(hour, planet):\n",
    "    '''\n",
    "    Raises an exception if the given `hour` is not in the list of hours for the given `planet`\n",
    "    '''\n",
    "    if hour not in hours[planet]:\n",
    "        raise Exception(hour + ' not an accepted hour for ' + planet + '. Try one of ' + str(hours[planet]))\n",
    "\n",
    "venus_spike_start_time = 9000 #in seconds, start of spike in 1700Å\n",
    "venus_spike_end_time = 15000 #in seconds, end of spike in 1700Å\n",
    "\n",
    "\n",
    "ingress_start_time = {\"venus\":4020,\"mercury\":5010} #in seconds, start of ingress\n",
    "ingress_end_time = {\"venus\":5150,\"mercury\":5280} #in seconds, end of ingress\n",
    "\n",
    "egress_start_time = {\"venus\":26100,\"mercury\":31440} #in seconds, start of egress\n",
    "egress_end_time = {\"venus\":27350,\"mercury\":31680} #in seconds, end of egress\n",
    "\n",
    "transit_start_time = ingress_end_time\n",
    "transit_end_time = egress_start_time\n",
    "\n",
    "def timeSinceStartOfTransit(planet, time):\n",
    "    '''\n",
    "    `planet` is the planet transit to use\n",
    "\n",
    "    Given a `time` since the start of the recorded period, \n",
    "        returns the difference between the given time \n",
    "        and when the transit actually started\n",
    "    '''\n",
    "    return time - transit_start_time[planet]\n",
    "\n",
    "def timeThroughTransit(planet, time):\n",
    "    '''\n",
    "    `planet` is the planet transit to use\n",
    "\n",
    "    Given a `time` since the start of the recorded period, \n",
    "        returns    the start time if the transit has yet to start,\n",
    "                   the end time   if the transit has already finished,\n",
    "                or the current time\n",
    "    '''\n",
    "    if time < transit_start_time[planet]:\n",
    "        return transit_start_time[planet]\n",
    "    if time < transit_end_time[planet]:\n",
    "        return time\n",
    "    return transit_end_time[planet]        \n",
    "\n",
    "def checkTimeBlock(planet, time_block, time):\n",
    "    '''\n",
    "    `planet` is the planet transit to use\n",
    "\n",
    "    `time_block` is an array of strings to limit the list to (see section \"Time\")\n",
    "   \n",
    "    `time` is the time to check\n",
    "\n",
    "    Checks if `time` is in the `time_block`\n",
    "    '''\n",
    "    if time_block == []:\n",
    "        return True\n",
    "    \n",
    "    doIAppend = False\n",
    "    \n",
    "    for block in time_block:\n",
    "        if type(block) in [int,float,long]:\n",
    "            if time < block:\n",
    "                doIAppend = True\n",
    "            else:\n",
    "                doIAppend = False\n",
    "            return doIAppend\n",
    "        \n",
    "    if \"transit\" in time_block:\n",
    "        if time > transit_start_time[planet] and time < transit_end_time[planet]:\n",
    "            doIAppend = True\n",
    "            \n",
    "    if \"not transit\" in time_block:\n",
    "        if time < ingress_start_time[planet] or time > egress_end_time[planet]:\n",
    "            doIAppend = True\n",
    "\n",
    "    if \"ingress\" in time_block:\n",
    "        if time > ingress_start_time[planet] and time < ingress_end_time[planet]:\n",
    "            doIAppend = True\n",
    "    if \"egress\" in time_block:\n",
    "        if time > egress_start_time[planet] and time < egress_end_time[planet]:\n",
    "            doIAppend = True\n",
    "\n",
    "    if \"not venus spike\" in time_block:\n",
    "        if time > venus_spike_start_time and time < venus_spike_end_time:\n",
    "            doIAppend = False\n",
    "    \n",
    "    return doIAppend\n",
    "\n",
    "def getDataFromTimeBlock(planet, time_block, data, times):\n",
    "    '''\n",
    "    `planet` is the planet transit to use\n",
    "\n",
    "    `time_block` is an array of strings to limit the list to (see section \"Time\")\n",
    "\n",
    "    `data` is an array of generic data points that correspond with the time points in `times`\n",
    "\n",
    "    `data` and `times` must be of same length    \n",
    "\n",
    "    Returns the data points that match the `time_block`\n",
    "    '''\n",
    "    good_data = []\n",
    "    for i, time in enumerate(times):\n",
    "        if checkTimeBlock(planet, time_block, time):\n",
    "            good_data.append(data[i])\n",
    "            \n",
    "    return good_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#data about Venus\n",
    "\n",
    "venus_radius = 3760.4 #miles\n",
    "sun_radius = 432168.6 #miles\n",
    "venus_semi_major_axis = 67237909. #miles\n",
    "\n",
    "#apparent venus radius in stellar radii\n",
    "#((tangent of angular radius of venus on june 5th 2012)*(1 astronomical unit - radius of sun - radius of earth) / (radius of sun))\n",
    "venus_apparant_radius = 0.023 #stellar radii\n",
    "\n",
    "venus_longitude_perihelion = 131.53298 #degrees\n",
    "\n",
    "venus_orbital_period_days = 224.701 #days\n",
    "venus_orbital_period = venus_orbital_period_days * 24 * 60 * 60 #seconds\n",
    "\n",
    "venus_eccentricity = 0.0067\n",
    "venus_orbital_inclination = 3.39 #degrees\n",
    "\n",
    "venus_time_inferior_conjunc = (transit_start_time[\"venus\"] + transit_end_time[\"venus\"] ) / 2 #seconds\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "'''\n",
    "Functions that are of the form nth * x^n + n-1th * x^(n-1) + … + one * x + zero\n",
    "\n",
    "They work for `func_one` and `func_two` for graphLightCurve\n",
    "'''\n",
    "\n",
    "def cubic(x,three,two,one,zero):\n",
    "    return np.multiply(three,np.power(x,3)) + np.multiply(two,np.square(x)) + np.multiply(one,x) + zero\n",
    "\n",
    "def quadratic(x,two,one,zero):\n",
    "    return np.multiply(two,np.square(x)) + np.multiply(one,x) + zero\n",
    "\n",
    "def linear(x,one,zero):\n",
    "    return np.multiply(one,x) + zero\n",
    "\n",
    "def quartic(x,four,three,two,one,zero):\n",
    "    return np.multiply(four,np.power(x,4)) + np.multiply(three,np.power(x,3)) + np.multiply(two,np.square(x)) + np.multiply(one,x) + zero\n",
    "\n",
    "def quintic(x,five,four,three,two,one,zero):\n",
    "    return np.multiply(five,np.power(x,5)) + np.multiply(four,np.power(x,4)) + np.multiply(three,np.power(x,3)) + np.multiply(two,np.square(x)) + np.multiply(one,x) + zero"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "wavelengths = [\"0094\",\"0131\",\"0171\",\"0193\",\"0211\",\"0304\",\"0335\",\"1600\",\"1700\",\"cont\"]\n",
    "#\"cont\" = HMI intensity continuum, all other are AIA wavelengths\n",
    "\n",
    "def checkIfInWavelengths(wavelength):\n",
    "    '''\n",
    "    Raises an exception if the given `wavelength` \n",
    "        is not in the list of wavelengths\n",
    "    '''\n",
    "    if wavelength not in wavelengths:\n",
    "        raise Exception('Wavelength ' + wavelength + \n",
    "                        ' not an accepted wavelength. Try one of ' \n",
    "                        + str(wavelengths))\n",
    "        \n",
    "def getWavelengthFromFile(filepath):\n",
    "    '''\n",
    "    `filepath` is a path to a FITS file\n",
    "    '''\n",
    "    if \"intensity\" in filepath:\n",
    "        return \"cont\"\n",
    "    else:\n",
    "        return filepath[-9:-5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def getPlanetFromFile(filepath):\n",
    "    '''\n",
    "    `filepath` is a path to a FITS file\n",
    "\n",
    "    Returns either \"venus\" or \"mercury\"\n",
    "    '''\n",
    "    if \"/2012/\" in filepath:\n",
    "        return \"venus\"\n",
    "    elif \"/2016/\" in filepath:\n",
    "        return \"mercury\"\n",
    "    else:\n",
    "        raise Exception(\"Can't parse filepath \" + filepath + \" to get a planet\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def readData(planet, wavelength):\n",
    "    '''\n",
    "    `planet` is the planet transit to use\n",
    "\n",
    "    `wavelength` is the wavelength to use\n",
    "\n",
    "    Returns the saved data for that wavelength in list form, sorted by time\n",
    "    '''\n",
    "    checkIfInWavelengths(wavelength)\n",
    "    \n",
    "    data = sorted(readFromFile(personal_data_directory + planet + 'data' + wavelength + '.txt').items(), key=lambda tup: tup[1][1])\n",
    "        \n",
    "    good_data = []\n",
    "    \n",
    "    for filepath, (count, time, pixels, count_in_split, pixels_in_split, count_out_of_split, pixels_out_of_split) in data:\n",
    "        if np.isfinite(count):\n",
    "            good_data.append((filepath, (count, time, pixels, count_in_split, pixels_in_split, count_out_of_split, \n",
    "                                         pixels_out_of_split)))\n",
    "    \n",
    "    return good_data\n",
    "\n",
    "def writeToFile(filepath, data):\n",
    "    '''\n",
    "    `filepath` is the filepath to write to\n",
    "\n",
    "    `data` is the data to save\n",
    "\n",
    "    Returns `data`\n",
    "    '''\n",
    "    with open(filepath, 'a') as target:\n",
    "        target.truncate(0)\n",
    "        target.write(str(data))\n",
    "        \n",
    "    return data\n",
    "        \n",
    "def readFromFile(filepath):\n",
    "    '''\n",
    "    `filepath` is the filepath to read from\n",
    "    \n",
    "    Returns the file, evaluated. \n",
    "    \n",
    "    This will turn, for example, a dictionary in the file \n",
    "        into a dictionary in python\n",
    "    '''\n",
    "    while not os.path.exists(filepath):\n",
    "        print(\"Waiting for file \" + filepath + \" to exist\")\n",
    "        time.sleep(1)\n",
    "    \n",
    "    try:\n",
    "        with open(filepath, 'r') as target:\n",
    "            s = target.read()\n",
    "            return eval(s)\n",
    "    except:\n",
    "        return {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def showFancyImage(filepath, subplot=None, new_figure=True):\n",
    "    '''\n",
    "    `filepath` is a path to a FITS file\n",
    "    \n",
    "    `subplot` is the argument passed into matplotlib's subplot function. If `subplot == None` then subplot is not used.\n",
    "\n",
    "    If `new_figure` is True, it creates a new matplotlib figure for the image\n",
    "\n",
    "    It makes a sunpy map of the given file, and puts the map on screen\n",
    "    \n",
    "    Returns the map\n",
    "    '''\n",
    "    \n",
    "    if new_figure:\n",
    "        plt.figure()\n",
    "        \n",
    "    if subplot != None:\n",
    "        plt.subplot(subplot)\n",
    "\n",
    "    smap = sunpy.map.Map(filepath)\n",
    "    smap.plot()\n",
    "    return smap\n",
    "\n",
    "def showArray(array, subplot=None, new_figure=True, color_map='Greys'):\n",
    "    '''\n",
    "    Graphs the given `array` using the colors from `color_map` \n",
    "    \n",
    "    `subplot` is the argument passed into matplotlib's subplot function. If `subplot == None` then subplot is not used.\n",
    "\n",
    "    If `new_figure` is True, it creates a new matplotlib figure for the graphed array\n",
    "    \n",
    "    Returns the array\n",
    "    '''\n",
    "    \n",
    "    if new_figure:\n",
    "        plt.figure()\n",
    "        \n",
    "    if subplot != None:\n",
    "        plt.subplot(subplot)\n",
    "        \n",
    "    plt.matshow(array, cmap=color_map)\n",
    "    \n",
    "    return array\n",
    "    \n",
    "\n",
    "def showHistogramOfPixelIntensities(array, boxes=50, subplot=None, new_figure=True):\n",
    "    '''\n",
    "    Shows a histogram of elements of the given `array`\n",
    "\n",
    "    `boxes` is the number of separate bars to sort the elements of the array into\n",
    "    \n",
    "    `subplot` is the argument passed into matplotlib's subplot function. If `subplot == None` then subplot is not used.\n",
    "\n",
    "    If `new_figure` is True, it creates a new matplotlib figure for the histogram\n",
    "\n",
    "    Note: uses lots of ram\n",
    "        \n",
    "    Returns the array, with all nan elements set to zero\n",
    "    '''\n",
    "            \n",
    "    if new_figure:\n",
    "        plt.figure()\n",
    "        \n",
    "    if subplot != None:\n",
    "        plt.subplot(subplot)\n",
    "\n",
    "    array = array[~np.isnan(array)]\n",
    "    \n",
    "    plt.hist(array, boxes)\n",
    "    plt.show()\n",
    "    \n",
    "    return array\n",
    "\n",
    "def showImage(filepath, split=no_split, rotate=no_rotate, replace_with_zero=False, subplot=None, new_figure=True):\n",
    "    '''\n",
    "    `filepath` is a path to a FITS file\n",
    "\n",
    "    `split` is the section of the array that is cared about (see section \"Split and Rotate\")\n",
    "\n",
    "    `rotate` is the number of degrees to rotate the array by (counter clockwise)\n",
    "\n",
    "    If `replace_with_zero` is true, it will replace the removed elements of the array with zeros, otherwise it graphs a truncated array.\n",
    "\n",
    "    `subplot` is the argument passed into matplotlib's subplot function. If `subplot == None` then subplot is not used.\n",
    "\n",
    "    If `new_figure` is True, it creates a new matplotlib figure for the graphed array\n",
    "\n",
    "    Graphs the array from the filepath, applying the split and rotate\n",
    "    '''\n",
    "            \n",
    "    if new_figure:\n",
    "        plt.figure()\n",
    "        \n",
    "    if subplot != None:\n",
    "        plt.subplot(subplot)\n",
    "\n",
    "    array = getRotatedAndSplitArrayFromFitsFile(filepath, split, rotate, replace_with_zero=replace_with_zero)\n",
    "    \n",
    "    showArray(array)\n",
    "    \n",
    "    return array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def getArrayFromFitsFile(filepath):\n",
    "    '''\n",
    "    `filepath` is a path to a FITS file\n",
    "\n",
    "    If the file doesn't exist, wait for it to.\n",
    "\n",
    "    Returns the array from the given file\n",
    "    '''\n",
    "    \n",
    "    while not os.path.exists(filepath):\n",
    "        print(\"Waiting for file \" + filepath + \" to exist\")\n",
    "        time.sleep(1)\n",
    "\n",
    "    try:\n",
    "        return np.array(fits.read(filepath)[1][0])\n",
    "    except:\n",
    "        print(\"Couldn't get array from \" + filepath)\n",
    "        time.sleep(1)\n",
    "        return getArrayFromFitsFile(filepath)\n",
    "    \n",
    "def getHeaderFromFitsFile(filepath):\n",
    "    '''\n",
    "    `filepath` is a path to a FITS file\n",
    "\n",
    "    If the file doesn't exist, wait for it to.\n",
    "\n",
    "    Returns the header from the given file\n",
    "    '''\n",
    "    while not os.path.exists(filepath):\n",
    "        print(\"Waiting for file \" + filepath + \" to exist\")\n",
    "        time.sleep(1)\n",
    "\n",
    "    try:\n",
    "        return fits.get_header(filepath)\n",
    "    except:\n",
    "        print(\"Couldn't get header from \" + filepath)\n",
    "        time.sleep(1)\n",
    "        return getHeaderFromFitsFile(filepath)\n",
    "        \n",
    "def findFileByTimeSinceStartOfTransit(planet, time, wavelength):\n",
    "    '''\n",
    "    `planet` is the planet transit to use\n",
    "\n",
    "    `time` is the time since the start of the transit to get the file from\n",
    "\n",
    "    `wavelength` is the wavelength to use\n",
    "\n",
    "    Returns the filepath to the FITS file closest to that time\n",
    "    '''\n",
    "    \n",
    "    checkIfInWavelengths(wavelength)\n",
    "\n",
    "    file_and_time_list = readListOfFilesAndTimeSinceStartOfTransit(planet, wavelength)\n",
    "    filepath = None\n",
    "    plusminus = 1\n",
    "    \n",
    "    if file_and_time_list == None:\n",
    "        return None\n",
    "    \n",
    "    while filepath == None:\n",
    "        #Look for a file in the current time range \n",
    "        filepath = next((i for i, v in enumerate(file_and_time_list) if v[1] < time + plusminus \n",
    "                     and v[1] > time - plusminus), None)\n",
    "        if filepath == None: #If no file is found, expand the time range\n",
    "            plusminus *= 2\n",
    "            \n",
    "    return file_and_time_list[filepath][0]\n",
    "\n",
    "def getRotatedAndSplitArrayFromFitsFile(filepath, split=no_split, rotate=no_rotate, replace_with_zero=False):\n",
    "    '''\n",
    "    `filepath` is a path to a FITS file\n",
    "\n",
    "    `split` is the section of the array that is cared about (see section \"Split and Rotate\")\n",
    "\n",
    "    `rotate` is the number of degrees to rotate the array by (counter clockwise)\n",
    "\n",
    "    If `replace_with_zero` is true, it will replace the removed elements of the array with zeros, otherwise it returns a truncated array.\n",
    "\n",
    "    Finds the array from the given file given file, \n",
    "        then applies the given split and rotate parameters as needed\n",
    "        \n",
    "    Returns the array created\n",
    "    '''\n",
    "    \n",
    "    wavelength = getWavelengthFromFile(filepath)\n",
    "    \n",
    "    planet = getPlanetFromFile(filepath)\n",
    "    \n",
    "    if rotate == transit_rotate:\n",
    "        rotate = transit_rotate[planet]\n",
    "            \n",
    "    split = split[planet]\n",
    "        \n",
    "    #Get the array\n",
    "    full_array = getArrayFromFitsFile(filepath)\n",
    "    if full_array == None:\n",
    "        return []\n",
    "    \n",
    "    #Replace all the nans with zero\n",
    "    full_array = np.nan_to_num(full_array)\n",
    "    \n",
    "    #Flips the array so that the orientation is consistent for all wavelengths\n",
    "    \n",
    "    \n",
    "    if wavelength == \"cont\":\n",
    "        full_array = np.fliplr(full_array)\n",
    "        split = split[1]\n",
    "    else:\n",
    "        full_array = np.flipud(full_array)\n",
    "        split = split[0]\n",
    "            \n",
    "    #Rotates the image, if neeeded\n",
    "    if rotate != 0:\n",
    "        full_array = scipy.ndimage.interpolation.rotate(full_array,rotate, reshape=False)\n",
    "    \n",
    "    #Splits the array into the different sections\n",
    "    arrays = np.array_split(full_array, split[0])\n",
    "    \n",
    "    split_shape = arrays[0].shape\n",
    "    \n",
    "    #Goes through all the split arrays and isolates the wanted splits\n",
    "    good_arrays = []    \n",
    "    for i, array in enumerate(arrays):\n",
    "        if i in split[1]:\n",
    "            good_arrays.append(array)\n",
    "        else:\n",
    "            if replace_with_zero:\n",
    "                good_arrays.append(np.zeros(split_shape))\n",
    "         \n",
    "    #Combines the wanted splits into a single image\n",
    "    array = np.concatenate(good_arrays)\n",
    "    return array\n",
    "\n",
    "def getTotalData(filepath, split=no_split, rotate=no_rotate):\n",
    "    '''\n",
    "    `filepath` is a path to a FITS file\n",
    "\n",
    "    `split` is the section of the array that is cared about (see section \"Split and Rotate\")\n",
    "\n",
    "    `rotate` is the number of degrees to rotate the array by (counter clockwise)\n",
    "    \n",
    "    It ignores the negative items in the array and items with a value \n",
    "        greater than the mean + std dev of the wavelength as a whole\n",
    "    \n",
    "    Returns the sum of the image data in the fits file\n",
    "    '''\n",
    "    \n",
    "    wavelength = getWavelengthFromFile(filepath)\n",
    "        \n",
    "    planet = getPlanetFromFile(filepath)\n",
    "    \n",
    "    array = getRotatedAndSplitArrayFromFitsFile(filepath, split, rotate)\n",
    "    \n",
    "    #Remove the bottom of cut-off images\n",
    "    array[np.where(array == -32768)] = 0\n",
    "    #Cap pixels at the mean plus one standard deviation\n",
    "    array[np.where(array > mean_stddev_skew_median_dict[planet][wavelength][0] \n",
    "                   + mean_stddev_skew_median_dict[planet][wavelength][1])] = (\n",
    "        mean_stddev_skew_median_dict[planet][wavelength][0] \n",
    "        + mean_stddev_skew_median_dict[planet][wavelength][1])\n",
    "    \n",
    "    return np.nansum(array)\n",
    "\n",
    "def getNumberOfPixels(filepath, split=no_split, rotate=no_rotate):\n",
    "    '''\n",
    "    `filepath` is a path to a FITS file\n",
    "\n",
    "    `split` is the section of the array that is cared about (see section \"Split and Rotate\")\n",
    "\n",
    "    `rotate` is the number of degrees to rotate the array by (counter clockwise)\n",
    "    \n",
    "    Returns the number of non-nan pixels\n",
    "    '''\n",
    "    \n",
    "    return np.count_nonzero(~np.isnan(getRotatedAndSplitArrayFromFitsFile(filepath, split, rotate)))\n",
    "\n",
    "def getExpTime(filepath):\n",
    "    '''\n",
    "    `filepath` is a path to a FITS file\n",
    "\n",
    "    Note: the HMI continuum doesn't have an exposure time, so returns a 1\n",
    "\n",
    "    Returns the saved exposure time of the file. This may be inaccurate, depending on the processing already done on the file\n",
    "    '''\n",
    "    if getWavelengthFromFile(filepath) == \"cont\":\n",
    "        return 1.0\n",
    "    return getHeaderFromFitsFile(filepath)[1][\"EXPTIME\"]\n",
    "\n",
    "def getActualExpTime(filepath):\n",
    "    '''\n",
    "    `filepath` is a path to a FITS file\n",
    "\n",
    "    Returns the actual exposure time of the file, based on the time stamps in the file header\n",
    "    '''\n",
    "    return (sunpy.time.parse_time(getHeaderFromFitsFile(filepath)[1][\"T_OBS\"]) - \n",
    "            sunpy.time.parse_time(getHeaderFromFitsFile(filepath)[1][\"DATE-OBS\"])).microseconds / 1000. / 1000. * 2.0\n",
    "\n",
    "def getCountsPerSecond(filepath, split=no_split, rotate=no_rotate):\n",
    "    '''\n",
    "    `filepath` is a path to a FITS file\n",
    "\n",
    "    `split` is the section of the array that is cared about (see section \"Split and Rotate\")\n",
    "\n",
    "    `rotate` is the number of degrees to rotate the array by (counter clockwise)\n",
    "\n",
    "    Returns the number of counts per second from the file, applying the split and rotation angle\n",
    "\n",
    "    '''\n",
    "    \n",
    "    return getTotalData(filepath, split=split, rotate=rotate)/getExpTime(filepath)\n",
    "\n",
    "def getTimeSinceStartOfTransit(planet, filepath):\n",
    "    '''\n",
    "    `filepath` is a path to a FITS file\n",
    "\n",
    "    Returns the time in seconds since the start of the recorded transit period\n",
    "    '''\n",
    "    \n",
    "    planet = getPlanetFromFile(filepath)\n",
    "    \n",
    "    if planet == \"venus\":\n",
    "        start_time = datetime.datetime(2012,6,5,21,0,0,0)\n",
    "    else:\n",
    "        start_time = datetime.datetime(2016,5,9,10,0,0,0)\n",
    "        \n",
    "    current_time = sunpy.time.parse_time(getHeaderFromFitsFile(filepath)[1]['T_OBS'])\n",
    "\n",
    "    return (current_time - start_time).seconds\n",
    "\n",
    "def getDistanceToSunFromSDO(filepath):\n",
    "    '''\n",
    "    `filepath` is a path to a FITS file\n",
    "\n",
    "    Returns distance of the sun to the SDO in meters (DSUN_OBS)\n",
    "    '''\n",
    "    \n",
    "    head = getHeaderFromFitsFile(filepath)[1]\n",
    "    \n",
    "    return head[\"DSUN_OBS\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def getDataMeanStdDevSkewMedian(filepath):\n",
    "    '''\n",
    "    `filepath` is a path to a FITS file\n",
    "\n",
    "    Returns tuple of (data mean, \n",
    "                      standard deviation, \n",
    "                      data skew, \n",
    "                      data median)\n",
    "    '''\n",
    "    \n",
    "    head = getHeaderFromFitsFile(filepath)[1]\n",
    "    \n",
    "    return (head[\"DATAMEAN\"],head[\"DATARMS\"],head[\"DATASKEW\"],head[\"DATAMEDN\"])\n",
    "\n",
    "def getDataMean(filepath):\n",
    "    '''\n",
    "    `filepath` is a path to a FITS file\n",
    "\n",
    "    Returns the data mean from the file header\n",
    "    '''\n",
    "    \n",
    "    return getDataMeanStdDevSkewMedian(filepath)[0]\n",
    "\n",
    "def getDataRMS(filepath):\n",
    "    '''\n",
    "    `filepath` is a path to a FITS file\n",
    "\n",
    "    Returns the data rms from the file header\n",
    "    '''\n",
    "    \n",
    "    return getDataMeanStdDevSkewMedian(filepath)[1]\n",
    "\n",
    "def getDataSkew(filepath):\n",
    "    '''\n",
    "    `filepath` is a path to a FITS file\n",
    "\n",
    "    Returns the data skew from the file header\n",
    "    '''\n",
    "    \n",
    "    return getDataMeanStdDevSkewMedian(filepath)[2]\n",
    "\n",
    "def getDataMedian(filepath):\n",
    "    '''\n",
    "    `filepath` is a path to a FITS file\n",
    "\n",
    "    Returns the data median from the file header\n",
    "    '''\n",
    "    \n",
    "    return getDataMeanStdDevSkewMedian(filepath)[3]\n",
    "\n",
    "def getMeanStddevSkewMedian(planet, wavelength):\n",
    "    '''\n",
    "    `planet` is the planet transit to use\n",
    "\n",
    "    `wavelength` is the wavelength to use\n",
    "    \n",
    "    Returns a tuple of (average data mean, \n",
    "                        average standard deviation, \n",
    "                        average data skew, \n",
    "                        average data median)\n",
    "        for the wavelength\n",
    "    '''\n",
    "    total_mean = 0\n",
    "    total_dev = 0\n",
    "    total_median = 0\n",
    "    total_skew = 0\n",
    "    list_of_files = readListOfFiles(planet, wavelength)\n",
    "    num_files = len(list_of_files)\n",
    "\n",
    "    if num_files == 0:\n",
    "        return (0,0,0,0)\n",
    "    \n",
    "    for filepath in list_of_files:\n",
    "        tup = getDataMeanStdDevSkewMedian(filepath)\n",
    "        total_mean += tup[0]\n",
    "        total_dev += tup[1]\n",
    "        total_skew += tup[2]\n",
    "        total_median += tup[3]\n",
    "\n",
    "    return (total_mean/num_files, total_dev/num_files, total_skew/num_files, total_median/num_files)\n",
    "\n",
    "def getDictOfMeanStddevSkewMedian(planet):\n",
    "    '''\n",
    "    `planet` is the planet transit to use\n",
    "\n",
    "    Returns a dictionary of {\n",
    "                             wavelength : (average data mean, \n",
    "                                           average standard deviation, \n",
    "                                           average data skew, \n",
    "                                           average data median)\n",
    "                            }\n",
    "        for all wavelengths\n",
    "    '''\n",
    "    wave_dict = {}\n",
    "    for wavelength in wavelengths:\n",
    "        wave_dict[wavelength] = getMeanStddevSkewMedian(planet, wavelength)\n",
    "        \n",
    "    return wave_dict\n",
    "\n",
    "def saveMeanStddevSkewMedianDict(planet):\n",
    "    '''\n",
    "    `planet` is the planet transit to use\n",
    "\n",
    "    Saves a dictionary of {\n",
    "                           wavelength : (average data mean, \n",
    "                                         average standard deviation, \n",
    "                                         average data skew, \n",
    "                                         average data median)\n",
    "                          }\n",
    "        to `personal_data_directory``planet`mean_stddev_skew_median.txt\n",
    "        for all wavelengths\n",
    "    '''\n",
    "    writeToFile(personal_data_directory + planet + \"mean_stddev_skew_median.txt\", getDictOfMeanStddevSkewMedian(planet))\n",
    "    \n",
    "def readMeanStddevSkewMedianDict(planet):\n",
    "    '''\n",
    "    `planet` is the planet transit to use\n",
    "\n",
    "    Returns the saved dictionary of {\n",
    "                                     wavelength : (average data mean, \n",
    "                                                   average standard deviation, \n",
    "                                                   average data skew, \n",
    "                                                   average data median)\n",
    "                                    }\n",
    "        from `personal_data_directory``planet`mean_stddev_skew_median.txt\n",
    "        for all wavelengths\n",
    "    '''\n",
    "    return readFromFile(personal_data_directory + planet + \"mean_stddev_skew_median.txt\")\n",
    "\n",
    "mean_stddev_skew_median_dict = {\"venus\": readMeanStddevSkewMedianDict(\"venus\"), \n",
    "                                \"mercury\": readMeanStddevSkewMedianDict(\"mercury\")}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def readListOfFiles(planet, wavelength, time_block=[]):\n",
    "    '''\n",
    "    `planet` is the planet transit to use\n",
    "\n",
    "    `wavelength` is the wavelength to use\n",
    "\n",
    "    `time_block` is an array of strings to limit the list to (see section \"Time\")\n",
    "    \n",
    "    Reads the saved data for that wavelength and gets the list of files from that\n",
    "    '''\n",
    "    checkIfInWavelengths(wavelength)\n",
    "    \n",
    "    file_list = []\n",
    "    for filepath, (count, time, pixels, count_in_split, pixels_in_split, count_out_of_split, \n",
    "                   pixels_out_of_split) in readData(planet, wavelength):\n",
    "        if checkTimeBlock(planet, time_block, time):\n",
    "            file_list.append(filepath)\n",
    "    return file_list\n",
    "\n",
    "def readListOfFilesAndTimeSinceStartOfTransit(planet, wavelength, time_block=[]):\n",
    "    '''\n",
    "    `planet` is the planet transit to use\n",
    "\n",
    "    `wavelength` is the wavelength to use\n",
    "\n",
    "    `time_block` is an array of strings to limit the list to (see section \"Time\")\n",
    "    \n",
    "    Reads the saved data for that wavelength and gets the list of (file, time) from that\n",
    "    '''\n",
    "    checkIfInWavelengths(wavelength)\n",
    "    \n",
    "    file_and_time_list = []\n",
    "    for filepath, (count, time, pixels, count_in_split, pixels_in_split, count_out_of_split, \n",
    "                   pixels_out_of_split) in readData(planet, wavelength):\n",
    "        if checkTimeBlock(planet, time_block, time):\n",
    "            file_and_time_list.append((filepath, time))\n",
    "    return file_and_time_list\n",
    "\n",
    "def readExposureTimeOfFiles(planet, wavelength, time_block=[]):\n",
    "    '''\n",
    "    `planet` is the planet transit to use\n",
    "\n",
    "    `wavelength` is the wavelength to use\n",
    "\n",
    "    `time_block` is an array of strings to limit the list to (see section \"Time\")\n",
    "    \n",
    "    Reads the saved data for that wavelength and gets the list of counts per second from that\n",
    "    '''\n",
    "    \n",
    "    checkIfInWavelengths(wavelength)\n",
    "        \n",
    "    exptime_list = []\n",
    "    for filepath, (count, time, pixels, count_in_split, pixels_in_split, count_out_of_split, \n",
    "                   pixels_out_of_split) in readData(planet, wavelength):\n",
    "                        \n",
    "        if checkTimeBlock(planet, time_block, time):\n",
    "            exptime_list.append(getExpTime(filepath))\n",
    "        \n",
    "    return exptime_list\n",
    "\n",
    "def readActualExposureTimeOfFiles(planet, wavelength, time_block=[]):\n",
    "    '''\n",
    "    `planet` is the planet transit to use\n",
    "\n",
    "    `wavelength` is the wavelength to use\n",
    "\n",
    "    `time_block` is an array of strings to limit the list to (see section \"Time\")\n",
    "    \n",
    "    Reads the saved data for that wavelength and gets the list of counts per second from that\n",
    "    '''\n",
    "    \n",
    "    checkIfInWavelengths(wavelength)\n",
    "        \n",
    "    exptime_list = []\n",
    "    for filepath, (count, time, pixels, count_in_split, pixels_in_split, count_out_of_split, \n",
    "                   pixels_out_of_split) in readData(planet, wavelength):\n",
    "                        \n",
    "        if checkTimeBlock(planet, time_block, time):\n",
    "            exptime_list.append(getActualExpTime(filepath))\n",
    "        \n",
    "    return exptime_list\n",
    "\n",
    "def readDistancefSunOfFiles(planet, wavelength, time_block=[]):\n",
    "    '''\n",
    "    `planet` is the planet transit to use\n",
    "\n",
    "    `wavelength` is the wavelength to use\n",
    "\n",
    "    `time_block` is an array of strings to limit the list to (see section \"Time\")\n",
    "    \n",
    "    Returns a list of the distance to the sun for all the files in the given wavelength\n",
    "    '''\n",
    "\n",
    "    checkIfInWavelengths(wavelength)\n",
    "    distance_list = []\n",
    "    for filepath, (count, time, pixels, count_in_split, pixels_in_split, count_out_of_split, \n",
    "                   pixels_out_of_split) in readData(planet, wavelength):\n",
    "        \n",
    "        if checkTimeBlock(planet, time_block, time):\n",
    "            distance_list.append(getDistanceToSunFromSDO(filepath))\n",
    "\n",
    "    return distance_list\n",
    "\n",
    "def readDataMeansOfFiles(planet, wavelength, time_block=[]):\n",
    "    '''\n",
    "    `planet` is the planet transit to use\n",
    "\n",
    "    `wavelength` is the wavelength to use\n",
    "\n",
    "    `time_block` is an array of strings to limit the list to (see section \"Time\")\n",
    "    \n",
    "    Returns a list of the data means for all the files in the given wavelength\n",
    "    '''\n",
    "\n",
    "    checkIfInWavelengths(wavelength)\n",
    "    mean_list = []\n",
    "    for filepath, (count, time, pixels, count_in_split, pixels_in_split, count_out_of_split, \n",
    "                   pixels_out_of_split) in readData(planet, wavelength):\n",
    "        if checkTimeBlock(planet, time_block, time):\n",
    "            mean_list.append(getDataMean(filepath))\n",
    "        \n",
    "    return mean_list\n",
    "\n",
    "def readTimeSinceStartOfTransitOfFiles(planet, wavelength, time_block=[]):\n",
    "    '''\n",
    "    `planet` is the planet transit to use\n",
    "\n",
    "    `wavelength` is the wavelength to use\n",
    "\n",
    "    `time_block` is an array of strings to limit the list to (see section \"Time\")\n",
    "    \n",
    "    Reads the saved data for that wavelength and gets the list of times from that\n",
    "    '''\n",
    "    checkIfInWavelengths(wavelength)\n",
    "    time_list = []\n",
    "    for filepath, (count, time, pixels, count_in_split, pixels_in_split, count_out_of_split, \n",
    "                   pixels_out_of_split) in readData(planet, wavelength):\n",
    "        \n",
    "        if checkTimeBlock(planet, time_block, time):\n",
    "            time_list.append(time)\n",
    "           \n",
    "    return time_list\n",
    "def getNumberOfPixelsOfFiles(planet, wavelength, split=no_split, rotate=no_rotate, time_block=[]):\n",
    "    '''\n",
    "    `planet` is the planet transit to use\n",
    "\n",
    "    `wavelength` is the wavelength to use\n",
    "\n",
    "    `split` is the section of the array that is cared about (see section \"Split and Rotate\")\n",
    "\n",
    "    `rotate` is the number of degrees to rotate the array by (counter clockwise)\n",
    "\n",
    "    `time_block` is an array of strings to limit the list to (see section \"Time\")\n",
    "    \n",
    "    Returns a list of the number of pixels for all the files in the given wavelength\n",
    "    '''\n",
    "    \n",
    "    checkIfInWavelengths(wavelength)\n",
    "    \n",
    "    pixel_count_list = []\n",
    "    for filepath, (count, time, pixels, count_in_split, pixels_in_split, count_out_of_split, \n",
    "                   pixels_out_of_split) in readData(planet, wavelength):\n",
    "                    \n",
    "        if checkTimeBlock(planet, time_block, time):\n",
    "            if split == no_split:\n",
    "                thingToAppend = pixels\n",
    "            elif split == inside_split and rotate == transit_rotate:\n",
    "                thingToAppend = pixels_in_split\n",
    "            elif split == outside_split and rotate == transit_rotate:\n",
    "                thingToAppend = pixels_out_of_split\n",
    "            else:\n",
    "                thingToAppend = getNumberOfPixels(filepath, split, rotate)\n",
    "            \n",
    "            pixel_count_list.append(thingToAppend)\n",
    "                    \n",
    "    return pixel_count_list\n",
    "\n",
    "def readCountsPerSecondOfFiles(planet, wavelength, split=no_split, rotate=no_rotate, time_block=[]):\n",
    "    '''\n",
    "    `planet` is the planet transit to use\n",
    "\n",
    "    `wavelength` is the wavelength to use\n",
    "\n",
    "    `split` is the section of the array that is cared about (see section \"Split and Rotate\")\n",
    "\n",
    "    `rotate` is the number of degrees to rotate the array by (counter clockwise)\n",
    "\n",
    "    `time_block` is an array of strings to limit the list to (see section \"Time\")\n",
    "    \n",
    "    Reads the saved data for that wavelength and gets the list of counts per second from that\n",
    "    '''\n",
    "    \n",
    "    checkIfInWavelengths(wavelength)\n",
    "        \n",
    "    count_list = []\n",
    "    for filepath, (count, time, pixels, count_in_split, pixels_in_split, count_out_of_split, \n",
    "                   pixels_out_of_split) in readData(planet, wavelength):\n",
    "                        \n",
    "        if checkTimeBlock(planet, time_block, time):\n",
    "            if split == no_split:\n",
    "                thingToAppend = count\n",
    "            elif split == inside_split and rotate == transit_rotate:\n",
    "                thingToAppend = count_in_split\n",
    "            elif split == outside_split and rotate == transit_rotate:\n",
    "                thingToAppend = count_out_of_split\n",
    "            else:\n",
    "                thingToAppend = getCountsPerSecond(filepath, split, rotate)\n",
    "            count_list.append(thingToAppend)\n",
    "        \n",
    "    return count_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def findGoodFiles(planet, hour, wavelength):\n",
    "    '''\n",
    "    `planet` is the planet transit to use\n",
    "\n",
    "    `hour` is what hour of the transit to get the data from, from \"00\" to \"23\"\n",
    "\n",
    "    `wavelength` is the wavelength to use\n",
    "        \n",
    "    Returns a list of non cutoff files\n",
    "    '''\n",
    "    \n",
    "    checkIfInHours(hour, planet)\n",
    "    \n",
    "    if planet == \"venus\":\n",
    "        year = \"2012\"\n",
    "        month = \"06\"\n",
    "        if hour in [\"21\",\"22\",\"23\"]:\n",
    "            day = \"05\"\n",
    "        else:\n",
    "            day = \"06\"\n",
    "    else:\n",
    "        year = \"2016\"\n",
    "        month = \"05\"\n",
    "        day = \"09\"\n",
    "    \n",
    "    checkIfInWavelengths(wavelength)\n",
    "    \n",
    "    if wavelength == \"cont\":\n",
    "        #Goes to the directory for the given hour and day\n",
    "        \n",
    "        os.chdir(hmi_data_directory + year + '/'  + month + '/' + day \n",
    "                                                       + '/H' + hour  + '00')\n",
    "        #Gets a list populated with all the files from the directory\n",
    "        files = filter(os.path.isfile, os.listdir( os.curdir ) )\n",
    "        \n",
    "        #Make sure that there were files on that day\n",
    "        if files == []:\n",
    "            print \"No files found on day: \" + day + \" hour: \" + hour\n",
    "            \n",
    "        good_files = []\n",
    "        \n",
    "        for filepath in files:\n",
    "            good_files.append(hmi_data_directory + year + '/' + month + '/' + day + '/H' + hour + '00/' + filepath)\n",
    "\n",
    "        return good_files\n",
    "    else:\n",
    "        #Goes to the directory for the given hour and day\n",
    "        os.chdir(aia_data_directory + year + '/' + month + '/' + day + '/H' + hour + '00')\n",
    "        #Gets a list populated with all the files from the directory\n",
    "        files = filter(os.path.isfile, os.listdir( os.curdir ) )\n",
    "\n",
    "        #Looks for all the files that are for the requested wavelength\n",
    "        correct_wavelength_files = []\n",
    "        for filepath in files:\n",
    "            if wavelength + \".fits\" in filepath:\n",
    "                correct_wavelength_files.append(filepath)\n",
    "\n",
    "        #Make sure that there were files on that day\n",
    "        if correct_wavelength_files == []:\n",
    "            print \"No files found on day: \" + day + \" hour: \" + hour\n",
    "            return []\n",
    "\n",
    "        good_files = []\n",
    "\n",
    "        #Goes through all the files for the wavelength and puts non-cutoff ones \n",
    "        #     in an array to be returned\n",
    "        for i,filepath in enumerate(correct_wavelength_files):        \n",
    "            if wavelength == \"1600\": \n",
    "                #All the 1600 wavelength images are cut off, \n",
    "                #     so we'll work with all of the cut off ones\n",
    "                if getHeaderFromFitsFile(filepath)[1]['MISSVALS'] != 0:\n",
    "                    good_files.append(aia_data_directory + year + '/' + month + '/' + day + '/H' + hour + '00/' + filepath)\n",
    "            else:\n",
    "                #Check if the sum of the data array is greater than zero\n",
    "                #This removes all cutoff images, \n",
    "                #     since the black pixels are all equal to -2^15\n",
    "                if getHeaderFromFitsFile(filepath)[1]['MISSVALS'] == 0:\n",
    "                    good_files.append(aia_data_directory + year + '/' + month + '/' + day + '/H' + hour + '00/' + filepath)\n",
    "\n",
    "        return good_files\n",
    "    \n",
    "def getAllGoodFiles(planet, wavelength):\n",
    "    '''\n",
    "    `planet` is the planet transit to use\n",
    "\n",
    "    `wavelength` is the wavelength to use\n",
    "\n",
    "    Returns a list of all the non cutoff files\n",
    "    '''\n",
    "    checkIfInWavelengths(wavelength)\n",
    "    \n",
    "    return list(map(lambda hour: findGoodFiles(planet,hour,wavelength), hours[planet]))\n",
    "    \n",
    "def getDictOfAllGoodFilesAllThings(planet, wavelength):\n",
    "    '''\n",
    "    `planet` is the planet transit to use\n",
    "\n",
    "    `wavelength` is the wavelength to use\n",
    "    \n",
    "    Returns a dictionary of             \n",
    "        filepath, (count, time, pixels, count_in_split, pixels_in_split, count_out_of_split, pixels_out_of_split)\n",
    "    '''\n",
    "    checkIfInWavelengths(wavelength)\n",
    "    \n",
    "    good_files = readListOfFiles(planet, wavelength)\n",
    "    if good_files == []:\n",
    "        good_files = getAllGoodFiles(planet, wavelength)\n",
    "            \n",
    "    d = {}\n",
    "    \n",
    "    for filepath in good_files:\n",
    "        d[filepath] = (getCountsPerSecond(filepath),getTimeSinceStartOfTransit(planet, filepath),\n",
    "                       getNumberOfPixels (filepath, no_split,      no_rotate),\n",
    "                       getCountsPerSecond(filepath, inside_split,  transit_rotate),\n",
    "                       getNumberOfPixels (filepath, inside_split,  transit_rotate),\n",
    "                       getCountsPerSecond(filepath, outside_split, transit_rotate),\n",
    "                       getNumberOfPixels (filepath, outside_split, transit_rotate))\n",
    "    return d\n",
    "\n",
    "def saveDictData(planet, wavelength, directory=personal_data_directory):\n",
    "    '''\n",
    "    `planet` is the planet transit to use\n",
    "\n",
    "    `wavelength` is the wavelength to use\n",
    "\n",
    "    `directory` is the directory to save the datafile to\n",
    "    \n",
    "    Saves a dictionary of \n",
    "        all the good files for the wavelength : \n",
    "            (their counts per second, their time since start of transit)\n",
    "        to the data file for that wavelength\n",
    "        \n",
    "    The saved data file is named `planet`data`wavelength`.txt\n",
    "\n",
    "    Returns the dictionary\n",
    "    '''\n",
    "    \n",
    "    checkIfInWavelengths(wavelength)\n",
    "    \n",
    "    return writeToFile(directory + planet + 'data' + wavelength + '.txt', \n",
    "                       getDictOfAllGoodFilesAllThings(planet, wavelength))\n",
    "    \n",
    "def saveAllWavelengthData(planet, directory=personal_data_directory):\n",
    "    '''\n",
    "    `planet` is the planet transit to use\n",
    "\n",
    "    `directory` is the directory to save the datafile to\n",
    "\n",
    "    For all wavelengths, saves a dictionary of \n",
    "        all the good files for the wavelength : \n",
    "            (their counts per second, their time since start of transit)\n",
    "        to the data file for that wavelength\n",
    "\n",
    "        Each saved data file is named `planet`data`wavelength`.txt\n",
    "    '''\n",
    "        \n",
    "    for wavelength in wavelengths:\n",
    "        saveDictData(planet, wavelength, directory)\n",
    "        print \"Saved data for\" + wavelength\n",
    "            \n",
    "def saveAllWavelengthDataMultiThreaded(planet):\n",
    "    '''\n",
    "    `planet` is the planet transit to use\n",
    "\n",
    "    For all wavelengths, saves a dictionary of \n",
    "        all the good files for the wavelength : \n",
    "            (their counts per second, their time since start of transit)\n",
    "        to the data file for that wavelength\n",
    "\n",
    "        Each saved data file is named `planet`data`wavelength`.txt\n",
    "\n",
    "    Does this in multiple threads (see section \"Multithreading\")\n",
    "\n",
    "    Due to limitations of multithreading, the directory to save data to is the default one from saveDictData\n",
    "    '''\n",
    "\n",
    "    #It's so repetitive because putting stuff in a %job does weird stuff with the arguments\n",
    "    if planet == \"venus\":\n",
    "        %job [saveDictData(\"venus\", \"0094\")]\n",
    "        %job [saveDictData(\"venus\", \"0131\")]\n",
    "        %job [saveDictData(\"venus\", \"0171\")]\n",
    "        %job [saveDictData(\"venus\", \"0193\")]\n",
    "        %job [saveDictData(\"venus\", \"0211\")]\n",
    "        %job [saveDictData(\"venus\", \"0304\")]\n",
    "        %job [saveDictData(\"venus\", \"0335\")]\n",
    "        %job [saveDictData(\"venus\", \"1600\")]\n",
    "        %job [saveDictData(\"venus\", \"1700\")]\n",
    "        %job [saveDictData(\"venus\", \"cont\")]\n",
    "    else:\n",
    "        %job [saveDictData(\"mercury\", \"0094\")]\n",
    "        %job [saveDictData(\"mercury\", \"0131\")]\n",
    "        %job [saveDictData(\"mercury\", \"0171\")]\n",
    "        %job [saveDictData(\"mercury\", \"0193\")]\n",
    "        %job [saveDictData(\"mercury\", \"0211\")]\n",
    "        %job [saveDictData(\"mercury\", \"0304\")]\n",
    "        %job [saveDictData(\"mercury\", \"0335\")]\n",
    "        %job [saveDictData(\"mercury\", \"1600\")]\n",
    "        %job [saveDictData(\"mercury\", \"1700\")]\n",
    "        %job [saveDictData(\"mercury\", \"cont\")]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def graphLightCurve(planet, wavelength, split=no_split, rotate=no_rotate, popt_one=[], func_one=None, popt_two=[], \n",
    "                    func_two=None, \n",
    "                    time_block=[], \n",
    "                    show_events=False, label=\"\", wavelength_name=True, new_figure=True, scale_to_one=True, shift=0, \n",
    "                    scale_to_one_based_on=\"max\", \n",
    "                    fontsize=10, numbersize=10, show_graph=True, subplot=None):\n",
    "    '''\n",
    "    `planet` is the planet transit to use\n",
    "    \n",
    "    `wavelength` is the wavelength to use\n",
    "\n",
    "    `split` is the section of the array that is cared about (see section \"Split and Rotate\")\n",
    "\n",
    "    `rotate` is the number of degrees to rotate the array by (counter clockwise)\n",
    "    \n",
    "    `func_one` is a function (number, parameters) -> number that is applied to the light curve before graphing it\n",
    "\n",
    "    `popt_one` is an array of optimized parameters for `func_one`\n",
    "\n",
    "    `func_two` is a function (number, parameters) -> number that is applied to the light curve before graphing it\n",
    "\n",
    "    `popt_one` is an array of optimized parameters for `func_two `\n",
    "\n",
    "    `time_block` is an array of strings to limit the list to (see section \"Time\")\n",
    "    \n",
    "    If `show_events` is True, the graph of the light curve includes lines for various events that happen during the transit\n",
    "        Shows the ingress and egress with black lines\n",
    "        Shows midnight (venus transit only) with a green line\n",
    "\n",
    "    `label` is text to add to the end of the data name that goes in the legend\n",
    "\n",
    "    If `wavelength_name` is True, the data name that goes in the legend includes the name of the wavelength\n",
    "    \n",
    "    If `new_figure` is True, it creates a new matplotlib figure for the light curve\n",
    "\n",
    "    If `scale_to_one` is True, it scales the light curve to be on a scale of 0 to 1\n",
    "\n",
    "    `scale_to_one_based_on` decides how to scale the light curve, if `scale_to_one` is True\n",
    "        \"max\"         scales so that the maximum element is 1\n",
    "        \"first\"       scales so that the first element is 1\n",
    "        \"last\"        scales so that the last element is 1\n",
    "        \"not transit\" scales so that the average not transit element is 1\n",
    "\n",
    "    `shift` is added to every datapoint in the array\n",
    "    \n",
    "    `fontsize` is the size of the text in the graphed light curve (doesn't include the tick mark labels)\n",
    "    \n",
    "    `numbersize` is the size of the tick mark labels in the graphed light curve\n",
    "\n",
    "    If `show_graph` is True, it graphs the light curve\n",
    "    \n",
    "    `subplot` is the argument passed into matplotlib's subplot function. If `subplot == None` then subplot is not used.\n",
    "    \n",
    "    Returns a tuple of the array of times and the array of data values that were graphed\n",
    "    '''\n",
    "    \n",
    "    def scaleDataToOutOfOne(planet, data, based_on=\"max\", times=None):\n",
    "        '''\n",
    "        `planet` is the planet transit to use\n",
    "\n",
    "        `data` is the array of data to scale\n",
    "\n",
    "        Scales the given data to be on a scale of 0 to 1. \n",
    "\n",
    "        If `based_on` is not given or is set to \"max\", it scales to 0 to 1\n",
    "\n",
    "        If `based_on` is \"first\", it scales so that the first element is 1\n",
    "\n",
    "        If `based_on` is \"last\", it scales so that the last element is 1\n",
    "\n",
    "        If `based_on` is \"not transit\", it scales so that the average not transit element is 1\n",
    "            If this is selected, you must also pass in an array of time elements into `times`\n",
    "        '''\n",
    "\n",
    "        #Change to be out of 1\n",
    "        if based_on==\"max\":\n",
    "            max_value = max(data)\n",
    "        elif based_on==\"first\":\n",
    "            max_value = data[0]\n",
    "        elif based_on==\"last\":\n",
    "            max_value = data[-1]\n",
    "        elif based_on==\"not transit\":\n",
    "            if time==None:\n",
    "                raise Exception(\"Time parameter needs to be set for scaleDataToOutOfOne with parameter \\\"not transit\\\"\")\n",
    "            good_data = getDataFromTimeBlock(planet, [\"not transit\"], data, times)\n",
    "            max_value = np.average(good_data)\n",
    "        else:\n",
    "            raise Exception(\"Couldn't scale to out of one because based_on parameter \" \n",
    "                            + based_on + \" is invalid for scaleDataToOutOfOne\")\n",
    "        for i, data_point in enumerate(data):\n",
    "            data[i] = data_point / max_value\n",
    "\n",
    "        return data    \n",
    "    \n",
    "    checkIfInWavelengths(wavelength)\n",
    "    \n",
    "    #Gets all of the sum data points, then divides them by the number of pixels to have averages\n",
    "    pts = np.divide(readCountsPerSecondOfFiles(planet, wavelength, split, rotate, time_block=time_block),\n",
    "                    getNumberOfPixelsOfFiles(planet, wavelength, split, rotate, time_block=time_block))\n",
    "    \n",
    "    #Gets the list of times\n",
    "    times = readTimeSinceStartOfTransitOfFiles(planet, wavelength, time_block=time_block)\n",
    "    \n",
    "    #Calculate the numbers of sections that make up the split and select either the HMI or AIA split\n",
    "    if split != None:\n",
    "        split = split[planet]\n",
    "\n",
    "        if wavelength == \"cont\":\n",
    "            split = split[1]\n",
    "        else:\n",
    "            split = split[0]\n",
    "        split_sections = split[0] * 1.0 / len(split[1]) \n",
    "    else:\n",
    "        split_sections = 1\n",
    "\n",
    "    if len(popt_one) == 0 and len(popt_two) == 0:\n",
    "        if scale_to_one:\n",
    "            pts = scaleDataToOutOfOne(planet, pts, based_on=scale_to_one_based_on, times=times)\n",
    "            \n",
    "        for i, data in enumerate(pts):\n",
    "            #The division needs to be done in the last step.\n",
    "            #If there is another curve fit, do it there, else do it here.\n",
    "            pts[i] = pts[i]/split_sections + ((split_sections - 1.0)/ split_sections)\n",
    "    else:\n",
    "        #Subtracts the first curve fit from the data\n",
    "        if len(popt_one) != 0:\n",
    "            for i, data in enumerate(pts):\n",
    "                #The division needs to be done in the last step.\n",
    "                #If there is another curve fit, do it there, else do it here.\n",
    "                if len(popt_two) != 0:\n",
    "                    pts[i] = pts[i] - func_one(timeThroughTransit(planet, times[i]),*popt_one) + popt_one[-1]\n",
    "                else:\n",
    "                    pts[i] = (pts[i] - func_one(timeThroughTransit(planet, times[i]),*popt_one))/split_sections+popt_one[-1]\n",
    "\n",
    "        #Subtracts the second curve fit from the data\n",
    "        if len(popt_two) != 0:\n",
    "            for i, data in enumerate(pts):\n",
    "                pts[i] = (pts[i] - func_two(timeThroughTransit(planet, times[i]),*popt_two))/split_sections + popt_two[-1]\n",
    "        if scale_to_one:\n",
    "            pts = scaleDataToOutOfOne(planet, pts, based_on=scale_to_one_based_on, times=times)\n",
    "\n",
    "    #apply shift\n",
    "    pts = np.add(pts, shift)\n",
    "\n",
    "    if show_graph:\n",
    "        if wavelength_name:\n",
    "            #Name the graph based on the wavelength\n",
    "            if wavelength == \"cont\":\n",
    "                data_name = \"HMI Continuum\"\n",
    "            else:\n",
    "                data_name = \"AIA \" + wavelength + \"Å\"  \n",
    "        else:\n",
    "            data_name = \"\"\n",
    "        data_name += \" \" + label\n",
    "        \n",
    "        #Prepare the canvas\n",
    "        if new_figure:\n",
    "            fig = plt.figure()\n",
    "        else:\n",
    "            fig = plt.gcf()\n",
    "        if subplot == None:\n",
    "            axes = plt.gca()\n",
    "        else:\n",
    "            axes = fig.add_subplot(subplot)\n",
    "            \n",
    "        #Graph the light curve\n",
    "        light_curve = LightCurve.create({data_name: pts}, index = times)    \n",
    "        light_curve.plot(fontsize=numbersize)\n",
    "\n",
    "        if new_figure:\n",
    "            plt.xlabel('Time (seconds)', fontsize=fontsize)\n",
    "            plt.ylabel('Brightness (percent)', fontsize=fontsize)\n",
    "\n",
    "            if planet == \"venus\":\n",
    "                fig.suptitle('Percent Brightness During the 2012 Venus Transit', fontsize=fontsize)\n",
    "            else:\n",
    "                fig.suptitle('Percent Brightness During the 2016 Mercury Transit', fontsize=fontsize)\n",
    "\n",
    "        plt.legend(loc=2,prop={'size':fontsize})\n",
    "\n",
    "        #Graph the different activities\n",
    "        if show_events:\n",
    "            #Ingress\n",
    "            plt.plot([ingress_start_time[planet], ingress_start_time[planet]], [axes.get_ylim()[0], axes.get_ylim()[1]], 'k-')\n",
    "            plt.plot([ingress_end_time[planet], ingress_end_time[planet]], [axes.get_ylim()[0], axes.get_ylim()[1]], 'k-')\n",
    "\n",
    "            #Egress\n",
    "            plt.plot([egress_start_time[planet], egress_start_time[planet]], [axes.get_ylim()[0], axes.get_ylim()[1]], 'k-')\n",
    "            plt.plot([egress_end_time[planet], egress_end_time[planet]], [axes.get_ylim()[0], axes.get_ylim()[1]], 'k-')\n",
    "\n",
    "            if planet == \"venus\":\n",
    "                #Midnight\n",
    "                plt.plot([10800, 10800], [axes.get_ylim()[0], axes.get_ylim()[1]], 'g-')\n",
    "    \n",
    "    return (times, pts)\n",
    "\n",
    "def graphAmountCausedByDistance(planet, wavelength):\n",
    "    '''\n",
    "    `planet` is the planet transit to use\n",
    "    \n",
    "    `wavelength` is the wavelength to use\n",
    "\n",
    "    Graphs the light curve and a line representing how much of that was caused by the SDO moving in respect to the Sun\n",
    "    '''\n",
    "    \n",
    "    distance_data = readDistancefSunOfFiles(planet, wavelength)\n",
    "    times = readTimeSinceStartOfTransitOfFiles(planet, wavelength)\n",
    "    pts = np.divide(readCountsPerSecondOfFiles(planet, wavelength), getNumberOfPixelsOfFiles(planet, wavelength))\n",
    "    max_value = max(pts)\n",
    "    for i, data in enumerate(pts):\n",
    "        pts[i] = data / max_value\n",
    "    graphLightCurve(wavelength, scale_to_one=True)\n",
    "    scaled_data = (np.divide(distance_data, max(distance_data)))\n",
    "    scaled_data = np.add(scaled_data,max(scaled_data)-min(scaled_data))\n",
    "    scaled_data = np.multiply(scaled_data, pts[0])\n",
    "    light_curve = LightCurve.create({\"Amount caused by distance\": scaled_data}, index = times)    \n",
    "    light_curve.plot()\n",
    "    \n",
    "def graphLightCurveAdjusted(planet, wavelength, show_events=False, use_primary_curve_fit=True, use_secondary_curve_fit=True,\n",
    "                         label=\"\", wavelength_name=True, new_figure=True, scale_to_one=True, shift=0,\n",
    "                         force_primary_curve_fit=False, scale_to_one_based_on=\"max\", fontsize=10, numbersize=10, show_graph=True):\n",
    "    '''\n",
    "    Parameters that are also in graphLightCurve behave in the same way\n",
    "\n",
    "    If `use_primary_curve_fit` is True, uses a primary curve fit to adjust the light curve\n",
    "\n",
    "    If `use_secondary_curve_fit` is True, uses a secondary linear curve fit to adjust the light curve\n",
    "    \n",
    "    Returns a tuple of the array of times and the array of data values that were graphed\n",
    "    '''\n",
    "    checkIfInWavelengths(wavelength)\n",
    "    \n",
    "    if planet == \"venus\":\n",
    "        if wavelength in [\"1700\",\"1600\"] and not force_primary_curve_fit:\n",
    "            use_primary_curve_fit = False\n",
    "\n",
    "        if use_primary_curve_fit:\n",
    "            if wavelength in [\"cont\"]:\n",
    "                split_for_primary = outside_split\n",
    "            else:\n",
    "                split_for_primary = inside_split\n",
    "\n",
    "            #Get the times and data for the outside split\n",
    "\n",
    "            time_block = [\"transit\"]\n",
    "            if wavelength == \"0304\":\n",
    "                time_block.append(\"not venus spike\")\n",
    "\n",
    "            times = readTimeSinceStartOfTransitOfFiles(planet, wavelength, time_block=time_block)\n",
    "            data = readCountsPerSecondOfFiles(planet, wavelength, time_block=time_block,\n",
    "                                              split=split_for_primary, rotate=transit_rotate)\n",
    "\n",
    "            num_pixels = getNumberOfPixelsOfFiles(planet, wavelength, time_block=time_block,\n",
    "                                                  split=split_for_primary, rotate=transit_rotate)\n",
    "            \n",
    "            #Adjust the data to be averages instead of totals\n",
    "            data = np.divide(data, num_pixels)\n",
    "            \n",
    "            if wavelength == \"0304\":\n",
    "                func_one = quadratic\n",
    "            else:\n",
    "                func_one = cubic\n",
    "                \n",
    "            #Do curve fitting\n",
    "            popt_one, pcov = curve_fit(func_one, times, data)\n",
    "\n",
    "            if split_for_primary == outside_split:\n",
    "                if wavelength == \"cont\":\n",
    "                    num = 1\n",
    "                else:\n",
    "                    num = 0\n",
    "                popt_one = np.divide(popt_one,inside_split[planet][num][0] * 1.0 / len(inside_split[planet][num][1]))\n",
    "            \n",
    "        else:\n",
    "            popt_one=[]\n",
    "            func_one=None\n",
    "\n",
    "        if use_secondary_curve_fit:\n",
    "\n",
    "            #Gets all of the sum data points, then divides them by the number of pixels to have averages\n",
    "            data = np.divide(readCountsPerSecondOfFiles(planet, wavelength, split=inside_split, rotate=transit_rotate),\n",
    "                             getNumberOfPixelsOfFiles(planet, wavelength, split=inside_split, rotate=transit_rotate))\n",
    "\n",
    "            #Gets the list of times\n",
    "            times = readTimeSinceStartOfTransitOfFiles(planet, wavelength)\n",
    "\n",
    "            #Subtracts the curve fit from the data\n",
    "            if len(popt_one) != 0:\n",
    "                for i, data_point in enumerate(data):\n",
    "                    data[i] = data[i] - func_one(timeThroughTransit(planet,times[i]),*popt_one) + popt_one[-1]\n",
    "\n",
    "            #Get rid of transit points\n",
    "            new_times = []\n",
    "            new_data = []\n",
    "\n",
    "            for i, time in enumerate(times):\n",
    "                if time < ingress_start_time[planet] or time > egress_end_time[planet]:\n",
    "                    new_times.append(time)\n",
    "                    new_data.append(data[i])\n",
    "\n",
    "            func_two = linear\n",
    "\n",
    "            #Do curve fitting\n",
    "            popt_two, pcov = curve_fit(func_two, new_times, new_data)\n",
    "        else:\n",
    "            popt_two=[]\n",
    "            func_two=None\n",
    "\n",
    "    else:\n",
    "        if wavelength in [\"cont\",\"1700\"]:\n",
    "            if use_primary_curve_fit:\n",
    "                \n",
    "                #Gets all of the sum data points, then divides them by the number of pixels to have averages\n",
    "                data = np.divide(readCountsPerSecondOfFiles(planet, wavelength, split=inside_split, rotate=transit_rotate),\n",
    "                                 getNumberOfPixelsOfFiles(planet, wavelength, split=inside_split, rotate=transit_rotate))\n",
    "\n",
    "                #Gets the list of times\n",
    "                times = readTimeSinceStartOfTransitOfFiles(planet, wavelength)\n",
    "\n",
    "                #Get rid of transit points\n",
    "                new_times = []\n",
    "                new_data = []\n",
    "\n",
    "                for i, time in enumerate(times):\n",
    "                    if checkTimeBlock(planet,[\"not transit\"],time):\n",
    "                        new_times.append(time)\n",
    "                        new_data.append(data[i])\n",
    "\n",
    "                func_one = linear\n",
    "\n",
    "                #Do curve fitting\n",
    "                popt_one, pcov = curve_fit(func_one, new_times, new_data)\n",
    "            else:\n",
    "                popt_one=[]\n",
    "                func_one=None\n",
    "            if use_secondary_curve_fit:\n",
    "                \n",
    "                #Gets all of the sum data points, then divides them by the number of pixels to have averages\n",
    "                data = np.divide(readCountsPerSecondOfFiles(planet, wavelength, split=inside_split, rotate=transit_rotate),\n",
    "                                 getNumberOfPixelsOfFiles(planet, wavelength, split=inside_split, rotate=transit_rotate))\n",
    "\n",
    "                #Gets the list of times\n",
    "                times = readTimeSinceStartOfTransitOfFiles(planet, wavelength)\n",
    "\n",
    "                \n",
    "                #Subtracts the curve fit from the data\n",
    "                if len(popt_one) != 0:\n",
    "                    for i, data_point in enumerate(data):\n",
    "                        data[i] = data[i] - func_one(timeThroughTransit(planet,times[i]),*popt_one) + popt_one[-1]\n",
    "\n",
    "                #Get rid of transit points\n",
    "                new_times = []\n",
    "                new_data = []\n",
    "\n",
    "                for i, time in enumerate(times):\n",
    "                    if checkTimeBlock(planet,[\"transit\"],time):\n",
    "                        new_times.append(time)\n",
    "                        new_data.append(data[i])\n",
    "\n",
    "                func_two = linear\n",
    "\n",
    "                #Do curve fitting\n",
    "                popt_two, pcov = curve_fit(func_two, new_times, new_data)\n",
    "            else:\n",
    "                popt_two=[]\n",
    "                func_two=None\n",
    "        else:\n",
    "            print \"graph adjusted transit doesn't work for mercury wavelength \" + wavelength + \"Å\"\n",
    "            popt_one=[]\n",
    "            func_one=None\n",
    "            popt_two=[]\n",
    "            func_two=None\n",
    "    \n",
    "    #Actually graph the light curve, using the optimal values from the curve fit. \n",
    "    #It only uses the proper split\n",
    "    return graphLightCurve(planet, wavelength, popt_one=popt_one, func_one=func_one, popt_two=popt_two, func_two=func_two, \n",
    "                split=inside_split, rotate=transit_rotate, show_events=show_events, label=label,\n",
    "                wavelength_name=wavelength_name, new_figure=new_figure, scale_to_one=scale_to_one, shift=shift,\n",
    "                scale_to_one_based_on=scale_to_one_based_on, fontsize=fontsize, numbersize=numbersize, show_graph=show_graph)\n",
    "    \n",
    "def graphAllLightCurves(planet, split=no_split, rotate=no_rotate, time_block=[], show_events=False, label=\"\", \n",
    "                        wavelength_name=True, \n",
    "                        new_figure=True, scale_to_one=True, remove=[], shift_up=False, all_new_figures=False, \n",
    "                        scale_to_one_based_on=\"max\", \n",
    "                        reverse=False, fontsize=10, numbersize=10, show_graph=True, shift_amount=0.0003):\n",
    "    '''\n",
    "    Parameters that are also in graphLightCurve behave in the same way\n",
    "    \n",
    "    If `shift_up` is True, graphs the light curves with `shift_amount` gap between them\n",
    "\n",
    "    If `all_new_figures` is True, graphs each light curve in a new figure, \n",
    "        otherwise it graphs all the light curves in a single image\n",
    "\n",
    "    `remove` is an array of wavelengths not to graph\n",
    "    '''\n",
    "    \n",
    "    for wavelength in remove:\n",
    "        checkIfInWavelengths(wavelength)\n",
    "\n",
    "    shift = 0\n",
    "    \n",
    "    if reverse:\n",
    "        r_wavelengths = wavelengths[::-1]\n",
    "    else:\n",
    "        r_wavelengths = wavelengths\n",
    "        \n",
    "    for wavelength in r_wavelengths:\n",
    "        if wavelength not in remove:\n",
    "            if shift_up:\n",
    "                shift = shift + shift_amount\n",
    "\n",
    "            graphLightCurve(planet, wavelength, new_figure=new_figure, show_events=show_events,\n",
    "                            wavelength_name=wavelength_name, label=label,\n",
    "                            scale_to_one=scale_to_one, time_block=time_block, split=split, rotate=rotate,\n",
    "                           scale_to_one_based_on=scale_to_one_based_on, fontsize=fontsize, numbersize=numbersize, show_graph=show_graph)\n",
    "            new_figure=all_new_figures\n",
    "            show_events=all_new_figures\n",
    "            \n",
    "\n",
    "def graphAllLightCurvesAdjusted(planet, use_primary_curve_fit=True, use_secondary_curve_fit=True, show_events=False, \n",
    "                                label=\"\", \n",
    "                                wavelength_name=True, new_figure=True, scale_to_one=True, remove=[], shift_up=False, \n",
    "                                all_new_figures=False,\n",
    "                                scale_to_one_based_on=\"max\", reverse=False, fontsize=10, numbersize=10, show_graph=True, shift_amount=0.0003):\n",
    "    '''\n",
    "    Parameters that are also in graphLightCurve, graphAllLightCurves, or graphLightCurveAdjusted behave in the same way\n",
    "    '''\n",
    "    \n",
    "    for wavelength in remove:\n",
    "        checkIfInWavelengths(wavelength)\n",
    "    \n",
    "    shift = 0\n",
    "    \n",
    "    if reverse:\n",
    "        r_wavelengths = wavelengths[::-1]\n",
    "    else:\n",
    "        r_wavelengths = wavelengths\n",
    "    \n",
    "    \n",
    "    for wavelength in r_wavelengths:\n",
    "        if wavelength not in remove:\n",
    "            if shift_up:\n",
    "                shift = shift + shift_amount\n",
    "            graphLightCurveAdjusted(planet, wavelength, new_figure=new_figure, show_events=show_events,\n",
    "                                 use_primary_curve_fit=use_primary_curve_fit, \n",
    "                                 use_secondary_curve_fit=use_secondary_curve_fit,\n",
    "                                 wavelength_name=wavelength_name, label=label,\n",
    "                                 scale_to_one=scale_to_one, shift=shift,\n",
    "                                scale_to_one_based_on=scale_to_one_based_on, fontsize=fontsize, numbersize=numbersize, show_graph=show_graph)\n",
    "            new_figure=all_new_figures\n",
    "            show_events=all_new_figures\n",
    "    \n",
    "def graphWavelengthAndLimbDarkening(planet, wavelength, \n",
    "                                    limb_darkening_model=\"quadratic\", limb_darkening_parameters=None, \n",
    "                                    depth=None, orbital_period_divider=11.3, semi_major_axis_const=14,\n",
    "                                    new_figure=True, show_wavelength=True, show_graph=True):\n",
    "    '''\n",
    "    Parameters that are also in graphLightCurve behave in the same way\n",
    "\n",
    "    `depth` increases the magnitude of the predicted light curve change\n",
    "\n",
    "    `semi_major_axis_const` and `orbital_period_divider` affect the predicted light curve in strange ways\n",
    "\n",
    "    Graphs the adjusted light curve and a predicted light curve \n",
    "        based on the `limb_darkening_model` and `limb_darkening_parameters`\n",
    "\n",
    "    If `limb_darkening_model` is \"quadratic\" and `limb_darkening_parameters` is None, \n",
    "        it uses parameters from Allen's Astrophysical Quantities\n",
    "    '''\n",
    "    \n",
    "    if planet == \"venus\":\n",
    "        if depth==None:\n",
    "            if wavelength in [\"1600\",\"1700\"]:\n",
    "                depth = 1.16\n",
    "            elif wavelength == \"cont\":\n",
    "                depth = 1.45\n",
    "            else:\n",
    "                depth = 1.4\n",
    "        params = batman.TransitParams()\n",
    "        params.t0 = venus_time_inferior_conjunc      #time of inferior conjunction\n",
    "        params.per = venus_orbital_period / orbital_period_divider         #orbital period\n",
    "        params.rp = venus_apparant_radius*depth       #planet radius (in units of stellar radii)\n",
    "        params.a = semi_major_axis_const #semi-major axis (in units of stellar radii)\n",
    "        params.inc = 90 - venus_orbital_inclination       #orbital inclination (in degrees)\n",
    "        params.ecc = venus_eccentricity              #eccentricity\n",
    "        params.w = venus_longitude_perihelion        #longitude of periastron (in degrees)\n",
    "        params.limb_dark = limb_darkening_model        #limb darkening model\n",
    "        if limb_darkening_parameters == None and limb_darkening_model == \"quadratic\":       #limb darkening coefficients\n",
    "            #These constants are from Allen's Astrophysical Quantities\n",
    "            if wavelength == \"cont\":\n",
    "                params.u = [0.88, -0.23]\n",
    "            else:\n",
    "                params.u = [0.12, 0.33]\n",
    "        else:\n",
    "            params.u = limb_darkening_parameters\n",
    "            \n",
    "        if params.u == None:\n",
    "            params.u = []\n",
    "            \n",
    "        t = np.linspace(0, 32400, 10000)  #times at which to calculate light curve\n",
    "        m = batman.TransitModel(params, t)    #initializes model\n",
    "        flux = m.light_curve(params) #creates light curve\n",
    "        if new_figure:\n",
    "            plt.figure()\n",
    "        plt.plot(t, flux,label='Generated') #plots generated light curve\n",
    "        if show_wavelength:\n",
    "            #plots actual light curve\n",
    "            graphLightCurveAdjusted(planet,wavelength,new_figure=False, scale_to_one_based_on=\"not transit\", show_graph=show_graph) \n",
    "    else:\n",
    "        raise Exception(\"graphWavelengthAndLimbDarkening not yet implemented for mercury\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def derotateWavelength(planet, wavelength, time_one=0, time_two=30000, file_one=None, file_two=None):\n",
    "    '''\n",
    "    `planet` is the planet transit to use\n",
    "    \n",
    "    `wavelength` is the wavelength to use\n",
    "\n",
    "    `time_one` is the time to derotate to\n",
    "\n",
    "    `time_two` is the time to get an image from to derotate from\n",
    "\n",
    "    `file_one` is the file to derotate to\n",
    "\n",
    "    `file_two` is the file to get an image from to derotate from\n",
    "\n",
    "    Doesn't work for HMI Continuum\n",
    "    \n",
    "    Returns the derotated image\n",
    "    '''\n",
    "    list_of_maps = []\n",
    "    if file_one == None:\n",
    "        list_of_maps.append(sunpy.map.Map(findFileByTimeSinceStartOfTransit(planet, time_one,wavelength)))  \n",
    "    else:\n",
    "        list_of_maps.append(sunpy.map.Map(file_one))  \n",
    "    if file_two == None:\n",
    "        list_of_maps.append(sunpy.map.Map(findFileByTimeSinceStartOfTransit(planet, time_two,wavelength)))  \n",
    "    else:\n",
    "        list_of_maps.append(sunpy.map.Map(file_two))  \n",
    "\n",
    "    mapcube = sunpy.map.Map(list_of_maps, cube=True)\n",
    "    derotated = mapcube_solar_derotate(mapcube)\n",
    "    return derotated.maps[-1]\n",
    "        \n",
    "def showDiffBetweenTwoTimes(planet, wavelength, file_one=None, file_two=None, time_one=None, time_two=None, \n",
    "                            one_special=False, two_special=False):\n",
    "    '''\n",
    "    `planet` is the planet transit to use\n",
    "    \n",
    "    `wavelength` is the wavelength to use\n",
    "\n",
    "    `time_one` is the time to get an image from to difference to\n",
    "\n",
    "    `time_two` is the time to get an image from to difference from\n",
    "\n",
    "    `file_one` is the file to get an image from to difference to\n",
    "\n",
    "    `file_two` is the file to get an image from to difference from\n",
    "\n",
    "    If `one_special` is True, the array is stored in the first section of the first file\n",
    "\n",
    "    If `two_special` is True, the array is stored in the first section of the second file\n",
    "    \n",
    "    Returns the difference image\n",
    "    '''\n",
    "    \n",
    "    checkIfInWavelengths(wavelength)\n",
    "\n",
    "    #Gets the two arrays\n",
    "    if time_one != None:\n",
    "        file_one = findFileByTimeSinceStartOfTransit(planet, time_one,wavelength)\n",
    "    if time_two != None:\n",
    "        file_two = findFileByTimeSinceStartOfTransit(planet, time_two,wavelength)\n",
    "\n",
    "    if one_special == False:\n",
    "        image_one = getArrayFromFitsFile(file_one)\n",
    "    else:\n",
    "        image_one = np.fliplr(np.flipud(fits.read(file_one)[0][0]))\n",
    "    if two_special == False:\n",
    "        image_two = getArrayFromFitsFile(file_two)\n",
    "    else:\n",
    "        image_two = np.fliplr(np.flipud(fits.read(file_two)[0][0]))\n",
    "             \n",
    "            \n",
    "    #Finds the differences between the two arrays\n",
    "    diff_image = np.subtract(image_two, image_one)\n",
    "    \n",
    "    #Flips the array so that the orientation is consistent for all wavelengths\n",
    "    if wavelength == \"cont\":\n",
    "        diff_image = np.fliplr(diff_image)\n",
    "    else:\n",
    "        diff_image = np.flipud(diff_image)\n",
    "        \n",
    "    showArray(diff_image)\n",
    "    \n",
    "    return diff_image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def getFourierFrequencyData(planet, wavelength, show_graph=True, show_light_curve=False, seconds_between_points=24):\n",
    "    '''\n",
    "    `planet` is the planet transit to use\n",
    "    \n",
    "    `wavelength` is the wavelength to use\n",
    "    \n",
    "    If `show_graph` is True, it graphs the frequency data\n",
    "    \n",
    "    If `show_light_curve` is True, it graphs the light curve that produced the frequency data\n",
    "\n",
    "    `seconds_between_points` is the number of seconds between data points used for interpolation. \n",
    "    Decreasing this number increases the accuracy of the function but increases the processing time it takes to run the function.\n",
    "\n",
    "    Returns the result of a Fourier transform done on the planet and wavelength data, \n",
    "    the data's linear spacing and interpolated points, and the data's average\n",
    "    '''\n",
    "    \n",
    "    times, data  = graphLightCurveAdjusted(planet,wavelength,show_graph=show_light_curve)\n",
    "\n",
    "    times = np.array(times,np.float64)\n",
    "    data  = np.array(data, np.float64)\n",
    "    \n",
    "    #Change data to have a mean of zero\n",
    "    #If this is not done, the Fourier transform just produced a large spike centered on zero with nothing else\n",
    "    data_average = np.average(data)\n",
    "    data  = np.subtract(data,data_average)\n",
    "\n",
    "    # number of signal points\n",
    "    N = int((max(times) - min(times)) / seconds_between_points)\n",
    "\n",
    "    #Interpolate the data, since the Fourier transform only accepts evenly spaced data\n",
    "    x = np.linspace(min(times), max(times), N)\n",
    "    y = interp1d(times, data)(x)\n",
    "    yf = fft(y)\n",
    "    xf = fftfreq(N, seconds_between_points)\n",
    "    \n",
    "    if show_graph:\n",
    "        plt.figure()\n",
    "        plt.plot(fftshift(xf), np.abs(fftshift(yf)))\n",
    "        plt.grid()\n",
    "        plt.show()\n",
    "    return (xf, yf, x, y, data_average)\n",
    "\n",
    "def getLowPassedData(planet, wavelength, show_graph=True, show_light_curve=True, seconds_between_points=24):\n",
    "    '''\n",
    "    `planet` is the planet transit to use\n",
    "    \n",
    "    `wavelength` is the wavelength to use\n",
    "    \n",
    "    If `show_graph` is True, it graphs the low passed data\n",
    "    \n",
    "    If `show_light_curve` is True, it graphs the light curve\n",
    "\n",
    "    `seconds_between_points` is the number of seconds between data points used for interpolation. \n",
    "    Decreasing this number increases the accuracy of the function but increases the processing time it takes to run the function.\n",
    "\n",
    "    Returns the planet and wavelength data with a low pass filter applied\n",
    "    \n",
    "    Source for this is https://tomroelandts.com/articles/how-to-create-a-simple-low-pass-filter\n",
    "    '''\n",
    "    \n",
    "    times, data  = graphLightCurveAdjusted(planet,wavelength,show_graph=show_light_curve)\n",
    "\n",
    "    times = np.array(times,np.float64)\n",
    "    data  = np.array(data, np.float64)\n",
    "    data  = np.subtract(data,np.average(data))\n",
    "\n",
    "    # number of signal points\n",
    "    N = int((max(times) - min(times)) / seconds_between_points)\n",
    "\n",
    "    #Interpolate the data, since the Fourier transform only accepts evenly spaced data\n",
    "    x = np.linspace(min(times), max(times), N)\n",
    "    y = interp1d(times, data)(x)\n",
    "\n",
    "    fc = 0.03  # Cutoff frequency as a fraction of the sampling rate (in (0, 0.5)).\n",
    "    b = 0.02  # Transition band, as a fraction of the sampling rate (in (0, 0.5)).\n",
    "    N = int(np.ceil((4 / b)))\n",
    "    if not N % 2: N += 1  # Make sure that N is odd.\n",
    "    n = np.arange(N)\n",
    "\n",
    "    # Compute sinc filter.\n",
    "    h = np.sinc(2 * fc * (n - (N - 1) / 2.))\n",
    "\n",
    "    # Compute Blackman window.\n",
    "    w = np.blackman(N)\n",
    "\n",
    "    # Multiply sinc filter with window.\n",
    "    h = h * w\n",
    "\n",
    "    # Normalize to get unity gain.\n",
    "    h = h / np.sum(h)\n",
    "    y = np.convolve(y, h, mode='same')\n",
    "    y = np.add(y, data_average)\n",
    "    \n",
    "    if show_graph:\n",
    "        plt.plot(x,y)\n",
    "        \n",
    "    return (x,y)\n",
    "\n",
    "def getLowSquashedData(planet, wavelength, data_limit=0.0001, show_graph=True, show_light_curve=False, \n",
    "                       show_pre_change_frequency_graph=False, \n",
    "                       show_post_change_frequency_graph=False, seconds_between_points=24, slope_adjust=1.15):\n",
    "    '''\n",
    "    `planet` is the planet transit to use\n",
    "    \n",
    "    `wavelength` is the wavelength to use\n",
    "    \n",
    "    `data_limit` is the frequency value above which all frequency data is squashed\n",
    "    \n",
    "    If `show_graph` is True, it graphs the processed data\n",
    "    \n",
    "    If `show_light_curve` is True, it graphs the light curve that produced the frequency data\n",
    "\n",
    "    If `show_pre_change_frequency_graph` is True, it graphs the unsquashed frequency data\n",
    "\n",
    "    If `show_post_change_frequency_graph` is True, it graphs the squashed frequency data\n",
    "    \n",
    "    `seconds_between_points` is the number of seconds between data points used for interpolation. \n",
    "    Decreasing this number increases the accuracy of the function but increases the processing time it takes to run the function.\n",
    "\n",
    "    `slope_adjust` is a multiplier for the squashing. It helps to remove any ringing\n",
    "\n",
    "    Returns the planet and wavelength data with the high frequency data squashed\n",
    "    '''\n",
    "\n",
    "    #The default value for `slope_adjust` was found by trial, error, and subjectivity\n",
    "    \n",
    "    xf, yf, x, y, data_average = getFourierFrequencyData(planet,wavelength,\n",
    "                                                         show_graph=show_pre_change_frequency_graph,\n",
    "                                                         show_light_curve=show_light_curve, seconds_between_points=seconds_between_points)\n",
    "\n",
    "    #Squash yf values\n",
    "    good_yf = []\n",
    "    for i, item in enumerate(xf):\n",
    "        if abs(item) < data_limit :\n",
    "            good_yf.append(yf[i])\n",
    "        elif abs(item) < data_limit*2: #This is to prevent ringing from happening, which dirties the data\n",
    "            good_yf.append(yf[i]*(2 - slope_adjust * abs(item) / data_limit))\n",
    "        else:\n",
    "            good_yf.append(0)\n",
    "    \n",
    "    if show_post_change_frequency_graph:\n",
    "        plt.figure()\n",
    "        plt.plot(fftshift(xf), np.abs(fftshift(good_yf)))\n",
    "        plt.grid()\n",
    "        plt.show()\n",
    "    \n",
    "    #to get it back on the original scale, since before Fourier transformations, the data is changed to have a mean of 0\n",
    "    data = np.add(ifft(good_yf),data_average)\n",
    "    \n",
    "    if show_graph:\n",
    "        plt.plot(x,data)\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}