Plot real cell and Brillouin zone.

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)