The Schrodinger equation
$-i\hbar \frac{\partial}{\partial t} \Psi(\mathbf{x}, t) = \hat{H} \Psi(\mathbf{x}, t)$
describes the time-evolution of the wavefunction $\Psi(\mathbf{x}, t)$ under the influence of the Hamiltonian $\hat{H}$. Here we are interested in the eigenstates and spectrum of a given Hamiltonian, e.g. such that
$\hat{H} \Psi_n(\mathbf{x}) = E_n \Psi_n(\mathbf{x})$
In general, the Hamiltonian for a particle in 3D space takes the form
$\displaystyle \hat{H} = K + U = \left[-\frac{\hbar^2}{2m}\nabla^2 + U(\mathbf{x})\right]$
Here we focus on 2D problems, and write the Hamiltonian on the form
$\displaystyle \hat{H}\Psi_n(\mathbf{x}, t) = -\hbar^2\left(k_{11}\frac{\partial^2}{\partial x_1^2} + k_{12}\frac{\partial^2}{\partial x_1x_2} + k_{22}\frac{\partial^2}{\partial x_2^2}\right) \Psi_n(x_1, x_2, t) + U(x_1, x_2)\Psi_n(x_1, x_2, t)$
In problems with diagonalized kinetic term, we have $k_{12}=0$.
The corresponding eigenvalue problem becomes
$\displaystyle \left[-\hbar^2\left(k_{11}\frac{\partial^2}{\partial x_1^2} + k_{12}\frac{\partial^2}{\partial x_1x_2} + k_{22}\frac{\partial^2}{\partial x_2^2}\right) + U(x_1, x_2)\right] \Psi_n(x_1, x_2) = E_n \Psi_n(x_1, x_2)$
To proceed further we need to pick a basis with which we can write $\Psi_n(x_1, x_2)$ as a vector and the PDE as a matrix eigenvalue problem.
If $U(x_1, x_2)$ is periodic with periods $T_{x_1}$ and $T_{x_2}$ in the $x_1$ and $x_2$ directions, respectively, we can write the wavefunctions as Fourier series
$\displaystyle \Psi(x_1, x_2) = \sum_{n_1}\sum_{n_2} \psi_{n_1n_2} e^{i2\pi n_1x_1/T_{x_1}}e^{i2\pi n_2x_2/T_{x_2}}$
$\displaystyle U(x_1, x_2) = \sum_{n_1}\sum_{n_2} u_{n_1n_2} e^{i2\pi n_1x_1/T_{x_1}}e^{i2\pi n_2x_2/T_{x_2}}$
Substituting this expression into the eigenvalue problem gives
$\displaystyle -\hbar^2\left(k_{11}\frac{\partial^2}{\partial x_1^2} + k_{12}\frac{\partial^2}{\partial x_1x_2} + k_{22}\frac{\partial^2}{\partial x_2^2}\right) \sum_{n_1}\sum_{n_2} \psi_{n_1n_2} e^{i2\pi n_1x_1/T_{x_1}}e^{i2\pi n_2x_2/T_{x_2}} $
$\displaystyle = -\hbar^2 \sum_{n_1}\sum_{n_2} \psi_{n_1n_2} \left(k_{11}\left(\frac{2\pi i n_1}{T_{x_1}}\right)^2 + k_{12}\frac{-(2\pi)^2 n_1n_2}{T_{x_1}T_{x_2}} + k_{22}\left(\frac{2\pi i n_2}{T_{x_2}}\right)^2 \right) e^{i2\pi n_1x_1/T_{x_1}}e^{i2\pi n_2x_2/T_{x_2}} $
$\displaystyle = \sum_{n_1}\sum_{n_2}\left[-\hbar^2 \sum_{m_1}\sum_{m_2} \psi_{m_1m_2} \delta_{n_1,m_1} \delta_{n_2,m_2} \left(k_{11}\left(\frac{2\pi i m_1}{T_{x_1}}\right)^2 + k_{12}\frac{-(2\pi)^2 m_1m_2}{T_{x_1}T_{x_2}} + k_{22}\left(\frac{2\pi i m_2}{T_{x_2}}\right)^2 \right)\right] e^{i2\pi n_1x_1/T_{x_1}}e^{i2\pi n_2x_2/T_{x_2}} $
$\displaystyle = \sum_{n_1}\sum_{n_2} \left[\sum_{m_1}\sum_{m_2} K_{n_1, n_2}^{m_1, m_2} \psi_{m_1m_2}\right] e^{i2\pi n_1x_1/T_{x_1}}e^{i2\pi n_2x_2/T_{x_2}} $
where
$\displaystyle K_{n_1, n_2}^{m_1, m_2} = \delta_{n_1,m_1} \delta_{n_2,m_2} \left(k_{11}\left(\frac{2\pi m_1}{T_{x_1}}\right)^2 + k_{12}\frac{(2\pi)^2 m_1m_2}{T_{x_1}T_{x_2}} + k_{22}\left(\frac{2\pi m_2}{T_{x_2}}\right)^2\right)$
and where we have absorbed the $\hbar^2$ into the definitions of $k_{11}, k_{12}, k_{22}$.
$\displaystyle \left[\sum_{m_1}\sum_{m_2} u_{m_1m_2} e^{i2\pi m_1x_1/T_{x_1}}e^{i2\pi m_2x_2/T_{x_2}}\right] \sum_{n_1}\sum_{n_2} \psi_{n_1n_2} e^{i2\pi n_1x_1/T_{x_1}}e^{i2\pi n_2x_2/T_{x_2}}$
$\displaystyle = \sum_{n_1}\sum_{n_2} \left[\sum_{m_1}\sum_{m_2} u_{m_1m_2} \psi_{n_1n_2} e^{i2\pi (n_1+m_1)x_1/T_{x_1}}e^{i2\pi (n_2+m_2)x_2/T_{x_2}} \right]$
shift variables $m_1, m_2$ such that $m_1 + n_1 \rightarrow \tilde{m}_1$ and $m_2 + n_2 \rightarrow \tilde{m}_2$
$\displaystyle = \sum_{n_1}\sum_{n_2} \left[\sum_{\tilde{m}_1}\sum_{\tilde{m}_2} u_{\tilde{m}_1-n_1,\tilde{m}_2-n_2} \psi_{n_1n_2} e^{i2\pi \tilde{m}_1x_1/T_{x_1}}e^{i2\pi \tilde{m}_2x_2/T_{x_2}} \right]$
$\displaystyle = \sum_{\tilde{m}_1}\sum_{\tilde{m}_2} \left[\sum_{n_1}\sum_{n_2} u_{\tilde{m}_1-n_1,\tilde{m}_2-n_2} \psi_{n_1n_2} \right] e^{i2\pi \tilde{m}_1x_1/T_{x_1}}e^{i2\pi \tilde{m}_2x_2/T_{x_2}}$
rename $n \rightarrow m$, $\tilde{m} \rightarrow n$
$\displaystyle = \sum_{n_1}\sum_{n_2} \left[\sum_{m_1}\sum_{m_2} u_{n_1-m_1,n_2-m_2} \psi_{m_1m_2} \right] e^{i2\pi n_1x_1/T_{x_1}}e^{i2\pi n_2x_2/T_{x_2}}$
$\displaystyle = \sum_{n_1}\sum_{n_2} \left[\sum_{m_1}\sum_{m_2} V_{n_1, n_2}^{m_1, m_2} \psi_{m_1m_2} \right] e^{i2\pi n_1x_1/T_{x_1}}e^{i2\pi n_2x_2/T_{x_2}}$
where
$\displaystyle V_{n_1, n_2}^{m_1, m_2} = u_{n_1-m_1,n_2-m_2}$
$\displaystyle E \sum_{n_1}\sum_{n_2} \psi_{n_1n_2} e^{i2\pi n_1x_1/T_{x_1}}e^{i2\pi n_2x_2/T_{x_2}} = \sum_{n_1}\sum_{n_2} \left[E \psi_{n_1n_2}\right] e^{i2\pi n_1x_1/T_{x_1}}e^{i2\pi n_2x_2/T_{x_2}}$
The eigenvalue problem then takes the form
$\displaystyle \sum_{n_1}\sum_{n_2} \left[ \sum_{m_1}\sum_{m_2} K_{n_1, n_2}^{m_1, m_2} \psi_{m_1m_2}
\sum_{m_1}\sum_{m_2} V_{n_1, n_2}^{m_1, m_2}\psi_{m_1m_2}
E \psi_{n_1n_2}\right] e^{i2\pi n_1x_1/T_{x_1}}e^{i2\pi n_2x_2/T_{x_2}} = 0$
and since this has to be true for all $n_1$ and $n_2$, we can identity
$\displaystyle \sum_{m_1}\sum_{m_2} K_{n_1, n_2}^{m_1, m_2} \psi_{m_1m_2} + \sum_{m_1}\sum_{m_2} V_{n_1, n_2}^{m_1, m_2}\psi_{m_1m_2} = E \psi_{n_1n_2}$
which in matrix form can be written
$\displaystyle \sum_{M}[K_{NM} + V_{NM}]\Psi_{M} = \sum_{M}H_{NM}\Psi_{M} = E \Psi_{N}$
or simply
$\displaystyle H\Psi = E \Psi$
This matrix eigenvalue problem can be solved with a standard solver. We therefore only need to assemble the matrix representation $H$ for Hamiltonian $\hat{H}$, by adding its kinetic and potential contributions.
%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cm as cm
from wavefunction.wavefunction2d import *
$\displaystyle K_{n_1, n_2}^{m_1, m_2} = \delta_{n_1,m_1} \delta_{n_2,m_2} \left(k_{11}\left(\frac{2\pi m_1}{T_{x_1}}\right)^2 + k_{12}\frac{(2\pi)^2 m_1m_2}{T_{x_1}T_{x_2}} + k_{22}\left(\frac{2\pi m_2}{T_{x_2}}\right)^2\right)$
# pick a truncation of the fourier series: n1 = [-L1, ..., L1], n2 = [-L2, ..., L2]
L1 = 3
L2 = 3
# pick periods for the coordinates
Tx1 = Tx2 = 2 * np.pi
#
k11, k12, k22 = 0.0000125, 0.0001, 0.0000125
help(assemble_K)
Help on function assemble_K in module wavefunction.wavefunction2d: assemble_K(L1, L2, k11, k12, k22, Tx1, Tx2, sparse=False) Assemble the matrix representation of the kinetic energy contribution to the Hamiltonian.
K = assemble_K(L1, L2, k11, k12, k22, Tx1, Tx2)
The potential contribution to the matrix representation of the Hamiltonian can be calculated analytically or numerically. Here we first focus on the analytical calculation for a flux qubit potential (period with period $2\pi$), and then discuss how it could be evaluated completely numerically.
The flux qubit potential we consider here is
$\displaystyle U(x_1, x_2)/E_J = U_0 - \epsilon_1 \cos(x_1) - \epsilon_2 \cos(x_2) - \epsilon_{12} \cos(2\pi f + x_1 - x_2)$
To obtain the Fourier series expansion of $U(x_1, x_2)$, we first write the $\cos$ expressions as exponential functions
$\displaystyle U(x_1, x_2)/E_J = U_0
Now, by comparing with the Fourier series
$\displaystyle U(x_1, x_2)/E_J = \sum_{n_1}\sum_{n_2} u_{n_1n_2} e^{in_1x_1}e^{in_2x_2}$
we can identity
$\displaystyle u_{n_1n_2} = \alpha \delta_{0,n_1}\delta_{0,n_2}
\frac{\epsilon_{12}}{2} (e^{i2\pi f} \delta_{1,n_1}\delta_{-1,n_2} + e^{-i2\pi f}\delta_{-1,n_1}\delta_{+1,n_2}) $
args = {'U0': 0, 'eps1': 1, 'eps2': 1, 'eps12': 0.75, 'f': 0.5}
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):
U0, eps1, eps2, eps12, f = args['U0'], args['eps1'], args['eps2'], args['eps12'], args['f']
return U0 - eps1 * np.cos(x1) - eps2 * np.cos(x2) - eps12 * np.cos(2 * np.pi * f + x1 - x2)
Z = np.real(U_flux_qubit(X1, X2, args).T)
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)
$\displaystyle u_{n_1n_2} = \alpha \delta_{0,n_1}\delta_{0,n_2}
\frac{\epsilon_{12}}{2} (e^{i2\pi f} \delta_{1,n_1}\delta_{-1,n_2} + e^{-i2\pi f}\delta_{-1,n_1}\delta_{1,n_2}) $
def assemble_u_flux_qubit(L1, L2, args, sparse=False):
U0, eps1, eps2, eps12, f = args['U0'], args['eps1'], args['eps2'], args['eps12'], args['f']
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] = U0 * delta(0, n1) * delta(0, n2) + \
- 0.5 * eps1 * (delta(1, n1) + delta(-1, n1)) * delta(0, n2) + \
- 0.5 * eps2 * (delta(1, n2) + delta(-1, n2)) * delta(0, n1)+ \
- 0.5 * eps12 * np.exp(+2j * np.pi * f) * delta(+1, n1) * delta(-1, n2) + \
- 0.5 * eps12 * np.exp(-2j * np.pi * f) * delta(-1, n1) * delta(+1, n2)
return u
def assemble_u_flux_qubit_direct(L1, L2, args, sparse=False):
U0, eps1, eps2, eps12, f = args['U0'], args['eps1'], args['eps2'], args['eps12'], args['f']
L1n = 2 * L1 + 1
L2n = 2 * L2 + 1
u = np.zeros((L1n, L2n), dtype=np.complex)
u[L1, L2] = U0
u[+1+L1, L2] = - 0.5 * eps1
u[-1+L1, L2] = - 0.5 * eps1
u[L1, +1+L2] = - 0.5 * eps2
u[L1, -1+L2] = - 0.5 * eps2
u[+1+L1, -1+L2] = -0.5 * eps12 * np.exp(+2j * np.pi * f)
u[-1+L1, +1+L2] = -0.5 * eps12 * np.exp(-2j * np.pi * f)
return np.real(u)
u2 = assemble_u_flux_qubit(L1, L2, args)
u = assemble_u_flux_qubit_direct(L1, L2, args)
abs(u - u2).max()
4.5924254968025742e-17
$\displaystyle U(x_1, x_2)/E_J = \sum_{n_1}\sum_{n_2} u_{n_1n_2} e^{in_1x_1}e^{in_2x_2}$
# show code of evalute_fourier_series
help(evalute_fourier_series)
Help on function evalute_fourier_series in module wavefunction.wavefunction2d: evalute_fourier_series(X1, X2, L1, L2, u)
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)
Given this Fourier series we can calcuate the potential contribution to the eigenvalue problem as
$\displaystyle V_{n_1, n_2}^{m_1, m_2} = u_{n_1-m_1,n_2-m_2}$
help(assemble_V)
Help on function assemble_V in module wavefunction.wavefunction2d: assemble_V(L1, L2, u, sparse=False) Assemble the matrix representation of the potential energy contribution to the Hamiltonian.
and specifically for the flux qubit potential
$\displaystyle V_{n_1, n_2}^{m_1, m_2} = \alpha \delta_{0,n_1-m_1}\delta_{0,n_2-m_2}
\frac{\epsilon_{12}}{2} (e^{i2\pi f} \delta_{1,n_1-m_1}\delta_{-1,n_2-m_2} + e^{-i2\pi f}\delta_{-1,n_1-m_1}\delta_{1,n_2-m_2}) $
def assemble_V_flux_qubit(L1, L2, args, sparse=False):
U0, eps1, eps2, eps12, f = args['U0'], args['eps1'], args['eps2'], args['eps12'], args['f']
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] = U0 * delta(m1, n1) * delta(m2, n2) + \
- 0.5 * eps1 * (delta(m1 + 1, n1) + delta(m1 - 1, n1)) * delta(m2, n2) + \
- 0.5 * eps2 * (delta(m2 + 1, n2) + delta(m2 - 1, n2)) * delta(m1, n1) + \
- 0.5 * eps12 * np.exp(+2j * np.pi * f) * delta(m1 + 1, n1) * delta(m2 - 1, n2) + \
- 0.5 * eps12 * np.exp(-2j * np.pi * f) * delta(m1 - 1, n1) * delta(m2 + 1, n2)
return V
V_full = assemble_V(L1, L2, u)
V = assemble_V_flux_qubit(L1, L2, args)
abs(V-V_full).max()
4.5924254968025742e-17
We now want to solve the eigenvalue problem
$H \Psi = E \Psi$
where $H$ is the matrix representation of the Hamiltonian assembled from $K$ and $V$ above.
H = K + V
help(solve_eigenproblem)
Help on function solve_eigenproblem in module wavefunction.wavefunction2d: solve_eigenproblem(H)
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, val.real * 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), squeeze=False)
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[1, 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)
e_vals = np.zeros((len(vals), len(f_vec)))
for f_idx, f in enumerate(f_vec):
K = assemble_K(L1, L2, k11, k12, k22, Tx1, Tx2)
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=(12,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[10, :].max())
ax.set_ylabel("Energy levels")
ax.set_xlabel(r'$f$');
np.diff(X1, axis=1)
array([[ 0.12693304, 0.12693304, 0.12693304, ..., 0.12693304, 0.12693304, 0.12693304], [ 0.12693304, 0.12693304, 0.12693304, ..., 0.12693304, 0.12693304, 0.12693304], [ 0.12693304, 0.12693304, 0.12693304, ..., 0.12693304, 0.12693304, 0.12693304], ..., [ 0.12693304, 0.12693304, 0.12693304, ..., 0.12693304, 0.12693304, 0.12693304], [ 0.12693304, 0.12693304, 0.12693304, ..., 0.12693304, 0.12693304, 0.12693304], [ 0.12693304, 0.12693304, 0.12693304, ..., 0.12693304, 0.12693304, 0.12693304]])
np.diff(X2, axis=0).shape
(99, 100)
psi.shape
(7, 7)
psi0 = wavefunction_normalize(X1, X2, evalute_fourier_series(X1, X2, L1, L2, convert_v2m(L1, L2, vecs[0])))
psi0.shape
(100, 100)
derivative(X1, X2, psi0, axis=1)
array([[ 0.13697142 +4.68239186e-16j, 0.27398563 +1.12761740e-15j, 0.26221110 +1.29208937e-15j, ..., 0.18616390 +1.85580292e-16j, 0.23197458 +4.51616884e-16j, 0.12595118 +3.50588618e-16j], [ 0.09394329 +4.41423641e-16j, 0.19033415 +1.01103714e-15j, 0.18697407 +1.27653495e-15j, ..., 0.11821312 +3.82424127e-16j, 0.15273278 +6.29149370e-16j, 0.08412528 +3.48633496e-16j], [ 0.05476196 +3.56142525e-16j, 0.11345187 +8.67829225e-16j, 0.11637742 +1.03946418e-15j, ..., 0.05918139 +5.02568593e-16j, 0.08247452 +6.43094710e-16j, 0.04670938 +3.58287962e-16j], ..., [ 0.22078104 +2.97143010e-16j, 0.43461480 +7.36421215e-16j, 0.40197856 +1.44495175e-15j, ..., 0.32779507 -4.02269418e-16j, 0.39250834 -7.50902914e-18j, 0.20960266 +2.24198156e-16j], [ 0.18048027 +3.65796991e-16j, 0.35779525 +1.19929923e-15j, 0.33600638 +1.50609670e-15j, ..., 0.25799529 -1.93089321e-17j, 0.31418404 +2.86415826e-16j, 0.16897932 +1.60907767e-16j], [ 0.13697142 +4.41228493e-16j, 0.27398563 +1.11026118e-15j, 0.26221110 +1.28672578e-15j, ..., 0.18616390 +2.61743301e-16j, 0.23197458 +4.96863790e-16j, 0.12595118 +3.28254262e-16j]])
inner_product(X1, X2, psi0, derivative(X1, X2, psi0, axis=0))
(7.5028596895989966e-16+8.0366056869616157e-15j)
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)
Z = X1 ** 2 + X2 ** 2
#Z = np.real(derivative(X1, X2, Z, axis=1))
fig, ax = plt.subplots()
p = ax.pcolor(X1/(2*np.pi), X2/(2*np.pi), Z, cmap=cm.RdBu, vmin=-Z.max(), vmax=Z.max())
cb = fig.colorbar(p, ax=ax)