Generalized grid with dsenum

kpoint.py

from itertools import product, combinations
from typing import List

import numpy as np
from pymatgen.core import Structure
from tqdm import tqdm

from dsenum.permutation_group import DerivativeMultiLatticeHash
from dsenum.superlattice import generate_all_superlattices, reduce_HNF_list_by_parent_lattice_symmetry
from dsenum.utils import get_symmetry_operations


class IrreducibleKPoints:

    def __init__(self, reciprocal_basis: np.ndarray, kpoint_stars: np.ndarray, rotations: np.ndarray):
        """
        Parameters
        ----------
        reciprocal_basis: (dim, dim)
            reciprocal_basis[i, :] is the i-th reciprocal basis vector
        kpoint_stars: (# of stars, dim)
        rotations: (# of rotation, dim, dim)
        """
        # TODO: Minkovski reduction
        self.reciprocal_basis = reciprocal_basis
        self.B = self.reciprocal_basis.T
        assert self.reciprocal_basis.shape[1] == self.dim

        self.kpoint_stars = kpoint_stars
        assert self.kpoint_stars.shape[1] == self.dim

        self.rotations = rotations
        assert rotations.shape[1] == self.dim
        assert rotations.shape[2] == self.dim

    @property
    def dim(self):
        return self.reciprocal_basis.shape[0]

    @property
    def num_stars(self):
        return self.kpoint_stars.shape[0]

    def max_distance(self):
        jimages = [np.array(offset) for offset in product([-1, 0, 1], repeat=self.dim)]

        ret = None

        # within a star
        for kp in self.kpoint_stars:
            for offset in jimages:
                for rotation in self.rotations:
                    reciprocal_coords = kp - np.dot(kp, rotation.T) + offset
                    dtmp = np.linalg.norm(np.dot(reciprocal_coords, self.B))
                    if np.isclose(dtmp, 0):
                        # identity operation case
                        continue
                    if (ret is None) or (dtmp > ret):
                        ret = dtmp

        # between two stars
        for k1, k2 in combinations(self.kpoint_stars, 2):
            for offset in jimages:
                for rotation in self.rotations:
                    reciprocal_coords = k1 - np.dot(k2, rotation.T) + offset
                    dtmp = np.linalg.norm(np.dot(reciprocal_coords, self.B))
                    if (ret is None) or (dtmp > ret):
                        ret = dtmp

        return ret


class GeneralizedGrid:

    def __init__(self, reciprocal_basis: np.ndarray, rotations: np.ndarray, transformation: np.ndarray):
        self.reciprocal_basis = reciprocal_basis
        assert self.reciprocal_basis.shape[1] == self.dim

        self.transformation = transformation
        assert self.transformation.shape[0] == self.dim
        assert self.transformation.shape[1] == self.dim

        self.rotations = rotations
        assert self.rotations.shape[1] == self.dim
        assert self.rotations.shape[2] == self.dim

    @property
    def dim(self):
        return self.reciprocal_basis.shape[0]

    def generate(self) -> IrreducibleKPoints:
        dmhash = DerivativeMultiLatticeHash(self.transformation, np.zeros((1, self.dim)))

        # lattice points in supercell

        factors = dmhash.get_all_factors()  # (index, dim)
        kpoints = [np.dot(np.array(factor).T, np.dot(np.linalg.inv(dmhash.snf), dmhash.left)) for factor in factors]
        kpoints = np.array([np.remainder(kp, 1) for kp in kpoints])

        stars = fold_kpoints_to_stars(kpoints, self.rotations)
        ikp = IrreducibleKPoints(self.reciprocal_basis, stars, self.rotations)
        return ikp

    @classmethod
    def with_structure(cls, structure: Structure, transformation: np.ndarray):
        rotations, _ = get_symmetry_operations(structure)
        reciprocal_basis = structure.lattice.reciprocal_lattice.matrix
        return cls(reciprocal_basis, rotations, transformation)


def fold_kpoints_to_stars(kpoints: np.ndarray, rotations: np.ndarray) -> np.ndarray:
    """
    Parameters
    ----------
    kpoints: (index, dim)
    rotations: (# of rotations, dim, dim)

    Returns
    -------
    stars: (# of stars, dim)
    """
    index = kpoints.shape[0]
    num_stars = 0
    stars = []
    visited = [-1 for _ in range(index)]
    for i, kp in enumerate(kpoints):
        if visited[i] != -1:
            continue
        stars.append(kp)
        for R in rotations:
            rotated = np.remainder(np.dot(kp, R.T), 1)
            # TODO: inefficient search
            for j in range(index):
                if visited[j] != -1:
                    continue
                if np.allclose(rotated, kpoints[j]):
                    visited[j] = num_stars
                    break
        num_stars += 1

    assert(np.all([(sm != -1) for sm in visited]))
    return np.array(stars)


def generate_distinct_sublattices_3d(index: int, rotations: np.ndarray):
    all_transformations = generate_all_superlattices(index)
    reduced = reduce_HNF_list_by_parent_lattice_symmetry(all_transformations, rotations)
    return reduced


if __name__ == '__main__':
    from collections import Counter
    matrix = np.array([[0, 0.5, 0.5], [0.5, 0, 0.5], [0.5, 0.5, 0]])
    species = ['Al']
    frac_coords = [[0., 0., 0.]]
    structure = Structure(matrix, species, frac_coords)

    rotations, _ = get_symmetry_operations(structure)
    index = 4

    hnfs = generate_distinct_sublattices_3d(index, rotations)
    counter = Counter()
    for hnf in tqdm(hnfs):
        gr = GeneralizedGrid.with_structure(structure, hnf)
        ikp = gr.generate()
        counter[ikp.max_distance()] += 1

    for key, cnt in counter.items():
        print(f"{key:.4f}: {cnt}")

    print("min:", min(counter.keys()))