# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
import sys
import numpy as np
from scipy.integrate import odeint
import matplotlib.pyplot as plt
from matplotlib.patches import Circle
import pandas as pd
import matplotlib

import ipywidgets as widgets
from ipywidgets import HBox, VBox
import numpy as np
import matplotlib.pyplot as plt
from IPython.display import display


#Need this line for widgets to work
# %matplotlib inline

matplotlib.rcParams['mathtext.fontset'] = 'stix'
matplotlib.rcParams['font.family'] = 'STIXGeneral'
matplotlib.rcParams['font.size'] = 14
# -

# # Load in $\zeta_{\ell_1, \ell_2, \ell_3}^{b_1, b_2,b_3}$

# +
#Load in data from 4pcf calculation
zeta = np.load('4pcf_coefficients_normed.npy')

print(np.shape(zeta))

index_dict = {}

for l_1 in range(0,ell_max+1, 1):
    for l_2 in range(0,ell_max+1, 1):
        for l_3 in range(0,ell_max+1, 1):
#         for l_3 in range(np.abs(l_1 - l_2), min(l_1 + l_2, ell_max)+1):
            if l_3 >= np.abs(l_1 - l_2) and l_3 < (min(l_1 + l_2, ell_max) + 1): 
                index_dict[str(l_1) + str(l_2) + str(l_3)] = ("$\\ell_1, \ell_2, \ell_3$ = " + str(l_1) + ',' +
                                                    str(l_2) + ',' + str(l_3))
            else:
                index_dict[str(l_1) + str(l_2) + str(l_3)] = "Not Valid Combo"
                
# -

print(len(index_dict))

# +

ell_max=5
nbins = np.shape(zeta)[-1]

slider1 = widgets.IntSlider(
                            value=0,
                            min=0,
                            max=ell_max,
                            step=1,
                            description=r'\(\ell_1\)')

slider2 = widgets.IntSlider(
                            value=0,
                            min=0,
                            max=ell_max,
                            step=1,
                            description=r'\(\ell_2\)')

slider3 = widgets.IntSlider(
                            value=0,
                            min=0,
                            max=ell_max,
                            step=1,
                            description=r'\(\ell_3\)')

slider4 = widgets.IntSlider(
                            value=0,
                            min=0,
                            max=nbins-1,
                            step=1,
                            description=r'\(b_1\)')

slider5 = widgets.IntSlider(
                            value=0,
                            min=0,
                            max=nbins-1,
                            step=1,
                            description=r'\(b_2\)')

slider6 = widgets.IntSlider(
                            value=0,
                            min=0,
                            max=nbins-1,
                            step=1,
                            description=r'\(b_3\)')


def make_plot(l_1=0, l_2=0, l_3=0, b_1=0, b_2=0, b_3=0):
    xylabel_fontsize = 22
    sub_title_fontsize = 22
    sup_title_fontsize = 30
    text_fontsize = 18
    
    f, ((ax1,ax2,ax3)) = plt.subplots(1,3, figsize=(16,6))
    f.suptitle('4PCF coefficients $\\zeta^{b_1 b_2 b_3}_{\ell_1 \ell_2 \ell_3}$', fontsize=sup_title_fontsize)
    
    ax1.cla()
    ax1.imshow(zeta[l_1,l_2,l_3,b_1,:,:].real, origin='lower')
    ax1.text(0,nbins-1, index_dict[str(l_1) + str(l_2) + str(l_3)], c='white', fontsize=text_fontsize)
    ax1.set_title('$b_1$ = ' + str(b_1), fontsize=sub_title_fontsize)
    ax1.set_xticks(range(nbins))
    ax1.set_yticks(range(nbins))
    ax1.set_xlabel("$b_2$",fontsize=xylabel_fontsize)
    ax1.set_ylabel("$b_3$", fontsize=xylabel_fontsize)
    
    ax2.cla()
    ax2.imshow(zeta[l_1,l_2,l_3,:,b_2,:].real, origin='lower')
    ax2.text(0,nbins-1, index_dict[str(l_1) + str(l_2) + str(l_3)], c='white', fontsize=text_fontsize)
    ax2.set_title('$b_2$ = ' + str(b_2), fontsize=sub_title_fontsize)
    ax2.set_xticks(range(nbins))
    ax2.set_yticks(range(nbins))
    ax2.set_xlabel("$b_1$",fontsize=xylabel_fontsize)
    ax2.set_ylabel("$b_3$", fontsize=xylabel_fontsize)
    
    ax3.cla()
    ax3.imshow(zeta[l_1,l_2,l_3,:,:,b_3].real, origin='lower')
    ax3.text(0,nbins-1, index_dict[str(l_1) + str(l_2) + str(l_3)], c='white', fontsize=text_fontsize)
    ax3.set_title('$b_3$ = ' + str(b_3), fontsize=sub_title_fontsize)
    ax3.set_xticks(range(nbins))
    ax3.set_yticks(range(nbins))
    ax3.set_xlabel("$b_1$",fontsize=xylabel_fontsize)
    ax3.set_ylabel("$b_2$", fontsize=xylabel_fontsize)
    
