Fu-Kane Spin Pump

%matplotlib notebook

import numpy as np
from scipy import sparse
from scipy.sparse import bsr, csr_matrix
from scipy.sparse.bsr import bsr_matrix
from typing import List
import scipy.linalg as linalg
from typing import Tuple

class FuKaneModel:
    h_0: float
    delta_0: float
    t_0: float
    period: float
    sites: int

    def __init__(self, h_0: float, delta_0: float, t_0: float, period: float, sites: int):
        self.h_0 = h_0
        self.delta_0 = delta_0
        self.t_0 = t_0
        self.period = period
        self.sites = sites
        assert sites > 0 and sites % 2 == 0

    def delta(self, t: float) -> float:
        return self.delta_0 * np.cos(2 * np.pi * t / self.period)

    def h_st(self, t: float) -> float:
        return self.h_0 * np.sin(2 * np.pi * t / self.period)

    def single_hamiltonian(self, t: float) -> np.ndarray:
        hamiltonian = np.zeros((2 * self.sites, 2 * self.sites))
        for i in range(0, self.sites):
            hamiltonian[2 * i, 2 * i]               = (-1) ** i * self.h_st(t)
            hamiltonian[2 * i + 1, 2 * i + 1]       = -(-1) ** i * self.h_st(t)
        for i in range(0, self.sites - 1):
            hamiltonian[2 * i, 2 * (i + 1)]         = 1 / 2 * (self.t_0 + self.delta(t) * (-1) ** i)
            hamiltonian[2 * i + 1, 2 * (i + 1) + 1] = 1 / 2 * (self.t_0 + self.delta(t) * (-1) ** i)
            hamiltonian[2 * (i + 1), 2 * i]         = np.conj(1 / 2 * (self.t_0 + self.delta(t) * (-1) ** i))
            hamiltonian[2 * (i + 1) + 1, 2 * i + 1] = np.conj(1 / 2 * (self.t_0 + self.delta(t) * (-1) ** i))
        return hamiltonian

    def single_exact_diagonalization(self, t: float) -> Tuple[np.ndarray, np.ndarray]:
        single_hamiltonian: np.ndarray = self.single_hamiltonian(t)
        energy_levels, vectors = linalg.eigh(single_hamiltonian)
        return energy_levels, vectors.T

    def bulk_levels(self, t: float) -> np.ndarray:
        k_points = np.linspace(-np.pi, np.pi, self.sites // 2 + 1)
        levels = np.sqrt(self.h_st(t) ** 2 + self.delta(t) ** 2 *
                         np.sin(k_points / 2) ** 2 + self.t_0 ** 2 * np.cos(k_points / 2) ** 2)
        return np.concatenate(
            (levels, -levels)
        )

h_0     = 0.5
delta_0 = 0.2
t_0     = 1
period  = 10
sites   = 100
model   = FuKaneModel(h_0, delta_0, t_0, period, sites)

steps = 201
duration = 2 * period
times = np.linspace(0, duration, steps)

exact_levels_and_vectors_by_time = [model.single_exact_diagonalization(t) for t in times]
exact_levels_by_time = np.array([levels_and_vectors[0] for levels_and_vectors in exact_levels_and_vectors_by_time])
exact_vectors_by_time = [levels_and_vectors[1] for levels_and_vectors in exact_levels_and_vectors_by_time]
occupied_edge_vector1_by_time = np.array([np.abs(exact_vectors[len(exact_vectors) // 2 - 2]) ** 2 for exact_vectors in exact_vectors_by_time])
occupied_edge_vector2_by_time = np.array([np.abs(exact_vectors[len(exact_vectors) // 2 - 1]) ** 2 for exact_vectors in exact_vectors_by_time])
unoccupied_edge_vector1_by_time = np.array([np.abs(exact_vectors[len(exact_vectors) // 2]) ** 2  for exact_vectors in exact_vectors_by_time])
unoccupied_edge_vector2_by_time = np.array([np.abs(exact_vectors[len(exact_vectors) // 2 + 1]) ** 2  for exact_vectors in exact_vectors_by_time])

bulk_levels_by_time = np.array([model.bulk_levels(t) for t in times])

import matplotlib
import matplotlib.pyplot as plt

exact_levels_fig, exact_levels_ax = plt.subplots(1)
for i in range(0, exact_levels_by_time.shape[1], 2):
    energy_levels = exact_levels_by_time[:, i]
    exact_levels_ax.scatter(times, energy_levels, s=3)
for i in range(1, exact_levels_by_time.shape[1], 2):
    energy_levels = exact_levels_by_time[:, i]
    exact_levels_ax.scatter(times, energy_levels, marker='_')
plt.show()

bulk_levels_fig, bulk_levels_ax = plt.subplots(1)
for i in range(bulk_levels_by_time.shape[1]):
    energy_levels = bulk_levels_by_time[:, i]
    bulk_levels_ax.scatter(times, energy_levels, s=1)
plt.show()

def plot_density(time_index):
    density_fig, density_ax = plt.subplots()
    density_ax.set_title(f'Edge states at time {time_index / steps * duration: .3}, period {duration}.')
    density_ax.set_xlabel('Relative position')
    density_ax.set_ylabel('Density summed on two sites')
    occupied1_x, occupied1_y = np.linspace(0, 1, sites)[::2], np.sum(occupied_edge_vector1_by_time[time_index][::2].reshape((-1, 2)) - occupied_edge_vector1_by_time[time_index][1::2].reshape((-1, 2)), axis=1)
    occupied2_x, occupied2_y = np.linspace(0, 1, sites)[::2], np.sum(occupied_edge_vector2_by_time[time_index][::2].reshape((-1, 2)) - occupied_edge_vector2_by_time[time_index][1::2].reshape((-1, 2)), axis=1)
    unoccupied1_x, unoccupied1_y = np.linspace(0, 1, sites)[::2], np.sum(unoccupied_edge_vector1_by_time[time_index][::2].reshape((-1, 2)) - unoccupied_edge_vector1_by_time[time_index][1::2].reshape((-1, 2)), axis=1)
    unoccupied2_x, unoccupied2_y = np.linspace(0, 1, sites)[::2], np.sum(unoccupied_edge_vector2_by_time[time_index][::2].reshape((-1, 2)) - unoccupied_edge_vector2_by_time[time_index][1::2].reshape((-1, 2)), axis=1)

    density_ax.fill_between(occupied1_x, occupied1_y, label='Occupied 1', color='blue')
    density_ax.plot(occupied1_x, occupied1_y, '--', color='blue')
    density_ax.fill_between(occupied2_x, occupied2_y, label='Occupied 2', color='blue')
    density_ax.plot(occupied2_x, occupied2_y, '--', color='blue')

    density_ax.fill_between(unoccupied1_x, unoccupied1_y, alpha=0.2, label='Unoccupied 1', color='orange')
    density_ax.plot(unoccupied1_x, unoccupied1_y, '--', alpha=0.2, color='orange')
    density_ax.fill_between(unoccupied2_x, unoccupied2_y, alpha=0.2, label='Unoccupied 2', color='orange')
    density_ax.plot(unoccupied2_x, unoccupied2_y, '--', alpha=0.2, color='orange')

    density_ax.legend()
    plt.show()

from ipywidgets import interact, interactive, fixed, interact_manual
import ipywidgets as widgets

interact(plot_density, time_index=widgets.IntSlider(
    value=0,
    min=0,
    max=steps,
    step=1)
)

<function __main__.plot_density(time_index)>

import matplotlib.animation as animation
from IPython import display

animation_fig, animation_ax = plt.subplots()

animation_plot_occupied1, = animation_ax.plot([], [], '--', color='blue')
animation_plot_occupied2, = animation_ax.plot([], [], '--', color='blue')
animation_plot_unoccupied1, = animation_ax.plot([], [], '--', alpha=0.2, color='orange')
animation_plot_unoccupied2, = animation_ax.plot([], [], '--', alpha=0.2, color='orange')

def animate_density(frame, *fargs):
    time_index = frame
    animation_ax.set_ylim([-1, 1])
    animation_ax.set_title(f'Edge states at time {time_index / steps * duration: .3}, period {duration}.')
    animation_ax.set_xlabel('Relative position')
    animation_ax.set_ylabel('Density summed on two sites')

    occupied1_x, occupied1_y = np.linspace(0, 1, sites)[::2], np.sum(occupied_edge_vector1_by_time[time_index][::2].reshape((-1, 2)) - occupied_edge_vector1_by_time[time_index][1::2].reshape((-1, 2)), axis=1)
    occupied2_x, occupied2_y = np.linspace(0, 1, sites)[::2], np.sum(occupied_edge_vector2_by_time[time_index][::2].reshape((-1, 2)) - occupied_edge_vector2_by_time[time_index][1::2].reshape((-1, 2)), axis=1)
    unoccupied1_x, unoccupied1_y = np.linspace(0, 1, sites)[::2], np.sum(unoccupied_edge_vector1_by_time[time_index][::2].reshape((-1, 2)) - unoccupied_edge_vector1_by_time[time_index][1::2].reshape((-1, 2)), axis=1)
    unoccupied2_x, unoccupied2_y = np.linspace(0, 1, sites)[::2], np.sum(unoccupied_edge_vector2_by_time[time_index][::2].reshape((-1, 2)) - unoccupied_edge_vector2_by_time[time_index][1::2].reshape((-1, 2)), axis=1)

    animation_ax.collections.clear()

    animation_plot_occupied1.set_xdata(occupied1_x)
    animation_plot_occupied1.set_ydata(occupied1_y)
    animation_plot_occupied2.set_xdata(occupied2_x)
    animation_plot_occupied2.set_ydata(occupied2_y)

    animation_fill_occupied1 = animation_ax.fill_between(occupied1_x, occupied1_y, label='Occupied 1', color='blue')
    animation_fill_occupied2 = animation_ax.fill_between(occupied2_x, occupied2_y, label='Occupied 2', color='blue')

    animation_plot_unoccupied1.set_xdata(occupied1_x)
    animation_plot_unoccupied1.set_ydata(occupied1_y)
    animation_plot_unoccupied2.set_xdata(occupied2_x)
    animation_plot_unoccupied2.set_ydata(occupied2_y)

    animation_fill_unoccupied1 = animation_ax.fill_between(unoccupied1_x, unoccupied1_y, alpha=0.2, label='Unoccupied 1', color='orange')
    animation_fill_unoccupied2 = animation_ax.fill_between(unoccupied2_x, unoccupied2_y, alpha=0.2, label='Unoccupied 1', color='orange')

    animation_ax.legend()

    return animation_fill_occupied1, animation_plot_occupied1, animation_fill_unoccupied1, animation_plot_unoccupied1, animation_fill_occupied2, animation_plot_occupied2, animation_fill_unoccupied2, animation_plot_unoccupied2

def animate_density_init():
    return animate_density(0)

density_animation = animation.FuncAnimation(animation_fig, animate_density, np.arange(steps), init_func=animate_density_init, interval=25, blit=True)
density_video = density_animation.to_html5_video()
display.display(display.HTML(density_video))