fermions.py

"""
Simulation of coupled Fermions by explicitly constructing the Fermion
operators as matrices and then diagonalizing the subsequent Hamiltonian
to find the energies and energy eigenvectors.

The primary reference for this is Chapter 4 of 
Introduction to Many Body Physics by Piers Coleman, pg 71 - 77 on the
Jordan-Wigner transformation.

P. Coleman, Simple examples of second quantization,
in Introduction to Many Body Physics, pg 71-94,
Cambridge University Press, 2015.

"""
import numpy as np
import matplotlib.pyplot as plt
import scipy.sparse as sparse
import matplotlib.animation as animation
from matplotlib.gridspec import GridSpec

from typing import List, Tuple, Union


# Maximum number of particles, or occupancy.
# The two shall be used interchangeably.
N = 10

# Fermionic creation operator for a single particle.
LADDER_UP = np.array([[0.0, 1.0], [0.0, 0.0]])

# This corresponds to the operator exp(-n i pi),
# where n counts the occupation.
STRAND = np.array([[-1.0, 0.0], [0.0, 1.0]])

I = np.array([[1.0, 0.0], [0.0, 1.0]]) # 2x2 identity matrix

# The state that corresponds to zero occupancy.
VACANT_STATE = np.array([1.0 if i == 0 else 0.0
                         for i in range(int(2**N))])[::-1]


def normalize(psi: np.ndarray) -> np.ndarray:
    return psi/np.sqrt(np.sum(np.abs(psi)**2))


def kron_power(base: Union[sparse.spmatrix, np.ndarray], 
               exponent: int) -> sparse.spmatrix:
    res = sparse.identity(base.shape[0]) if exponent == 0 else base.copy()
    for _ in range(exponent - 1):
        res = sparse.kron(res, base)
    return res


def get_ladder_operators(n: int) -> List[sparse.spmatrix]:
    """
    Returns a list of the coupled Fermionic creation operators for n particles.

    The way this works is that one first starts with ladder spin operators
    for a one-dimensional linear sequence of spins, where one now interprets
    spin down as corresponding to the state of zero occupancy, and spin up as the
    occupied state. However, the ladder spin operators do not follow
    Fermionic anticommuting relations; rather they commute. To fix this,
    one must use the Jordan-Wigner transformation, where each ladder spin
    operator in the sequence is matrix multiplied by an operator exp(-i n pi),
    where n counts the occupation of states for **only** the preceding ladder
    spin operators in the sequence [1].

    [1] P. Coleman, Simple examples of second quantization.
        in Introduction to Many Body Physics, pg 71-94,
        Cambridge University Press, 2015.

    """
    ladders_u = []
    for i in range(n):
        string = sparse.identity(1) if i == 0 else kron_power(STRAND, i)
        o1 = sparse.kron(sparse.identity(2**i) @ string, LADDER_UP)
        o2 = sparse.kron(o1, sparse.identity(2**(N-i-1)))
        ladders_u.append(o2)
    return ladders_u


def make_to_density_vector_matrix(n):
    n_mat = np.zeros([n, int(2**n)])
    val: int = 1
    for i in range(n):
        for j in range(int(2**n)):
            element = float(int((val&(~j)) > 0))
            # if element > 0.0: print(i, j, element)
            n_mat[i, j] = element
        val *= 2
    return np.flip(n_mat, axis=0)

n_mat = make_to_density_vector_matrix(N)
print(n_mat)
f_dag = get_ladder_operators(N)
f = [o.T for o in f_dag]

# Fermionic ladder operators that are decoupled when using them to construct
# the Hamiltonian
a_dag = [sum([f_dag[j]*np.exp(-2.0j*np.pi*j*k/N) for j in range(N)])/np.sqrt(N)
         for k in range(N)]
a = [sum([f[j]*np.exp(-2.0j*np.pi*j*k/N) for j in range(N)])/np.sqrt(N)
     for k in range(N)]

