Add an option to change the orbital order when folding. (#101)

Add an option to choose between different ways of sorting orbitals
in the `fold_model` method: Either by their index in the given
model (the previous default), or by the occurrence of keys in the
`orbital_labels` input (the new default).

The reason for this new feature is that the sort order can be
unpredictable when it is unclear which orbitals end up inside or
outside of the new unit cell.

Author: greschd

setup.cfg

@@ -56,8 +56,9 @@ dev =
     pre-commit
     pylint==2.6.0
     isort==5.5.1
-    mypy==0.782
+    mypy==0.812
     ruamel.yaml
+    typing-extensions
 
 [options.entry_points]
 console_scripts =

tbmodels/_tb_model.py

@@ -6,6 +6,8 @@
 Implements the Model class, which describes a tight-binding model.
 """
 
+from __future__ import annotations
+
 import re
 import os
 import copy
@@ -24,6 +26,8 @@
 from fsc.hdf5_io import subscribe_hdf5, HDF5Enabled
 
 if ty.TYPE_CHECKING:
+    # Replace with typing.Literal once Python 3.7 support is dropped.
+    from typing_extensions import Literal
     import symmetry_representation  # pylint: disable=unused-import
 
 from .kdotp import KdotpModel
@@ -334,7 +338,7 @@ def from_hop_list(
         hop_list: ty.Iterable[ty.Tuple[complex, int, int, ty.Tuple[int, ...]]] = (),
         size: ty.Optional[int] = None,
         **kwargs,
-    ) -> "Model":
+    ) -> Model:
         """
         Create a :class:`.Model` from a list of hopping terms.
 
@@ -529,7 +533,7 @@ def _mat_to_hr(R, mat):
     @classmethod
     def from_wannier_folder(
         cls, folder: str = ".", prefix: str = "wannier", **kwargs
-    ) -> "Model":
+    ) -> Model:
         """
         Create a :class:`.Model` instance from Wannier90 output files,
         given the folder containing the files and file prefix.
@@ -573,7 +577,7 @@ def from_wannier_files(  # pylint: disable=too-many-locals
         pos_kind: str = "wannier",
         distance_ratio_threshold: float = 3.0,
         **kwargs,
-    ) -> "Model":
+    ) -> Model:
         """
         Create a :class:`.Model` instance from Wannier90 output files.
 
@@ -988,7 +992,7 @@ def construct_kdotp(self, k: ty.Sequence[float], order: int):
     @classmethod
     def from_hdf5_file(  # pylint: disable=arguments-differ
         cls, hdf5_file: str, **kwargs
-    ) -> "Model":
+    ) -> Model:
         """
         Returns a :class:`.Model` instance read from a file in HDF5
         format.
@@ -1013,7 +1017,7 @@ def from_hdf5_file(  # pylint: disable=arguments-differ
     @classmethod
     def from_hdf5(  # pylint: disable=arguments-differ
         cls, hdf5_handle, **kwargs
-    ) -> "Model":
+    ) -> Model:
         # For compatibility with a development version which created a top-level
         # 'tb_model' attribute.
         try:
@@ -1349,10 +1353,10 @@ def _input_kwargs(self):
 
     def symmetrize(
         self,
-        symmetries: ty.Sequence["symmetry_representation.SymmetryOperation"],
+        symmetries: ty.Sequence[symmetry_representation.SymmetryOperation],
         full_group: bool = False,
         position_tolerance: float = 1e-5,
-    ) -> "Model":
+    ) -> Model:
         """
         Returns a model which is symmetrized w.r.t. the given
         symmetries. This is done by performing a group average over the
@@ -1404,7 +1408,7 @@ def symmetrize(
 
     def _apply_operation(  # pylint: disable=too-many-locals
         self, symmetry_operation, position_tolerance
-    ) -> "Model":
+    ) -> Model:
         """
         Helper function to apply a symmetry operation to the model.
         """
