#!/usr/bin/env python
# coding: utf-8

# GetDist iPython Plot Gallery
# ===============

# In[1]:


#Show plots inline, and load main getdist plot module and samples class
get_ipython().run_line_magic('matplotlib', 'inline')
from __future__ import print_function
from getdist import plots, MCSamples
import getdist, IPython
print('Version: ',getdist.__version__)


# In[2]:


#Get some random samples for demonstration:
#make random covariance, then independent samples from Gaussian
import numpy as np
ndim = 4
nsamp = 10000
np.random.seed(10)
A = np.random.rand(ndim,ndim)
cov = np.dot(A, A.T)
samps = np.random.multivariate_normal([0]*ndim, cov, size=nsamp)
A = np.random.rand(ndim,ndim)
cov = np.dot(A, A.T)
samps2 = np.random.multivariate_normal([0]*ndim, cov, size=nsamp)


# In[3]:


#Get the getdist MCSamples objects for the samples, specifying same parameter
#names and labels; if not specified weights are assumed to all be unity
names = ["x%s"%i for i in range(ndim)]
labels =  ["x_%s"%i for i in range(ndim)]
samples = MCSamples(samples=samps,names = names, labels = labels)
samples2 = MCSamples(samples=samps2,names = names, labels = labels)


# In[4]:


#Triangle plot
g = plots.getSubplotPlotter()
g.triangle_plot([samples, samples2], filled=True)


# In[5]:


#Here we are using inline plots, but if you wanted to export to file you'd just to e.g.
#g.export('output_file.pdf')


# In[6]:


#1D marginalized plot
g = plots.getSinglePlotter(width_inch=4)
g.plot_1d(samples, 'x2')


# In[7]:


#1D marginalized comparison plot
g = plots.getSinglePlotter(width_inch=3)
g.plot_1d([samples, samples2], 'x1')


# In[8]:


#1D normalized comparison plot
g = plots.getSinglePlotter(width_inch=4)
g.plot_1d([samples, samples2], 'x1', normalized=True)


# In[9]:


#2D line contour comparison plot with extra bands and markers
g = plots.getSinglePlotter()
g.plot_2d([samples, samples2], 'x1', 'x2')
g.add_x_marker(0)
g.add_y_bands(0, 1)


# In[10]:


#Filled 2D comparison plot with legend
g = plots.getSinglePlotter(width_inch=4, ratio=1)
g.plot_2d([samples, samples2], 'x1', 'x2', filled=True)
g.add_legend(['sim 1', 'sim 2'], colored_text=True);


# In[11]:


#Shaded 2D comparison plot
g = plots.getSinglePlotter(width_inch=4)
g.plot_2d([samples, samples2], 'x1', 'x2', shaded=True);


# In[12]:


#Customized 2D filled comparison plot
g = plots.getSinglePlotter(width_inch=6, ratio=3 / 5.)
g.settings.legend_fontsize = 10
g.plot_2d([samples, samples2], 'x1', 'x2', filled=True, 
    colors=['green', ('#F7BAA6', '#E03424')], lims=[-4, 7, -5, 5])
g.add_legend(['Sim ', 'Sim 2'], legend_loc='upper right');


# In[13]:


#Change the contours levels for marge stats and plots
#(note you need a lot of samples for 99% confidence contours to be accurate)
g = plots.getSinglePlotter()
samples.updateSettings({'contours': [0.68, 0.95, 0.99]})
g.settings.num_plot_contours = 3
g.plot_2d(samples, 'x1', 'x2');


# In[14]:


#2D scatter (3D) plot
g = plots.getSinglePlotter(width_inch=5)
g.plot_3d(samples, ['x1', 'x2', 'x3'])


# In[15]:


#Multiple 1D subplots
g = plots.getSubplotPlotter(width_inch=5)
g.settings.axes_fontsize = 9
g.settings.legend_fontsize = 10
g.plots_1d(samples, ['x0', 'x1', 'x2', 'x3'], nx=2);


# In[16]:


#Multiple 2D subplots
g = plots.getSubplotPlotter(subplot_size=2.5)
g.plots_2d(samples, param_pairs=[['x0', 'x1'], ['x2', 'x3']], 
           nx=2, filled=True);


# In[17]:


#Customized triangle plot
g = plots.getSubplotPlotter()
g.settings.figure_legend_frame = False
g.settings.alpha_filled_add=0.4
legends = ['Simulation', 'Simulation 2']
g.triangle_plot([samples, samples2], ['x0', 'x1', 'x2'], 
    filled_compare=True, 
    legend_labels=legends, 
    legend_loc='upper right', 
    line_args=[{'ls':'--', 'color':'green'},
               {'lw':2, 'color':'darkblue'}], 
    contour_colors=['green','darkblue'])


