%matplotlib inline

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cm as cm

from wavefunction import *
from wavefunction.wavefunction2d import *

args = {'Ec': 1/80.0, 'Ej': 1.0, 'alpha': 0.8, 'gamma': 0.0, 'f': 0.50}

x1 = np.linspace(-2*np.pi, 2*np.pi, 100)
x2 = np.linspace(-2*np.pi, 2*np.pi, 100)
X1, X2 = np.meshgrid(x1, x2)

def U_flux_qubit(x1, x2, args):
    globals().update(args)
    return 2 + alpha - np.cos(x1) - np.cos(x2) - alpha * np.cos(2 * np.pi * f + x1 - x2)

Z = np.real(U_flux_qubit(X1, X2, args))
fig, ax = plt.subplots()
p = ax.pcolor(X1/(2*np.pi), X2/(2*np.pi), Z, cmap=cm.RdBu, vmin=Z.min(), vmax=Z.max())
ax.set_xlabel(r'$x_1$', fontsize=16)
ax.set_ylabel(r'$x_2$', fontsize=16)
cb = fig.colorbar(p, ax=ax)

# pick a truncation of the fourier series: n1 = [-L1, ..., L1], n2 = [-L2, ..., L2]
L1 = 5
L2 = 5

# pick periods for the coordinates
Tx1 = Tx2 = 2 * np.pi

#
k11 = k22 = 4 * (Ec / Ej) * (1 + alpha + gamma) / ((1 + gamma) * (1 + 2 * alpha + gamma))
k12 = 4 * (Ec / Ej) * 2 * alpha / ((1 + gamma) * (1 + 2 * alpha + gamma))

K = assemble_K(L1, L2, k11, k12, k22, Tx1, Tx2)

def assemble_u_flux_qubit(L1, L2, args, sparse=False):

    globals().update(args)
    
    L1n = 2 * L1 + 1
    L2n = 2 * L2 + 1
    
    u = np.zeros((L1n, L2n), dtype=np.complex)
    
    for n1 in np.arange(-L1, L1+1):
        N1 = n1 + L1
        for n2 in np.arange(-L2, L2+1):
            N2 = n2 + L2
            u[N1, N2] = (2 + alpha) * delta(0, n1) * delta(0, n2) + \
                        - 0.5 * (delta(1, n1) + delta(-1, n1)) * delta(0, n2) + \
                        - 0.5 * (delta(1, n2) + delta(-1, n2)) * delta(0, n1)+ \
                        - 0.5 * alpha * np.exp(+2j * np.pi * f) * delta(+1, n1) * delta(-1, n2) + \
                        - 0.5 * alpha * np.exp(-2j * np.pi * f) * delta(-1, n1) * delta(+1, n2)                    
    return u

u = assemble_u_flux_qubit(L1, L2, args)

Z = np.real(evalute_fourier_series(X1, X2, L1, L2, u))
fig, ax = plt.subplots()
p = ax.pcolor(X1/(2*np.pi), X2/(2*np.pi), Z, cmap=cm.RdBu, vmin=Z.min(), vmax=Z.max())
cb = fig.colorbar(p, ax=ax)

V = assemble_V(L1, L2, u)

def assemble_V_flux_qubit(L1, L2, args, sparse=False):
    
    globals().update(args)
    
    L1n = 2 * L1 + 1
    L2n = 2 * L2 + 1
    
    V = np.zeros((L1n*L1n, L2n*L2n), dtype=np.complex)
    
    for n1 in np.arange(-L1, L1+1):
        for n2 in np.arange(-L1, L1+1):
            N = index_m2v(L1, n1, n2)
            for m1 in np.arange(-L2, L2+1):
                for m2 in np.arange(-L2, L2+1):
                    M = index_m2v(L2, m1, m2)

                    V[N,M] = (2 + alpha) * delta(m1, n1) * delta(m2, n2) + \
                             - 0.5 * (delta(m1 + 1, n1) + delta(m1 - 1, n1)) * delta(m2, n2) + \
                             - 0.5 * (delta(m2 + 1, n2) + delta(m2 - 1, n2)) * delta(m1, n1) + \
                             - 0.5 * alpha * np.exp(+2j * np.pi * f) * delta(m1 + 1, n1) * delta(m2 - 1, n2) + \
                             - 0.5 * alpha * np.exp(-2j * np.pi * f) * delta(m1 - 1, n1) * delta(m2 + 1, n2)

    return V