# Explicitly show the matrices that correspond to our fermionic operators
# for i, op in enumerate(f_dag):
#     print(i, op)
# for j, op in enumerate(f):
#     print(j, op)

# Sanity check to ensure that our ladder operators satisfy all the
# correct commutation relations.
# for i, fi_dag in enumerate(f_dag):
#     for j, fj in enumerate(f):
#         print(i, j, '\n', fi_dag @ fj + fj @ fi_dag)
# for i, fi_dag in enumerate(f_dag):
#     for j, fj_dag in enumerate(f_dag):
#         print(i, j, '\n', fi_dag @ fj_dag + fj_dag @ fi_dag)
# for i, fi in enumerate(f):
#     for j, fj in enumerate(f):
#         print(i, j, '\n', fi @ fj + fj @ fi)

# Construct the Hamiltonian
alpha, beta, gamma = 1.0, 2.0, 0.0
H = sum([-alpha*(f_dag[(i+1)%N] @ f[i] + f_dag[(i-1)] @ f[i])
         + beta*f_dag[i] @ f[i]
         + gamma*f_dag[i] @ f[i] @ f_dag[i] @ f[i] for i in range(N)]
        ) # + sparse.identity(int(2**N))

# Numerically solve for the energies and their associated eigenvectors
energies, eigenstates = np.linalg.eigh(H.toarray())

# List of singly occupied states
single_particle_states_list = [f_dag[i] @ VACANT_STATE for i in range(N)]

# List of doubly occupied states
two_particle_states_list = [f_dag[i] @ f_dag[j] @ VACANT_STATE
                            for i in range(N) for j in range(N)]

single_particle_states = np.array(single_particle_states_list)
two_particle_states = np.array(two_particle_states_list)

# Make the initial state
# psi0 = (f_dag[3] + f_dag[3] @ f_dag[2]) @ VACANT_STATE
psi0 = (a_dag[0] @ a_dag[1] + f_dag[6]
         + a_dag[0] @ a_dag[6] @ a_dag[7]) @ VACANT_STATE
# psi0 = np.exp(-np.random.rand(2**N)*2.0j*np.pi)*np.random.rand(2**N)
# psi0 = normalize(psi0)
# psi0 = (a_dag[1] + a_dag[-1] + a_dag[0] @ a_dag[-1]
#          + a_dag[0] @ a_dag[1]) @ VACANT_STATE
# psi0 = (a_dag[-3] + a_dag[2] @ a_dag[1] @ a_dag[0]) @ VACANT_STATE
psi0_e = np.conj(eigenstates.T) @ psi0
# psi0_e[0], psi0_e[1] = 0.0, 0.0

# Show the energy levels for the entire system
# plt.scatter([i for i in range(int(2**N))], np.sort(energies))
# plt.show()
# plt.close()

# Show the sum of each single particle state in the original basis
# plt.scatter([i for i in range(int(2**N))], sum(single_particle_states))
# plt.show()
# plt.close()

# Show psi0 in terms of energy eigenstates
# plt.scatter([i for i in range(int(2**N))],  np.real(psi0_e))
# plt.scatter([i for i in range(int(2**N))],  np.imag(psi0_e))
# plt.show()
# plt.close()

fig = plt.figure()
grid_spec = GridSpec(2, 2, figure=fig)
axes = [fig.add_subplot(grid_spec[0, 0]),
        fig.add_subplot(grid_spec[0, 1]),
        fig.add_subplot(grid_spec[1, 0]),
        fig.add_subplot(grid_spec[1, 1])]
re_, = axes[0].plot(np.real(np.conj(single_particle_states) @ psi0))
im_, = axes[0].plot(np.imag(np.conj(single_particle_states) @ psi0))
abs_, = axes[0].plot(np.abs(np.conj(single_particle_states) @ psi0),
                     color='black')
