Pylab

All posts tagged Pylab

A preliminary investigation showed a project a colleague and I were considering probably isn’t worth doing. But for that investigation, I took a few hours to make a rather complicated plot using pylab, so I thought I’d share how I did that.

First, here’s the plot:

Tidal decay timescales for members of multi-planet systems.

Tidal decay timescales for members of multi-planet systems.

The plot shows the timescales for tidal decay of members of multi-planet systems. Unfortunately, the x-axis labels aren’t legible unless you zoom in, but if you do, you can see the font colors match up with the corresponding line colors.

Below is the ipython notebook I used to generate the plot. The excel spreadsheet with the data is here.

<span class="c">#Show the plot inline</span>
<span class="o">%</span><span class="k">matplotlib</span> <span class="n">inline</span>

<span class="c">#load in the required modules</span>
<span class="kn">import</span> <span class="nn">pandas</span> <span class="kn">as</span> <span class="nn">pd</span>
<span class="kn">import</span> <span class="nn">pylab</span> <span class="kn">as</span> <span class="nn">pl</span>
<span class="kn">import</span> <span class="nn">itertools</span> <span class="kn">as</span> <span class="nn">it</span>
<span class="kn">import</span> <span class="nn">numpy</span> <span class="kn">as</span> <span class="nn">np</span>

<span class="c"># using the ExcelFile class</span>
<span class="n">xls</span> <span class="o">=</span> <span class="n">pd</span><span class="o">.</span><span class="n">ExcelFile</span><span class="p">(</span><span class="s">'exoplanet-archive_2015Mar25.xlsx'</span><span class="p">)</span>
<span class="n">data</span> <span class="o">=</span> <span class="n">xls</span><span class="o">.</span><span class="n">parse</span><span class="p">(</span><span class="s">'obj of interest'</span><span class="p">,</span> <span class="n">index_col</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span>
<span class="n">data</span> <span class="o">=</span> <span class="n">data</span><span class="p">[</span><span class="n">pd</span><span class="o">.</span><span class="n">notnull</span><span class="p">(</span><span class="n">data</span><span class="p">[</span><span class="s">'a/(da/dt)Qs=1e6 (Gyrs)'</span><span class="p">])]</span>

<span class="c">#Make a nice, big figure</span>
<span class="n">fig</span> <span class="o">=</span> <span class="n">pl</span><span class="o">.</span><span class="n">figure</span><span class="p">(</span><span class="n">figsize</span><span class="o">=</span><span class="p">(</span><span class="mi">15</span><span class="p">,</span><span class="mi">15</span><span class="p">))</span>
<span class="n">ax</span> <span class="o">=</span> <span class="n">fig</span><span class="o">.</span><span class="n">add_subplot</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>

<span class="c">#Make a list, indexing the dataframe labels</span>
<span class="n">indices</span> <span class="o">=</span> <span class="nb">range</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="nb">set</span><span class="p">(</span><span class="n">data</span><span class="o">.</span><span class="n">index</span><span class="p">)))</span>

<span class="c">#Make a list with the indices as the entries</span>
<span class="n">labels</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="nb">set</span><span class="p">(</span><span class="n">data</span><span class="o">.</span><span class="n">index</span><span class="p">))</span>
<span class="c">#For concision, drop "Kepler" wherever it's found</span>
<span class="n">labels</span> <span class="o">=</span> <span class="p">[</span><span class="n">w</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span><span class="s">'Kepler-'</span><span class="p">,</span> <span class="s">''</span><span class="p">)</span> <span class="k">for</span> <span class="n">w</span> <span class="ow">in</span> <span class="n">labels</span><span class="p">]</span>

<span class="c">#Make a cycle of line and text colors, blue, green, red, etc.</span>
<span class="n">colors</span> <span class="o">=</span> <span class="n">it</span><span class="o">.</span><span class="n">cycle</span><span class="p">([</span><span class="s">'b'</span><span class="p">,</span> <span class="s">'g'</span><span class="p">,</span> <span class="s">'r'</span><span class="p">,</span> <span class="s">'c'</span><span class="p">,</span> <span class="s">'m'</span><span class="p">,</span> <span class="s">'y'</span><span class="p">,</span> <span class="s">'k'</span><span class="p">])</span>