V_fq = assemble_V_flux_qubit(L1, L2, args)

abs(V - V_fq).max()

H = K + V

vals, vecs = solve_eigenproblem(H)

fig, ax = plt.subplots()

U_x = np.real(U_flux_qubit(x1, -x1, args))

for val in vals:
    ax.plot(x1, np.real(val) * np.ones_like(x1), 'k')

ax.plot(x1, U_x, label="potential", lw='2')
    
ax.axis('tight')
ax.set_ylim(min(vals[0], U_x.min()), U_x.max())
ax.set_ylabel(r'$U(x_1, x_2=x_1)$', fontsize=16)
ax.set_xlabel(r'$x_1$', fontsize=16);

Nstates = 4

fig, axes = plt.subplots(Nstates, 3, figsize=(16, 5 * Nstates))

for n in range(Nstates):

    psi = convert_v2m(L1, L2, vecs[n])
    
    Z = np.real(evalute_fourier_series(X1, X2, L1, L2, u))
    PSI = evalute_fourier_series(X1, X2, L1, L2, psi)
    
    p = axes[n, 0].pcolor(X1/(2*np.pi), X2/(2*np.pi), np.real(PSI),
                          cmap=cm.RdBu, vmin=-abs(PSI.real).max(), vmax=abs(PSI.real).max())
    c = axes[n, 0].contour(X1/(2*np.pi), X2/(2*np.pi), Z, cmap=cm.RdBu, vmin=Z.min(), vmax=Z.max())
    cb = fig.colorbar(p, ax=axes[n, 0])

    p = axes[n, 1].pcolor(X1/(2*np.pi), X2/(2*np.pi), np.imag(PSI),
                          cmap=cm.RdBu, vmin=-abs(PSI.imag).max(), vmax=abs(PSI.imag).max())
    c = axes[n, 1].contour(X1/(2*np.pi), X2/(2*np.pi), Z, cmap=cm.RdBu, vmin=Z.min(), vmax=Z.max())
    cb = fig.colorbar(p, ax=axes[n, 1])
    
    p = axes[n, 2].pcolor(X1/(2*np.pi), X2/(2*np.pi), np.abs(PSI), cmap=cm.Blues, vmin=0, vmax=abs(PSI).max())
    c = axes[n, 2].contour(X1/(2*np.pi), X2/(2*np.pi), Z, cmap=cm.RdBu, vmin=Z.min(), vmax=Z.max())
    cb = fig.colorbar(p, ax=axes[n, 2])

    axes[n, 0].set_ylabel("%d state" % n);

axes[0, 0].set_title(r"$\mathrm{re}\;\Psi(\varphi_p, \varphi_m)$", fontsize=16);
axes[0, 1].set_title(r"$\mathrm{im}\;\Psi(\varphi_p, \varphi_m)$", fontsize=16);
axes[0, 2].set_title(r"$|\Psi(\varphi_p, \varphi_m)|^2$", fontsize=16);

f_vec = np.linspace(0.45, 0.55, 50)
e_vals = np.zeros((len(vals), len(f_vec)))

K = assemble_K(L1, L2, k11, k12, k22, Tx1, Tx2)

for f_idx, f in enumerate(f_vec):
    args['f'] = f
    #u = assemble_u_flux_qubit(L1, L2, args)
    #V = assemble_V(L1, L2, u)
    V = assemble_V_flux_qubit(L1, L2, args)
    H = K + V
    vals, vecs = solve_eigenproblem(H)
    
    e_vals[:, f_idx] = np.real(vals)

fig, ax = plt.subplots(1, 1, figsize=(10,6))

for n in range(len(vals)):
    ax.plot(f_vec, e_vals[n, :])

ax.axis('tight')

ax.set_ylim(e_vals[0, :].min(), e_vals[20, :].max())

ax.set_ylabel("Energy levels")
ax.set_xlabel(r'$f$');

args = {'Ec': 1/80.0, 'Ej': 1.0, 'alpha': 0.8, 'gamma': 0.0, 'f': 0.50}

x_p = np.linspace(-2*np.pi, 2*np.pi, 100)
x_m = np.linspace(-2*np.pi, 2*np.pi, 100)
X1, X2 = np.meshgrid(x_p, x_m)

def U_flux_qubit(x_p, x_m, args):
    globals().update(args)
    return 2 + alpha - 2 * np.cos(x_p) * np.cos(x_m) - alpha * np.cos(2 * np.pi * f + 2 * x_m)