axes[0].set_title('Singly Occupied States ($f_i^{\dagger}|0>$)')
axes[0].set_ylabel('Amplitude')
axes[0].set_xlabel('i ($f_i^{\dagger}|0>$)')
axes[0].set_xlim(0, N-1)
axes[0].set_ylim(-0.75, 0.75)
e_re, = axes[3].plot([i for i in range(int(2**N))],  np.real(psi0_e),
                     linewidth=1.0, label='Real')
e_im, = axes[3].plot([i for i in range(int(2**N))],  np.imag(psi0_e),
                     linewidth=1.0, label='Imag.')
e_abs, = axes[3].plot([i for i in range(int(2**N))],
                      np.abs(psi0_e), color='black', linewidth=1.0,
                      label='Abs.')
axes[3].set_title('Energy States for Entire System')
axes[3].set_ylabel('Amplitude')
axes[3].set_xlabel('Enumeration from lowest to highest energy')
axes[3].set_ylim(-1.0, 1.0)
axes[3].set_xlim(0, 2**N-1)
axes[3].legend(loc='lower right')
axes[2].set_title('Total Occupation Density')
axes[2].set_ylabel('Arbitrary Units')
axes[2].set_xlabel('index i')
density = axes[2].bar([i for i in range(N)], n_mat @ np.abs(psi0)**2)
# print(density)
# print(density.__dict__)
bg_im = axes[1].imshow(np.zeros([N, N]), cmap='Greys_r')
im = axes[1].imshow(np.angle(np.reshape((np.conj(two_particle_states)
                                         @ psi0), (N, N))),
                    cmap='hsv', interpolation='nearest',
                    alpha=np.abs(np.reshape((np.conj(two_particle_states)
                                             @ psi0),
                                            (N, N)))**2)
axes[1].set_title('Doubly Occupied States '
                  + '($\propto f_i^{\dagger}f_j^{\dagger}|0>$)')
axes[1].set_ylabel('$f_i^{\dagger}$')
axes[1].set_xlabel('$f_j^{\dagger}$')
# axes[2].set_xlim(0, N-1)
# axes[2].set_ylim(0, N-1)
plt.tight_layout()


data = {'psi': psi0_e, 't': 0.0}

def animation_func(*args):
    data['t'] += 0.1 # - 1.0j*0.001
    data['psi'] = psi0_e*np.exp(-1.0j*energies*data['t'])
    # data['psi'] = normalize(data['psi'])
    psi = eigenstates @ data['psi']
    psi_sp = np.conj(single_particle_states) @ psi
    re_.set_ydata(np.real(psi_sp))
    im_.set_ydata(np.imag(psi_sp))
    abs_.set_ydata(np.abs(psi_sp))
    e_re.set_ydata(np.real(data['psi']))
    e_im.set_ydata(np.imag(data['psi']))
    e_abs.set_ydata(np.abs(data['psi']))
    density.datavalues = n_mat @ np.abs(psi)**2
    # https://stackoverflow.com/a/42143866
    for i, d in enumerate(density.datavalues):
        density[i].set_height(d)
    im.set_data(np.angle(np.reshape((np.conj(two_particle_states)
                                     @ psi), (N, N))))
    abs_val = np.abs(np.reshape((np.conj(two_particle_states)
                                 @ psi), (N, N)))**2
    abs_val_show = 20.0*abs_val
    im.set_alpha(np.where(abs_val_show > 1.0, 1.0, abs_val_show))
    # print(np.sum(np.abs(psi)**2))
    # print(np.sum(np.abs(psi_sp)**2))
    return re_, im_, abs_, e_re, e_im, e_abs, bg_im, *density.patches, im

# for i in range(N):
#     plt.scatter([i for i in range(int(2**N))],
#                 f_dag[i] @ GROUND, label=f'{i}')
# plt.legend()

data['ani'] = animation.FuncAnimation(fig, animation_func,
                                      blit=True, interval=1000//60)


plt.show()
plt.close()