1D Shockley Surface States, with a Symmetrical System

import numpy as np
import matplotlib.pyplot as plt
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

from scipy import fft

N = 100
M = 1
m = 1
K = 1
k = 5
a = 1

Matrix = np.arange((2*N)*(2*N)).reshape(2*N, 2*N)

for i in range(2*N):
    for j in range(2*N):
        if i == j == 0:
            Matrix[i, j] = (K+k) / M
        elif i == j == 2*N-1:
            Matrix[i, j] = (K+k) / M
        elif i%2 == 0:
            if i == j:
                Matrix[i, j] = (K+k) / m
            elif i == j-1:
                Matrix[i, j] = -K / m
            elif i == j+1:
                Matrix[i, j] = -k / m 
            else:
                Matrix[i, j] = 0
        elif i%2 == 1:
            if i == j:
                Matrix[i, j] = (K+k) / M
            elif i == j-1:
                Matrix[i, j] = -k / M 
            elif i == j+1:
                Matrix[i, j] = -K / M
            else:
                Matrix[i, j] = 0

print(Matrix)

values, vectors = np.linalg.eigh(Matrix)
vectors = np.transpose(vectors)

idx = np.argsort(values)
values_sorted = values
vectors_sorted = vectors

n=np.argmin(values_sorted)
print(n, values_sorted[n])
plt.plot(vectors_sorted[n])
plt.title('Lowest Frequency Standing Wave Solution')
plt.xlabel(r'$n$')
plt.ylabel(r'$u$')

# plt.savefig('./fig/Shockley-parity/lowest_solution.svg')

def dstfreq(N, a):
    return np.array([(i+1)/(2*a*(N+1)) for i in range(N)])

q = []

for i in range(2*N):
    peaks = np.abs(fft.dst(vectors_sorted[i][1::2], type=1))
    freqs = dstfreq(N, 2*a) * 2

    if i == N // 2 + 3:
        plt.plot(freqs, peaks, '-')
        # print(freqs)
        # print(peaks)
    q.append(2*np.pi*freqs[np.argmax(peaks)])

plt.title('DFT Spectra of a Normal Standing Wave')
plt.xlabel(r'$q$')
plt.ylabel('Amplitude')

# plt.savefig('./fig/Shockley-parity/dft_spectra_normal.svg')
# plt.plot(vectors_sorted[np.argmin(q)])

q0 = np.array([j * np.pi / (a * (N+1)) for j in range(1, N)])

plt.plot(q, np.sqrt(values_sorted), '.', color='blue', label='Actual dispersion', markersize=5)
# plt.plot(q0, np.sqrt(4*K/M)*np.abs(np.sin(q0*a/2.0)), color='red', label='Perfect dispersion', alpha=0.8)
plt.legend()

plt.title('Dispersion Relation')
plt.xlabel(r'$q$')
plt.ylabel(r'$\omega$')

# plt.savefig('./fig/Shockley-parity/dispersion.svg')

n_surface = [a for a in values_sorted if (a < 1.5**2) and (a > 3.0**2)]
print(n_surface)

n = [i for v,i in enumerate(values_sorted) if (v > 1.5**2) and (v < 2.8**2)]
print(n)

n_max = N
n_maxes = np.argsort(values_sorted)[2*N-2:]
print(n_maxes)

print(n_max, q[n_max], np.sqrt(values_sorted[n_max]))

plt.plot(vectors_sorted[n_max], color='blue', label='Surface state localized right')
plt.legend()
plt.title('Localized Surface State')
plt.xlabel(r'$n$')
plt.ylabel(r'$u$')

plt.plot(vectors_sorted[n_max-1], color='red', label='Surface state localized left')
plt.legend()
plt.title('Localized Surface State')
plt.xlabel(r'$n$')
plt.ylabel(r'$u$')

# plt.savefig('./fig/Shockley-parity/localized_waves.svg')

peaks = np.abs(fft.dst(vectors_sorted[n_max][1::2], type=1))
freqs = dstfreq(N, a)
plt.plot(freqs, peaks, '-', color='blue', label='Surface state localized right')
plt.plot(freqs, peaks, '-', color='red', label='Surface state localized left')


plt.title('DFT Spectra of Localized Surface States')
plt.xlabel(r'$q$')
plt.ylabel('Amplitude')
plt.legend()

# plt.savefig('./fig/Shockley-parity/dft_spectra.svg')

plt.hlines(values_sorted[:n_max], linewidth=0.2, xmin=0, xmax=1, color='blue', label='Lower bulk band')
plt.hlines(values_sorted[n_max+1:], linewidth=0.2, xmin=0, xmax=1, color='red', label='Upper bulk band')
plt.hlines(values_sorted[n_max], linewidth=0.8, xmin=0, xmax=1, color='purple', label='Surface state')
plt.legend()
plt.title('Energy Bands')
plt.ylabel(r'$E/\hbar$')

# plt.savefig('./fig/Shockley-parity/energy_bands.svg')