#     f.tight_layout()



out = widgets.interactive_output(make_plot, {'l_1':slider1,
                                     'l_2':slider2, 'l_3':slider3, 'b_1':slider4,
                                            'b_2':slider5, 'b_3':slider6,})

# we can add vertical and horizontal boxes for styling!
box1 = widgets.VBox([slider1, slider4])
ui1 = widgets.HBox([box1])

box2 = widgets.VBox([slider2,slider5])
ui2 = widgets.HBox([box2])

box3 = widgets.VBox([slider3, slider6])
ui3 = widgets.HBox([box3])

ui = widgets.HBox([ui1, ui2, ui3])



display(ui,out)


# +
"""If you prefer %matplotlib notebook settings, use this version. They do the same thing.
(just remember to switch to the notebook setting in the first cell)"""

# ell_max=5
# nbins = np.shape(zeta)[-1]

# slider1 = widgets.IntSlider(
#                             value=0,
#                             min=0,
#                             max=ell_max,
#                             step=1,
#                             description=r'\(\ell_1\)')

# slider2 = widgets.IntSlider(
#                             value=0,
#                             min=0,
#                             max=ell_max,
#                             step=1,
#                             description=r'\(\ell_2\)')

# slider3 = widgets.IntSlider(
#                             value=0,
#                             min=0,
#                             max=ell_max,
#                             step=1,
#                             description=r'\(\ell_3\)')

# slider4 = widgets.IntSlider(
#                             value=0,
#                             min=0,
#                             max=nbins-1,
#                             step=1,
#                             description=r'\(b_1\)')

# slider5 = widgets.IntSlider(
#                             value=0,
#                             min=0,
#                             max=nbins-1,
#                             step=1,
#                             description=r'\(b_2\)')

# slider6 = widgets.IntSlider(
#                             value=0,
#                             min=0,
#                             max=nbins-1,
#                             step=1,
#                             description=r'\(b_3\)')


# def make_plot(l_1=0, l_2=0, l_3=0, b_1=0, b_2=0, b_3=0):
    
#     xylabel_fontsize = 22
#     sub_title_fontsize = 22
#     sup_title_fontsize = 30
#     text_fontsize = 18
    
#     ax1.cla()
#     ax1.imshow(zeta[l_1,l_2,l_3,b_1,:,:].real, origin='lower')
#     ax1.text(0,nbins-1, index_dict[str(l_1) + str(l_2) + str(l_3)], c='white', fontsize=text_fontsize)
#     ax1.set_title('$b_1$ = ' + str(b_1), fontsize=sub_title_fontsize)
#     ax1.set_xticks(range(nbins))
#     ax1.set_yticks(range(nbins))
#     ax1.set_xlabel("$b_2$",fontsize=xylabel_fontsize)
#     ax1.set_ylabel("$b_3$", fontsize=xylabel_fontsize)
    