@@ -1473,7 +1477,7 @@ def _apply_operation(  # pylint: disable=too-many-locals
 
         return Model(**co.ChainMap(dict(hop=new_hop), self._input_kwargs))
 
-    def slice_orbitals(self, slice_idx: ty.List[int]) -> "Model":
+    def slice_orbitals(self, slice_idx: ty.List[int]) -> Model:
         """
         Returns a new model with only the orbitals as given in the
         ``slice_idx``. This can also be used to re-order the orbitals.
@@ -1491,7 +1495,7 @@ def slice_orbitals(self, slice_idx: ty.List[int]) -> "Model":
         return Model(**co.ChainMap(dict(hop=new_hop, pos=new_pos), self._input_kwargs))
 
     @classmethod
-    def join_models(cls, *models: "Model") -> "Model":
+    def join_models(cls, *models: Model) -> Model:
         """
         Creates a tight-binding model which contains all orbitals of the
         given input models. The orbitals are ordered by model, such that
@@ -1559,7 +1563,7 @@ def change_unit_cell(  # pylint: disable=too-many-branches
         uc: ty.Optional[ty.Sequence[ty.Sequence[float]]] = None,
         offset: ty.Sequence[float] = (0, 0, 0),
         cartesian: bool = False,
-    ) -> "Model":
+    ) -> Model:
         """Return a model with a different unit cell of the same volume.
 
         Creates a model with a changed unit cell - with a different
@@ -1646,7 +1650,7 @@ def change_unit_cell(  # pylint: disable=too-many-branches
 
     def supercell(  # pylint: disable=too-many-locals
         self, size: ty.Sequence[int]
-    ) -> "Model":
+    ) -> Model:
         """Generate a model for a supercell of the current unit cell.
 
         Parameters
@@ -1738,7 +1742,8 @@ def fold_model(  # pylint: disable=too-many-locals,too-many-branches,too-many-st
         check_uc_volume: bool = True,
         uc_volume_tolerance: float = 1e-6,
         check_orbital_ratio: bool = True,
-    ) -> "Model":
+        order_by: Literal["label", "index"] = "label",
+    ) -> Model:
         """
         Returns a model with a smaller unit cell. Orbitals which are related
         by a lattice vector of the new unit cell are "folded" into a single
@@ -1789,6 +1794,15 @@ def fold_model(  # pylint: disable=too-many-locals,too-many-branches,too-many-st
             should be checked to be the same as the initial ratio. If
             this is set to False, the resulting model will always
             have ``occ=None``.
+        order_by :
+            Determines how the orbitals in the new model are ordered.
+            For ``order_by="index"``, the orbitals are ordered exactly the
+            same as in the original model. Note that this order will
+            depend on which orbitals end up inside the unit cell.
+            For ``order_by="label"``, the orbitals are ordered
+            according to the **occurrence** (not the value) of the
+            labels in the ``orbital_labels`` input. Orbitals with the
+            same label are again ordered by index.
         """
         if len(orbital_labels) != self.size:
             raise ValueError(
@@ -1863,6 +1877,24 @@ def fold_model(  # pylint: disable=too-many-locals,too-many-branches,too-many-st
                 axis=-1,
             )
         ).flatten()
+        if order_by == "label":
+            idx = 0
+            orbital_sort_idx = {}
+            for label in orbital_labels:
+                if label not in orbital_sort_idx:
+                    orbital_sort_idx[label] = idx
+                    idx += 1
+            in_uc_sort_idx = [
+                orbital_sort_idx[orbital_labels[i]] for i in in_uc_indices
+            ]
+            in_uc_indices = in_uc_indices[
+                np.argsort(in_uc_sort_idx, kind="mergesort")  # need stable sorting
+            ]
+        else:
+            if order_by != "index":
+                raise ValueError(
+                    f"Invalid input '{order_by}' for 'order_by', must be either 'label' or 'index'."
+                )
         if target_indices is not None:
             if not np.all(target_indices == in_uc_indices):
                 raise ValueError(
@@ -2001,7 +2033,7 @@ def get_matching_idx_and_offset(pos_reduced, orbital_label):
             )
         )
 
