{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from astropy.io import fits\n",
    "from scipy.interpolate import interp1d\n",
    "import math as m\n",
    "import scipy.constants as p\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import timeit\n",
    "import sys\n",
    "sys.path.append(\"..\")\n",
    "import physics\n",
    "sys.path.append(\"../..\")\n",
    "from spec.spectrum import *\n",
    "#import Electrons_v1\n",
    "#%run 'Electrons_v1.ipynb'\n",
    "import xsecs\n",
    "np.set_printoptions(precision=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Test spectrum\n",
    "eng1 = np.logspace(0.,3.,num = 500)\n",
    "dNdE1 = np.logspace(0.,5.,num = 500)\n",
    "\n",
    "eng2   = np.logspace(0.,3.,num=2)\n",
    "dNdE2 = np.logspace(0.,5.,num=2)\n",
    "\n",
    "testMat = np.array([[1,2],[2,1]])\n",
    "\n",
    "rs = 2000.\n",
    "\n",
    "specObj1 = Spectrum(eng1,dNdE1,rs)\n",
    "specObj2 = Spectrum(eng2,dNdE2,rs)\n",
    "\n",
    "ion_frac=0.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 96.3758]\n",
      "[ 0.]\n"
     ]
    }
   ],
   "source": [
    "eng1 = np.logspace(0.,3.,num = 500)\n",
    "dNdE1 = np.logspace(0.,5.,num = 500)\n",
    "rs=1000\n",
    "E_spec = Spectrum(eng1,dNdE1,rs)\n",
    "photon_spec = Spectrum(E_spec.eng,np.zeros(len(E_spec.eng)),rs)\n",
    "TM=np.zeros([len(photon_spec.eng)+2,len(E_spec.eng)]) #possibly make pandas\n",
    "Transfer_Matrix=pd.DataFrame(TM) #or np.zeros\n",
    "#print(Transfer_Matrix)\n",
    "print(np.array([E_spec.eng[330]]))\n",
    "print(xsecs.ionize_cs(np.array([E_spec.eng[100]]),np.ones(1)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "27970736981.1\n",
      "27970742355.1\n",
      "38196427.8074\n",
      "38196417.9184\n"
     ]
    }
   ],
   "source": [
    "eng1 = np.logspace(0.,3.,num = 500)\n",
    "dNdE1 = np.logspace(0.,5.,num = 500)\n",
    "rs=1000\n",
    "ion_frac=[0.5,0.5,0.5]\n",
    "E_spec = Spectrum(eng1,dNdE1,rs)\n",
    "\n",
    "\n",
    "shift = Spectrum((E_spec.eng+xsecs.heating_dE(E_spec.eng, ion_frac[0], rs, nH=physics.nH)), E_spec.dNdE, rs)\n",
    "shift.rebin(E_spec.eng)\n",
    "print(shift.toteng())\n",
    "print(E_spec.toteng())\n",
    "\n",
    "print(shift.totN())\n",
    "print(E_spec.totN())\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.07166764803\n",
      "0.0166493283539\n",
      "0.0\n",
      "0.0166493283539\n"
     ]
    }
   ],
   "source": [
    "        print(E_spec.eng[5])\n",
    "        dE = xsecs.heating_dE(E_spec.eng[5], 0.5, rs, nH=physics.nH)\n",
    "    \n",
    "        delta_dNdE = np.zeros(len(E_spec.dNdE))\n",
    "        np.put(delta_dNdE, 5, E_spec.dNdE[5])\n",
    "        delta = Spectrum(E_spec.eng, delta_dNdE, rs)\n",
    "        print(delta.totN())\n",
    "   \n",
    "        shift_delta_eng = E_spec.eng+dE\n",
    "        delta.shift_eng(shift_delta_eng)\n",
    "        \n",
    "        delta.rebin(E_spec.eng)\n",
    "        print(delta.totN())\n",
    "        \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def electron_low_E(E_spec, rs, ion_frac=[0.5,0.5,0.5], den=[1,1,1], dt=1, all_outputs=False): #update preset density values\n",
    "    #V1\n",
    "    #E_spec=electron spectrum class object, rs= current redshift, \n",
    "    #ion_frac= H_frac, He_frac, Hep_frac, density= H_den, He_den, Hep_den\n",
    "    #all_outputs controls if function returns transfer matrix (False) or transfer matrix, photon spectrum, ionization fraction, heating loss (True)\n",
    "    \n",
    "    #Initialize photon spectrum (ask about length)\n",
    "    photon_spec = Spectrum(E_spec.eng,np.zeros(len(E_spec.eng)),rs)\n",
    "    \n",
    "    #initialize transfer matrix\n",
    "    #columns match E_spec\n",
    "    #rows are 0:heating, 1:ionization, 2->len(photon_spec)+1:photons\n",
    "    TM=np.zeros([len(photon_spec.eng)+2,len(E_spec.eng)]) #possibly make pandas\n",
    "    Transfer_Matrix=pd.DataFrame(TM) #read more pd\n",
    "    \n",
    "    #check density calculation...\n",
    "    #insure ion_frac is correct for den[2]\n",
    "    den[0]=(physics.nH*(1-ion_frac[0]))/(rs)**3\n",
    "    den[1]=(physics.nHe*(1-ion_frac[1])*(1-ion_frac[2]))/(rs)**3\n",
    "    den[2]=(physics.nHe*(1-ion_frac[1])*(ion_frac[2]))/(rs)**3 #unsure about how ionfrac works for helium\n",
    "    \n",
    "    \n",
    "    #HEATING\n",
    "    heat_vector=np.zeros(len(E_spec.eng))\n",
    "    \n",
    "    \n",
    "    #main loop to fill transfer matrix\n",
    "        #change zeros if heating becomes cross section\n",
    "    for n in range(len(E_spec.eng)):\n",
    "        \n",
    "        \n",
    "        #Find xsecs with Electrons_v1\n",
    "        xsec=np.zeros(9)\n",
    "        #ionization\n",
    "        xsec[0] = xsecs.ionize_cs(np.array([E_spec.eng[n]]),np.ones(1)) #H ion\n",
    "        xsec[1] = xsecs.ionize_cs(np.array([E_spec.eng[n]]),2*np.ones(1)) #He ion\n",
    "        xsec[2] = xsecs.ionize_cs(np.array([E_spec.eng[n]]),3*np.ones(1)) #Hep ion\n",
    "        #excitation\n",
    "        #xsec[3] = excite_cs(np.array([E_spec.eng[n]]),1,1) #H-s excite\n",
    "        #xsec[4] = excite_cs(np.array([E_spec.eng[n]]),2,1) #He-s excite\n",
    "        #xsec[5] = excite_cs(np.array([E_spec.eng[n]]),3,1) #Hep-s excite\n",
    "        #xsec[6] = excite_cs(np.array([E_spec.eng[n]]),1,2) #H-p excite\n",
    "        #xsec[7] = excite_cs(np.array([E_spec.eng[n]]),2,2) #He-p excite\n",
    "        #xsec[8] = excite_cs(np.array([E_spec.eng[n]]),3,2) #Hep-p excite\n",
    "            #xsec[9] = thermalize_cs(numpy.array([E_spec.eng[n]])) #heating\n",
    "            \n",
    "        #Find probabilities from xsecs\n",
    "        #find velocity to calculate rate\n",
    "        v=np.sqrt(2*np.array([E_spec.eng[n]])/physics.me) #velocity\n",
    "        \n",
    "        #multipliers\n",
    "        mult=np.zeros(3)\n",
    "        mult[0]= v*den[0] #H\n",
    "        mult[1]= v*den[1] #He\n",
    "        mult[2]= v*den[2] #Hep\n",
    "            #mult[3]= #heating (v*electron_density)?\n",
    "        prob=np.zeros(9)\n",
    "        #ionization\n",
    "        ###possibly remove\n",
    "        prob[0] = xsec[0]*mult[0] #H ion\n",
    "        prob[1] = xsec[1]*mult[1] #He ion\n",
    "        prob[2] = xsec[2]*mult[2] #Hep ion\n",
    "        #excitation\n",
    "        prob[3] = xsec[3]*mult[0] #H-s excite\n",
    "        prob[4] = xsec[4]*mult[1] #He-s excite\n",
    "        prob[5] = xsec[5]*mult[2] #Hep-s excite\n",
    "        prob[6] = xsec[6]*mult[0] #H-p excite\n",
    "        prob[7] = xsec[7]*mult[1] #He-p excite\n",
    "        prob[8] = xsec[8]*mult[2] #Hep-p excite\n",
    "            #prob[9] = xsec[9]*mult[3] #heating\n",
    "        \n",
    "        \n",
    "        #Call HEAT related things\n",
    "        dE = xsecs.heating_dE(E_spec.eng[n], 0.5, rs, nH=physics.nH)\n",
    "        delta_dNdE = np.zeros(len(E_spec.dNdE))\n",
    "        np.put(delta_dNdE, n, E_spec.dNdE[n])\n",
    "        delta = Spectrum(E_spec.eng, delta_dNdE, rs)\n",
    "        shift_delta_eng = E_spec.eng+dE\n",
    "        delta.shift_eng(shift_delta_eng)\n",
    "        delta.rebin(E_spec.eng)\n",
    "        \n",
    "        heating_frac = delta.eng[n]/E_spec.eng[n]\n",
    "        \n",
    "        \n",
    "        prob_s=np.zeros((3,n))\n",
    "        #singly diff ionization prob from xsec\n",
    "        for i in range(n):\n",
    "            prob_s[0,i]=mult[0]*ionize_s_cs(E_spec.eng[n],E_spec.eng[i],1) #H\n",
    "            prob_s[1,i]=mult[1]*ionize_s_cs(E_spec.eng[n],E_spec.eng[i],2) #He\n",
    "            prob_s[2,i]=mult[2]*ionize_s_cs(E_spec.eng[n],E_spec.eng[i],3) #Hep\n",
    "        #sum probabilities from singly diff \n",
    "        prob_atom_ionize_s = sum(prob_s, axis=1) #vector giving prob sum for each element\n",
    "        \n",
    "        print(prob[0]-prob_atom_ionize_s[0])\n",
    "        print(prob[1]-prob_atom_ionize_s[1])\n",
    "        print(prob[2]-prob_atom_ionize_s[2])\n",
    "        \n",
    "        \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mult=array([1,2,3])\n",
    "prob_s=numpy.zeros((3,5))\n",
    "for i in range(5):\n",
    "    prob_s[0,i]=mult[0]*ionize_s_cs(E_spec.eng[5],E_spec.eng[i],1) #H\n",
    "    prob_s[1,i]=mult[1]*ionize_s_cs(E_spec.eng[5],E_spec.eng[i],2) #He\n",
    "    prob_s[2,i]=mult[2]*ionize_s_cs(E_spec.eng[5],E_spec.eng[i],3) #Hep\n",
    "#sum probabilities from singly diff \n",
    "prob_atom_ionize_s = sum(prob_s, axis=1)\n",
    "print(prob_s)\n",
    "print(prob_atom_ionize_s)\n",
    "\n",
    "prob_s[0,i]=mult[0]*ionize_s_cs(E_spec.eng[5],E_spec.eng[i],1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "sum() takes no keyword arguments",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-22-c44e78eb067e>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0melectron_low_E\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mE_spec\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mrs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mion_frac\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0.5\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m0.5\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m0.5\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mden\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mall_outputs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m<ipython-input-21-a033e467121e>\u001b[0m in \u001b[0;36melectron_low_E\u001b[0;34m(E_spec, rs, ion_frac, den, dt, all_outputs)\u001b[0m\n\u001b[1;32m     88\u001b[0m             \u001b[0mprob_s\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mmult\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0mionize_s_cs\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mE_spec\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0meng\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mn\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mE_spec\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0meng\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m3\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;31m#Hep\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     89\u001b[0m         \u001b[0;31m#sum probabilities from singly diff\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 90\u001b[0;31m         \u001b[0mprob_atom_ionize_s\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0msum\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mprob_s\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0maxis\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;31m#vector giving prob sum for each element\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     91\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     92\u001b[0m         \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mprob\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0mprob_atom_ionize_s\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mTypeError\u001b[0m: sum() takes no keyword arguments"
     ]
    }
   ],
   "source": [
    "electron_low_E(E_spec, rs, ion_frac=[0.5,0.5,0.5], den=[1,1,1], all_outputs=False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
