Fix plotting bug

fermisurfaces · Sep 30, 2024 · 40f405f · 40f405f
1 parent b39b8b2
commit 40f405f
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/src/ifermi/plot.py b/src/ifermi/plot.py
@@ -1341,8 +1341,14 @@ def get_isosurface_colors(
             raise ValueError(
                 "colors dict must have same number of spin channels as spins to plot"
             )
+        if isinstance(fermi_object, FermiSurface):
+            return [
+                colors[s]
+                for s in spins
+                for _ in range(fermi_object.n_surfaces_per_spin[s])
+            ]
         return [
-            colors[s] for s in spins for _ in range(fermi_object.n_surfaces_per_spin[s])
+            colors[s] for s in spins for _ in range(fermi_object.n_lines_per_spin[s])
         ]
 
     elif isinstance(colors, str):