-    def __add__(self, model: "Model") -> "Model":
+    def __add__(self, model: Model) -> Model:
         """
         Adds two models together by adding their hopping terms.
         """
@@ -2059,19 +2091,19 @@ def __add__(self, model: "Model") -> "Model":
         # -------------------
         return Model(**co.ChainMap(dict(hop=new_hop), self._input_kwargs))
 
-    def __sub__(self, model: "Model") -> "Model":
+    def __sub__(self, model: Model) -> Model:
         """
         Substracts one model from another by substracting all hopping terms.
         """
         return self + -model
 
-    def __neg__(self) -> "Model":
+    def __neg__(self) -> Model:
         """
         Changes the sign of all hopping terms.
         """
         return -1 * self
 
-    def __mul__(self, x: float) -> "Model":
+    def __mul__(self, x: float) -> Model:
         """
         Multiplies hopping terms by x.
         """
@@ -2081,13 +2113,13 @@ def __mul__(self, x: float) -> "Model":
 
         return Model(**co.ChainMap(dict(hop=new_hop), self._input_kwargs))
 
-    def __rmul__(self, x: float) -> "Model":
+    def __rmul__(self, x: float) -> Model:
         """
         Multiplies hopping terms by x.
         """
         return self.__mul__(x)
 
-    def __truediv__(self, x: float) -> "Model":
+    def __truediv__(self, x: float) -> Model:
         """
         Divides hopping terms by x.
         """

tests/test_folding.py

@@ -202,3 +202,56 @@ def test_new_unit_cell_check(get_model):
             orbital_labels=["a", "b"] * 2,
         )
     assert "new unit cell is not contained" in str(excinfo.value)
+
+
+def test_fold_by_label_order(get_model, models_close):
+    """
+    Test that the ``order_by`` input creates the expected result.
+    """
+    model = get_model(
+        0.1,
+        0.4,
+        uc=[[1, 0, 0.5], [0.1, 0.4, 0.0], [0.0, 0.0, 1.2]],
+        pos=[[0.1, 0.05, 0.2], [0.62, 0.3, 0.7]],
+    )
+    supercell_model = model.supercell(size=(1, 1, 3))
+
+    # switch the order of the two target orbitals
+    supercell_model_resliced = supercell_model.slice_orbitals(
+        slice_idx=[0, 1, 3, 2, 4, 5]
+    )
+    orbital_labels = ["a", "b", "b", "a", "a", "b"]
+    folded_model_by_label = supercell_model_resliced.fold_model(
+        new_unit_cell=model.uc,
+        unit_cell_offset=supercell_model_resliced.uc.T @ [0, 0, 1 / 3],
+        orbital_labels=orbital_labels,
+        target_indices=[3, 2],
+        order_by="label",
+    )
+    folded_model_by_index = supercell_model_resliced.fold_model(
+        new_unit_cell=model.uc,
+        unit_cell_offset=supercell_model_resliced.uc.T @ [0, 0, 1 / 3],
+        orbital_labels=orbital_labels,
+        target_indices=[2, 3],
+        order_by="index",
+    )
+    assert models_close(model.slice_orbitals([1, 0]), folded_model_by_index)
+    assert models_close(model, folded_model_by_label)
+
+
+def test_invalid_order_by(get_model):
+    """Test that passing an invalid ``order_by`` raises."""
+    model = get_model(
+        0.1,
+        0.3,
+        uc=np.diag([1, 2, 3]),
+        pos=[[0.5, 0.5, 0.5], [0.5, 0.5, 0.5]],
+    )
+    supercell_model = model.supercell(size=(1, 1, 2))
+    with pytest.raises(ValueError) as excinfo:
+        supercell_model.fold_model(
+            new_unit_cell=model.uc, orbital_labels=["a", "b"] * 2, order_by="spam"
+        )
+    assert "Invalid input" in str(excinfo.value)
+    assert "order_by" in str(excinfo.value)
+    assert "spam" in str(excinfo.value)