%matplotlib inline

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

from wavefunction.wavefunction2d import *

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

globals().update(args)

# 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 = 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 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)

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)

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)

U1 = U_flux_qubit(X1, X2, args)

U2 = evalute_fourier_series(X1, X2, L1, L2, u)

fig, axes = plt.subplots(1, 2, figsize=(10, 4))

# reconstructed
Z = np.real(U2)
p = axes[0].pcolor(X1/(2*np.pi), X2/(2*np.pi), Z, cmap=cm.RdBu, vmin=Z.min(), vmax=Z.max())
cb = fig.colorbar(p, ax=axes[0])
axes[0].set_xlabel(r'$\varphi_1$', fontsize=16)
axes[0].set_ylabel(r'$\varphi_2$', fontsize=16)
axes[0].set_title('Potential from Fourier series')

# original
Z = np.real(U1)
p = axes[1].pcolor(X1/(2*np.pi), X2/(2*np.pi), Z, cmap=cm.RdBu, vmin=Z.min(), vmax=Z.max())
axes[1].set_xlabel(r'$\varphi_1$', fontsize=16)
axes[1].set_ylabel(r'$\varphi_2$', fontsize=16)
cb = fig.colorbar(p, ax=axes[1])
axes[1].set_title('Potential directly form expression');

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=(12, 3 * 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])
    cb.set_clim(-5, 5)

    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])
    cb.set_clim(-5, 5)
    
    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.47, 0.50, 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[6, :].max())

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

args['f'] = 0.49
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)

# current operator for junction 1
I1 = np.sin(X1)

# current operator for junction 2
I2 = np.sin(X2)

# ground and excited state wave functions
psi0 = wavefunction_normalize(X1, X2, evalute_fourier_series(X1, X2, L1, L2, convert_v2m(L1, L2, vecs[0])))
psi1 = wavefunction_normalize(X1, X2, evalute_fourier_series(X1, X2, L1, L2, convert_v2m(L1, L2, vecs[1])))

expectation_value(X1, X2, I1, psi0).real  # I1/Ic

expectation_value(X1, X2, I1, psi1).real  # I2/Ic

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

# 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)

    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)

phi_p = np.linspace(-1*np.pi, 1*np.pi, 100)
phi_m = np.linspace(-1*np.pi, 1*np.pi, 100)
PHI_P, PHI_M = np.meshgrid(phi_p, phi_m)

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

U1 = U_flux_qubit(PHI_P, PHI_M, args)

U2 = evalute_fourier_series(PHI_P, PHI_M, L1, L2, u)

fig, axes = plt.subplots(1, 2, figsize=(10, 4))

# reconstructed
Z = np.real(U2)
p = axes[0].pcolor(PHI_P/(2*np.pi), PHI_M/(2*np.pi), Z, cmap=cm.RdBu, vmin=Z.min(), vmax=Z.max())
cb = fig.colorbar(p, ax=axes[0])
axes[0].axis('tight')
axes[0].set_xlabel(r'$\varphi_p$', fontsize=16)
axes[0].set_ylabel(r'$\varphi_m$', fontsize=16)
axes[0].set_title('Potential from Fourier series')

# original
Z = np.real(U1)
p = axes[1].pcolor(PHI_P/(2*np.pi), PHI_M/(2*np.pi), Z, cmap=cm.RdBu, vmin=Z.min(), vmax=Z.max())
cb = fig.colorbar(p, ax=axes[1])
axes[1].set_xlabel(r'$\varphi_p$', fontsize=16)
axes[1].set_ylabel(r'$\varphi_m$', fontsize=16)
axes[1].axis('tight')
axes[1].set_title('Potential directly form expression');

V_full = 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(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 = 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, phi_m, args))

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

ax.plot(phi_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 = 4

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

for n in range(Nstates):

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

    p = axes[n, 1].pcolor(PHI_P/(2*np.pi), PHI_M/(2*np.pi), np.imag(PSI),
                          cmap=cm.RdBu, vmin=-abs(PSI.imag).max(), vmax=abs(PSI.imag).max())
    c = axes[n, 1].contour(PHI_P/(2*np.pi), PHI_M/(2*np.pi), Z, cmap=cm.RdBu, vmin=Z.min(), vmax=Z.max())
    cb = fig.colorbar(p, ax=axes[n, 1])
    cb.set_clim(-5, 5)
    
    p = axes[n, 2].pcolor(PHI_P/(2*np.pi), PHI_M/(2*np.pi), np.abs(PSI), cmap=cm.Blues, vmin=0, vmax=abs(PSI).max())
    c = axes[n, 2].contour(PHI_P/(2*np.pi), PHI_M/(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.47, 0.5, 25)

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
    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[12, :].max())

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

args['f'] = 0.49
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)

# current operator for junction 1
I1 = np.sin(PHI_P + PHI_M)

# current operator for junction 2
I2 = np.sin(PHI_P - PHI_M)

# ground and excited state wave functions
psi0 = wavefunction_normalize(PHI_P, PHI_M, evalute_fourier_series(PHI_P, PHI_M, L1, L2, convert_v2m(L1, L2, vecs[0])))
psi1 = wavefunction_normalize(PHI_P, PHI_M, evalute_fourier_series(PHI_P, PHI_M, L1, L2, convert_v2m(L1, L2, vecs[1])))

expectation_value(PHI_P, PHI_M, I1, psi0).real  # I1/Ic

expectation_value(PHI_P, PHI_M, I2, psi0).real  # I2/Ic

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