#     ax2.cla()
#     ax2.imshow(zeta[l_1,l_2,l_3,:,b_2,:].real, origin='lower')
#     ax2.text(0,nbins-1, index_dict[str(l_1) + str(l_2) + str(l_3)], c='white', fontsize=text_fontsize)
#     ax2.set_title('$b_2$ = ' + str(b_2), fontsize=sub_title_fontsize)
#     ax2.set_xticks(range(nbins))
#     ax2.set_yticks(range(nbins))
#     ax2.set_xlabel("$b_1$",fontsize=xylabel_fontsize)
#     ax2.set_ylabel("$b_3$", fontsize=xylabel_fontsize)
    
#     ax3.cla()
#     ax3.imshow(zeta[l_1,l_2,l_3,:,:,b_3].real, origin='lower')
#     ax3.text(0,nbins-1, index_dict[str(l_1) + str(l_2) + str(l_3)], c='white', fontsize=text_fontsize)
#     ax3.set_title('$b_3$ = ' + str(b_3), fontsize=sub_title_fontsize)
#     ax3.set_xticks(range(nbins))
#     ax3.set_yticks(range(nbins))
#     ax3.set_xlabel("$b_1$",fontsize=xylabel_fontsize)
#     ax3.set_ylabel("$b_2$", fontsize=xylabel_fontsize)



# out = widgets.interactive_output(make_plot, {'l_1':slider1,
#                                      'l_2':slider2, 'l_3':slider3, 'b_1':slider4,
#                                             'b_2':slider5, 'b_3':slider6,})



# f, ((ax1,ax2,ax3)) = plt.subplots(1,3, figsize=(11,4))
# f.suptitle('4PCF coefficients $\\zeta^{b_1 b_2 b_3}_{\ell_1 \ell_2 \ell_3}$', fontsize=22)



# # we can add vertical and horizontal boxes for styling!
# box1 = widgets.VBox([slider1])
# ui1 = widgets.HBox([box1])

# box2 = widgets.VBox([slider2])
# ui2 = widgets.HBox([box2])

# box3 = widgets.VBox([slider3])
# ui3 = widgets.HBox([box3])

# box4 = widgets.VBox([slider4])
# ui4 = widgets.HBox([box4])

# box5 = widgets.VBox([slider5])
# ui5 = widgets.HBox([box5])

# box6 = widgets.VBox([slider6])
# ui6 = widgets.HBox([box6])

# display(ui1, out)
# display(ui2, out)
# display(ui3, out)
# print("This slider is for left plot")
# display(ui4, out)
# print("This slider is for middle plot")
# display(ui5, out)
# print("This slider is for right plot")
# display(ui6, out)
# make_plot()
# -

# # One can visualize all the plots side by side with the following cells

# +
# plt.figure(figsize=(15, 70))
# plt.suptitle('4PCF coefficients $\\zeta^{b_1 b_2 b_3}_{\ell_1 \ell_2 \ell_3}$ for $b_3 = 0$', fontsize='large')

# i = 0
# for l_1 in range(0,ell_max+1, 1):
#     for l_2 in range(0,ell_max+1, 1):
#         for l_3 in range(np.abs(l_1 - l_2), min(l_1 + l_2, ell_max)+1): 
#             plt.subplot(23, 5, (i)+1) 
#             plt.text(0,nbins - 1, index_dict[str(l_1) + str(l_2) + str(l_3)], c='white', fontsize=18)
#             plt.imshow(zeta[:,:,0,l_1,l_2,l_3].real, origin='lower')
#             plt.xticks(range(nbins))
#             plt.yticks(range(nbins))
#             if i == 0:
#                 plt.ylabel("$b_2$", fontsize=20)
#                 plt.title("$b_1$", fontsize = 20)
#             i+=1
            

# plt.tight_layout(rect=[0, 0.03, 1, 0.98])
# plt.show()

# print("\n\n\n")

# plt.figure(figsize=(15, 70))
# plt.suptitle('4PCF coefficients $\\zeta^{b_1 b_2 b_3}_{\ell_1 \ell_2 \ell_3}$ for $b_3 = 1$', fontsize='large')

