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

# ## Plotting options 

# In[1]:


import numpy as np
import matplotlib.pyplot as plt
get_ipython().run_line_magic('matplotlib', 'inline')
import matplotlib as mpl
import sys
label_size = 18
mpl.rcParams['xtick.labelsize'] = label_size 
mpl.rcParams['ytick.labelsize'] = label_size 
from matplotlib import rc
rc('font',**{'family':'sans-serif','sans-serif':['Helvetica']})
## for Palatino and other serif fonts use:
#rc('font',**{'family':'serif','serif':['Palatino']})
rc('text', usetex=True)


# ### Domain

# In[2]:


ψ1, ψ2 = np.linspace(0, 4*np.pi, 500), np.linspace(0, 4*np.pi,500)


# In[3]:


X,Y = np.meshgrid(ψ1, ψ2)


# ### Function, potential energy for a tunable $\alpha$-qubit

# In[4]:


def tunable_alpha_qubit(β, ψ1, ψ2, fα, f):
    return 2 + 2*β - np.cos(ψ1) - np.cos(ψ2) - 2*β*np.cos(np.pi*fα)*np.cos(2*np.pi*f + np.pi*fα - ψ1- ψ2)


# In[5]:


Z = tunable_alpha_qubit(.61, X, Y, 0.2, .1).T


# ### Density plot 

# In[6]:


plt.figure(figsize=(6,4))
plt.imshow(Z, interpolation = 'nearest', origin="lower", extent=[0,4*np.pi,0,4*np.pi],
           cmap = "inferno")
cbar = plt.colorbar()
cbar.ax.tick_params(labelsize = 16)
cbar.outline.set_visible(False)
plt.xticks([0,np.pi, 2*np.pi,3*np.pi, 4*np.pi], [r"$0$", r"$\pi$", r"$2\pi$", r"$3\pi$", r"$4\pi$"])
plt.yticks([0,np.pi, 2*np.pi,3*np.pi, 4*np.pi],[r"$0$", r"$\pi$", r"$2\pi$", r"$3\pi$", r"$4\pi$"])
plt.xlabel('$\psi_1$', fontsize = 22)
plt.ylabel('$\psi_2$', fontsize = 22)
plt.savefig('alphaqubitD' + '.png', format='png', dpi= 150, bbox_inches='tight');


# ### Packages, version control

# In[3]:


print(f"Python version: {sys.version}")
print(" ")
print("numpy == ", np.__version__)
print("matplotlib == ", mpl.__version__)