Feat: exciton wavefunction (#185)

* Read and plot exciton charge density
* Shift the cell if view.center=True
* Add tests for new features

---------

Co-authored-by: Alexey Tal <[email protected]>

src/py4vasp/_calculation/__init__.py

@@ -40,7 +40,7 @@
     "_stoichiometry",
 )
 GROUPS = {
-    "exciton": ("eigenvector",),
+    "exciton": ("density", "eigenvector"),
     "phonon": ("band", "dos"),
 }
 

src/py4vasp/_calculation/exciton_density.py

@@ -0,0 +1,137 @@
+# Copyright © VASP Software GmbH,
+# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0)
+import numpy as np
+
+from py4vasp import _config, exception
+from py4vasp._calculation import _stoichiometry, base, structure
+from py4vasp._third_party import view
+from py4vasp._util import import_, index, select
+
+pretty = import_.optional("IPython.lib.pretty")
+
+
+_DEFAULT_SELECTION = "1"
+
+
+class ExcitonDensity(base.Refinery, structure.Mixin, view.Mixin):
+    """This class accesses exciton charge densities of VASP.
+
+    The exciton charge densities can be calculated via the BSE/TDHF algorithm in
+    VASP. With this class you can extract these charge densities.
+    """
+
+    @base.data_access
+    def __str__(self):
+        _raise_error_if_no_data(self._raw_data.exciton_charge)
+        grid = self._raw_data.exciton_charge.shape[1:]
+        raw_stoichiometry = self._raw_data.structure.stoichiometry
+        stoichiometry = _stoichiometry.Stoichiometry.from_data(raw_stoichiometry)
+        return f"""exciton charge density:
+    structure: {pretty.pretty(stoichiometry)}
+    grid: {grid[2]}, {grid[1]}, {grid[0]}
+    excitons: {len(self._raw_data.exciton_charge)}"""
+
+    @base.data_access
+    def to_dict(self):
+        """Read the exciton density into a dictionary.
+
+        Returns
+        -------
+        dict
+            Contains the supercell structure information as well as the exciton
+            charge density represented on a grid in the supercell.
+        """
+        _raise_error_if_no_data(self._raw_data.exciton_charge)
+        result = {"structure": self._structure.read()}
+        result.update({"charge": self.to_numpy()})
+        return result
+
+    @base.data_access
+    def to_numpy(self):
+        """Convert the exciton charge density to a numpy array.
+
+        Returns
+        -------
+        np.ndarray
+            Charge density of all excitons.
+        """
+        return np.moveaxis(self._raw_data.exciton_charge, 0, -1).T
+
+    @base.data_access
+    def to_view(self, selection=None, supercell=None, center=False, **user_options):
+        """Plot the selected exciton density as a 3d isosurface within the structure.
+
+        Parameters
+        ----------
+        selection : str
+            Can be exciton index or a combination, i.e., "1" or "1+2+3"
+
+        supercell : int or np.ndarray
+            If present the data is replicated the specified number of times along each
+            direction.
+
+        center : bool
+            Shift the origin of the unit cell to the center. This is helpful if you
+            the exciton is at the corner of the cell.
+
+        user_options
+            Further arguments with keyword that get directly passed on to the
+            visualizer. Most importantly, you can set isolevel to adjust the
+            value at which the isosurface is drawn.
+
+        Returns
+        -------
+        View
+            Visualize an isosurface of the exciton density within the 3d structure.
+
+        Examples
+        --------
+        >>> calc = py4vasp.Calculation.from_path(".")
+        Plot an isosurface of the first exciton charge density
+        >>> calc.exciton.density.plot()
+        Plot an isosurface of the third exciton charge density
+        >>> calc.exciton.density.plot("3")
+        Plot an isosurface of the sum of first and second exciton charge
+        densities
+        >>> calc.exciton.density.plot("1+2")
+        """
+        _raise_error_if_no_data(self._raw_data.exciton_charge)
+        selection = selection or _DEFAULT_SELECTION
+        viewer = self._structure.plot(supercell)
+        map_ = self._create_map()
+        selector = index.Selector({0: map_}, self._raw_data.exciton_charge)
+        tree = select.Tree.from_selection(selection)
+        viewer.grid_scalars = [
+            self._grid_quantity(selector, selection, map_, user_options)
+            for selection in tree.selections()
+        ]
+        if center:
+            viewer.shift = (0.5, 0.5, 0.5)
+        return viewer
+
+    def _create_map(self):
+        num_excitons = self._raw_data.exciton_charge.shape[0]
+        return {str(choice + 1): choice for choice in range(num_excitons)}
+
+    def _grid_quantity(self, selector, selection, map_, user_options):
+        return view.GridQuantity(
+            quantity=(selector[selection].T)[np.newaxis],
+            label=selector.label(selection),
+            isosurfaces=self._isosurfaces(**user_options),
+        )
+
+    def _isosurfaces(self, isolevel=0.8, color=None, opacity=0.6):
+        color = color or _config.VASP_COLORS["cyan"]
+        return [view.Isosurface(isolevel, color, opacity)]
+
+
+def _raise_error_if_no_data(data):
+    if data.is_none():
+        raise exception.NoData(
+            "Exciton charge density was not found. Note that in order to calculate the"
+            "exciton charge density the number of eigenvectors has to be selected with"
+            "the tag NBSEEIG and the position of the hole or the electron has to be"
+            "provided with the tag BSEHOLE or BSEELECTRON, correspondingly. The exciton"
+            "density is written to vaspout.h5 if the tags LCHARGH5=T or LH5=T are set"
+            "in the INCAR file, otherwise the charge density is written to CHG.XXX files."
+        )

