feat: partial charges and stm (#144)

* implement partial charges
* implement to_stm for producing STM pictures of partial charges
* use current units of nA
* constant current outputs correct heights over the surface layer
* make STM settings readable to user
* use consistent selection logic
* allow for non-orthogonal z-axis

---------
Co-authored-by: Martin Schlipf <[email protected]>

core/pyproject.toml

@@ -9,7 +9,9 @@ authors = [
     "Orest Dubay <[email protected]>",
     "Jonathan Lahnsteiner <[email protected]>",
     "Eisuke Kawashima <[email protected]>",
-    "Sudarshan Vijay <[email protected]>"
+    "Sudarshan Vijay <[email protected]>",
+    "Marie-Therese Huebsch <[email protected]>",
+    "Michael Wolloch <[email protected]>",
 ]
 license = "Apache-2.0"
 readme = "README.md"

pyproject.toml

@@ -9,7 +9,9 @@ authors = [
     "Orest Dubay <[email protected]>",
     "Jonathan Lahnsteiner <[email protected]>",
     "Eisuke Kawashima <[email protected]>",
-    "Sudarshan Vijay <[email protected]>"
+    "Sudarshan Vijay <[email protected]>",
+    "Marie-Therese Huebsch <[email protected]>",
+    "Michael Wolloch <[email protected]>",
 ]
 license = "Apache-2.0"
 readme = "README.md"
@@ -30,6 +32,7 @@ ase = ">=3.22.1"
 plotly = ">=5.9.0"
 kaleido = ">=0.2.1,<0.2.1.post1 || >0.2.1.post1"
 ipython = ">=8.12"
+scipy = ">=1.12.0"
 
 [tool.poetry.group.dev.dependencies]
 pytest = ">=7.1.2"

src/py4vasp/_raw/data.py

@@ -343,6 +343,24 @@ class PairCorrelation:
     "Describes which indices correspond to which element pair."
 
 
+@dataclasses.dataclass
+class PartialCharge:
+    """Electronic partial charge and magnetization density on the fine Fourier grid
+
+    Possibly not only split by spin, but also by band and kpoint."""
+
+    structure: Structure
+    "The atomic structure to represent the densities."
+    partial_charge: VaspData
+    "The data of electronic charge and magnetization density."
+    kpoints: VaspData
+    "The kpoints at which the partial charge is evaluated."
+    bands: VaspData
+    "The bands at which the partial charge is evaluated."
+    grid: VaspData
+    "The fine FFT grid at which the partial charge is evaluated."
+
+
 @dataclasses.dataclass
 class PhononBand:
     """The band structure of the phonons.

src/py4vasp/_raw/definition.py

@@ -395,6 +395,16 @@ def selections(quantity):
     labels=f"{group}/labels",
 )
 #
+schema.add(
+    raw.PartialCharge,
+    required=raw.Version(6, 5),
+    structure=Link("structure", DEFAULT_SOURCE),
+    partial_charge="results/partial_charges/parchg",
+    bands="results/partial_charges/bands",
+    kpoints="results/partial_charges/kpoints",
+    grid="results/partial_charges/grid",
+)
+#
 group = "results/phonons"
 schema.add(
     raw.PhononBand,

src/py4vasp/_third_party/graph/contour.py

@@ -37,9 +37,10 @@ class Contour(trace.Trace):
     the dimensions should be the ones of the grid, if the data is 3d the first dimension
     should be a 2 for a vector in the plane of the grid and the other two dimensions
     should be the grid."""
-    lattice: Lattice
-    """2 vectors spanning the plane in which the data is represented. Each vector should
-    have two components, so remove any element normal to the plane."""
+    lattice: Plane
+    """Lattice plane in which the data is represented spanned by 2 vectors.
+    Each vector should have two components, so remove any element normal to
+    the plane. Can be generated with the 'plane' function in py4vasp._util.slicing."""
     label: str
     "Assign a label to the visualization that may be used to identify one among multiple plots."
     isolevels: bool = False

src/py4vasp/calculation/__init__.py

@@ -62,6 +62,7 @@ class provides a more flexible interface with which you can determine the source
     "magnetism",
     "OSZICAR",
     "pair_correlation",
+    "partial_charge",
     "phonon_band",
     "phonon_dos",
     "piezoelectric_tensor",