forked from zhangyan31415/continuum-model-twist-tmds
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathplot_quantum_geo.py
97 lines (75 loc) · 3.44 KB
/
plot_quantum_geo.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
import numpy as np
import matplotlib.pyplot as plt
import model_geometric as mg
nk_mesh = mg.nk_mesh
kmesh = mg.kmesh
s = mg.s
G1 = mg.model.bM1
G2 = mg.model.bM2
m_k1_vec = (G2+G1)/3
Oxy = np.load("TrG.npy")
print(Oxy.shape)
Oxy1 = Oxy[:,-1]
Oxy2 = Oxy[:,-2]
Oxy3 = Oxy[:,-3]
print(np.sum(Oxy1)*s/(2*np.pi))
print(np.sum(Oxy2)*s/(2*np.pi))
print(np.sum(Oxy3)*s/(2*np.pi))
def draw_hexagon(ax, corner):
# Calculate the length of the side of the hexagon
# Assuming the hexagon is regular
side_length = np.linalg.norm(corner)
# Generate the angles for the hexagon corners
angles = np.linspace(0, 2 * np.pi, 7) # 7 points to close the hexagon
# Calculate the coordinates of the hexagon corners
hexagon_points = np.array([[corner[0] * np.cos(angle) - corner[1] * np.sin(angle),
corner[0] * np.sin(angle) + corner[1] * np.cos(angle)]
for angle in angles])
# Plot the hexagon
ax.plot(hexagon_points[:, 0], hexagon_points[:, 1], ls='-', lw=0.5,c='white')
fig, axs = plt.subplots(1, 3, figsize=(10,4), layout='constrained', sharey=True)
axs[0].set_aspect('equal')
axs[1].set_aspect('equal')
axs[2].set_aspect('equal')
# pcm1 = axs[0].pcolormesh(KX, KY, Oxy1, shading='gouraud', cmap='RdBu')
# pcm2 = axs[1].pcolormesh(KX, KY, Oxy2, shading='gouraud', cmap='RdBu')
# pcm3 = axs[2].pcolormesh(KX, KY, Oxy3, shading='gouraud', cmap='RdBu')
# real_space = [[i, j] for i in range(-1,1,1) for j in range(-1,1,1)]
# for pair in real_space:
# nk = nk_mesh
# kmesh_now = kmesh+pair[0]*G1+pair[1]*G2
# KX = np.zeros((nk, nk))
# KY = np.zeros((nk, nk))
# for i in range(nk):
# for j in range(nk):
# KX[i, j] = kmesh_now[i*nk+j, 0]
# KY[i, j] = kmesh_now[i*nk+j, 1]
# pcm1 = axs[0].pcolormesh(KX, KY, Oxy1.reshape((nk,nk)),shading='gouraud',cmap='RdBu')
# pcm2 = axs[1].pcolormesh(KX, KY, Oxy2.reshape((nk,nk)),shading='gouraud',cmap='RdBu')
# pcm3 = axs[2].pcolormesh(KX, KY, Oxy3.reshape((nk,nk)),shading='gouraud',cmap='RdBu')
# Assuming nk, kmesh, G1, G2, Oxy1, Oxy2, Oxy3 are defined
nk = nk_mesh # Example: nk_mesh = 10
real_space = np.array([[i, j] for i in range(-1, 1) for j in range(-1, 1)])
# Initialize large grids
KX = np.zeros((2 * nk, 2 * nk))
KY = np.zeros((2 * nk, 2 * nk))
Oxy1_large = np.zeros((2 * nk, 2 * nk))
Oxy2_large = np.zeros((2 * nk, 2 * nk))
Oxy3_large = np.zeros((2 * nk, 2 * nk))
# Fill the large grids
for idx, (i_shift, j_shift) in enumerate(real_space):
kmesh_now = kmesh + i_shift * G1 + j_shift * G2
KX[idx // 2 * nk:(idx // 2 + 1) * nk, idx % 2 * nk:(idx % 2 + 1) * nk] = kmesh_now[:, 0].reshape(nk, nk)
KY[idx // 2 * nk:(idx // 2 + 1) * nk, idx % 2 * nk:(idx % 2 + 1) * nk] = kmesh_now[:, 1].reshape(nk, nk)
Oxy1_large[idx // 2 * nk:(idx // 2 + 1) * nk, idx % 2 * nk:(idx % 2 + 1) * nk] = Oxy1.reshape(nk, nk)
Oxy2_large[idx // 2 * nk:(idx // 2 + 1) * nk, idx % 2 * nk:(idx % 2 + 1) * nk] = Oxy2.reshape(nk, nk)
Oxy3_large[idx // 2 * nk:(idx // 2 + 1) * nk, idx % 2 * nk:(idx % 2 + 1) * nk] = Oxy3.reshape(nk, nk)
# Plot the larger grids
pcm1 = axs[0].pcolormesh(KX, KY, Oxy1_large, shading='gouraud',cmap='RdBu')
pcm2 = axs[1].pcolormesh(KX, KY, Oxy2_large, shading='gouraud',cmap='RdBu')
pcm3 = axs[2].pcolormesh(KX, KY, Oxy3_large, shading='gouraud',cmap='RdBu')
draw_hexagon(axs[0], m_k1_vec)
draw_hexagon(axs[1], m_k1_vec)
draw_hexagon(axs[2], m_k1_vec)
plt.savefig("quantum.pdf")
plt.show()