src/py4vasp/_raw/data.py

@@ -229,6 +229,15 @@ class Energy:
     "Energy specified by labels for all iteration steps."
 
 
+@dataclasses.dataclass
+class ExcitonDensity:
+    "The exciton charge density on the real space grid."
+    structure: Structure
+    "The atomic structure to represent the densities."
+    exciton_charge: VaspData
+    "The data of exciton charge density."
+
+
 @dataclasses.dataclass
 class ExcitonEigenvector:
     """Contains the BSE data required to produce a plot of eigenvector contributions."""

src/py4vasp/_raw/definition.py

@@ -88,6 +88,13 @@ def selections(quantity):
     scale="results/phonons/primitive/scale",
     lattice_vectors="results/phonons/primitive/lattice_vectors",
 )
+schema.add(
+    raw.Cell,
+    name="exciton",
+    required=raw.Version(6, 5),
+    scale="results/supercell/scale",
+    lattice_vectors="results/supercell/lattice_vectors",
+)
 #
 schema.add(
     raw.CONTCAR,
@@ -302,6 +309,13 @@ def selections(quantity):
 )
 #
 group = "results/linear_response"
+schema.add(
+    raw.ExcitonDensity,
+    required=raw.Version(6, 5),
+    file="vaspout.h5",
+    structure=Link("structure", "exciton"),
+    exciton_charge=f"{group}/exciton_charge",
+)
 schema.add(
     raw.ExcitonEigenvector,
     required=raw.Version(6, 4),
@@ -489,6 +503,13 @@ def selections(quantity):
     ion_types="results/phonons/primitive/ion_types",
     number_ion_types="results/phonons/primitive/number_ion_types",
 )
+schema.add(
+    raw.Stoichiometry,
+    name="exciton",
+    required=raw.Version(6, 5),
+    ion_types="results/supercell/ion_types",
+    number_ion_types="results/supercell/number_ion_types",
+)
 #
 schema.add(
     raw.Stress,
@@ -510,6 +531,14 @@ def selections(quantity):
     cell=Link("cell", "final"),
     positions="results/positions/position_ions",
 )