# i = 0
# for l_1 in range(0,ell_max+1, 1):
#     for l_2 in range(0,ell_max+1, 1):
#         for l_3 in range(np.abs(l_1 - l_2), min(l_1 + l_2, ell_max)+1): 
#             plt.subplot(23, 5, (i)+1) 
#             plt.text(0,nbins - 1, index_dict[str(l_1) + str(l_2) + str(l_3)], c='white', fontsize=18)
#             plt.imshow(zeta[:,:,1,l_1,l_2,l_3].real, origin='lower')
#             plt.xticks(range(nbins))
#             plt.yticks(range(nbins))
#             if i == 0:
#                 plt.ylabel("$b_2$", fontsize=20)
#                 plt.title("$b_1$", fontsize = 20)
#             i+=1
            

# plt.tight_layout(rect=[0, 0.03, 1, 0.98])
# plt.show()   

# print("\n\n\n")

# plt.figure(figsize=(15, 70))
# plt.suptitle('4PCF coefficients $\\zeta^{b_1 b_2 b_3}_{\ell_1 \ell_2 \ell_3}$ for $b_3 = 2$', fontsize='large')

# i = 0
# for l_1 in range(0,ell_max+1, 1):
#     for l_2 in range(0,ell_max+1, 1):
#         for l_3 in range(np.abs(l_1 - l_2), min(l_1 + l_2, ell_max)+1): 
#             plt.subplot(23, 5, (i)+1) 
#             plt.text(0,nbins - 1, index_dict[str(l_1) + str(l_2) + str(l_3)], c='white', fontsize=18)
#             plt.imshow(zeta[:,:,2,l_1,l_2,l_3].real, origin='lower')
#             plt.xticks(range(nbins))
#             plt.yticks(range(nbins))
#             if i == 0:
#                 plt.ylabel("$b_2$", fontsize=20)
#                 plt.title("$b_1$", fontsize = 20)
#             i+=1
            

# plt.tight_layout(rect=[0, 0.03, 1, 0.98])
# plt.show()  

# print("\n\n\n")

# plt.figure(figsize=(15, 70))
# plt.suptitle('4PCF coefficients $\\zeta^{b_1 b_2 b_3}_{\ell_1 \ell_2 \ell_3}$ for $b_3 = 3$', fontsize='large')

# i = 0
# for l_1 in range(0,ell_max+1, 1):
#     for l_2 in range(0,ell_max+1, 1):
#         for l_3 in range(np.abs(l_1 - l_2), min(l_1 + l_2, ell_max)+1): 
#             plt.subplot(23, 5, (i)+1) 
#             plt.text(0,nbins - 1, index_dict[str(l_1) + str(l_2) + str(l_3)], c='white', fontsize=18)
#             plt.imshow(zeta[:,:,3,l_1,l_2,l_3].real, origin='lower')
#             plt.xticks(range(nbins))
#             plt.yticks(range(nbins))
#             if i == 0:
#                 plt.ylabel("$b_2$", fontsize=20)
#                 plt.title("$b_1$", fontsize = 20)
#             i+=1
            

# plt.tight_layout(rect=[0, 0.03, 1, 0.98])
# plt.show()  

# +
# plt.figure(figsize=(15, 70))
# plt.suptitle('4PCF coefficients $\\zeta^{b_1 b_2 b_3}_{\ell_1 \ell_2 \ell_3}$ for $b_2 = 0$', fontsize='large')

# i = 0
# for l_1 in range(0,ell_max+1, 1):
#     for l_2 in range(0,ell_max+1, 1):
#         for l_3 in range(np.abs(l_1 - l_2), min(l_1 + l_2, ell_max)+1): 
#             plt.subplot(23, 5, (i)+1) 
#             plt.text(0,nbins - 1, index_dict[str(l_1) + str(l_2) + str(l_3)], c='white', fontsize=18)
#             plt.imshow(zeta[:,0,:,l_1,l_2,l_3].real, origin='lower')
#             plt.xticks(range(nbins))
#             plt.yticks(range(nbins))
#             if i == 0:
#                 plt.ylabel("$b_3$", fontsize=20)
#                 plt.title("$b_1$", fontsize = 20)
#             i+=1
            