<span class="c">#Since each member of the multi-system should be plotted with the same x-value,</span>
<span class="c">#  I need to generate a new list of all the same value with as many entries</span>
<span class="c">#  as members. That's what "i" is for.</span>
<span class="n">i</span> <span class="o">=</span> <span class="mi">0</span>
<span class="k">for</span> <span class="n">unq</span> <span class="ow">in</span> <span class="nb">set</span><span class="p">(</span><span class="n">data</span><span class="o">.</span><span class="n">index</span><span class="p">):</span>

    <span class="c">#Retrieve the decay timescales calculated in the spreadsheet</span>
    <span class="n">taus</span> <span class="o">=</span> <span class="n">data</span><span class="o">.</span><span class="n">loc</span><span class="p">[</span><span class="n">unq</span><span class="p">,</span> <span class="s">'a/(da/dt)Qs=1e6 (Gyrs)'</span><span class="p">]</span>

    <span class="c">#Generate the list of all the same x-value</span>
    <span class="n">idx</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">ones_like</span><span class="p">(</span><span class="n">taus</span><span class="p">)</span><span class="o">*</span><span class="n">i</span>

    <span class="c">#Make the scatter plot points with the current color</span>
    <span class="n">ax</span><span class="o">.</span><span class="n">semilogy</span><span class="p">(</span><span class="n">idx</span><span class="p">,</span> <span class="n">taus</span><span class="p">,</span> <span class="n">marker</span><span class="o">=</span><span class="s">'o'</span><span class="p">,</span> <span class="n">color</span><span class="o">=</span><span class="n">cur_color</span><span class="p">)</span>

    
    <span class="c">#Get the next line color</span>
    <span class="n">cur_color</span> <span class="o">=</span> <span class="nb">next</span><span class="p">(</span><span class="n">colors</span><span class="p">)</span>
    
    <span class="c">#Next x-value</span>
    <span class="n">i</span> <span class="o">+=</span> <span class="mi">1</span>
    
<span class="c">#Give a little space to the left and right of the first and last x-values</span>
<span class="n">pl</span><span class="o">.</span><span class="n">xlim</span><span class="p">([</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="nb">len</span><span class="p">(</span><span class="nb">set</span><span class="p">(</span><span class="n">data</span><span class="o">.</span><span class="n">index</span><span class="p">))</span><span class="o">+</span><span class="mi">1</span><span class="p">])</span>

<span class="c">#Switch out the x-values with the system names</span>
<span class="n">pl</span><span class="o">.</span><span class="n">xticks</span><span class="p">(</span><span class="n">indices</span><span class="p">,</span> <span class="n">labels</span><span class="p">,</span> <span class="n">rotation</span><span class="o">=</span><span class="s">'vertical'</span><span class="p">,</span> <span class="n">size</span><span class="o">=</span><span class="s">'small'</span><span class="p">,</span> <span class="n">ha</span><span class="o">=</span><span class="s">'center'</span><span class="p">)</span>

<span class="c">#Increase the size of the y-axis label font</span>
<span class="n">pl</span><span class="o">.</span><span class="n">yticks</span><span class="p">(</span><span class="n">size</span><span class="o">=</span><span class="mi">36</span><span class="p">)</span>
<span class="c">#Label the y-axis</span>
<span class="n">pl</span><span class="o">.</span><span class="n">ylabel</span><span class="p">(</span><span class="s">'$a/</span><span class="se">\\</span><span class="s">left(</span><span class="se">\\</span><span class="s">frac{da}{dt}</span><span class="se">\\</span><span class="s">right)_{Q_{</span><span class="se">\\</span><span class="s">rm s} = 10^6}$ (Gyrs)'</span><span class="p">,</span> <span class="n">fontsize</span><span class="o">=</span><span class="mi">36</span><span class="p">)</span>

<span class="c">#Reset the colors cycle</span>
<span class="n">colors</span> <span class="o">=</span> <span class="n">it</span><span class="o">.</span><span class="n">cycle</span><span class="p">([</span><span class="s">'b'</span><span class="p">,</span> <span class="s">'g'</span><span class="p">,</span> <span class="s">'r'</span><span class="p">,</span> <span class="s">'c'</span><span class="p">,</span> <span class="s">'m'</span><span class="p">,</span> <span class="s">'y'</span><span class="p">,</span> <span class="s">'k'</span><span class="p">])</span>

<span class="c">#Set a new color for each x-axis label</span>
<span class="k">for</span> <span class="n">tick</span> <span class="ow">in</span> <span class="n">ax</span><span class="o">.</span><span class="n">xaxis</span><span class="o">.</span><span class="n">get_ticklabels</span><span class="p">():</span>
    <span class="n">tick</span><span class="o">.</span><span class="n">set_color</span><span class="p">(</span><span class="n">cur_color</span><span class="p">)</span>
    <span class="n">cur_color</span> <span class="o">=</span> <span class="nb">next</span><span class="p">(</span><span class="n">colors</span><span class="p">)</span>

<span class="n">pl</span><span class="o">.</span><span class="n">savefig</span><span class="p">(</span><span class="s">'Comparing multi-planet system a_dadt.png'</span><span class="p">,</span> <span class="n">bbox_inches</span><span class="o">=</span><span class="s">'tight'</span><span class="p">,</span> <span class="n">orientation</span><span class="o">=</span><span class="s">'landscape'</span><span class="p">,</span> <span class="n">dpi</span><span class="o">=</span><span class="mi">250</span><span class="p">)</span>