+schema.add(
+    raw.Structure,
+    name="exciton",
+    required=raw.Version(6, 5),
+    cell=Link("cell", "exciton"),
+    stoichiometry=Link("stoichiometry", "exciton"),
+    positions="results/supercell/position_ions",
+)
 #
 schema.add(raw.System, system="input/incar/SYSTEM")
 #

src/py4vasp/_third_party/view/view.py

@@ -114,6 +114,10 @@ class View:
     """Defines if the axes is shown in the viewer"""
     show_axes_at: Sequence[float] = None
     """Defines where the axis is shown, defaults to the origin"""
+    shift: npt.ArrayLike = None
+    """Defines the shift of the origin"""
+    camera: str = "orthographic"
+    """Defines the camera view type (orthographic or perspective)"""
 
     def __post_init__(self):
         self._verify()
@@ -133,6 +137,7 @@ def to_ngl(self):
         trajectory = [self._create_atoms(i) for i in self._iterate_trajectory_frames()]
         ngl_trajectory = nglview.ASETrajectory(trajectory)
         widget = nglview.NGLWidget(ngl_trajectory)
+        widget.camera = self.camera
         if self.grid_scalars:
             self._show_isosurface(widget, trajectory)
         if self.ion_arrows:
@@ -191,10 +196,10 @@ def _raise_error_if_any_shape_is_incorrect(self):
 
     def _create_atoms(self, step):
         symbols = "".join(self.elements[step])
-        atoms = ase.Atoms(symbols)
-        atoms.cell = self.lattice_vectors[step]
-        atoms.set_scaled_positions(self.positions[step])
-        atoms.set_pbc(True)
+        atoms = ase.Atoms(symbols, cell=self.lattice_vectors[step], pbc=True)
+        shift = np.zeros(3) if self.shift is None else self.shift
+        atoms.set_scaled_positions(np.add(self.positions[step], shift))
+        atoms.wrap()
         atoms = atoms.repeat(self.supercell)
         return atoms
 
@@ -226,9 +231,11 @@ def _show_isosurface(self, widget, trajectory):
         for grid_scalar in self.grid_scalars:
             if not grid_scalar.isosurfaces:
                 continue
+            quantity = grid_scalar.quantity[step]
+            quantity = self._shift_quantity(quantity)
+            quantity = self._repeat_isosurface(quantity)
             atoms = trajectory[step]
             self._set_atoms_in_standard_form(atoms)
-            quantity = self._repeat_isosurface(grid_scalar.quantity[step])
             with tempfile.TemporaryDirectory() as tmp:
                 filename = os.path.join(tmp, CUBE_FILENAME)
                 ase_cube.write_cube(open(filename, "w"), atoms=atoms, data=quantity)
@@ -241,6 +248,12 @@ def _show_isosurface(self, widget, trajectory):
                 }
                 component.add_surface(**isosurface_options)
 
+    def _shift_quantity(self, quantity):
+        if self.shift is None:
+            return quantity
+        shift_indices = np.round(quantity.shape * self.shift).astype(np.int32)
+        return np.roll(quantity, shift_indices, axis=(0, 1, 2))
+
     def _show_arrows_at_atoms(self, widget, trajectory):
         step = 0
         for _arrows in self.ion_arrows:

tests/calculation/test_exciton_density.py