# plt.tight_layout(rect=[0, 0.03, 1, 0.98])
# plt.show()

# print("\n\n\n")

# plt.figure(figsize=(15, 70))
# plt.suptitle('4PCF coefficients $\\zeta^{b_1 b_2 b_3}_{\ell_1 \ell_2 \ell_3}$ for $b_2 = 1$', fontsize='large')

# i = 0
# for l_1 in range(0,ell_max+1, 1):
#     for l_2 in range(0,ell_max+1, 1):
#         for l_3 in range(np.abs(l_1 - l_2), min(l_1 + l_2, ell_max)+1): 
#             plt.subplot(23, 5, (i)+1) 
#             plt.text(0,nbins - 1, index_dict[str(l_1) + str(l_2) + str(l_3)], c='white', fontsize=18)
#             plt.imshow(zeta[:,1,:,l_1,l_2,l_3].real, origin='lower')
#             plt.xticks(range(nbins))
#             plt.yticks(range(nbins))
#             if i == 0:
#                 plt.ylabel("$b_3$", fontsize=20)
#                 plt.title("$b_1$", fontsize = 20)
#             i+=1
            

# plt.tight_layout(rect=[0, 0.03, 1, 0.98])
# plt.show()   

# print("\n\n\n")

# plt.figure(figsize=(15, 70))
# plt.suptitle('4PCF coefficients $\\zeta^{b_1 b_2 b_3}_{\ell_1 \ell_2 \ell_3}$ for $b_2 = 2$', fontsize='large')

# i = 0
# for l_1 in range(0,ell_max+1, 1):
#     for l_2 in range(0,ell_max+1, 1):
#         for l_3 in range(np.abs(l_1 - l_2), min(l_1 + l_2, ell_max)+1): 
#             plt.subplot(23, 5, (i)+1) 
#             plt.text(0,nbins - 1, index_dict[str(l_1) + str(l_2) + str(l_3)], c='white', fontsize=18)
#             plt.imshow(zeta[:,2,:,l_1,l_2,l_3].real, origin='lower')
#             plt.xticks(range(nbins))
#             plt.yticks(range(nbins))
#             if i == 0:
#                 plt.ylabel("$b_3$", fontsize=20)
#                 plt.title("$b_1$", fontsize = 20)
#             i+=1
            

# plt.tight_layout(rect=[0, 0.03, 1, 0.98])
# plt.show()  

# print("\n\n\n")

# plt.figure(figsize=(15, 70))
# plt.suptitle('4PCF coefficients $\\zeta^{b_1 b_2 b_3}_{\ell_1 \ell_2 \ell_3}$ for $b_2 = 3$', fontsize='large')

# i = 0
# for l_1 in range(0,ell_max+1, 1):
#     for l_2 in range(0,ell_max+1, 1):
#         for l_3 in range(np.abs(l_1 - l_2), min(l_1 + l_2, ell_max)+1): 
#             plt.subplot(23, 5, (i)+1) 
#             plt.text(0,nbins - 1, index_dict[str(l_1) + str(l_2) + str(l_3)], c='white', fontsize=18)
#             plt.imshow(zeta[:,3,:,l_1,l_2,l_3].real, origin='lower')
#             plt.xticks(range(nbins))
#             plt.yticks(range(nbins))
#             if i == 0:
#                 plt.ylabel("$b_3$", fontsize=20)
#                 plt.title("$b_1$", fontsize = 20)
#             i+=1
            

# plt.tight_layout(rect=[0, 0.03, 1, 0.98])
# plt.show()  

# +
# plt.figure(figsize=(15, 70))
# plt.suptitle('4PCF coefficients $\\zeta^{b_1 b_2 b_3}_{\ell_1 \ell_2 \ell_3}$ for $b_1 = 0$', fontsize='large')

