Skip to content

Instantly share code, notes, and snippets.

@Ionizing

Ionizing/gridbz.py

Last active July 26, 2024 09:04
Show Gist options
  • Save Ionizing/e6ff28f7abb13bf6c5e923a2c4d2f19e to your computer and use it in GitHub Desktop.
Save Ionizing/e6ff28f7abb13bf6c5e923a2c4d2f19e to your computer and use it in GitHub Desktop.
Download ZIP
Plot real cell and Brillouin zone.
Raw
gridbz.py
#!/usr/bin/env python3
import numpy as np
from numpy.typing import NDArray
from scipy.spatial import Voronoi
import matplotlib.pyplot as plt
from matplotlib.axes import Axes
class RealCell:
def __init__(self,
cell: NDArray,
transform_matrix: NDArray):
assert cell.shape == (2, 2)
acell_prm = cell
acell_sup = transform_matrix @ acell_prm
self.acell_prm = acell_prm
self.acell_sup = acell_sup
bcell_prm = np.linalg.inv(acell_prm).T
bcell_sup = np.linalg.inv(acell_sup).T
self.bcell_prm = bcell_prm
self.bcell_sup = bcell_sup
@staticmethod
def generate_reduced_voronoi(cell: NDArray):
points = np.tensordot(
cell, np.mgrid[-1:2, -1:2], axes=(0, 0)
).reshape(2, -1).T
vor = Voronoi(points)
bz_vertices = []
for pid, rid in zip(vor.ridge_points, vor.ridge_vertices):
if 4 in pid:
bz_vertices += rid
bz_vertices = list(set(bz_vertices))
return np.array(vor.vertices[bz_vertices])
@staticmethod
def generate_polygon_edges(V: NDArray):
assert len(V.shape) == 2 and V.shape[1] == 2
centroid = np.mean(V, axis=0)
sorting_idx = np.argsort(np.arctan2((V-centroid)[:,0],
(V-centroid)[:,1]))
V = V[sorting_idx,:]
V = np.vstack((V, V[0,:]))
edges = []
for (left, right) in zip(V[:-1], V[1:]):
edges.append([left, right])
return np.asarray(edges)
@staticmethod
def irange(N: int):
start = -(N // 2)
end = start + N
return np.mgrid[start:end]
@staticmethod
def generate_polygon_edges_grid(E: NDArray, cell: NDArray,
*,
ngrid=(1,1)):
edges_total = np.zeros((0, 2, 2))
for i in RealCell.irange(ngrid[0]):
for j in RealCell.irange(ngrid[1]):
xy = [i, j] @ cell
edges_shift = E[:,:,:] + xy[None, None, :]
edges_total = np.vstack((edges_total, edges_shift))
return edges_total
@staticmethod
def edges_sort_dedup_x(E: NDArray):
assert len(E.shape) == 3 and \
E.shape[1:] == (2, 2)
E = E.copy()
for iseg, (left, right) in enumerate(E):
if left[0] > right[0]:
E[iseg, :, :] = np.array([right, left])
sorting_idx = np.argsort(E[:, 0, 0])
E = E[sorting_idx, :, :]
return np.unique(E, axis=0)
@staticmethod
def edges_sort_dedup_y(E: NDArray):
assert len(E.shape) == 3 and \
E.shape[1:] == (2, 2)
E = E.copy()
for iseg, (left, right) in enumerate(E):
if left[1] > right[1]:
E[iseg, :, :] = np.array([right, left])
sorting_idx = np.argsort(E[:, 0, 1])
E = E[sorting_idx, :, :]
return np.unique(E, axis=0)
@staticmethod
def edges_sort_dedup(E: NDArray,
*,
decimals=6):
assert len(E.shape) == 3 and \
E.shape[1:] == (2, 2)
E = np.round(E, decimals=decimals)
E = RealCell.edges_sort_dedup_x(E)
E = RealCell.edges_sort_dedup_y(E)
return E
def generate_bz_primitive_edges(self):
cell = self.bcell_prm
V = RealCell.generate_reduced_voronoi(cell)
E = RealCell.generate_polygon_edges(V)
return E
def generate_bz_supercell_edges(self,
*,
ngrid=(1,1)):
cell = self.bcell_sup
V = RealCell.generate_reduced_voronoi(cell)
E = RealCell.generate_polygon_edges(V)
E = RealCell.generate_polygon_edges_grid(E, cell, ngrid=ngrid)
E = RealCell.edges_sort_dedup(E)
return E
@staticmethod
def plot_edges(ax: Axes, edges: NDArray,
*,
ls: str,
lw: float,
color: str):
assert len(edges.shape) == 3 and \
edges.shape[1:] == (2, 2)
for edge in edges:
ax.plot(edge[:,0], edge[:,1], ls=ls, color=color, lw=lw)
if "__main__" == __name__:
cell = np.array([[ 2.7628944419999999,-1.5951578500000001],
[ 0.0000000000000000, 3.1903157000000002]]);
transform_matrix = np.array([[4, 2],
[0, 3]])
rc = RealCell(cell, transform_matrix)
E_bzprm = rc.generate_bz_primitive_edges()
E_bzsup = rc.generate_bz_supercell_edges(ngrid=(5, 5))
fig, ax = plt.subplots(figsize=(6, 6))
RealCell.plot_edges(ax, E_bzprm, color="b", ls= "-", lw=1.0)
RealCell.plot_edges(ax, E_bzsup, color="k", ls="--", lw=1.0)
ax.set_aspect("equal")
ax.axis("off")
fig.tight_layout(pad=0.5)
fig.savefig("supbz.png", dpi=400)
fig.savefig("supbz.pdf")
Raw
supbz.py
#!/usr/bin/env python3
import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
from scipy.spatial import Voronoi
def get_reduced_voronoi(cell):
assert cell.shape == (2,2)
points = np.tensordot(
cell, np.mgrid[-1:2, -1:2], axes=(0,0)
).reshape(2,-1).T
vor = Voronoi(points)
bz_facets = []
bz_ridges = []
bz_vertices = []
for pid, rid in zip(vor.ridge_points, vor.ridge_vertices):
if 4 in pid:
bz_ridges.append(vor.vertices[np.r_[rid, [rid[0]]]])
bz_facets.append(vor.vertices[rid])
bz_vertices += rid
bz_vertices = list(set(bz_vertices))
return (vor.vertices[bz_vertices], # V
bz_ridges, # E
bz_facets) # F
def plot_reduced_voronoi(cell, ax, lc, lw, ls, reci=False):
# icell = np.linalg.inv(cell).T
icell = cell
V, E, F = get_reduced_voronoi(icell)
# l1, l2 = np.linalg.norm(icell, axis=1)
xmin = min(ax.get_xlim()[0], icell[:,0].min())
xmax = max(ax.get_xlim()[1], icell[:,0].max())
ymin = min(ax.get_ylim()[0], icell[:,1].min())
ymax = max(ax.get_ylim()[1], icell[:,1].max())
for xx in E:
ax.plot(xx[:,0], xx[:,1], color=lc, lw=lw, ls=ls)
xmin = min(xmin, xx[:,0].min())
xmax = max(xmax, xx[:,0].max())
ymin = min(ymin, xx[:,1].min())
ymax = max(ymax, xx[:,1].max())
arrowprops = dict(arrowstyle="->",
shrinkA=0,
shrinkB=0,
color=lc)
s = "b" if reci else "a"
print(icell, tuple(icell[0,:]), tuple(icell[1,:]))
ax.text(*icell[0,:], s+"1", color=lc)
ax.annotate("", xy=tuple(icell[0,:]), xytext=(0,0), arrowprops=arrowprops)
ax.text(*icell[1,:], s+"2", color=lc)
ax.annotate("", xy=tuple(icell[1,:]), xytext=(0,0), arrowprops=arrowprops)
ax.set_xlim(xmin-0.02, xmax+0.02)
ax.set_ylim(ymin-0.02, ymax+0.02)
if '__main__' == __name__:
# 2D system only
# cell_primitive = np.array([[ 3.1903157000000002, 0.0000000000000000],
# [-1.5951578500000001, 2.7628944419999999]]);
cell_primitive = np.array([[ 2.7628944419999999,-1.5951578500000001],
[ 0.0000000000000000, 3.1903157000000002]]);
transform_matrix = np.array([[4, 2],
[0, 3]])
cell_supercell = transform_matrix @ cell_primitive
icell_primitive = np.linalg.inv(cell_primitive).T
icell_supercell = np.linalg.inv(cell_supercell).T
fig, axs = plt.subplots(ncols=2, figsize=(8, 4), dpi=400)
ax = axs[0]
ax.set_xlim(0,0.00001)
ax.set_ylim(0,0.00001)
plot_reduced_voronoi(cell_primitive, ax, lc="k", lw=1.0, ls='-', reci=False)
plot_reduced_voronoi(cell_supercell, ax, lc="b", lw=1.0, ls='-', reci=False)
ax.set_aspect('equal')
ax.axis('off')
ax = axs[1]
ax.set_xlim(0,0.00001)
ax.set_ylim(0,0.00001)
plot_reduced_voronoi(icell_primitive, ax, lc="k", lw=1.0, ls='-', reci=True)
plot_reduced_voronoi(icell_supercell, ax, lc="b", lw=1.0, ls='-', reci=True)
ax.set_aspect('equal')
ax.axis('off')
fig.tight_layout(pad=0.0)
fig.savefig("bz.png", dpi=400)
@Ionizing
Copy link
Author

bz

@Ionizing
Copy link
Author

Example output of gridbz.py

supbz

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment