Exciton wavefunction

src/py4vasp/_raw/data.py

@@ -213,6 +213,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 Fatband:
     """Contains the BSE data required to produce a fatband plot."""

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,
@@ -286,6 +293,15 @@ 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",
+)
+#
+group = "results/linear_response"
 schema.add(
     raw.Fatband,
     required=raw.Version(6, 4),
@@ -489,6 +505,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"),
+    topology=Link("topology", "exciton"),
+    positions="results/supercell/position_ions",
+)
 #
 schema.add(raw.System, system="input/incar/SYSTEM")
 #
@@ -504,6 +528,13 @@ def selections(quantity):
     ion_types="results/phonons/primitive/ion_types",
     number_ion_types="results/phonons/primitive/number_ion_types",
 )
+schema.add(
+    raw.Topology,
+    name="exciton",
+    required=raw.Version(6, 5),
+    ion_types="results/supercell/ion_types",
+    number_ion_types="results/supercell/number_ion_types",
+)
 #
 schema.add(
     raw.Velocity,

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"""
 
     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:
@@ -195,6 +200,11 @@ def _create_atoms(self, step):
         atoms.cell = self.lattice_vectors[step]
         atoms.set_scaled_positions(self.positions[step])
         atoms.set_pbc(True)
+        if(self.shift):
+            translation = np.sum(atoms.cell,axis=0)
+            translation = [translation[i]*self.shift[i] for i in range(3)]
+            atoms.positions += translation
+            atoms.wrap()
         atoms = atoms.repeat(self.supercell)
         return atoms
 
@@ -226,9 +236,12 @@ def _show_isosurface(self, widget, trajectory):
         for grid_scalar in self.grid_scalars:
             if not grid_scalar.isosurfaces:
                 continue
+            quantity = grid_scalar.quantity[step]
+            if(self.shift):
+                quantity = np.roll(quantity,[int(quantity.shape[i]*self.shift[i]) for i in range(3)] ,axis=(0,1,2))
+            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)

src/py4vasp/calculation/__init__.py

@@ -56,6 +56,7 @@ class provides a more flexible interface with which you can determine the source
     "dos",
     "elastic_modulus",
     "energy",
+    "exciton_density",
     "fatband",
     "force",
     "force_constant",

src/py4vasp/calculation/_exciton_density.py

@@ -0,0 +1,141 @@
+# 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, calculation, exception
+from py4vasp._third_party import view
+from py4vasp._util import documentation, import_, index, select
+from py4vasp.calculation import _base, _structure
+
+pretty = import_.optional("IPython.lib.pretty")
+
+
+_INTERNAL = "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:]
+        excitons = self._raw_data.exciton_charge.shape[0]
+        topology = calculation.topology.from_data(self._raw_data.structure.topology)
+        return f"""exciton charge density:
+    structure: {pretty.pretty(topology)}
+    grid: {grid[2]}, {grid[1]}, {grid[0]}
+    excitons : {excitons}"""
+
+    @_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.
+
+        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 _INTERNAL
+        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)
+        selections = tree.selections()
+        viewer.grid_scalars = [
+            self._grid_quantity(selector, selection, map_, user_options)
+            for selection in selections
+        ]
+        if center: 
+            viewer.shift = (0.5,0.5,0.5)
+        return viewer
+
+    def _create_map(self):
+        excitons=self._raw_data.exciton_charge.shape[0]
+        map_ = {
+            str(choice): choice-1
+            for choice in range(1,excitons+1)
+        }
+        return map_
+
+    def _grid_quantity(self, selector, selection, map_, user_options):
+        component_label = selector.label(selection)
+        return view.GridQuantity(
+            quantity=(selector[selection].T)[np.newaxis],
+            label=self._label(component_label),
+            isosurfaces=self._isosurfaces(**user_options),
+        )
+
+    def _label(self, component_label):
+        if self._selection:
+            return f"{self._selection}({component_label})"
+        else:
+            return component_label
+
+    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 the exciton density is"
+            "written to vaspout.h5 if the tags LCHARGH5=T or LH5=T are set in"
+            "the INCAR file"
+        )