# i = 0
# for l_1 in range(0,ell_max+1, 1):
#     for l_2 in range(0,ell_max+1, 1):
#         for l_3 in range(np.abs(l_1 - l_2), min(l_1 + l_2, ell_max)+1): 
#             plt.subplot(23, 5, (i)+1) 
#             plt.text(0,nbins - 1, index_dict[str(l_1) + str(l_2) + str(l_3)], c='white', fontsize=18)
#             plt.imshow(zeta[0,:,:,l_1,l_2,l_3].real, origin='lower')
#             plt.xticks(range(nbins))
#             plt.yticks(range(nbins))
#             if i == 0:
#                 plt.ylabel("$b_3$", fontsize=20)
#                 plt.title("$b_2$", fontsize = 20)
#             i+=1
            

# plt.tight_layout(rect=[0, 0.03, 1, 0.98])
# plt.show()

# print("\n\n\n")

# plt.figure(figsize=(15, 70))
# plt.suptitle('4PCF coefficients $\\zeta^{b_1 b_2 b_3}_{\ell_1 \ell_2 \ell_3}$ for $b_1 = 1$', fontsize='large')

# i = 0
# for l_1 in range(0,ell_max+1, 1):
#     for l_2 in range(0,ell_max+1, 1):
#         for l_3 in range(np.abs(l_1 - l_2), min(l_1 + l_2, ell_max)+1): 
#             plt.subplot(23, 5, (i)+1) 
#             plt.text(0,nbins - 1, index_dict[str(l_1) + str(l_2) + str(l_3)], c='white', fontsize=18)
#             plt.imshow(zeta[1,:,:,l_1,l_2,l_3].real, origin='lower')
#             plt.xticks(range(nbins))
#             plt.yticks(range(nbins))
#             if i == 0:
#                 plt.ylabel("$b_3$", fontsize=20)
#                 plt.title("$b_2$", fontsize = 20)
#             i+=1
            

# plt.tight_layout(rect=[0, 0.03, 1, 0.98])
# plt.show()   

# print("\n\n\n")

# plt.figure(figsize=(15, 70))
# plt.suptitle('4PCF coefficients $\\zeta^{b_1 b_2 b_3}_{\ell_1 \ell_2 \ell_3}$ for $b_1 = 2$', fontsize='large')

# i = 0
# for l_1 in range(0,ell_max+1, 1):
#     for l_2 in range(0,ell_max+1, 1):
#         for l_3 in range(np.abs(l_1 - l_2), min(l_1 + l_2, ell_max)+1): 
#             plt.subplot(23, 5, (i)+1) 
#             plt.text(0,nbins - 1, index_dict[str(l_1) + str(l_2) + str(l_3)], c='white', fontsize=18)
#             plt.imshow(zeta[2,:,:,l_1,l_2,l_3].real, origin='lower')
#             plt.xticks(range(nbins))
#             plt.yticks(range(nbins))
#             if i == 0:
#                 plt.ylabel("$b_3$", fontsize=20)
#                 plt.title("$b_2$", fontsize = 20)
#             i+=1
            

# plt.tight_layout(rect=[0, 0.03, 1, 0.98])
# plt.show()  

# print("\n\n\n")

# plt.figure(figsize=(15, 70))
# plt.suptitle('4PCF coefficients $\\zeta^{b_1 b_2 b_3}_{\ell_1 \ell_2 \ell_3}$ for $b_1 = 3$', fontsize='large')

# i = 0
# for l_1 in range(0,ell_max+1, 1):
#     for l_2 in range(0,ell_max+1, 1):
#         for l_3 in range(np.abs(l_1 - l_2), min(l_1 + l_2, ell_max)+1): 
#             plt.subplot(23, 5, (i)+1) 
#             plt.text(0,nbins - 1, index_dict[str(l_1) + str(l_2) + str(l_3)], c='white', fontsize=18)
#             plt.imshow(zeta[3,:,:,l_1,l_2,l_3].real, origin='lower')
#             plt.xticks(range(nbins))
#             plt.yticks(range(nbins))
#             if i == 0:
#                 plt.ylabel("$b_3$", fontsize=20)
#                 plt.title("$b_2$", fontsize = 20)
#             i+=1
            

# plt.tight_layout(rect=[0, 0.03, 1, 0.98])
# plt.show()  

StarcoderdataJupyterScriptsDedupFiltered