draw_2D_BZs.py ダウンロード/コピー
draw_2D_BZs.py
draw_2D_BZs.py
1"""
2概要: 2次元のブリルアンゾーンと関連する格子点をプロットするスクリプト。
3詳細説明: このスクリプトは、指定された数のブリルアンゾーン（BZ）を計算し、
4それらを色分けしてプロットします。また、BZの境界線、逆格子点、および
5カスタムカラーバー（凡例）も描画します。
6物理的な単位系は、2π/a を単位とする逆格子空間で表現されます。
7
8関連リンク: :doc:`draw_2D_BZs_usage`
9"""
10import numpy as np
11import matplotlib.pyplot as plt
12from matplotlib.colors import ListedColormap
13import matplotlib.patches as patches
14from mpl_toolkits.axes_grid1 import make_axes_locatable
15import matplotlib.patheffects as patheffects # 修正1: ここを追加
16import argparse
17
18def draw_custom_colorbar(fig, ax, colors, max_zone):
19 """
20 概要: メインのAxesの右側に、自作のカラーバー（凡例）を描画する。
21 詳細説明: この関数は、ブリルアンゾーンの番号に対応する色と番号を表示する
22 カスタムカラーバーを作成し、メインプロットの右側に配置します。
23 各色領域の中央に、対応するブリルアンゾーンの番号が白字と黒い縁取りで表示されます。
24
25 :param fig: matplotlib.figure.Figure: メインのFigureオブジェクト。
26 :param ax: matplotlib.axes.Axes: メインのAxesオブジェクト。このAxesの右側にカラーバーが追加されます。
27 :param colors: list[tuple]: BZ番号に対応する色のリスト。リストのインデックスはBZ番号に相当します。
28 :param max_zone: int: 描画する最大のブリルアンゾーン番号。
29 :returns: None: この関数は何も返しません。
30 """
31 divider = make_axes_locatable(ax)
32 cax = divider.append_axes("right", size="5%", pad=0.1)
33
34 cax.set_xlim(0, 1)
35 cax.set_ylim(0, max_zone)
36
37 cax.set_xticks([])
38 cax.set_yticks([])
39 for spine in cax.spines.values():
40 spine.set_visible(False)
41
42 for i in range(max_zone):
43 bz_number = i + 1
44
45 # 四角形の描画
46 rect = patches.Rectangle((0, i), 1, 1,
47 linewidth=0.5, edgecolor='black',
48 facecolor=colors[bz_number])
49 cax.add_patch(rect)
50
51 # テキスト（番号）を中央に配置
52 # 修正1対応: patheffectsを正しく呼び出し
53 cax.text(0.5, i + 0.5, str(bz_number),
54 ha='center', va='center', fontsize=12, fontweight='bold',
55 color='white',
56 path_effects=[patheffects.withStroke(linewidth=2, foreground='black')])
57
58 cax.set_title("BZ", fontsize=10)
59
60def plot_brillouin_zones(max_zone, limit = None):
61 """
62 概要: 2次元正方格子のブリルアンゾーンを計算し、可視化する。
63 詳細説明: 指定された数のブリルアンゾーン（BZ）を計算し、各ゾーンを異なる色で描画します。
64 プロットには、ブリルアンゾーンの境界線（垂直二等分線）と逆格子点も含まれます。
65 また、ゾーン番号を示すカスタムカラーバーが右側に表示されます。
66 計算は、各点 (X, Y) がどの逆格子点 G=(i, j) に最も近いか（ただし G=(0,0) を除く）
67 を判断することで行われます。最も近い逆格子点が k 空間原点 (0,0) であれば第1BZ、
68 最も近い逆格子点が G1 であれば、その点と原点、および G1 と原点の中間地点を結ぶ線で区切られます。
69 一般的な第 n BZ は、原点に最も近い n 個の逆格子点よりも、原点に近い領域です。
70
71 :param max_zone: int: 描画する最大のブリルアンゾーン番号。
72 :param limit: float or None: プロットの表示範囲（-limitからlimitまで）。単位は 2π/a。
73 Noneの場合、`np.sqrt(max_zone) + 1.0` を基に自動的に計算されます。
74 :returns: None: この関数は何も返しません。
75 """
76 resolution = 1000
77 if limit is None: limit = np.sqrt(max_zone) + 1.0
78
79 x = np.linspace(-limit, limit, resolution)
80 y = np.linspace(-limit, limit, resolution)
81 X, Y = np.meshgrid(x, y)
82
83 dist_origin_sq = X**2 + Y**2
84 zone_index = np.ones_like(X, dtype=int)
85
86 search_range = int(np.ceil(limit + 1))
87 lines_to_draw = []
88
89 # --- ゾーン計算 ---
90 for i in range(-search_range, search_range + 1):
91 for j in range(-search_range, search_range + 1):
92 if i == 0 and j == 0:
93 continue
94
95 dist_g_sq = (X - i)**2 + (Y - j)**2
96 zone_index += (dist_g_sq < dist_origin_sq)
97
98 rhs = 0.5 * (i**2 + j**2)
99 if rhs / np.sqrt(i**2 + j**2) < limit * 1.5:
100 lines_to_draw.append((i, j, rhs))
101
102 # --- 描画準備 ---
103 plot_data = np.where(zone_index <= max_zone, zone_index, 0)
104
105 base_cmap = plt.get_cmap('tab10')
106 colors = [base_cmap(i % 10) for i in range(max_zone + 1)]
107 colors[0] = (1, 1, 1, 1) # ゾーン0 (範囲外) は白
108
109 custom_cmap = ListedColormap(colors)
110
111 fig, ax = plt.subplots(figsize=(10, 8))
112
113 ax.imshow(plot_data, extent=[-limit, limit, -limit, limit],
114 origin='lower', cmap=custom_cmap, interpolation='nearest')
115
116 x_range = np.linspace(-limit, limit, 100)
117 for (i, j, rhs) in lines_to_draw:
118 if j == 0: # 垂直線の場合
119 x_val = rhs / i
120 if -limit <= x_val <= limit:
121 ax.vlines(x_val, -limit, limit, colors='black', linestyles='-.', linewidth=0.5)
122 else: # それ以外の場合
123 y_vals = (rhs - i * x_range) / j
124 ax.plot(x_range, y_vals, color='black', linestyle='-.', linewidth=0.5)
125
126 grid_points_x = []
127 grid_points_y = []
128 for i in range(-search_range, search_range + 1):
129 for j in range(-search_range, search_range + 1):
130 grid_points_x.append(i)
131 grid_points_y.append(j)
132 ax.scatter(grid_points_x, grid_points_y, c='black', s=15, zorder=10) # 逆格子点をプロット
133
134 ax.set_xlim(-limit, limit)
135 ax.set_ylim(-limit, limit)
136 ax.set_aspect('equal')
137 ax.set_title(f'Square Lattice Brillouin Zones (1-{max_zone})')
138
139 # 修正2: raw string (r'...') を使用して警告を回避
140 ax.set_xlabel(r'$k_x / (2\pi/a)$')
141 ax.set_ylabel(r'$k_y / (2\pi/a)$')
142
143 draw_custom_colorbar(fig, ax, colors, max_zone)
144
145 plt.tight_layout()
146 plt.show()
147
148if __name__ == "__main__":
149 parser = argparse.ArgumentParser(description='Draw Brillouin Zones with Custom Legend.')
150 parser.add_argument('max_zone', type=int, nargs='?', default=10,
151 help='The maximum Brillouin zone number to draw (default: 5)')
152 parser.add_argument('limit', type=float, nargs='?', default=None,
153 help='The maximum plot range in 2pi/a (default: None)')
154 args = parser.parse_args()
155
156 plot_brillouin_zones(args.max_zone, args.limit)