Z = np.real(U_flux_qubit(X1, X2, args))
fig, ax = plt.subplots()
p = ax.pcolor(X1/(2*np.pi), X2/(2*np.pi), Z, cmap=cm.RdBu, vmin=Z.min(), vmax=Z.max())
ax.set_xlabel(r'$x_p$', fontsize=16)
ax.set_ylabel(r'$x_m$', fontsize=16)
cb = fig.colorbar(p, ax=ax)

# pick a truncation of the fourier series: n1 = [-L1, ..., L1], n2 = [-L2, ..., L2]
L1 = 10
L2 = 10

# pick periods for the coordinates
Tx1 = Tx2 = 2 * np.pi

#
k11, k12, k22 = 2 / (1 + gamma) * Ec / Ej, 0.0, 2 / (1 + 2 * alpha + gamma) * Ec / Ej

K = assemble_K(L1, L2, k11, k12, k22, Tx1, Tx2)

def assemble_u_flux_qubit(L1, L2, args, sparse=False):

    globals().update(args)
    
    L1n = 2 * L1 + 1
    L2n = 2 * L2 + 1
    
    u = np.zeros((L1n, L2n), dtype=np.complex)
    
    for n1 in np.arange(-L1, L1+1):
        N1 = n1 + L1
        for n2 in np.arange(-L2, L2+1):
            N2 = n2 + L2
            u[N1, N2] = (2 + alpha) * delta(0, n1) * delta(0, n2) + \
                        - 0.5 * (delta(+1, n1) * delta(+1, n2) + delta(+1, n1) * delta(-1, n2) +
                                 delta(-1, n1) * delta(+1, n2) + delta(-1, n1) * delta(-1, n2)) + \
                        - 0.5 * alpha * np.exp(+2j * np.pi * f) * delta(0, n1) * delta(+2, n2) + \
                        - 0.5 * alpha * np.exp(-2j * np.pi * f) * delta(0, n1) * delta(-2, n2)  
            
    return u

def assemble_u_flux_qubit_direct(L1, L2, args, sparse=False):

    globals().update(args)
    
    L1n = 2 * L1 + 1
    L2n = 2 * L2 + 1
    
    u = np.zeros((L1n, L2n), dtype=np.complex)

    u[L1, L2] = 2 + alpha

    u[+1+L1, +1+L2] = - 0.5
    u[-1+L1, +1+L2] = - 0.5
    u[+1+L1, -1+L2] = - 0.5
    u[-1+L1, -1+L2] = - 0.5

    u[L1, +2+L2] = -0.5 * alpha * np.exp(+2j * np.pi * f)
    u[L1, -2+L2] = -0.5 * alpha * np.exp(-2j * np.pi * f)

    return u

u = assemble_u_flux_qubit(L1, L2, args)

u_direct = assemble_u_flux_qubit_direct(L1, L2, args)

abs(u - u_direct).max() == 0

Z = np.real(evalute_fourier_series(X1, X2, L1, L2, u))
fig, ax = plt.subplots()
p = ax.pcolor(X1/(2*np.pi), X2/(2*np.pi), Z, cmap=cm.RdBu, vmin=Z.min(), vmax=Z.max())
ax.set_xlabel(r'$x_p$', fontsize=16)
ax.set_ylabel(r'$x_m$', fontsize=16)
cb = fig.colorbar(p, ax=ax)

def assemble_V_flux_qubit(L1, L2, args, sparse=False):
    
    globals().update(args)
    
    L1n = 2 * L1 + 1
    L2n = 2 * L2 + 1
    
    V = np.zeros((L1n*L1n, L2n*L2n), dtype=np.complex)
    
    for n1 in np.arange(-L1, L1+1):
        for n2 in np.arange(-L1, L1+1):
            N = index_m2v(L1, n1, n2)
            for m1 in np.arange(-L2, L2+1):
                for m2 in np.arange(-L2, L2+1):
                    M = index_m2v(L2, m1, m2)

                    V[N,M] = (2 + alpha) * delta(m1, n1) * delta(m2, n2) + \
                             - 0.5 * (delta(m1 + 1, n1) * delta(m2 + 1, n2) + 
                                      delta(m1 + 1, n1) * delta(m2 - 1, n2) + 
                                      delta(m1 - 1, n1) * delta(m2 + 1, n2) + 
                                      delta(m1 - 1, n1) * delta(m2 - 1, n2)) + \
                             - 0.5 * alpha * np.exp(+2j * np.pi * f) * delta(m1, n1) * delta(m2 + 2, n2) + \
                             - 0.5 * alpha * np.exp(-2j * np.pi * f) * delta(m1, n1) * delta(m2 - 2, n2)

    return V

