feat: allow changing STM settings via argument of to_stm method (#158)

* made STM_settings a keyword argument to the to_stm method, so the settings can be changed by the user.
* added stm_settings property
* skip calling optional scipy in conftest.py

---------

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

src/py4vasp/calculation/_partial_charge.py

@@ -112,6 +112,7 @@ def to_stm(
         tip_height: float = 2.0,
         current: float = 1.0,
         supercell: Union[int, np.array] = 2,
+        stm_settings: STM_settings = STM_settings(),
     ) -> Graph:
         """Generate STM image data from the partial charge density.
 
@@ -129,6 +130,9 @@ def to_stm(
             Only used in "constant_current" mode.
         supercell : int | np.array
             The supercell to be used for plotting the STM. The default is 2.
+        stm_settings : STM_settings
+            Settings for the STM simulation concerning smoothening parameters
+            and interpolation. The default is STM_settings().
 
         Returns
         -------
@@ -143,8 +147,9 @@ def to_stm(
             if index > 0:
                 message = "Selecting more than one STM is not implemented."
                 raise exception.NotImplemented(message)
-            contour = self._make_contour(selection, tip_height, current)
+            contour = self._make_contour(selection, tip_height, current, stm_settings)
         contour.supercell = self._parse_supercell(supercell)
+        contour.settings = stm_settings
         return Graph(series=contour, title=contour.label)
 
     def _parse_supercell(self, supercell):
@@ -156,16 +161,18 @@ def _parse_supercell(self, supercell):
         The supercell is used to multiply the x and y directions of the lattice."""
         raise exception.IncorrectUsage(message)
 
-    def _make_contour(self, selection, tip_height, current):
+    def _make_contour(self, selection, tip_height, current, stm_settings):
         self._raise_error_if_tip_too_far_away(tip_height)
         mode = self._parse_mode(selection)
         spin = self._parse_spin(selection)
         self._raise_error_if_selection_not_understood(selection, mode, spin)
-        smoothed_charge = self._get_stm_data(spin)
+        smoothed_charge = self._get_stm_data(spin, stm_settings)
         if mode == "constant_height" or mode is None:
-            return self._constant_height_stm(smoothed_charge, tip_height, spin)
+            return self._constant_height_stm(
+                smoothed_charge, tip_height, spin, stm_settings
+            )
         current = current * 1e-09  # convert nA to A
-        return self._constant_current_stm(smoothed_charge, current, spin)
+        return self._constant_current_stm(smoothed_charge, current, spin, stm_settings)
 
     def _parse_mode(self, selection):
         for mode, aliases in _STM_MODES.items():
@@ -185,14 +192,14 @@ def _raise_error_if_selection_not_understood(self, selection, mode, spin):
             message = f"STM mode '{selection}' was parsed as mode='{mode}' and spin='{spin}' which could not be used. Please use 'constant_height' or 'constant_current' as mode and 'up', 'down', or 'total' as spin."
             raise exception.IncorrectUsage(message)
 
-    def _constant_current_stm(self, smoothed_charge, current, spin):
+    def _constant_current_stm(self, smoothed_charge, current, spin, stm_settings):
         z_start = _min_of_z_charge(
-            self._get_stm_data(spin),
-            sigma=self.stm_settings.sigma_z,
-            truncate=self.stm_settings.truncate,
+            self._get_stm_data(spin, stm_settings),
+            sigma=stm_settings.sigma_z,
+            truncate=stm_settings.truncate,
         )
         grid = self.grid()
-        z_step = 1 / self.stm_settings.interpolation_factor
+        z_step = 1 / stm_settings.interpolation_factor
         # roll smoothed charge so that we are not bothered by the boundary of the
         # unit cell if the slab is not centered. z_start is now the first index
         smoothed_charge = np.roll(smoothed_charge, -z_start, axis=2)
@@ -205,9 +212,9 @@ def _constant_current_stm(self, smoothed_charge, current, spin):
         label = f"STM of {topology} for {spin_label} at constant current={current*1e9:.2f} nA"
         return Contour(data=scan, lattice=self._get_stm_plane(), label=label)
 
-    def _constant_height_stm(self, smoothed_charge, tip_height, spin):
+    def _constant_height_stm(self, smoothed_charge, tip_height, spin, stm_settings):
         zz = self._z_index_for_height(tip_height + self._get_highest_z_coord())
-        height_scan = smoothed_charge[:, :, zz] * self.stm_settings.enhancement_factor
+        height_scan = smoothed_charge[:, :, zz] * stm_settings.enhancement_factor
         spin_label = "both spin channels" if spin == "total" else f"spin {spin}"
         topology = self._topology()
         label = f"STM of {topology} for {spin_label} at constant height={float(tip_height):.2f} Angstrom"
@@ -250,27 +257,27 @@ def _raise_error_if_tip_too_far_away(self, tip_height):
             You would be sampling the bottom of your slab, which is not supported."""
             raise exception.IncorrectUsage(message)
 
-    def _get_stm_data(self, spin):
+    def _get_stm_data(self, spin, stm_settings):
         if 0 not in self.bands() or 0 not in self.kpoints():
             massage = """Simulated STM images are only supported for non-separated bands and k-points.
             Please set LSEPK and LSEPB to .FALSE. in the INCAR file."""
             raise exception.NotImplemented(massage)
         chg = self._correct_units(self.to_numpy(spin, band=0, kpoint=0))
-        return self._smooth_stm_data(chg)
+        return self._smooth_stm_data(chg, stm_settings)
 
     def _correct_units(self, charge_data):
         grid_volume = np.prod(self.grid())
         cell_volume = self._structure.volume()
         return charge_data / (grid_volume * cell_volume)
 
-    def _smooth_stm_data(self, data):
+    def _smooth_stm_data(self, data, stm_settings):
         sigma = (
-            self.stm_settings.sigma_xy,
-            self.stm_settings.sigma_xy,
-            self.stm_settings.sigma_z,
+            stm_settings.sigma_xy,
+            stm_settings.sigma_xy,
+            stm_settings.sigma_z,
         )
         return ndimage.gaussian_filter(
-            data, sigma=sigma, truncate=self.stm_settings.truncate, mode="wrap"
+            data, sigma=sigma, truncate=stm_settings.truncate, mode="wrap"
         )
 
     def _get_stm_plane(self):

tests/calculation/test_partial_charge.py

@@ -8,6 +8,7 @@
 
 from py4vasp import calculation
 from py4vasp._util.slicing import plane
+from py4vasp.calculation._partial_charge import STM_settings
 from py4vasp.exception import IncorrectUsage, NoData, NotImplemented
 
 
@@ -102,7 +103,7 @@ def make_reference_partial_charge(raw_data, selection):
     return parchg
 
 
-def test_read(PartialCharge, Assert):
+def test_read(PartialCharge, Assert, not_core):
     actual = PartialCharge.read()
     expected = PartialCharge.ref
     Assert.allclose(actual["bands"], expected.bands)
@@ -113,31 +114,31 @@ def test_read(PartialCharge, Assert):
     Assert.same_structure(actual["structure"], expected.structure.read())
 
 
-def test_topology(PartialCharge):
+def test_topology(PartialCharge, not_core):
     actual = PartialCharge._topology()
     expected = str(PartialCharge.ref.structure._topology())
     assert actual == expected
 
 
-def test_bands(PartialCharge, Assert):
+def test_bands(PartialCharge, Assert, not_core):
     actual = PartialCharge.bands()
     expected = PartialCharge.ref.bands
     Assert.allclose(actual, expected)
 
 
-def test_kpoints(PartialCharge, Assert):
+def test_kpoints(PartialCharge, Assert, not_core):
     actual = PartialCharge.kpoints()
     expected = PartialCharge.ref.kpoints
     Assert.allclose(actual, expected)
 
 
-def test_grid(PartialCharge, Assert):
+def test_grid(PartialCharge, Assert, not_core):
     actual = PartialCharge.grid()
     expected = PartialCharge.ref.grid
     Assert.allclose(actual, expected)
 
 
-def test_non_split_to_numpy(PolarizedNonSplitPartialCharge, Assert):
+def test_non_split_to_numpy(PolarizedNonSplitPartialCharge, Assert, not_core):
     actual = PolarizedNonSplitPartialCharge.to_numpy("total")
     expected = PolarizedNonSplitPartialCharge.ref.partial_charge
     Assert.allclose(actual, expected[0, 0, 0].T)
@@ -149,7 +150,7 @@ def test_non_split_to_numpy(PolarizedNonSplitPartialCharge, Assert):
     Assert.allclose(actual, 0.5 * (expected[0, 0, 0].T - expected[0, 0, 1].T))
 
 
-def test_split_to_numpy(PolarizedAllSplitPartialCharge, Assert):
+def test_split_to_numpy(PolarizedAllSplitPartialCharge, Assert, not_core):
     bands = PolarizedAllSplitPartialCharge.ref.bands
     kpoints = PolarizedAllSplitPartialCharge.ref.kpoints
     for band_index, band in enumerate(bands):
@@ -175,28 +176,30 @@ def test_split_to_numpy(PolarizedAllSplitPartialCharge, Assert):
     assert msg in str(excinfo.value)
 
 
-def test_non_polarized_to_numpy(NonSplitPartialCharge, spin, Assert):
+def test_non_polarized_to_numpy(NonSplitPartialCharge, spin, Assert, not_core):
     actual = NonSplitPartialCharge.to_numpy(selection=spin)
     expected = NonSplitPartialCharge.ref.partial_charge
     Assert.allclose(actual, np.asarray(expected).T[:, :, :, 0, 0, 0])
 
 
-def test_split_bands_to_numpy(NonPolarizedBandSplitPartialCharge, spin, Assert):
+def test_split_bands_to_numpy(
+    NonPolarizedBandSplitPartialCharge, spin, Assert, not_core
+):
     bands = NonPolarizedBandSplitPartialCharge.ref.bands
     for band_index, band in enumerate(bands):
         actual = NonPolarizedBandSplitPartialCharge.to_numpy(spin, band=band)
     expected = NonPolarizedBandSplitPartialCharge.ref.partial_charge
     Assert.allclose(actual, np.asarray(expected).T[:, :, :, 0, band_index, 0])
 
 
-def test_to_stm_split(PolarizedAllSplitPartialCharge):
+def test_to_stm_split(PolarizedAllSplitPartialCharge, not_core):
     msg = "set LSEPK and LSEPB to .FALSE. in the INCAR file."
     with pytest.raises(NotImplemented) as excinfo:
         PolarizedAllSplitPartialCharge.to_stm(selection="constant_current")
     assert msg in str(excinfo.value)
 
 
-def test_to_stm_nonsplit_tip_to_high(NonSplitPartialCharge):
+def test_to_stm_nonsplit_tip_to_high(NonSplitPartialCharge, not_core):
     actual = NonSplitPartialCharge
     tip_height = 8.4
     error = f"""The tip position at {tip_height:.2f} is above half of the
@@ -208,27 +211,28 @@ def test_to_stm_nonsplit_tip_to_high(NonSplitPartialCharge):
 
 def test_to_stm_nonsplit_not_orthogonal_no_vacuum(
     PolarizedNonSplitPartialChargeSr2TiO4,
+    not_core,
 ):
     msg = "The vacuum region in your cell is too small for STM simulations."
     with pytest.raises(IncorrectUsage) as excinfo:
         PolarizedNonSplitPartialChargeSr2TiO4.to_stm()
     assert msg in str(excinfo.value)
 
 
-def test_to_stm_wrong_spin_nonsplit(PolarizedNonSplitPartialCharge):
+def test_to_stm_wrong_spin_nonsplit(PolarizedNonSplitPartialCharge, not_core):
     msg = "'up', 'down', or 'total'"
     with pytest.raises(IncorrectUsage) as excinfo:
         PolarizedNonSplitPartialCharge.to_stm(selection="all")
     assert msg in str(excinfo.value)
 
 
-def test_to_stm_wrong_mode(PolarizedNonSplitPartialCharge):
+def test_to_stm_wrong_mode(PolarizedNonSplitPartialCharge, not_core):
     with pytest.raises(IncorrectUsage) as excinfo:
         PolarizedNonSplitPartialCharge.to_stm(selection="stm")
     assert "STM mode" in str(excinfo.value)
 
 
-def test_wrong_vacuum_direction(NonSplitPartialChargeNi_100):
+def test_wrong_vacuum_direction(NonSplitPartialChargeNi_100, not_core):
     msg = """The vacuum region in your cell is not located along
         the third lattice vector."""
     with pytest.raises(NotImplemented) as excinfo:
@@ -311,7 +315,7 @@ def test_to_stm_nonsplit_constant_current_non_ortho(
     assert f"{current:.2f}" in actual.title
 
 
-def test_stm_default_settings(PolarizedNonSplitPartialCharge):
+def test_stm_settings(PolarizedNonSplitPartialCharge, not_core):
     actual = dataclasses.asdict(PolarizedNonSplitPartialCharge.stm_settings)
     defaults = {
         "sigma_xy": 4.0,
@@ -321,6 +325,49 @@ def test_stm_default_settings(PolarizedNonSplitPartialCharge):
         "interpolation_factor": 10,
     }
     assert actual == defaults
+    modified = STM_settings(
+        sigma_xy=2.0,
+        sigma_z=2.0,
+        truncate=1.0,
+        enhancement_factor=500,
+        interpolation_factor=5,
+    )
+    graph = PolarizedNonSplitPartialCharge.to_stm(stm_settings=modified)
+    assert graph.series.settings == modified
+
+
+def test_smoothening_change(PolarizedNonSplitPartialCharge, not_core):
+    mod_settings = STM_settings(sigma_xy=2.0, sigma_z=2.0, truncate=1.0)
+    data = PolarizedNonSplitPartialCharge.to_numpy("total", band=0, kpoint=0)
+    default_smoothed_density = PolarizedNonSplitPartialCharge._smooth_stm_data(
+        data=data, stm_settings=STM_settings()
+    )
+    new_smoothed_density = PolarizedNonSplitPartialCharge._smooth_stm_data(
+        data=data, stm_settings=mod_settings
+    )
+    assert not np.allclose(default_smoothed_density, new_smoothed_density)
+
+
+def test_enhancement_setting_change(PolarizedNonSplitPartialCharge, Assert, not_core):
+    enhance_settings = STM_settings(
+        enhancement_factor=STM_settings().enhancement_factor / 2.0
+    )
+    graph_def = PolarizedNonSplitPartialCharge.to_stm("constant_height")
+    graph_less_enhanced = PolarizedNonSplitPartialCharge.to_stm(
+        "constant_height", stm_settings=enhance_settings
+    )
+    Assert.allclose(graph_def.series.data, graph_less_enhanced.series.data * 2)
+
+
+def test_interpolation_setting_change(PolarizedNonSplitPartialCharge, not_core):
+    interp_settings = STM_settings(
+        interpolation_factor=STM_settings().interpolation_factor / 4.0
+    )
+    graph_def = PolarizedNonSplitPartialCharge.to_stm("constant_current", current=1)
+    graph_less_interp_points = PolarizedNonSplitPartialCharge.to_stm(
+        "constant_current", current=1, stm_settings=interp_settings
+    )
+    assert not np.allclose(graph_def.series.data, graph_less_interp_points.series.data)
 
 
 def test_factory_methods(raw_data, check_factory_methods):

tests/conftest.py

@@ -9,6 +9,9 @@
 from numpy.testing import assert_array_almost_equal_nulp
 
 from py4vasp import exception, raw
+from py4vasp._util import import_
+
+stats = import_.optional("scipy.stats")
 
 number_steps = 4
 number_atoms = 7
@@ -26,26 +29,25 @@
 
 
 @pytest.fixture(scope="session")
-def is_core():
+def only_core():
+    if not _is_core():
+        pytest.skip("This test checks py4vasp-core functionality not used by py4vasp.")
+
+
+@pytest.fixture(scope="session")
+def not_core():
+    if _is_core():
+        pytest.skip("This test requires features not present in py4vasp-core.")
+
+
+def _is_core():
     try:
         importlib.metadata.distribution("py4vasp-core")
         return True
     except importlib.metadata.PackageNotFoundError:
         return False
 
 
-@pytest.fixture()
-def only_core(is_core):
-    if not is_core:
-        pytest.skip("This test checks py4vasp-core functionality not used by py4vasp.")
-
-
-@pytest.fixture()
-def not_core(is_core):
-    if is_core:
-        pytest.skip("This test requires features not present in py4vasp-core.")
-
-
 class _Assert:
     @staticmethod
     def allclose(actual, desired, tolerance=1):
@@ -710,14 +712,22 @@ def _partial_charge(selection):
     else:
         spin_dimension = 1
     grid = raw.VaspData(tuple(reversed(grid_dim)))
-    random_charge = raw.VaspData(
-        np.random.rand(len(kpoints), len(bands), spin_dimension, *grid_dim)
-    )
+    gaussian_charge = np.zeros((len(kpoints), len(bands), spin_dimension, *grid_dim))
+    if not _is_core():
+        cov = grid_dim[0] / 10  # standard deviation
+        z = np.arange(grid_dim[0])  # z range
+        for gy in range(grid_dim[1]):
+            for gx in range(grid_dim[2]):
+                m = int(grid_dim[0] / 2) + gy / 10 + gx / 10
+                val = stats.multivariate_normal(mean=m, cov=cov).pdf(z)
+                # Fill the gaussian_charge array
+                gaussian_charge[:, :, :, :, gy, gx] = val
+    gaussian_charge = raw.VaspData(gaussian_charge)
     return raw.PartialCharge(
         structure=structure,
         bands=bands,
         kpoints=kpoints,
-        partial_charge=random_charge,
+        partial_charge=gaussian_charge,
         grid=grid,
     )