From fbb1e1ef0b606ac4598b64ad15cec6e4df6aafe6 Mon Sep 17 00:00:00 2001
From: Bonan Zhu <zhubonan@outlook.com>
Date: Mon, 26 Aug 2024 21:51:53 +0800
Subject: [PATCH] Improve handling of discontinues paths

Combine kpoint labels when plotting the unfolded effective band
structure.
---
 easyunfold/__init__.py |  2 +-
 easyunfold/plotting.py |  2 +-
 easyunfold/unfold.py   | 26 +++++++++++++++++++++++++-
 3 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/easyunfold/__init__.py b/easyunfold/__init__.py
index 64637c2..add4e56 100644
--- a/easyunfold/__init__.py
+++ b/easyunfold/__init__.py
@@ -2,4 +2,4 @@
 Collection of code for band unfolding
 """
 
-__version__ = '0.3.6'
+__version__ = '0.3.7'
diff --git a/easyunfold/plotting.py b/easyunfold/plotting.py
index 1c2817a..93c1f0b 100644
--- a/easyunfold/plotting.py
+++ b/easyunfold/plotting.py
@@ -372,7 +372,7 @@ def _plot_spectral_function_rgba(
     def _add_kpoint_labels(self, ax: plt.Axes, x_is_kidx=False):
         """Add labels to the k-points for a given axes"""
         # Label the kpoints
-        labels = self.unfold.kpoint_labels
+        labels = self.unfold.get_combined_kpoint_labels()
         kdist = self.unfold.get_kpoint_distances()
 
         # Explicit label indices
diff --git a/easyunfold/unfold.py b/easyunfold/unfold.py
index eb24db8..f664d63 100644
--- a/easyunfold/unfold.py
+++ b/easyunfold/unfold.py
@@ -615,7 +615,7 @@ def as_dict(self) -> dict:
             output[key] = getattr(self, key)
         return output
 
-    def get_kpoint_distances(self):
+    def get_kpoint_distances(self, hide_discontinuities=True):
         """
         Distances between the kpoints along the path in the reciprocal space.
         This does not take account of the breaking of the path.
@@ -623,14 +623,38 @@ def get_kpoint_distances(self):
         The reciprocal lattice vectors includes the $2\\pi$ factor, e.g. `np.linalg.inv(L).T * 2 * np.pi`.
         :::
         """
+        # Check for
         kpts = self.kpts_pc
         pc_latt = self.pc_latt
         kpts_path = kpts @ np.linalg.inv(pc_latt).T * np.pi * 2  # Kpoint path in the reciprocal space
         dists = np.cumsum(np.linalg.norm(np.diff(kpts_path, axis=0), axis=-1))
         dists = np.append([0], dists)
 
+        if hide_discontinuities:
+            last_idx = -2
+            for idx, _ in self.kpoint_labels:
+                if idx - last_idx == 1:
+                    # This label is directly adjacent to the previous one - this is a discontinuity
+                    shift = dists[idx] - dists[idx - 1]
+                    # Shift the distances beyond
+                    dists[idx:] -= shift
+                last_idx = idx
+
         return dists
 
+    def get_combined_kpoint_labels(self):
+        """Get kpoints label with discontinuities combined into a single label"""
+        last_entry = [-2, None]
+        comnbined_labels = []
+        for idx, name in self.kpoint_labels:
+            if idx - last_entry[0] == 1:
+                comnbined_labels.append([last_entry[0], last_entry[1] + '|' + name])
+            else:
+                comnbined_labels.append([idx, name])
+            last_entry = [idx, name]
+
+        return comnbined_labels
+
 
 def LorentzSmearing(x, x0, sigma=0.02):
     r"""