@@ -0,0 +1,116 @@
+# Copyright © VASP Software GmbH,
+# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0)
+import types
+
+import numpy as np
+import pytest
+
+from py4vasp import _config, exception, raw
+from py4vasp._calculation.exciton_density import ExcitonDensity
+from py4vasp._calculation.structure import Structure
+
+
+@pytest.fixture
+def exciton_density(raw_data):
+    raw_density = raw_data.exciton_density()
+    density = ExcitonDensity.from_data(raw_density)
+    density.ref = types.SimpleNamespace()
+    density.ref.structure = Structure.from_data(raw_density.structure)
+    expected_charge = [component.T for component in raw_density.exciton_charge]
+    density.ref.density = np.array(expected_charge)
+    print(density.ref.density.shape)
+    return density
+
+
+@pytest.fixture
+def empty_density(raw_data):
+    raw_density = raw.ExcitonDensity(
+        raw_data.structure("Sr2TiO4"), exciton_charge=raw.VaspData(None)
+    )
+    return ExcitonDensity.from_data(raw_density)
+
+
+def test_read(exciton_density, Assert):
+    actual = exciton_density.read()
+    actual_structure = actual.pop("structure")
+    Assert.same_structure(actual_structure, exciton_density.ref.structure.read())
+    Assert.allclose(actual["charge"], exciton_density.ref.density)
+
+
+def test_missing_data(empty_density):
+    with pytest.raises(exception.NoData):
+        empty_density.read()
+
+
+def test_to_numpy(exciton_density, Assert):
+    actual = exciton_density.to_numpy()
+    Assert.allclose(actual, exciton_density.ref.density)
+
+
+@pytest.mark.parametrize("selection, indices", [(None, 0), ("2", 1), ("1, 3", (0, 2))])
+def test_plot_selection(exciton_density, selection, indices, Assert):
+    indices = np.atleast_1d(indices)
+    if selection is None:
+        view = exciton_density.plot()
+    else:
+        view = exciton_density.plot(selection)
+    Assert.same_structure_view(view, exciton_density.ref.structure.plot())
+    assert len(view.grid_scalars) == len(indices)
+    for grid_scalar, index in zip(view.grid_scalars, indices):
+        selected_exciton = exciton_density.ref.density[index]
+        assert grid_scalar.label == str(index + 1)
+        assert grid_scalar.quantity.ndim == 4
+        Assert.allclose(grid_scalar.quantity, selected_exciton)
+        assert len(grid_scalar.isosurfaces) == 1
+        isosurface = grid_scalar.isosurfaces[0]
+        assert isosurface.isolevel == 0.8
+        assert isosurface.color == _config.VASP_COLORS["cyan"]
+        assert isosurface.opacity == 0.6
+
+
+def test_plot_addition(exciton_density, Assert):
+    view = exciton_density.plot("1 + 3")
+    assert len(view.grid_scalars) == 1
+    grid_scalar = view.grid_scalars[0]
+    selected_exciton = exciton_density.ref.density[0] + exciton_density.ref.density[2]
+    assert grid_scalar.label == "1 + 3"
+    Assert.allclose(grid_scalar.quantity, selected_exciton)
+
+
+def test_plot_centered(exciton_density, Assert):
+    view = exciton_density.plot()
+    assert view.shift is None
+    view = exciton_density.plot(center=True)
+    Assert.allclose(view.shift, 0.5)
+
+
+@pytest.mark.parametrize("supercell", (2, (3, 1, 2)))
+def test_plot_supercell(exciton_density, supercell, Assert):
+    view = exciton_density.plot(supercell=supercell)
+    Assert.allclose(view.supercell, supercell)
+
+
+def test_plot_user_options(exciton_density):
+    view = exciton_density.plot(isolevel=0.4, color="red", opacity=0.5)
+    assert len(view.grid_scalars) == 1
+    grid_scalar = view.grid_scalars[0]
+    assert len(grid_scalar.isosurfaces) == 1
+    isosurface = grid_scalar.isosurfaces[0]
+    assert isosurface.isolevel == 0.4
+    assert isosurface.color == "red"
+    assert isosurface.opacity == 0.5
+
+
+def test_print(exciton_density, format_):
+    actual, _ = format_(exciton_density)
+    expected_text = """\
+exciton charge density:
+    structure: Sr2TiO4
+    grid: 10, 12, 14
+    excitons: 3"""
+    assert actual == {"text/plain": expected_text}
+
+
+def test_factory_methods(raw_data, check_factory_methods):
+    data = raw_data.exciton_density()
+    check_factory_methods(ExcitonDensity, data)