structure-info.py

import json
import spglib
import numpy as np


def irreducible_kpoints(fname_sirius_json="sirius.json"):
    sirius_config = json.load(open(fname_sirius_json, "r"))
    pos = sirius_config["unit_cell"]["atoms"]
    C = np.array(sirius_config["unit_cell"]["lattice_vectors"])
    positions = list(pos.values())
    indicator = [np.ones(len(x)) * i for i, x in enumerate(positions)]
    indicator = np.hstack(indicator)

    # assume atom pos given in a.u.
    pos = np.vstack(positions)
    rpos = np.linalg.solve(C.T, pos[:, :3].T).T
    cell = (C, rpos, indicator)
    mesh = sirius_config["parameters"]["ngridk"]
    if (rpos > 1).any() or (rpos < 0).any():
        print("WARNING: atoms not inside unit cell")

    # Gamma centre mesh
    mapping, grid = spglib.get_ir_reciprocal_mesh(mesh, cell, is_shift=[0, 0, 0])

    # Irreducible k-points
    print("ngridk              :", mesh)
    print("Number of ir-kpoints: %d" % len(np.unique(mapping)))
    # print(grid[np.unique(mapping)] / np.array(mesh, dtype=float))


if __name__ == "__main__":

    irreducible_kpoints()

    # count electrons
    sirius_json = json.load(open("sirius.json", "r"))

    ne = {}

    for atom_type, upf in sirius_json["unit_cell"]["atom_files"].items():
        ne[atom_type] = json.load(open(upf, 'r'))["pseudo_potential"]["header"]["z_valence"]

    print('-------------------------------')
    z_valence = 0
    for atom_type in sirius_json['unit_cell']['atoms']:
        coords = sirius_json['unit_cell']['atoms'][atom_type]
        ne_loc = ne[atom_type]
        natoms_type = len(coords)
        zloc = natoms_type * ne_loc
        print(atom_type, '*', natoms_type, '=', ne_loc, 'electrons')
        z_valence += zloc

    print('-------------------------------')
    print('total: ', z_valence, 'electrons')
    print('-------------------------------')