# In[18]:


#3D (scatter) triangle plot
g = plots.getSubplotPlotter(width_inch=6)
g.settings.axes_fontsize = 9
g.settings.legend_fontsize = 11
g.triangle_plot([samples, samples2], ['x1', 'x2', 'x3'], 
                plot_3d_with_param='x0', legend_labels=legends)


# In[19]:


#You can reference g.subplots for manual tweaking, 
#e.g. let's add a vertical axis line in the first column 
for ax in g.subplots[:,0]:
    ax.axvline(0, color='gray', ls='--')
IPython.display.display(g.fig)


# In[20]:


#Rectangle 2D comparison plots
g = plots.getSubplotPlotter()
g.settings.figure_legend_frame = False
g.rectangle_plot(['x0', 'x1'], ['x2', 'x3'], 
    roots=[samples, samples2], filled=True, 
    plot_texts=[['Test Label', None], ['Test 2', None]]);


# In[21]:


#Example of how to handle boundaries (samples are restricted to x0 >-0.5)
cut_samps = samps[samps[:,0]>-0.5,:]
cut_samples = MCSamples(samples=cut_samps, names = names, labels = labels, 
                        ranges={'x0':(-0.5, None)})
g = plots.getSubplotPlotter()
g.plots_1d(cut_samples,nx=4);
g = plots.getSinglePlotter(width_inch=4, ratio=1)
g.plot_2d(cut_samples, 'x0', 'x1', filled=True);


# In[22]:


#Add and plot a new derived parameter
#getParms gets p so that p.x0, p.x1.. are numpy vectors of sample values
p = samples.getParams() 
samples.addDerived((5+p.x2)** 2, name='z', label='z_d')
g = plots.getSubplotPlotter()
g.plots_2d(samples,'x1',['x2','z'], nx=2);


# In[23]:


#Many other things you can do besides plot, e.g. get latex
print(cut_samples.getInlineLatex('x0',limit=2))
print(samples.getInlineLatex('x0',limit=2))


# In[24]:


print(samples.getInlineLatex('x1',limit=1))


# In[25]:


print(samples.getTable().tableTex())


# In[26]:


print(samples.PCA(['x1','x2']))


# In[27]:


stats = cut_samples.getMargeStats()
lims0 = stats.parWithName('x0').limits
lims1 = stats.parWithName('x1').limits
for conf, lim0, lim1 in zip(samples.contours,lims0, lims1):
    print('x0 %s%% lower: %.3f upper: %.3f (%s)'%(conf, lim0.lower, lim0.upper, lim0.limitType()))
    print('x1 %s%% lower: %.3f upper: %.3f (%s)'%(conf, lim1.lower, lim1.upper, lim1.limitType()))
       


# In[28]:


#Save to file
import tempfile, os
tempdir = os.path.join(tempfile.gettempdir(),'testchaindir')
os.mkdirs(tempdir)
rootname = os.path.join(tempdir, 'testchain')
samples.saveAsText(rootname)


# In[ ]:


#Load from file
from getdist import loadMCSamples
readsamps = loadMCSamples(rootname)


# In[ ]:


#Make plots from chain files, loading automatically as needed
g = plots.getSinglePlotter(chain_dir=tempdir, width_inch=4)
g.plot_2d('testchain', 'x1', 'x2', shaded=True);


# In[ ]:


#Custom settings for all loaded chains can be set as follows
g = plots.getSinglePlotter(chain_dir=tempdir, 
            analysis_settings={'ignore_rows': 0.2, 'contours':[0.2, 0.4, 0.6, 0.8]});
g.settings.num_plot_contours = 4
g.plot_2d('testchain', 'x1', 'x2', filled=False);


# In[ ]:


#Chains can be loaded by searching in multiple directories, give list as chain_dir
#(not chain names must be unique)

#make second test chain in new temp dir
temp2 = os.path.join(tempdir,'chaindir2')
cut_samples.saveAsText(os.path.join(temp2, 'testchain2'), make_dirs=True)
#Plot from chain files
g = plots.getSinglePlotter(chain_dir=[tempdir, temp2])
g.plot_2d(['testchain','testchain2'], 'x1', 'x2', filled=True);


# In[ ]:


#cleanup test files
import shutil
shutil.rmtree(tempdir)