From 661b9efed3824b36e67c559f4120b5ecfd231bce Mon Sep 17 00:00:00 2001 From: Michael Wolloch Date: Tue, 18 Jun 2024 20:54:46 +0200 Subject: [PATCH] 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 --- poetry.lock | 2 +- src/py4vasp/calculation/_partial_charge.py | 45 +++++++------ tests/calculation/test_partial_charge.py | 77 +++++++++++++++++----- tests/conftest.py | 44 ++++++++----- 4 files changed, 116 insertions(+), 52 deletions(-) diff --git a/poetry.lock b/poetry.lock index 19fe5d15..f3481526 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 1.8.2 and should not be changed by hand. +# This file is automatically @generated by Poetry 1.8.3 and should not be changed by hand. [[package]] name = "alabaster" diff --git a/src/py4vasp/calculation/_partial_charge.py b/src/py4vasp/calculation/_partial_charge.py index 6841a879..7834bb49 100644 --- a/src/py4vasp/calculation/_partial_charge.py +++ b/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): diff --git a/tests/calculation/test_partial_charge.py b/tests/calculation/test_partial_charge.py index 4d8770a8..53009a3f 100644 --- a/tests/calculation/test_partial_charge.py +++ b/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,13 +176,15 @@ 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) @@ -189,14 +192,14 @@ def test_split_bands_to_numpy(NonPolarizedBandSplitPartialCharge, spin, Assert): 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,6 +211,7 @@ 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: @@ -215,20 +219,20 @@ def test_to_stm_nonsplit_not_orthogonal_no_vacuum( 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): diff --git a/tests/conftest.py b/tests/conftest.py index 1458b8bb..e58db1b0 100644 --- a/tests/conftest.py +++ b/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,7 +29,18 @@ @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 @@ -34,18 +48,6 @@ def is_core(): 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, )