Source code for icet.core.sublattices

from copy import deepcopy
from icet.core.structure import Structure
from typing import Iterator
from ase import Atoms
import copy
from itertools import product
from string import ascii_uppercase
import numpy as np
from icet.tools.geometry import chemical_symbols_to_numbers




class Sublattice:
    """
    This class stores and provides information about a specific
    sublattice. A sublattice is always supercell specific since
    it contains lattice indices.

    Note
    ----
    As a user you will usually not interact directly with objects of this type.

    Parameters
    ----------
    chemical_symbols
        The allowed species on this sublattice.
    indices
        The lattice indices the sublattice consists of.
    symbol
        String used to mark the sublattice.
    """

    def __init__(self,
                 chemical_symbols: list[str],
                 indices: list[int],
                 symbol: str):
        self._chemical_symbols = chemical_symbols
        self._indices = indices
        self._symbol = symbol
        self._numbers = chemical_symbols_to_numbers(chemical_symbols)

    @property
    def chemical_symbols(self):
        return copy.deepcopy(self._chemical_symbols)

    @property
    def atomic_numbers(self):
        return self._numbers.copy()

    @property
    def indices(self):
        return self._indices.copy()

    @property
    def symbol(self):
        """Symbol representation of sublattice, i.e. A, B, C, etc."""
        return self._symbol





class Sublattices:
    """
    This class stores and provides information about the sublattices
    of a structure.

    Note
    ----
    As a user you will usually not interact directly with objects of this type.

    Parameters
    ----------
    allowed_species
        List of the allowed species on each site of the primitve
        structure. For example this can be the chemical_symbols from
        a cluster space.
    primitive_structure
        The primitive structure the allowed species reference to.
    structure
        The structure that the sublattices will be based on.
    fractional_position_tolerance
        Tolerance applied when comparing positions in fractional coordinates.
    """

    def __init__(self,
                 allowed_species: list[list[str]],
                 primitive_structure: Atoms,
                 structure: Atoms,
                 fractional_position_tolerance: float):
        self._structure = structure
        # sorted unique sites, this basically decides A, B, C... sublattices
        active_lattices = sorted(set([tuple(sorted(symbols))
                                      for symbols in allowed_species if len(symbols) > 1]))
        inactive_lattices = sorted(
            set([tuple(sorted(symbols)) for symbols in allowed_species if len(symbols) == 1]))
        self._allowed_species = active_lattices + inactive_lattices

        n = int(np.sqrt(len(self._allowed_species))) + 1
        symbols = [''.join(p) for r in range(1, n+1) for p in product(ascii_uppercase, repeat=r)]

        cpp_prim_structure = Structure.from_atoms(primitive_structure)
        self._sublattices = []
        sublattice_to_indices = [[] for _ in range(len(self._allowed_species))]
        for index, position in enumerate(structure.positions):
            lattice_site = cpp_prim_structure.find_lattice_site_by_position(
                position=position, fractional_position_tolerance=fractional_position_tolerance)

            # Get allowed species on this site
            species = allowed_species[lattice_site.index]

            # Get what sublattice those species correspond to
            sublattice = self._allowed_species.index(tuple(sorted(species)))

            sublattice_to_indices[sublattice].append(index)

        for symbol, species, indices in zip(symbols, self._allowed_species, sublattice_to_indices):
            sublattice = Sublattice(chemical_symbols=species, indices=indices, symbol=symbol)
            self._sublattices.append(sublattice)

        # Map lattice index to sublattice index
        self._index_to_sublattice = {}
        for k, sublattice in enumerate(self):
            for index in sublattice.indices:
                self._index_to_sublattice[index] = k

    def __getitem__(self, key: int) -> Sublattice:
        """ Returns a sublattice according to key. """
        return self._sublattices[key]

    def __len__(self):
        """ Returns number of sublattices. """
        return len(self._sublattices)

    def __iter__(self) -> Iterator[Sublattice]:
        """ Generator over sublattices. """
        yield from self._sublattices



    def get_sublattice_index_from_site_index(self, index: int) -> int:
        """ Returns the index of the sublattice the
        index in the structure belongs to.

        Parameters
        ----------
        index
            Index of site in the structure.
        """
        return self._index_to_sublattice[index]


    @property
    def allowed_species(self) -> list[list[str]]:
        """Lists of the allowed species on each sublattice in order. """
        return deepcopy(self._allowed_species)



    def get_sublattice_sites(self, index: int) -> list[int]:
        """Returns the sites that belong to the sublattice with the
        corresponding index.

        Parameters
        ----------
        index
            Index of the sublattice.
        """
        return self[index].indices




    def get_allowed_symbols_on_site(self, index: int) -> list[str]:
        """Returns the allowed symbols on the site.

        Parameters
        -----------
        index
            Lattice site index.
        """
        return self[self._index_to_sublattice[index]].chemical_symbols




    def get_allowed_numbers_on_site(self, index: int) -> list[int]:
        """Returns the allowed atomic numbers on the site.

        Parameters
        -----------
        index
            Lattice site index.
        """
        return self[self._index_to_sublattice[index]].atomic_numbers


    @property
    def active_sublattices(self) -> list[Sublattice]:
        """ List of active sublattices. """
        return [sl for sl in self if len(sl.chemical_symbols) > 1]

    @property
    def inactive_sublattices(self) -> list[Sublattice]:
        """ List of inactive sublattices. """
        return [sl for sl in self if len(sl.chemical_symbols) == 1]



    def assert_occupation_is_allowed(self, chemical_symbols: list[str]):
        """ Asserts that the current occupation obeys the sublattices. """
        if len(chemical_symbols) != len(self._structure):
            raise ValueError(f'Length of input chemical symbols ({len(chemical_symbols)}) does not'
                             f' match length of supercell ({len(self._structure)}')
        for sl in self:
            for i in sl.indices:
                if not chemical_symbols[i] in sl.chemical_symbols:
                    msg = ('Occupations of structure not compatible with the sublattice.'
                           ' Site {} with occupation {} not allowed on sublattice {}'
                           .format(i, chemical_symbols[i], sl.chemical_symbols))
                    raise ValueError(msg)