From 372275e44ed0af948da2bac4d299f1d97c5cd67c Mon Sep 17 00:00:00 2001
From: mt-huebsch <42400095+mt-huebsch@users.noreply.github.com>
Date: Fri, 22 Dec 2023 16:44:54 +0100
Subject: [PATCH] feat: plotting ncl densities (#120)

* add kinetic energy density
* return density in different format for kinetic energy density

* plot ncl magnetization
* make plotting work for more densities
* fix plots for non standard-cell ystems
* allow multiple isosurfaces in single plot

* add to_numpy function

* add selections for source to all classes
* improve documentation for selections
* selection for quantity and components with synonyms for easy access
* improve documentation of Density.selections() to see available options

* implement aliases for density
* format sources that are aliases different in YAML output

---------

Co-authored-by: Marie-Therese Huebsch <marie-therese.huebsch@vasp.at>
Co-authored-by: Martin Schlipf <martin.schlipf@gmail.com>
---
 .gitignore                       |   1 +
 src/py4vasp/_data/base.py        |  24 ++-
 src/py4vasp/_data/density.py     | 284 ++++++++++++++++++++++++------
 src/py4vasp/_data/phonon_band.py |   4 +-
 src/py4vasp/_data/phonon_dos.py  |   6 +-
 src/py4vasp/_data/projector.py   |   2 +-
 src/py4vasp/_data/viewer3d.py    |   5 +-
 src/py4vasp/_raw/definition.py   |  10 ++
 src/py4vasp/_raw/schema.py       |  42 +++--
 src/py4vasp/_util/index.py       |  10 +-
 src/py4vasp/_util/select.py      |  32 +++-
 tests/data/test_base.py          |  17 ++
 tests/data/test_density.py       | 287 ++++++++++++++++++++++---------
 tests/data/test_dos.py           |   1 -
 tests/data/test_viewer3d.py      |   1 +
 tests/raw/conftest.py            |   5 +-
 tests/raw/test_schema.py         |  32 +++-
 tests/util/test_index.py         |  13 +-
 tests/util/test_select.py        |  24 ++-
 19 files changed, 623 insertions(+), 177 deletions(-)

diff --git a/.gitignore b/.gitignore
index 9a19cdfc..669452a8 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,4 @@
+*~
 # File open for editing in vi
 *.swp
 # autogenerated by dephell
diff --git a/src/py4vasp/_data/base.py b/src/py4vasp/_data/base.py
index 28474a27..037cd4f1 100644
--- a/src/py4vasp/_data/base.py
+++ b/src/py4vasp/_data/base.py
@@ -122,6 +122,10 @@ def path(self):
     def _raw_data(self):
         return self._data_context.data
 
+    @property
+    def _selection(self):
+        return self._data_context.selection
+
     @data_access
     def print(self):
         "Print a string representation of this instance."
@@ -131,6 +135,23 @@ def read(self, *args, **kwargs):
         "Convenient wrapper around to_dict. Check that function for examples and optional arguments."
         return self.to_dict(*args, **kwargs)
 
+    @data_access
+    def selections(self):
+        """Returns possible alternatives for this particular quantity VASP can produce.
+
+        The returned dictionary contains a single item with the name of the quantity
+        mapping to all possible selections. Each of these selection may be passed to
+        other functions of this quantity to select which output of VASP is used.
+
+        Returns
+        -------
+        dict
+            The key indicates this quantity and the values possible choices for arguments
+            to other functions of this quantity.
+        """
+        sources = list(raw.selections(self._data_context.quantity))
+        return {self._data_context.quantity: sources}
+
     def _repr_pretty_(self, p, cycle):
         p.text(str(self))
 
@@ -308,4 +329,5 @@ def __exit__(self, *_):
             self._stack.close()
 
     def set_selection(self, selection):
-        self.selection = selection
+        if self._counter == 0:
+            self.selection = selection
diff --git a/src/py4vasp/_data/density.py b/src/py4vasp/_data/density.py
index b13eb7e8..069a0ba2 100644
--- a/src/py4vasp/_data/density.py
+++ b/src/py4vasp/_data/density.py
@@ -2,9 +2,9 @@
 # Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0)
 import numpy as np
 
-from py4vasp import data, exception
+from py4vasp import _config, data, exception
 from py4vasp._data import base, structure
-from py4vasp._util import import_
+from py4vasp._util import documentation, import_, index, select
 
 pretty = import_.optional("IPython.lib.pretty")
 
@@ -12,11 +12,40 @@
 class _ViewerWrapper:
     def __init__(self, viewer):
         self._viewer = viewer
-        self._options = {"isolevel": 0.2, "color": "yellow", "opacity": 0.6}
+        self._options = {"isolevel": 0.2, "opacity": 0.6}
 
-    def show_isosurface(self, data, **options):
+    def show_isosurface(self, data, symmetric, **options):
         options = {**self._options, **options}
-        self._viewer.show_isosurface(data, **options)
+        if symmetric:
+            _raise_error_if_color_is_specified(**options)
+            self._viewer.show_isosurface(data, color=_config.VASP_BLUE, **options)
+            self._viewer.show_isosurface(-data, color=_config.VASP_RED, **options)
+        else:
+            self._viewer.show_isosurface(data, color=_config.VASP_CYAN, **options)
+
+
+def _raise_error_if_color_is_specified(**user_options):
+    if "color" in user_options:
+        msg = "Specifying the color of a magnetic isosurface is not implemented."
+        raise exception.NotImplemented(msg)
+
+
+_DEFAULT = 0
+_COMPONENTS = {
+    0: ["0", "unity", "sigma_0", "scalar"],
+    1: ["1", "sigma_x", "x", "sigma_1"],
+    2: ["2", "sigma_y", "y", "sigma_2"],
+    3: ["3", "sigma_z", "z", "sigma_3"],
+}
+_MAGNETIZATION = ("magnetization", "mag", "m")
+
+
+def _join_with_emphasis(data):
+    emph_data = [f"*{x}*" for x in data]
+    if len(data) < 3:
+        return " and ".join(emph_data)
+    emph_data.insert(-1, "and")
+    return ", ".join(emph_data)
 
 
 class Density(base.Refinery, structure.Mixin):
@@ -31,25 +60,103 @@ def __str__(self):
         _raise_error_if_no_data(self._raw_data.charge)
         grid = self._raw_data.charge.shape[1:]
         topology = data.Topology.from_data(self._raw_data.structure.topology)
-        if self.nonpolarized():
-            name = "nonpolarized"
-        elif self.collinear():
-            name = "collinear"
+        if self._selection == "tau":
+            name = "Kinetic energy"
+        elif self.is_nonpolarized():
+            name = "Nonpolarized"
+        elif self.is_collinear():
+            name = "Collinear"
         else:
-            name = "noncollinear"
+            name = "Noncollinear"
         return f"""{name} density:
     structure: {pretty.pretty(topology)}
     grid: {grid[2]}, {grid[1]}, {grid[0]}"""
 
+    @documentation.format(
+        component0=_join_with_emphasis(_COMPONENTS[0]),
+        component1=_join_with_emphasis(_COMPONENTS[1]),
+        component2=_join_with_emphasis(_COMPONENTS[2]),
+        component3=_join_with_emphasis(_COMPONENTS[3]),
+    )
+    @base.data_access
+    def selections(self):
+        """Returns possible densities VASP can produce along with all available components.
+
+        In the dictionary, the key *density* lists all different densities you can access
+        from the VASP output provided you set the relevant INCAR tags. You can combine
+        any of these with any possible choice from the key *component* to further
+        specify the particular output you will receive. If you do not specify a *density*
+        or a *component* the other routines will default to the electronic charge and
+        the 0-th component.
+
+        To nest density and component, please use parentheses, e.g. ``charge(1, 2)`` or
+        ``3(tau)``.
+
+        For convenience, py4vasp accepts the following aliases
+
+        electronic charge density
+            *charge*, *n*, *charge_density*, and *electronic_charge_density*
+
+        kinetic energy density
+            *tau*, *kinetic_energy*, and *kinetic_energy_density*
+
+        0th component
+            {component0}
+
+        1st component
+            {component1}
+
+        2nd component
+            {component2}
+
+        3rd component
+            {component3}
+
+        Returns
+        -------
+        dict
+            Possible densities and components to pass as selection in other functions on density.
+
+        Notes
+        -----
+        In the special case of collinear calculations, *magnetization*, *mag*, and *m*
+        are another alias for the 3rd component of the charge density.
+
+        Examples
+        --------
+        >>> calc = py4vasp.Calculation.from_path(".")
+        >>> calc.density.to_dict("n")
+        >>> calc.density.plot("magnetization")
+        Using synonyms and nesting
+        >>> calc.density.plot("n m(1,2) mag(sigma_z)")
+        """
+        sources = super().selections()
+        if self._raw_data.charge.is_none():
+            return sources
+        if self.is_nonpolarized():
+            components = [_COMPONENTS[0][_DEFAULT]]
+        elif self.is_collinear():
+            components = [_COMPONENTS[0][_DEFAULT], _COMPONENTS[3][_DEFAULT]]
+        else:
+            components = [_COMPONENTS[i][_DEFAULT] for i in range(4)]
+        return {**sources, "component": components}
+
     @base.data_access
     def to_dict(self):
-        """Read the electronic density into a dictionary.
+        """Read the density into a dictionary.
+
+        Parameters
+        ----------
+        selection : str
+            VASP computes different densities depending on the INCAR settings. With this
+            parameter, you can control which one of them is returned. Please use the
+            `selections` routine to get a list of all possible choices.
 
         Returns
         -------
         dict
             Contains the structure information as well as the density represented
-            of a grid in the unit cell.
+            on a grid in the unit cell.
         """
         _raise_error_if_no_data(self._raw_data.charge)
         result = {"structure": self._structure.read()}
@@ -57,22 +164,44 @@ def to_dict(self):
         return result
 
     def _read_density(self):
-        density = np.moveaxis(self._raw_data.charge, 0, -1).T
-        yield "charge", density[0]
-        if self.collinear():
-            yield "magnetization", density[1]
-        elif self.noncollinear():
-            yield "magnetization", density[1:]
+        density = self.to_numpy()
+        if self._selection:
+            yield self._selection, density
+        else:
+            yield "charge", density[0]
+            if self.is_collinear():
+                yield "magnetization", density[1]
+            elif self.is_noncollinear():
+                yield "magnetization", density[1:]
+
+    @base.data_access
+    def to_numpy(self):
+        """Convert the density to a numpy array.
+
+        The number of components is 1 for nonpolarized calculations, 2 for collinear
+        calculations, and 4 for noncollinear calculations. Each component is 3
+        dimensional according to the grid VASP uses for the FFTs.
+
+        Returns
+        -------
+        np.ndarray
+            All components of the selected density.
+        """
+        return np.moveaxis(self._raw_data.charge, 0, -1).T
 
     @base.data_access
-    def plot(self, selection="charge", **user_options):
+    def plot(self, selection="0", **user_options):
         """Plot the selected density as a 3d isosurface within the structure.
 
         Parameters
         ----------
         selection : str
             Can be either *charge* or *magnetization*, depending on which quantity
-            should be visualized.
+            should be visualized.  For a noncollinear calculation, the density has
+            4 components which can be represented in a 2x2 matrix. Specify the
+            component of the density in terms of the Pauli matrices: sigma_1,
+            sigma_2, sigma_3.
+
         user_options
             Further arguments with keyword that get directly passed on to the
             visualizer. Most importantly, you can set isolevel to adjust the
@@ -82,53 +211,101 @@ def plot(self, selection="charge", **user_options):
         -------
         Viewer3d
             Visualize an isosurface of the density within the 3d structure.
+
+        Examples
+        --------
+        >>> calc = py4vasp.Calculation.from_path(".")
+        Plot an isosurface of the electronic charge density
+        >>> calc.density.plot()
+        Plot isosurfaces for positive (blue) and negative (red) magnetization
+        of a spin-polarized calculation (ISPIN=2)
+        >>> calc.density.plot("m")
+        Plot the isosurface for the third component of a noncollinear magnetization
+        >>> calc.density.plot("m(3)")
         """
         _raise_error_if_no_data(self._raw_data.charge)
         viewer = self._structure.plot()
-        if selection == "charge":
-            self._plot_charge(_ViewerWrapper(viewer), **user_options)
-        elif selection == "magnetization":
-            self._plot_magnetism(_ViewerWrapper(viewer), **user_options)
-        else:
-            msg = f"'{selection}' is an unknown option, please use 'charge' or 'magnetization' instead."
-            raise exception.IncorrectUsage(msg)
+        wrapper = _ViewerWrapper(viewer)
+        map_ = self._create_map()
+        selector = index.Selector({0: map_}, self._raw_data.charge)
+        tree = select.Tree.from_selection(selection)
+        selections = self._filter_noncollinear_magnetization_from_selections(tree)
+        for selection in selections:
+            label = selector.label(selection)
+            symmetric = self._use_symmetric_isosurface(label, map_)
+            wrapper.show_isosurface(selector[selection], symmetric, **user_options)
         return viewer
 
+    def _filter_noncollinear_magnetization_from_selections(self, tree):
+        if self._selection or not self.is_noncollinear():
+            yield from tree.selections()
+        else:
+            filtered_selections = tree.selections(filter=set(_MAGNETIZATION))
+            for filtered, unfiltered in zip(filtered_selections, tree.selections()):
+                if filtered != unfiltered and len(filtered) != 1:
+                    _raise_component_not_specified_error(unfiltered)
+                yield filtered
+
+    def _create_map(self):
+        map_ = {
+            choice: self._index_component(component)
+            for component, choices in _COMPONENTS.items()
+            for choice in choices
+        }
+        self._add_magnetization_for_charge_and_collinear(map_)
+        return map_
+
+    def _index_component(self, component):
+        if self.is_collinear():
+            component = (0, 2, 3, 1)[component]
+        return component
+
+    def _add_magnetization_for_charge_and_collinear(self, map_):
+        if self._selection or not self.is_collinear():
+            return
+        for key in _MAGNETIZATION:
+            map_[key] = 1
+
+    def _use_symmetric_isosurface(self, label, map_):
+        component = map_.get(label, -1)
+        if component > 0 and self.is_nonpolarized():
+            _raise_is_nonpolarized_error()
+        if component > 1 and self.is_collinear():
+            _raise_is_collinear_error()
+        return component > 0
+
     @base.data_access
-    def nonpolarized(self):
+    def is_nonpolarized(self):
         "Returns whether the density is not spin polarized."
         return len(self._raw_data.charge) == 1
 
     @base.data_access
-    def collinear(self):
+    def is_collinear(self):
         "Returns whether the density has a collinear magnetization."
         return len(self._raw_data.charge) == 2
 
     @base.data_access
-    def noncollinear(self):
+    def is_noncollinear(self):
         "Returns whether the density has a noncollinear magnetization."
         return len(self._raw_data.charge) == 4
 
-    def _plot_charge(self, viewer, **user_options):
-        viewer.show_isosurface(self._raw_data.charge[0], **user_options)
+    @property
+    def _selection(self):
+        selection_map = {
+            "tau": "tau",
+            "kinetic_energy": "tau",
+            "kinetic_energy_density": "tau",
+        }
+        return selection_map.get(super()._selection)
 
-    def _plot_magnetism(self, viewer, **user_options):
-        if self.nonpolarized():
-            _raise_is_nonpolarized_error()
-        if self.collinear():
-            return self._plot_collinear_magnetism(viewer, **user_options)
-        if self.noncollinear():
-            return self._plot_noncollinear_magnetism(viewer, **user_options)
-
-    def _plot_collinear_magnetism(self, viewer, **user_options):
-        _raise_error_if_color_is_specified(**user_options)
-        magnetization = self._raw_data.charge[1]
-        viewer.show_isosurface(magnetization, color="blue", **user_options)
-        viewer.show_isosurface(-magnetization, color="red", **user_options)
 
-    def _plot_noncollinear_magnetism(self, viewer, **user_options):
-        magnetization = np.linalg.norm(self._raw_data.charge[1:], axis=0)
-        viewer.show_isosurface(magnetization, **user_options)
+def _raise_component_not_specified_error(selec_tuple):
+    msg = (
+        "Invalid selection: selection='"
+        + ", ".join(selec_tuple)
+        + "'. For a noncollinear calculation, the density has 4 components which can be represented in a 2x2 matrix. Specify the component of the density in terms of the Pauli matrices: sigma_1, sigma_2, sigma_3. E.g.: m(sigma_1)."
+    )
+    raise exception.IncorrectUsage(msg)
 
 
 def _raise_is_nonpolarized_error():
@@ -136,6 +313,11 @@ def _raise_is_nonpolarized_error():
     raise exception.NoData(msg)
 
 
+def _raise_is_collinear_error():
+    msg = "Density does not contain noncollinear magnetization. Please rerun VASP with LNONCOLLINEAR = T to obtain it."
+    raise exception.NoData(msg)
+
+
 def _raise_error_if_no_data(data):
     if data.is_none():
         raise exception.NoData(
@@ -145,9 +327,3 @@ def _raise_error_if_no_data(data):
             'common issue is when you create `Calculation.from_file("vaspout.h5")` '
             "because this will overwrite the default file behavior."
         )
-
-
-def _raise_error_if_color_is_specified(**user_options):
-    if "color" in user_options:
-        msg = "Specifying the color of a magnetic isosurface is not implemented."
-        raise exception.NotImplemented(msg)
diff --git a/src/py4vasp/_data/phonon_band.py b/src/py4vasp/_data/phonon_band.py
index fec2afd2..d543f73f 100644
--- a/src/py4vasp/_data/phonon_band.py
+++ b/src/py4vasp/_data/phonon_band.py
@@ -8,7 +8,7 @@
 from py4vasp._util import convert, documentation, index, select
 
 
-class PhononBand(base.Refinery, phonon.Mixin, graph.Mixin):
+class PhononBand(phonon.Mixin, base.Refinery, graph.Mixin):
     """The phonon band structure.
 
     Use this to examine the phonon band structure along a high-symmetry path in the
@@ -73,7 +73,7 @@ def _projections(self, selection, width):
         if not selection:
             return None
         maps = {2: self._init_atom_dict(), 3: self._init_direction_dict()}
-        selector = index.Selector(maps, np.abs(self._modes()))
+        selector = index.Selector(maps, np.abs(self._modes()), use_number_labels=True)
         tree = select.Tree.from_selection(selection)
         return {
             selector.label(selection): width * selector[selection]
diff --git a/src/py4vasp/_data/phonon_dos.py b/src/py4vasp/_data/phonon_dos.py
index 8d77b35e..b78a0d26 100644
--- a/src/py4vasp/_data/phonon_dos.py
+++ b/src/py4vasp/_data/phonon_dos.py
@@ -8,7 +8,7 @@
 from py4vasp._util import documentation, index, select
 
 
-class PhononDos(base.Refinery, phonon.Mixin, graph.Mixin):
+class PhononDos(phonon.Mixin, base.Refinery, graph.Mixin):
     """The phonon density of states (DOS).
 
     You can use this class to extract the phonon DOS data of a VASP
@@ -71,7 +71,9 @@ def _read_data(self, selection):
         if not selection:
             return {}
         maps = {0: self._init_atom_dict(), 1: self._init_direction_dict()}
-        selector = index.Selector(maps, self._raw_data.projections)
+        selector = index.Selector(
+            maps, self._raw_data.projections, use_number_labels=True
+        )
         tree = select.Tree.from_selection(selection)
         return {
             selector.label(selection): selector[selection]
diff --git a/src/py4vasp/_data/projector.py b/src/py4vasp/_data/projector.py
index c7f8760d..4f520503 100644
--- a/src/py4vasp/_data/projector.py
+++ b/src/py4vasp/_data/projector.py
@@ -161,7 +161,7 @@ def _make_selector(self, projections):
         maps = self.to_dict()
         maps = {1: maps["atom"], 2: maps["orbital"], 0: maps["spin"]}
         try:
-            return index.Selector(maps, projections)
+            return index.Selector(maps, projections, use_number_labels=True)
         except exception._Py4VaspInternalError as error:
             message = f"""Error reading the projections. Please make sure that the passed
                 projections has the right format, i.e., the indices correspond to spin,
diff --git a/src/py4vasp/_data/viewer3d.py b/src/py4vasp/_data/viewer3d.py
index e8dd084e..96e5d30d 100644
--- a/src/py4vasp/_data/viewer3d.py
+++ b/src/py4vasp/_data/viewer3d.py
@@ -67,15 +67,16 @@ def from_structure(cls, structure, supercell=None):
             If present the cell is extended by the specified factor along each axis.
         """
         ase = structure.to_ase(supercell)
+        res = cls(nglview.show_ase(ase))
         # ngl works with the standard form, so we need to store the positions in the same format
         standard_cell, transformation = ase.cell.standard_form()
-        ase.set_cell(standard_cell)
-        res = cls(nglview.show_ase(ase))
+        ase.set_cell(standard_cell, scale_atoms=True)
         res._lengths = tuple(ase.cell.lengths())
         res._angles = tuple(ase.cell.angles())
         res._positions = ase.positions
         res._number_cells = res._calculate_number_cells(supercell)
         res._transformation = transformation
+        # res._transformation = np.eye(3)
         return res
 
     def _calculate_number_cells(self, supercell):
diff --git a/src/py4vasp/_raw/definition.py b/src/py4vasp/_raw/definition.py
index 7261995f..c81dc332 100644
--- a/src/py4vasp/_raw/definition.py
+++ b/src/py4vasp/_raw/definition.py
@@ -101,10 +101,20 @@ def selections(quantity):
 )
 schema.add(
     raw.Density,
+    alias=["charge", "n", "charge_density", "electronic_charge_density"],
     file="vaspwave.h5",
     structure=Link("structure", DEFAULT_SOURCE),
     charge="charge/charge",
 )
+schema.add(
+    raw.Density,
+    name="tau",
+    required=raw.Version(6, 5),
+    alias=["kinetic_energy", "kinetic_energy_density"],
+    file="vaspwave.h5",
+    structure=Link("structure", DEFAULT_SOURCE),
+    charge="kinetic_energy_density/values",
+)
 #
 group = "results/linear_response"
 energies = "energies_dielectric_function"
diff --git a/src/py4vasp/_raw/schema.py b/src/py4vasp/_raw/schema.py
index 0c80f809..67c39b90 100644
--- a/src/py4vasp/_raw/schema.py
+++ b/src/py4vasp/_raw/schema.py
@@ -5,6 +5,8 @@
 import dataclasses
 import textwrap
 
+import numpy as np
+
 from py4vasp import exception
 from py4vasp._util import convert
 
@@ -15,7 +17,7 @@ def __init__(self, version):
         self._version = version
         self._verified = False
 
-    def add(self, cls, name="default", file=None, required=None, **kwargs):
+    def add(self, cls, name="default", alias=[], file=None, required=None, **kwargs):
         """Add a new quantity to the schema.
 
         The name of the quantity is deduced from the class you pass in, for example
@@ -44,14 +46,21 @@ def add(self, cls, name="default", file=None, required=None, **kwargs):
             You need to specify for all the fields of the class from where in the HDF5
             file they can be obtained.
         """
-        class_name = convert.quantity_name(cls.__name__)
-        self._sources.setdefault(class_name, {})
-        if name in self._sources[class_name]:
-            message = f"{class_name}/{name} already in the schema. Please choose a different name."
-            raise exception.IncorrectUsage(message)
-        self._sources[class_name][name] = Source(cls(**kwargs), file, required)
+        quantity = convert.quantity_name(cls.__name__)
+        self._sources.setdefault(quantity, {})
+        labels = [name] + list(np.atleast_1d(alias))
+        for label in labels:
+            self._raise_error_if_already_in_schema(quantity, label)
+            alias_for = name if label != name else None
+            source = Source(cls(**kwargs), file, required, alias_for)
+            self._sources[quantity][label] = source
         self._verified = False
 
+    def _raise_error_if_already_in_schema(self, quantity, label):
+        if label in self._sources[quantity]:
+            message = f"{quantity}/{label} already in the schema. Please choose a different name."
+            raise exception._Py4VaspInternalError(message)
+
     @property
     def sources(self):
         return self._sources
@@ -88,12 +97,14 @@ def _verify_source(self, key, source):
     def _verify_quantity_is_in_schema(self, key, field):
         message = f"""Verifying the schema failed in link resolution for {key}, because
     {field.quantity} is not defined in the schema."""
-        assert field.quantity in self._sources, message
+        if field.quantity not in self._sources:
+            raise exception._Py4VaspInternalError(message)
 
     def _verify_source_is_defined_for_quantity(self, key, field):
         message = f"""Verifying the schema failed in link resolution for {key}, because
     {field.source} is not a source defined in the schema for the quantity {field.quantity}."""
-        assert field.source in self._sources[field.quantity], message
+        if field.source not in self._sources[field.quantity]:
+            raise exception._Py4VaspInternalError(message)
 
     def __str__(self):
         version = _parse_version(self.version)
@@ -109,6 +120,7 @@ class Source:
     data: Any
     file: str = None
     required: Version = None
+    alias_for: str = None
 
 
 @dataclasses.dataclass
@@ -128,7 +140,8 @@ def _parse_version(version):
     return f"""version:
     major: {version.major}
     minor: {version.minor}
-    patch: {version.patch}"""
+    patch: {version.patch}
+"""
 
 
 def _parse_quantities(quantities):
@@ -136,12 +149,15 @@ def _parse_quantities(quantities):
         if name == "version":
             continue
         sources = (_parse_source(name, *source) for source in sources.items())
-        yield f"{name}:\n" + "\n".join(sources)
+        yield f"{name}:\n" + "\n".join(sources) + "\n"
 
 
 def _parse_source(quantity, source, specification):
-    specs = _parse_specification(specification)
-    return 4 * " " + f"{source}:  &{quantity}-{source}\n" + "\n".join(specs)
+    if specification.alias_for:
+        return 4 * " " + f"{source}: *{quantity}-{specification.alias_for}"
+    else:
+        specs = _parse_specification(specification)
+        return 4 * " " + f"{source}:  &{quantity}-{source}\n" + "\n".join(specs)
 
 
 def _parse_specification(specification):
diff --git a/src/py4vasp/_util/index.py b/src/py4vasp/_util/index.py
index 6e607aff..9f296167 100644
--- a/src/py4vasp/_util/index.py
+++ b/src/py4vasp/_util/index.py
@@ -53,15 +53,21 @@ class Selector:
         The function used to reduce over the dimensions listed in the map. If not
         specified a summation is performed. Note that the function must have an axis
         argument with the same meaning as `np.sum`.
+    use_number_labels :  bool
+        If set numbers will be replaced by the corresponding label of the slice. If you
+        have e.g. the label *A* corresponding to the first three elements and *1*
+        corresponds to the first element in total, setting this flag will label *1* as
+        *A_1* instead.
     """
 
-    def __init__(self, maps, data, *, reduction=np.sum):
+    def __init__(self, maps, data, *, reduction=np.sum, use_number_labels=False):
         self._data = raw.VaspData(data)
         self._axes = tuple(maps.keys())
         _raise_error_if_duplicate_keys(maps)
         if not self._data.is_none():
             _raise_error_if_map_out_of_bounds(maps.keys(), self._data.ndim)
         self._map = self._make_map(maps)
+        self._use_number_labels = use_number_labels
         self._number_labels = self._make_number_labels(maps)
         self._indices = self._make_default_indices(maps, self._data.ndim)
         self._reduction = reduction
@@ -83,6 +89,8 @@ def _make_number_labels(self, maps):
         }
 
     def _make_label(self, map_, number, index, size):
+        if not self._use_number_labels:
+            return number
         indices = range(size)
         index, *rest = list(indices[_make_slice(index)])
         message = f"Integer label {number} maps to more than a single index."
diff --git a/src/py4vasp/_util/select.py b/src/py4vasp/_util/select.py
index 5cd252b2..bb51eb12 100644
--- a/src/py4vasp/_util/select.py
+++ b/src/py4vasp/_util/select.py
@@ -96,7 +96,7 @@ def nodes(self):
     def __str__(self):
         return str(self._content)
 
-    def selections(self, selected=()):
+    def selections(self, selected=(), filter={}):
         """Core routine generating all user selections parsed.
 
         This will generate one selection at a time so it should be used in a loop or
@@ -108,6 +108,8 @@ def selections(self, selected=()):
             Prior selections obtained from a different source. These selections will
             be added to any additional selection parsed from the user input. If not
             set, if defaults to giving just the user selections.
+        filter : set
+            Remove any element found in the set from the resulting selection.
 
         Yields
         ------
@@ -115,21 +117,29 @@ def selections(self, selected=()):
             Each selection corresponds to one path from the root of the tree to one of
             its leaves.
         """
-        content = (self._content,) if self._content else ()
+        if self._content and self._content not in filter:
+            content = (self._content,)
+        else:
+            content = ()
         if not self._children:
             yield selected + content
         elif self._is_operation:
-            yield from self._operation_selections(selected)
+            yield from self._operation_selections(selected, filter)
         else:
             for child in self._children:
-                yield from child.selections(selected + content)
+                yield from child.selections(selected + content, filter)
 
-    def _operation_selections(self, selected):
-        left_operands = self._children[0].selections()
-        right_operands = self._children[1].selections()
+    def _operation_selections(self, selected, filter):
+        left_operands = self._get_operands(self._children[0], filter)
+        right_operands = self._get_operands(self._children[1], filter)
         for left_op, right_op in itertools.product(left_operands, right_operands):
             yield *selected, Operation(left_op, self._content.operator, right_op)
 
+    def _get_operands(self, child, filter):
+        for operand in child.selections(filter=filter):
+            child._raise_error_if_content_and_operator_are_incompatible(operand)
+            yield operand
+
     def to_mermaid(self):
         "Helper routine to visualize the Tree using Mermaid"
         return "\n".join(self._to_mermaid(root=True))
@@ -298,6 +308,12 @@ def _raise_error_if_operation_misses_right_hand_side(self):
         message = f"The operator {self._content} is not followed by an element."
         raise exception._Py4VaspInternalError(message)
 
+    def _raise_error_if_content_and_operator_are_incompatible(self, operand):
+        if bool(self._content) == bool(operand):
+            return
+        message = f"The operand `{operand}` has a qualitatively different behavior then the content `{self}`. This may occur when a filter would replace the last element."
+        raise exception.IncorrectUsage(message)
+
 
 @dataclasses.dataclass
 class Group:
@@ -307,6 +323,7 @@ class Group:
     separator: str
     "The string separating the members of the group."
     __str__ = lambda self: self.separator.join(self.group)
+    __hash__ = lambda self: hash(str(self))
 
     def __iadd__(self, character):
         self.group[-1] += character
@@ -318,6 +335,7 @@ class _Operator:
     operator: str
     _id: int
     __str__ = lambda self: f"_{self._id}_[{self.operator}]"
+    __hash__ = lambda self: hash(self.operator)
 
 
 @dataclasses.dataclass
diff --git a/tests/data/test_base.py b/tests/data/test_base.py
index 423aff5a..514c1eb7 100644
--- a/tests/data/test_base.py
+++ b/tests/data/test_base.py
@@ -68,6 +68,10 @@ def with_selection_argument(self, selection=DEFAULT_SELECTION):
     def selection_without_default(self, selection):
         return selection
 
+    @base.data_access
+    def selection_from_property(self):
+        return self._selection
+
     @base.data_access
     def __str__(self):
         return self._raw_data.content
@@ -361,3 +365,16 @@ def test_syntax_error_still_raised(mock_schema):
     example = Example.from_data(RAW_DATA)
     with pytest.raises(TypeError):
         example.read(1, 2)
+
+
+def test_selections(mock_schema):
+    example = Example.from_data(RAW_DATA)
+    assert example.selections() == {"example": ["default", "alternative"]}
+
+
+def test_selection_from_property(mock_access):
+    example = Example.from_data(RAW_DATA)
+    assert example.selection_from_property() is None
+    example = Example.from_path()
+    assert example.selection_from_property() is None
+    assert example.selection_from_property(SELECTION) == SELECTION
diff --git a/tests/data/test_density.py b/tests/data/test_density.py
index 5f9dd184..296aabb7 100644
--- a/tests/data/test_density.py
+++ b/tests/data/test_density.py
@@ -6,58 +6,100 @@
 import numpy as np
 import pytest
 
-from py4vasp import exception, raw
+from py4vasp import _config, exception, raw
 from py4vasp._data import viewer3d
 from py4vasp.data import Density, Structure
 
 
+@pytest.fixture(params=[None, "tau"])
+def density_source(request):
+    return request.param
+
+
 @pytest.fixture(params=["Sr2TiO4", "Fe3O4 collinear", "Fe3O4 noncollinear"])
-def reference_density(raw_data, request):
-    return make_reference_density(raw_data, request.param)
+def reference_density(raw_data, density_source, request):
+    return make_reference_density(raw_data, request.param, density_source)
+
+
+@pytest.fixture
+def nonpolarized_density(raw_data):
+    return make_reference_density(raw_data, "Sr2TiO4")
+
+
+@pytest.fixture
+def collinear_density(raw_data, density_source):
+    return make_reference_density(raw_data, "Fe3O4 collinear", density_source)
 
 
 @pytest.fixture
-def collinear_density(raw_data):
-    return make_reference_density(raw_data, "Fe3O4 collinear")
+def noncollinear_density(raw_data, density_source):
+    return make_reference_density(raw_data, "Fe3O4 noncollinear", density_source)
 
 
-def make_reference_density(raw_data, selection):
+@pytest.fixture
+def empty_density(raw_data):
+    raw_density = raw.Density(raw_data.structure("Sr2TiO4"), charge=raw.VaspData(None))
+    return Density.from_data(raw_density)
+
+
+@pytest.fixture
+def mock_viewer():
+    obj = viewer3d.Viewer3d
+    cm_init = patch.object(obj, "__init__", autospec=True, return_value=None)
+    cm_cell = patch.object(obj, "show_cell")
+    cm_surface = patch.object(obj, "show_isosurface")
+    with cm_init as init, cm_cell as cell, cm_surface as surface:
+        yield {"init": init, "cell": cell, "surface": surface}
+
+
+def make_reference_density(raw_data, selection, source=None):
     raw_density = raw_data.density(selection)
     density = Density.from_data(raw_density)
     density.ref = types.SimpleNamespace()
     density.ref.structure = Structure.from_data(raw_density.structure).read()
-    density.ref.output = get_expected_dict(raw_density.charge)
-    density.ref.string = get_expected_string(raw_density.charge)
+    density.ref.output = get_expected_dict(raw_density.charge, source)
+    density.ref.string = get_expected_string(selection, source)
+    density.ref.selections = get_expected_selections(raw_density.charge)
+    density._data_context.selection = source
+    density.ref.source = source or "charge"
     return density
 
 
-def get_expected_dict(charge):
-    if len(charge) == 1:  # nonpolarized
-        return {"charge": charge[0].T}
-    if len(charge) == 2:  # collinear
-        return {"charge": charge[0].T, "magnetization": charge[1].T}
-    # noncollinear
-    return {"charge": charge[0].T, "magnetization": np.moveaxis(charge[1:].T, -1, 0)}
-
-
-def get_expected_string(charge):
-    if len(charge) == 1:
-        return """\
-nonpolarized density:
-    structure: Sr2TiO4
-    grid: 10, 12, 14"""
-    elif len(charge) == 2:
-        return """\
-collinear density:
-    structure: Fe3O4
-    grid: 10, 12, 14"""
+def get_expected_dict(charge, source):
+    if source:
+        return {source: np.array([component.T for component in charge])}
+    else:
+        if len(charge) == 1:  # nonpolarized
+            return {"charge": charge[0].T}
+        if len(charge) == 2:  # collinear
+            return {"charge": charge[0].T, "magnetization": charge[1].T}
+        # noncollinear
+        magnetization = np.moveaxis(charge[1:].T, -1, 0)
+        return {"charge": charge[0].T, "magnetization": magnetization}
+
+
+def get_expected_string(selection, source):
+    structure, *density = selection.split()
+    if source == "tau":
+        density = "Kinetic energy"
+    elif not density:
+        density = "Nonpolarized"
     else:
-        return """\
-noncollinear density:
-    structure: Fe3O4
+        density = density[0].capitalize()
+    return f"""{density} density:
+    structure: {structure}
     grid: 10, 12, 14"""
 
 
+def get_expected_selections(charge):
+    result = {"density": list(raw.selections("density")), "component": ["0"]}
+    if len(charge) == 2:  # collinear
+        result["component"] += ["3"]
+    if len(charge) == 4:  # noncollinear
+        result["component"] += ["1", "2", "3"]
+    return result
+
+
 def test_read(reference_density, Assert):
     actual = reference_density.read()
     actual_structure = actual.pop("structure")
@@ -67,83 +109,164 @@ def test_read(reference_density, Assert):
         Assert.allclose(actual[key], reference_density.ref.output[key])
 
 
-def test_empty_density(raw_data):
-    raw_density = raw.Density(raw_data.structure("Sr2TiO4"), charge=raw.VaspData(None))
-    density = Density.from_data(raw_density)
+def test_empty_density(empty_density):
     with pytest.raises(exception.NoData):
-        density.read()
+        empty_density.read()
 
 
-def test_charge_plot(reference_density, Assert, not_core):
-    obj = viewer3d.Viewer3d
-    cm_init = patch.object(obj, "__init__", autospec=True, return_value=None)
-    cm_cell = patch.object(obj, "show_cell")
-    cm_surface = patch.object(obj, "show_isosurface")
-    with cm_init as init, cm_cell as cell, cm_surface as surface:
-        result = reference_density.plot()
-        assert isinstance(result, viewer3d.Viewer3d)
-        init.assert_called_once()
-        cell.assert_called_once()
-        surface.assert_called_once()
-        args, kwargs = surface.call_args
-    Assert.allclose(args[0], reference_density.ref.output["charge"].T)
-    assert kwargs == {"isolevel": 0.2, "color": "yellow", "opacity": 0.6}
-
-
-def test_magnetization_plot(reference_density, Assert, not_core):
-    if reference_density.nonpolarized():
-        check_accessing_spin_raises_error(reference_density)
+@pytest.mark.parametrize("selection", [None, "0", "unity", "sigma_0", "scalar"])
+def test_charge_plot(selection, reference_density, mock_viewer, Assert, not_core):
+    source = reference_density.ref.source
+    if source == "charge":
+        expected_density = reference_density.ref.output["charge"].T
+    else:
+        expected_density = reference_density.ref.output[source][0].T
+    if selection:
+        result = reference_density.plot(selection)
     else:
-        check_plotting_magnetization_density(reference_density, Assert)
+        result = reference_density.plot()
+    assert isinstance(result, viewer3d.Viewer3d)
+    mock_viewer["init"].assert_called_once()
+    mock_viewer["cell"].assert_called_once()
+    mock_viewer["surface"].assert_called_once()
+    args, kwargs = mock_viewer["surface"].call_args
+    Assert.allclose(args[0], expected_density)
+    assert kwargs == {"isolevel": 0.2, "color": _config.VASP_CYAN, "opacity": 0.6}
 
 
-def check_accessing_spin_raises_error(nonpolarized_density):
+def test_accessing_spin_raises_error(nonpolarized_density, not_core):
     with pytest.raises(exception.NoData):
-        nonpolarized_density.plot("magnetization")
+        nonpolarized_density.plot("3")
 
 
-def check_plotting_magnetization_density(polarized_density, Assert):
-    obj = viewer3d.Viewer3d
-    cm_init = patch.object(obj, "__init__", autospec=True, return_value=None)
-    cm_cell = patch.object(obj, "show_cell")
-    cm_surface = patch.object(obj, "show_isosurface")
-    with cm_init as init, cm_cell as cell, cm_surface as surface:
-        result = polarized_density.plot("magnetization", isolevel=0.1, smooth=1)
+@pytest.mark.parametrize(
+    "selection", ["3", "sigma_z", "z", "sigma_3", "magnetization", "mag", "m"]
+)
+def test_collinear_plot(selection, collinear_density, mock_viewer, Assert, not_core):
+    source = collinear_density.ref.source
+    if source == "charge":
+        expected_density = collinear_density.ref.output["magnetization"].T
+    else:
+        expected_density = collinear_density.ref.output[source][1].T
+        if selection in (
+            "magnetization",
+            "mag",
+            "m",
+        ):  # magnetization not allowed for tau
+            return
+    result = collinear_density.plot(selection, isolevel=0.1, smooth=1)
+    assert isinstance(result, viewer3d.Viewer3d)
+    calls = mock_viewer["surface"].call_args_list
+    check_magnetization_plot(expected_density, calls, Assert)
+
+
+def test_accessing_noncollinear_element_raises_error(collinear_density, not_core):
+    with pytest.raises(exception.NoData):
+        collinear_density.plot("1")
+
+
+@pytest.mark.parametrize(
+    "selections",
+    [
+        ("1", "2", "3"),
+        ("sigma_x", "sigma_y", "sigma_z"),
+        ("x", "y", "z"),
+        ("sigma_1", "sigma_2", "sigma_3"),
+        (
+            "m(1)",
+            "mag(2)",
+            "magnetization(3)",
+        ),  # the magnetization label should be ignored
+    ],
+)
+def test_plotting_noncollinear_density(
+    selections, noncollinear_density, mock_viewer, Assert, not_core
+):
+    source = noncollinear_density.ref.source
+    if source == "charge":
+        expected_density = noncollinear_density.ref.output["magnetization"]
+    else:
+        expected_density = noncollinear_density.ref.output[source][1:]
+        if "(" in selections[0]:  # magnetization not allowed for tau
+            return
+    for component, selection in enumerate(selections):
+        result = noncollinear_density.plot(selection, isolevel=0.1, smooth=1)
         assert isinstance(result, viewer3d.Viewer3d)
-        calls = surface.call_args_list
-    reference_magnetization = polarized_density.ref.output["magnetization"].T
-    if polarized_density.collinear():
-        check_collinear_plot(reference_magnetization, calls, Assert)
-    elif polarized_density.noncollinear():
-        check_noncollinear_plot(reference_magnetization, calls, Assert)
+        calls = mock_viewer["surface"].call_args_list
+        check_magnetization_plot(expected_density[component].T, calls, Assert)
+        mock_viewer["surface"].reset_mock()
 
 
-def check_collinear_plot(magnetization, calls, Assert):
+def check_magnetization_plot(magnetization, calls, Assert):
     assert len(calls) == 2
     args, kwargs = calls[0]
     Assert.allclose(args[0], magnetization)
-    assert kwargs == {"isolevel": 0.1, "color": "blue", "opacity": 0.6, "smooth": 1}
+    blue = _config.VASP_BLUE
+    assert kwargs == {"isolevel": 0.1, "color": blue, "opacity": 0.6, "smooth": 1}
     args, kwargs = calls[1]
     Assert.allclose(args[0], -magnetization)
-    assert kwargs == {"isolevel": 0.1, "color": "red", "opacity": 0.6, "smooth": 1}
+    red = _config.VASP_RED
+    assert kwargs == {"isolevel": 0.1, "color": red, "opacity": 0.6, "smooth": 1}
 
 
-def check_noncollinear_plot(magnetization, calls, Assert):
-    magnetization = np.linalg.norm(magnetization, axis=-1)
-    assert len(calls) == 1
-    args, kwargs = calls[0]
-    Assert.allclose(args[0], magnetization)
-    assert kwargs == {"isolevel": 0.1, "color": "yellow", "opacity": 0.6, "smooth": 1}
+def test_adding_components(noncollinear_density, mock_viewer, Assert, not_core):
+    source = noncollinear_density.ref.source
+    if source == "charge":
+        expected_density = noncollinear_density.ref.output["magnetization"]
+    else:
+        expected_density = noncollinear_density.ref.output[source][1:]
+    expected_density = expected_density[0] + expected_density[1]
+    result = noncollinear_density.plot("1 + 2", isolevel=0.4)
+    assert isinstance(result, viewer3d.Viewer3d)
+    mock_viewer["surface"].assert_called_once()
+    args, kwargs = mock_viewer["surface"].call_args
+    Assert.allclose(args[0], expected_density.T)
+    assert kwargs == {"isolevel": 0.4, "color": _config.VASP_CYAN, "opacity": 0.6}
+
+
+def test_to_numpy(reference_density, Assert):
+    source = reference_density.ref.source
+    if source == "charge":
+        if reference_density.is_nonpolarized():
+            expected_density = [reference_density.ref.output["charge"]]
+        elif reference_density.is_collinear():
+            expected_density = [
+                reference_density.ref.output["charge"],
+                reference_density.ref.output["magnetization"],
+            ]
+        else:
+            expected_density = [
+                reference_density.ref.output["charge"],
+                *reference_density.ref.output["magnetization"],
+            ]
+    else:
+        expected_density = reference_density.ref.output[source]
+    Assert.allclose(reference_density.to_numpy(), expected_density)
 
 
-def test_missing_element(reference_density, Assert, not_core):
+def test_selections(reference_density):
+    assert reference_density.selections() == reference_density.ref.selections
+
+
+def test_selections_empty_density(empty_density):
+    assert empty_density.selections() == {"density": list(raw.selections("density"))}
+
+
+def test_missing_element(reference_density, not_core):
     with pytest.raises(exception.IncorrectUsage):
         reference_density.plot("unknown tag")
 
 
-def test_color_specified_for_magnetism(collinear_density, Assert, not_core):
+def test_color_specified_for_sigma_z(collinear_density, not_core):
     with pytest.raises(exception.NotImplemented):
-        collinear_density.plot("magnetization", color="brown")
+        collinear_density.plot("3", color="brown")
+
+
+@pytest.mark.parametrize("selection", ("m", "mag", "magnetization"))
+def test_magnetization_without_component(selection, raw_data, not_core):
+    data = raw_data.density("Fe3O4 noncollinear")
+    with pytest.raises(exception.IncorrectUsage):
+        Density.from_data(data).plot(selection)
 
 
 def test_print(reference_density, format_):
diff --git a/tests/data/test_dos.py b/tests/data/test_dos.py
index a14f1d85..3b396dfe 100644
--- a/tests/data/test_dos.py
+++ b/tests/data/test_dos.py
@@ -203,7 +203,6 @@ def test_plot_combine_projectors(Fe3O4_projectors, Assert):
     subtraction_up = Fe3O4_projectors.ref.Fe_up - Fe3O4_projectors.ref.p_up
     subtraction_down = Fe3O4_projectors.ref.Fe_down - Fe3O4_projectors.ref.p_down
     names = [d.name for d in data]
-    print(data[names.index("Fe_up + O_d_down")])
     Assert.allclose(data[names.index("Fe_up + O_d_down")].y, addition)
     Assert.allclose(data[names.index("Fe_up - p_up")].y, subtraction_up)
     Assert.allclose(data[names.index("Fe_down - p_down")].y, -subtraction_down)
diff --git a/tests/data/test_viewer3d.py b/tests/data/test_viewer3d.py
index 9e034277..3c0686da 100644
--- a/tests/data/test_viewer3d.py
+++ b/tests/data/test_viewer3d.py
@@ -151,6 +151,7 @@ def test_serializable(not_core):
         json.json_clean(element)
 
 
+@pytest.mark.xfail(reason="Not compatible with density plotting")
 def test_nonstandard_form(nonstandard_form, Assert, assert_arrow_message):
     viewer = nonstandard_form
     Assert.allclose(viewer._positions, viewer.ref.positions)
diff --git a/tests/raw/conftest.py b/tests/raw/conftest.py
index 88d46f4b..18709fed 100644
--- a/tests/raw/conftest.py
+++ b/tests/raw/conftest.py
@@ -31,15 +31,16 @@ def complex_schema():
     schema.add(OptionalArgument, name=name, mandatory=only_mandatory.mandatory)
     schema.add(OptionalArgument, mandatory=both.mandatory, optional=both.optional)
     schema.add(WithLink, required=version, baz=pointer.baz, simple=pointer.simple)
-    schema.add(WithLength, num_data=length.num_data)
+    schema.add(WithLength, alias="alias_name", num_data=length.num_data)
     schema.add(Complex, opt=first.opt, link=first.link, length=first.length)
     schema.add(Complex, name=name, opt=second.opt, link=second.link)
+    alias_source = Source(length, alias_for="default")
     reference = {
         "version": {"default": Source(VERSION)},
         "simple": {"default": Source(simple, file=filename)},
         "optional_argument": {"default": Source(both), name: Source(only_mandatory)},
         "with_link": {"default": Source(pointer, required=version)},
-        "with_length": {"default": Source(length)},
+        "with_length": {"default": Source(length), "alias_name": alias_source},
         "complex": {"default": Source(first), name: Source(second)},
     }
     return schema, reference
diff --git a/tests/raw/test_schema.py b/tests/raw/test_schema.py
index 796dd0f6..7383c194 100644
--- a/tests/raw/test_schema.py
+++ b/tests/raw/test_schema.py
@@ -78,6 +78,24 @@ def test_length():
     assert remove_version(schema.sources) == reference
 
 
+def test_alias():
+    first = Simple("foo1", "bar1")
+    second = Simple("foo2", "bar2")
+    schema = Schema(VERSION)
+    schema.add(Simple, foo=first.foo, bar=first.bar, alias=["first", "other"])
+    schema.add(Simple, name="second", foo=second.foo, bar=second.bar, alias="more")
+    reference = {
+        "simple": {
+            "default": Source(first),
+            "first": Source(first, alias_for="default"),
+            "other": Source(first, alias_for="default"),
+            "second": Source(second),
+            "more": Source(second, alias_for="second"),
+        },
+    }
+    assert remove_version(schema.sources) == reference
+
+
 def remove_version(sources):
     version = sources.pop("version")
     assert version == {"default": Source(VERSION)}
@@ -102,25 +120,31 @@ def test_complex_str(complex_schema):
     major: major_dataset
     minor: minor_dataset
     patch: patch_dataset
+
 simple:
     default:  &simple-default
         file: other_file
         foo: foo_dataset
         bar: bar_dataset
+
 optional_argument:
     mandatory:  &optional_argument-mandatory
         mandatory: mandatory1
     default:  &optional_argument-default
         mandatory: mandatory2
         optional: optional
+
 with_link:
     default:  &with_link-default
         required: 1.2.3
         baz: baz_dataset
         simple: *simple-default
+
 with_length:
     default:  &with_length-default
         num_data: length(dataset)
+    alias_name: *with_length-default
+
 complex:
     default:  &complex-default
         opt: *optional_argument-default
@@ -128,7 +152,7 @@ def test_complex_str(complex_schema):
         length: *with_length-default
     mandatory:  &complex-mandatory
         opt: *optional_argument-mandatory
-        link: *with_link-default\
+        link: *with_link-default
 """
     assert str(schema) == reference
 
@@ -149,7 +173,7 @@ def test_missing_quantity():
 def test_adding_twice_error():
     schema = Schema(VERSION)
     schema.add(Simple, foo="foo1", bar="bar1")
-    with pytest.raises(exception.IncorrectUsage):
+    with pytest.raises(exception._Py4VaspInternalError):
         schema.add(Simple, foo="foo2", bar="bar2")
 
 
@@ -167,12 +191,12 @@ def test_incomplete_schema():
     schema.add(WithLink, baz=pointer.baz, simple=pointer.simple)
     # test missing quantity
     assert not schema.verified
-    with pytest.raises(AssertionError):
+    with pytest.raises(exception._Py4VaspInternalError):
         schema.verify()
     assert not schema.verified
     # test missing source
     schema.add(Simple, foo=target.foo, bar=target.bar)
     assert not schema.verified
-    with pytest.raises(AssertionError):
+    with pytest.raises(exception._Py4VaspInternalError):
         schema.verify()
     assert not schema.verified
diff --git a/tests/util/test_index.py b/tests/util/test_index.py
index 489e2927..00bfd655 100644
--- a/tests/util/test_index.py
+++ b/tests/util/test_index.py
@@ -303,7 +303,7 @@ def test_label(selection, label):
         2: {"x": 0, "y": 1, "z": 2, "u~v": 3},
         0: {"total": slice(0, 2), "up": 0, "down": 1},
     }
-    selector = index.Selector(map_, np.zeros((2, 8, 4, 0)))
+    selector = index.Selector(map_, np.zeros((2, 8, 4, 0)), use_number_labels=True)
     assert selector.label(selection) == label
 
 
@@ -326,18 +326,25 @@ def test_label_operations(selection, label):
         2: {"x": 0, "y": 1, "z": 2, "u~v": 3},
         0: {"total": slice(0, 2), "up": 0, "down": 1},
     }
-    selector = index.Selector(map_, np.zeros((2, 8, 4, 0)))
+    selector = index.Selector(map_, np.zeros((2, 8, 4, 0)), use_number_labels=True)
     assert selector.label(selection) == label
 
 
+def test_labels_without_number_labels():
+    map_ = {1: {"A": 0, "0": 0}}
+    selector = index.Selector(map_, np.zeros((10, 10)))
+    assert selector.label(("0",)) == "0"
+
+
 def test_error_when_duplicate_key():
     with pytest.raises(exception._Py4VaspInternalError):
         index.Selector({0: {"A": 1}, 1: {"A": 2}}, None)
 
 
 def test_error_when_numbers_are_longer_than_one():
+    map_ = {0: {"A": 1, "1": slice(0, 2)}}
     with pytest.raises(exception._Py4VaspInternalError):
-        index.Selector({0: {"A": 1, "1": slice(0, 2)}}, np.zeros(3))
+        index.Selector(map_, np.zeros(3), use_number_labels=True)
 
 
 @pytest.mark.parametrize(
diff --git a/tests/util/test_select.py b/tests/util/test_select.py
index 2054c615..3a320994 100644
--- a/tests/util/test_select.py
+++ b/tests/util/test_select.py
@@ -8,6 +8,7 @@
 
 def test_empty_tree():
     assert selections(None) == ((),)
+    assert selections("A", filter={"A"}) == ((),)
     assert graph(None) == "graph LR"
 
 
@@ -214,6 +215,25 @@ def test_complex_operation():
     assert graph(selection) == expected
 
 
+@pytest.mark.parametrize(
+    "unfiltered, filtered",
+    [
+        ("A(B) B(A)", "B B"),
+        ("A(B C)", "B C"),
+        ("A~B B~A", "A~B B~A"),
+        ("A(B) + C, B - C(A)", "B + C, B - C"),
+    ],
+)
+def test_selection_filter(unfiltered, filtered):
+    assert selections(unfiltered, filter={"A"}) == selections(filtered)
+
+
+@pytest.mark.parametrize("selection", ("A + B", "B + A", "A - B", "B - A"))
+def test_incorrect_combination(selection):
+    with pytest.raises(exception.IncorrectUsage):
+        selections(selection, filter={"A"})
+
+
 @pytest.mark.parametrize(
     "input, output",
     [
@@ -272,9 +292,9 @@ def test_default_constructor_raises_error(selection):
         select.Tree(selection)
 
 
-def selections(selection):
+def selections(selection, filter={}):
     tree = select.Tree.from_selection(selection)
-    return tuple(tree.selections())
+    return tuple(tree.selections(filter=filter))
 
 
 def graph(selection):