Plot real cell and Brillouin zone.

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")

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)

Example output of gridbz.py