Fix broken test

vasp-dev · Dec 22, 2023 · 925f7d0 · 925f7d0
1 parent 2f28d50
commit 925f7d0
Showing 1 changed file with 9 additions and 11 deletions.
diff --git a/src/py4vasp/_data/density.py b/src/py4vasp/_data/density.py
@@ -31,16 +31,6 @@ def _raise_error_if_color_is_specified(**user_options):
 
 
 _DEFAULT = 0
-_SELECTIONS = {
-    None: None,
-    "charge": None,
-    "n": None,
-    "charge_density": None,
-    "electronic_charge_density": None,
-    "tau": "tau",
-    "kinetic_energy": "tau",
-    "kinetic_energy_density": "tau",
-}
 _COMPONENTS = {
     0: ["0", "unity", "sigma_0", "scalar"],
     1: ["1", "sigma_x", "x", "sigma_1"],
@@ -107,6 +97,9 @@ def selections(self):
         electronic charge density
             *charge*, *n*, *charge_density*, and *electronic_charge_density*
 
+        kinetic energy density
+            *tau*, *kinetic_energy*, and *kinetic_energy_density*
+
         0th component
             {component0}
 
@@ -298,7 +291,12 @@ def is_noncollinear(self):
 
     @property
     def _selection(self):
-        return _SELECTIONS[super()._selection]
+        selection_map = {
+            "tau": "tau",
+            "kinetic_energy": "tau",
+            "kinetic_energy_density": "tau",
+        }
+        return selection_map.get(super()._selection)
 
 
 def _raise_component_not_specified_error(selec_tuple):