V_full = assemble_V(L1, L2, u)

V = assemble_V_flux_qubit(L1, L2, args)

abs(V-V_full).max()

H = K + V

vals, vecs = solve_eigenproblem(H)

fig, ax = plt.subplots()

U_x = np.real(U_flux_qubit(0, x_m, args))

for val in vals:
    ax.plot(x_m, np.real(val) * np.ones_like(x_m), 'k')

ax.plot(x_m, U_x, label="potential", lw='2')
    
ax.axis('tight')
ax.set_ylim(min(vals[0], U_x.min()), U_x.max())
ax.set_ylabel(r'$U(x_1, x_2=x_1)$', fontsize=16)
ax.set_xlabel(r'$x_1$', fontsize=16);

Nstates = 6

fig, axes = plt.subplots(Nstates, 3, figsize=(16, 5 * Nstates))

for n in range(Nstates):

    psi = convert_v2m(L1, L2, vecs[n])
    
    Z = np.real(evalute_fourier_series(X1, X2, L1, L2, u))
    PSI = evalute_fourier_series(X1, X2, L1, L2, psi)
    
    p = axes[n, 0].pcolor(X1/(2*np.pi), X2/(2*np.pi), np.real(PSI),
                          cmap=cm.RdBu, vmin=-abs(PSI.real).max(), vmax=abs(PSI.real).max())
    c = axes[n, 0].contour(X1/(2*np.pi), X2/(2*np.pi), Z, cmap=cm.RdBu, vmin=Z.min(), vmax=Z.max())
    cb = fig.colorbar(p, ax=axes[n, 0])

    p = axes[n, 1].pcolor(X1/(2*np.pi), X2/(2*np.pi), np.imag(PSI),
                          cmap=cm.RdBu, vmin=-abs(PSI.imag).max(), vmax=abs(PSI.imag).max())
    c = axes[n, 1].contour(X1/(2*np.pi), X2/(2*np.pi), Z, cmap=cm.RdBu, vmin=Z.min(), vmax=Z.max())
    cb = fig.colorbar(p, ax=axes[n, 1])
    
    p = axes[n, 2].pcolor(X1/(2*np.pi), X2/(2*np.pi), np.abs(PSI), cmap=cm.Blues, vmin=0, vmax=abs(PSI).max())
    c = axes[n, 2].contour(X1/(2*np.pi), X2/(2*np.pi), Z, cmap=cm.RdBu, vmin=Z.min(), vmax=Z.max())
    cb = fig.colorbar(p, ax=axes[n, 2])

    axes[n, 0].set_ylabel("%d state" % n);

axes[0, 0].set_title(r"$\mathrm{re}\;\Psi(\varphi_p, \varphi_m)$", fontsize=16);
axes[0, 1].set_title(r"$\mathrm{im}\;\Psi(\varphi_p, \varphi_m)$", fontsize=16);
axes[0, 2].set_title(r"$|\Psi(\varphi_p, \varphi_m)|^2$", fontsize=16);

f_vec = np.linspace(0.4, 0.6, 50)
f_vec = np.linspace(0.47, 0.5, 25)

f_vec = np.linspace(0.45, 0.55, 50)

e_vals = np.zeros((len(vals), len(f_vec)))

for f_idx, f in enumerate(f_vec):
    args['f'] = f
    K = assemble_K(L1, L2, k11, k12, k22, Tx1, Tx2)
    V = assemble_V_flux_qubit(L1, L2, args)
    H = K + V
    vals, vecs = solve_eigenproblem(H)
    
    e_vals[:, f_idx] = np.real(vals)

fig, ax = plt.subplots(1, 1, figsize=(10,6))

for n in range(len(vals)):
    ax.plot(f_vec, e_vals[n, :])

ax.axis('tight')

ax.set_ylim(e_vals[0, :].min(), e_vals[20, :].max())

ax.set_ylabel("Energy levels")
ax.set_xlabel(r'$f$');

%reload_ext version_information
%version_information numpy, scipy, matplotlib, wavefunction