Merge pull request #36 from martin-schlipf/master

Fixes for workshop (#133)

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "py4vasp"
-version = "0.3.0"
+version = "0.3.1"
 description = ""
 authors = [
     "Martin Schlipf <[email protected]>",

src/py4vasp/__init__.py

@@ -1,4 +1,5 @@
 from .calculation import Calculation
 from py4vasp._third_party.interactive import set_error_handling
 
+__version__ = "0.3.1"
 set_error_handling("Minimal")

src/py4vasp/data/dielectric_function.py

@@ -34,9 +34,20 @@ def _to_string(self):
         return f"""
 dielectric function:
     energies: [{energies[0]:0.2f}, {energies[-1]:0.2f}] {len(energies)} points
+    components: {", ".join(self._components())}
     directions: isotropic, xx, yy, zz, xy, yz, xz
         """.strip()
 
+    def _components(self):
+        components = ()
+        if self._raw_data.density_density is not None:
+            components += ("density",)
+        if self._raw_data.current_current is not None:
+            components += ("current",)
+        if self._raw_data.ion is not None:
+            components += ("ion",)
+        return components
+
     def _to_dict(self):
         f"""Read the data into a dictionary.
 
@@ -45,11 +56,12 @@ def _to_dict(self):
         dict
             Contains the energies at which the dielectric function was evaluated
             and the dielectric tensor (3x3 matrix) at these energies."""
+        data = self._raw_data
         return {
-            "energies": self._raw_data.energies[:],
-            "density_density": _convert.to_complex(self._raw_data.density_density[:]),
-            "current_current": _convert.to_complex(self._raw_data.current_current[:]),
-            "ion": _convert.to_complex(self._raw_data.ion[:]),
+            "energies": data.energies[:],
+            "density_density": _convert_to_complex_if_not_none(data.density_density),
+            "current_current": _convert_to_complex_if_not_none(data.current_current),
+            "ion": _convert_to_complex_if_not_none(data.ion),
         }
 
     def _to_plotly(self, selection=None):
@@ -68,30 +80,43 @@ def _to_plotly(self, selection=None):
             plotly figure containing the dielectric function for the selected
             directions and components."""
         selection = _default_selection_if_none(selection)
-        choices = _parse_selection(selection)
         data = self._to_dict()
+        choices = _parse_selection(selection, data)
         plots = [_make_plot(data, *choice) for choice in choices]
         default = {
             "xaxis": {"title": {"text": "Energy (eV)"}},
-            "yaxis": {"title": {"text": r"$\epsilon$"}},
+            "yaxis": {"title": {"text": "dielectric function ϵ"}},
         }
         return go.Figure(data=plots, layout=default)
 
 
+def _convert_to_complex_if_not_none(array):
+    if array is None:
+        return None
+    else:
+        return _convert.to_complex(array[:])
+
+
 def _default_selection_if_none(selection):
     return "isotropic" if selection is None else selection
 
 
+def _parse_selection(selection, data):
+    tree = _selection.SelectionTree.from_selection(selection)
+    yield from _parse_recursive(tree, _default_choice(data))
+
+
 class _Choice(typing.NamedTuple):
-    component: str = "density"
+    component: str
     direction: str = "isotropic"
     real_or_imag: str = _Selection.default
 
 
-def _parse_selection(selection):
-    tree = _selection.SelectionTree.from_selection(selection)
-    default_choice = _Choice()
-    yield from _parse_recursive(tree, default_choice)
+def _default_choice(data):
+    if data["density_density"] is not None:
+        return _Choice("density")
+    else:
+        return _Choice("ion")
 
 
 def _parse_recursive(tree, current_choice):
@@ -165,14 +190,10 @@ def _select_data_direction(tensor, direction):
 
 
 def _build_name(component, direction, real_or_imag):
-    subscript = "" if direction == "isotropic" else f"_{{{direction}}}"
-    if component == "density":
-        superscript = "^{dd}"
-    elif component == "current":
-        superscript = "^{jj}"
-    else:
-        superscript = "^{ion}"
-    return f"{real_or_imag[:2].capitalize()}($\\epsilon{superscript}{subscript}$)"
+    name = f"{real_or_imag[:2].capitalize()},{component}"
+    if direction != "isotropic":
+        name += f",{direction}"
+    return name
 
 
 def _scatter_plot(x, y, name):

src/py4vasp/data/dielectric_tensor.py

@@ -34,7 +34,8 @@ def _to_dict(self):
     def _to_string(self):
         data = self._to_dict()
         return f"""
-MACROSCOPIC STATIC DIELECTRIC TENSOR {_description(data["method"])}
+Macroscopic static dielectric tensor (dimensionless)
+  {_description(data["method"])}
 ------------------------------------------------------
 {_dielectric_tensor_string(data["clamped_ion"], "clamped-ion")}
 {_dielectric_tensor_string(data["relaxed_ion"], "relaxed-ion")}
@@ -55,11 +56,11 @@ def _dielectric_tensor_string(tensor, label):
 
 def _description(method):
     if method == "dft":
-        return "(including local field effects in DFT)"
+        return "including local field effects in DFT"
     elif method == "rpa":
-        return "(including local field effects in RPA (Hartree))"
+        return "including local field effects in RPA (Hartree)"
     elif method == "scf":
-        return "(including local field effects)"
+        return "including local field effects"
     elif method == "nscf":
-        return "(excluding local field effects)"
+        return "excluding local field effects"
     assert False  # unknown method

src/py4vasp/raw/file.py

@@ -348,14 +348,15 @@ def dielectric_function(self):
 
     def _read_dielectric_function(self):
         self._raise_error_if_closed()
-        group = f"results/linear_response"
+        group = "results/linear_response"
+        suffix = "_dielectric_function"
         if group not in self._h5f:
             return None
         return RawDielectricFunction(
-            energies=self._h5f[f"{group}/energies_dielectric_function"],
-            density_density=self._h5f[f"{group}/density_density_dielectric_function"],
-            current_current=self._h5f[f"{group}/current_current_dielectric_function"],
-            ion=self._h5f[f"{group}/ion_dielectric_function"],
+            energies=self._h5f[f"{group}/energies{suffix}"],
+            density_density=self._safe_get_key(f"{group}/density_density{suffix}"),
+            current_current=self._safe_get_key(f"{group}/current_current{suffix}"),
+            ion=self._safe_get_key(f"{group}/ion{suffix}"),
         )
 
     @property

tests/conftest.py

@@ -59,7 +59,12 @@ def density(selection):
 
     @staticmethod
     def dielectric_function(selection):
-        return _dielectric_function()
+        if selection == "electron":
+            return _electron_dielectric_function()
+        elif selection == "ion":
+            return _ion_dielectric_function()
+        else:
+            raise exception.NotImplemented()
 
     @staticmethod
     def dielectric_tensor(selection):
@@ -181,14 +186,24 @@ def _number_components(selection):
         raise exception.NotImplemented()
 
 
-def _dielectric_function():
-    shape = (3, axes, axes, number_points, complex_)
+def _electron_dielectric_function():
+    shape = (2, axes, axes, number_points, complex_)
     data = np.linspace(0, 1, np.prod(shape)).reshape(shape)
     return raw.RawDielectricFunction(
         energies=np.linspace(0, 1, number_points),
         density_density=data[0],
         current_current=data[1],
-        ion=data[2],
+        ion=None,
+    )
+
+
+def _ion_dielectric_function():
+    shape = (axes, axes, number_points, complex_)
+    return raw.RawDielectricFunction(
+        energies=np.linspace(0, 1, number_points),
+        density_density=None,
+        current_current=None,
+        ion=np.linspace(0, 1, np.prod(shape)).reshape(shape),
     )