From 1397ec75207ae78b72ba564666aa31a36e7166c2 Mon Sep 17 00:00:00 2001 From: Martin Schlipf Date: Thu, 5 Feb 2026 12:15:35 +0100 Subject: [PATCH 1/3] Implement cutoff for radial plots --- src/py4vasp/_calculation/effective_coulomb.py | 29 +++++++---- tests/calculation/test_effective_coulomb.py | 52 ++++++++++++++----- 2 files changed, 59 insertions(+), 22 deletions(-) diff --git a/src/py4vasp/_calculation/effective_coulomb.py b/src/py4vasp/_calculation/effective_coulomb.py index c40f1b53..f2302721 100644 --- a/src/py4vasp/_calculation/effective_coulomb.py +++ b/src/py4vasp/_calculation/effective_coulomb.py @@ -226,6 +226,7 @@ def to_graph( selection: str = "U J V", omega: None | EllipsisType | np.ndarray = None, radius: None | EllipsisType | np.ndarray = None, + radius_max: None | float = None, config: interpolate.AAAConfig = interpolate.AAAConfig(), ) -> graph.Graph: """Generate a graph representation of the effective Coulomb interaction. @@ -234,8 +235,8 @@ def to_graph( are provided: - If only omega is given: creates a frequency-dependent plot - - If only radius is given: creates a radial-dependent plot - - If both omega and radius are given: creates a frequency plot for all radii + - If only radius/radius_max is given: creates a radial-dependent plot + - If both omega and radius/radius_max are given: creates a frequency plot for all radii Parameters ---------- @@ -260,6 +261,8 @@ def to_graph( are used. You can also provide specific radii, then the data will be interpolated to the selected radii. + radius_max + config Configuration for the analytic continuation of the frequency-dependent data. Use this if you need to adjust the parameters of the analytic continuation. @@ -271,15 +274,16 @@ def to_graph( interaction data. """ tree = select.Tree.from_selection(selection) - plotter = self._make_plotter(omega, radius, config) + plotter = self._make_plotter(omega, radius, radius_max,config) potentials = self._get_effective_potentials(tree, plotter) series = plotter.make_all_series(potentials) return graph.Graph( series, xlabel=plotter.xlabel, ylabel="Coulomb potential (eV)" ) - def _make_plotter(self, omega, radius, config): - if omega is not None and radius is not None: + def _make_plotter(self, omega, radius, radius_max, config): + radius_set = (radius is not None or radius_max is not None) + if omega is not None and radius_set: if radius is not ...: raise exception.NotImplemented( "Interpolating radial data for frequency plots is not implemented." @@ -287,9 +291,9 @@ def _make_plotter(self, omega, radius, config): omega_in = self._read_frequencies().get("frequencies") positions = self._read_positions() return _OmegaPlotter(omega_in, omega, config, positions) - if radius is not None: + if radius_set: positions = self._read_positions() - return _RadialPlotter(positions, radius) + return _RadialPlotter(positions, radius, radius_max) else: omega_in = self._read_frequencies().get("frequencies") return _OmegaPlotter(omega_in, omega, config) @@ -464,10 +468,15 @@ def _make_one_series(self, potential, position=0, suffix=""): class _RadialPlotter: xlabel = "Radius (Å)" - def __init__(self, positions, radius_out): + def __init__(self, positions, radius_out, radius_max): if not positions: raise exception.DataMismatch("The output does not contain position data.") self.radius_in = self._transform_positions_to_radial(positions) + if radius_max: + self.mask = self.radius_in <= radius_max + self.radius_in = self.radius_in[self.mask] + else: + self.mask = slice(None) self.interpolate = radius_out is not None and radius_out is not ... self.radius_out = radius_out if self.interpolate else self.radius_in self.marker = None if self.interpolate else "*" @@ -484,9 +493,9 @@ def interpolate_screened_if_necessary(self, potential): return self._ohno_interpolation(potential) def _ohno_interpolation(self, potential): + potential = potential.real[..., self.mask] if not self.interpolate: - return potential.real - potential = potential.real + return potential if potential.ndim == 2: # if multiple frequencies are present, take only omega = 0 potential = potential[0] diff --git a/tests/calculation/test_effective_coulomb.py b/tests/calculation/test_effective_coulomb.py index fbd09034..5b8ea1a1 100644 --- a/tests/calculation/test_effective_coulomb.py +++ b/tests/calculation/test_effective_coulomb.py @@ -353,6 +353,16 @@ def check_radial_plot_raises_error(effective_coulomb): with pytest.raises(exception.DataMismatch): effective_coulomb.plot(radius=...) +def test_plot_radial_with_cutoff(nonpolarized_crpar, Assert): + effective_coulomb = nonpolarized_crpar + graph = effective_coulomb.plot("U", radius_max=10) + assert len(graph) == 1 + series = graph[0] + mask = effective_coulomb.ref.radial_data["radius"] < 10 + Assert.allclose(series.x, effective_coulomb.ref.radial_data["radius"][mask]) + Assert.allclose(series.y, effective_coulomb.ref.radial_data["screened U"][0,mask]) + assert series.label == "screened U" + @pytest.mark.parametrize("selection", ["total", "up~up", "down~down", "up~down"]) def test_plot_radial_selected_spin(collinear_crpa, selection, Assert): @@ -384,38 +394,56 @@ def test_plot_radial_invalid_selection(nonpolarized_crpar, selection): def test_plot_radial_interpolation(nonpolarized_crpar, not_core, Assert): radial_data = nonpolarized_crpar.ref.radial_data - radius = np.linspace(0.0, np.max(radial_data["radius"]), 30) + radius_in = radial_data["radius"] + radius_out = np.linspace(0.0, np.max(radius_in), 30) ref_graph = nonpolarized_crpar.plot(radius=...) - graph = nonpolarized_crpar.plot(radius=radius) + graph = nonpolarized_crpar.plot(radius=radius_out) assert len(graph) == len(ref_graph) assert graph.xlabel == "Radius (Å)" assert graph.ylabel == "Coulomb potential (eV)" for series, ref_series in zip(graph, ref_graph): - Assert.allclose(series.x, radius) - expected_interpolated_values = interpolate_r(radial_data, ref_series.y, radius) - Assert.allclose(series.y, expected_interpolated_values, tolerance=100) + Assert.allclose(series.x, radius_out) + interpolated_values = interpolate_r(radius_in, ref_series.y, radius_out) + Assert.allclose(series.y, interpolated_values, tolerance=100) + assert series.label == ref_series.label + assert series.marker is None + + +def test_plot_radial_interpolation_with_cutoff(nonpolarized_crpar, not_core, Assert): + radial_data = nonpolarized_crpar.ref.radial_data + radius_in = radial_data["radius"][radial_data["radius"] < 10] + radius_out = np.linspace(0.0, np.max(radius_in), 30) + ref_graph = nonpolarized_crpar.plot(radius=..., radius_max=10) + graph = nonpolarized_crpar.plot(radius=radius_out, radius_max=10) + assert len(graph) == len(ref_graph) + assert graph.xlabel == "Radius (Å)" + assert graph.ylabel == "Coulomb potential (eV)" + for series, ref_series in zip(graph, ref_graph): + Assert.allclose(series.x, radius_out) + interpolated_values = interpolate_r(radius_in, ref_series.y, radius_out) + Assert.allclose(series.y, interpolated_values, tolerance=100) assert series.label == ref_series.label assert series.marker is None def test_plot_radial_interpolation_spin_selection(collinear_crpa, not_core, Assert): effective_coulomb = collinear_crpa - radial_data = effective_coulomb.ref.radial_data - radius = np.linspace(0, 10) + radius_in = effective_coulomb.ref.radial_data["radius"] + radius_out = np.linspace(0, 10) ref_graph = effective_coulomb.plot("up~down(U V)", radius=...) - graph = effective_coulomb.plot("up~down(U V)", radius=radius) + graph = effective_coulomb.plot("up~down(U V)", radius=radius_out) expected_labels = ["screened up~down_U", "bare up~down_V"] assert len(graph) == len(ref_graph) for series, ref_series, label in zip(graph, ref_graph, expected_labels): - Assert.allclose(series.x, radius) - Assert.allclose(series.y, interpolate_r(radial_data, ref_series.y, radius)) + Assert.allclose(series.x, radius_out) + Assert.allclose(series.y, interpolate_r(radius_in, ref_series.y, radius_out)) assert series.label == label -def interpolate_r(radial_data, potential, radius): +def interpolate_r(radius_in, potential, radius_out): U0 = potential[0] interpolation = numeric.interpolate_with_function( - EffectiveCoulomb.ohno_potential, radial_data["radius"], potential / U0, radius + EffectiveCoulomb.ohno_potential, radius_in, potential / U0, radius_out ) return U0 * interpolation From ff96efa9255f66c3c12ccd6618dcf007e7d2dc86 Mon Sep 17 00:00:00 2001 From: Martin Schlipf Date: Thu, 5 Feb 2026 12:39:08 +0100 Subject: [PATCH 2/3] Implement cutoff for frequency plot with multiple radii --- src/py4vasp/_calculation/effective_coulomb.py | 57 +++++++++++-------- tests/calculation/test_effective_coulomb.py | 23 +++++++- 2 files changed, 55 insertions(+), 25 deletions(-) diff --git a/src/py4vasp/_calculation/effective_coulomb.py b/src/py4vasp/_calculation/effective_coulomb.py index f2302721..a0146c5c 100644 --- a/src/py4vasp/_calculation/effective_coulomb.py +++ b/src/py4vasp/_calculation/effective_coulomb.py @@ -274,7 +274,7 @@ def to_graph( interaction data. """ tree = select.Tree.from_selection(selection) - plotter = self._make_plotter(omega, radius, radius_max,config) + plotter = self._make_plotter(omega, radius, radius_max, config) potentials = self._get_effective_potentials(tree, plotter) series = plotter.make_all_series(potentials) return graph.Graph( @@ -282,15 +282,15 @@ def to_graph( ) def _make_plotter(self, omega, radius, radius_max, config): - radius_set = (radius is not None or radius_max is not None) + radius_set = radius is not None or radius_max is not None if omega is not None and radius_set: - if radius is not ...: + if radius is not ... and radius is not None: raise exception.NotImplemented( "Interpolating radial data for frequency plots is not implemented." ) omega_in = self._read_frequencies().get("frequencies") positions = self._read_positions() - return _OmegaPlotter(omega_in, omega, config, positions) + return _OmegaPlotter(omega_in, omega, config, positions, radius_max) if radius_set: positions = self._read_positions() return _RadialPlotter(positions, radius, radius_max) @@ -419,7 +419,7 @@ class _CoulombPotential: class _OmegaPlotter: - def __init__(self, omega_in, omega_out, config, positions=None): + def __init__(self, omega_in, omega_out, config, positions=None, radius_max=None): self.omega_in = omega_in self.interpolate = omega_out is not None and omega_out is not ... if omega_in is None: @@ -429,7 +429,9 @@ def __init__(self, omega_in, omega_out, config, positions=None): self.omega_out = omega_out if self.interpolate else omega_in self.xlabel = "ω (eV)" if self.interpolate else "Im(ω) (eV)" self.config = config - self.positions = positions + if positions is not None: + self.positions = positions["positions"] + _, self.mask = transform_positions_to_radial(positions, radius_max) def interpolate_bare_if_necessary(self, potential): num_omega = len(self.omega_out) @@ -446,22 +448,29 @@ def interpolate_screened_if_necessary(self, potential): ).T def make_all_series(self, potentials): - if self.positions: + if hasattr(self, "positions"): return [ - self._make_one_series(potential, position=i, suffix=f" @ {position}") + self._make_one_series(potential, position=position) for potential in potentials - for i, position in enumerate(self.positions["positions"]) + for position in self.positions[self.mask] ] else: return [self._make_one_series(potential) for potential in potentials] - def _make_one_series(self, potential, position=0, suffix=""): + def _make_one_series(self, potential, position=None): + if position is None: + position_index = 0 + suffix = "" + else: + lookup = np.all(self.positions == position, axis=1) + position_index = np.squeeze(np.where(lookup)) + suffix = f" @ {position}" omega = self.omega_out.real if self.interpolate else self.omega_out.imag label = f"{potential.component} {potential.selector_label}{suffix}" if potential.strength.ndim == 1: strength = potential.strength else: - strength = potential.strength[:, position] + strength = potential.strength[:, position_index] return graph.Series(omega, strength.real, label=label) @@ -469,23 +478,11 @@ class _RadialPlotter: xlabel = "Radius (Å)" def __init__(self, positions, radius_out, radius_max): - if not positions: - raise exception.DataMismatch("The output does not contain position data.") - self.radius_in = self._transform_positions_to_radial(positions) - if radius_max: - self.mask = self.radius_in <= radius_max - self.radius_in = self.radius_in[self.mask] - else: - self.mask = slice(None) + self.radius_in, self.mask = transform_positions_to_radial(positions, radius_max) self.interpolate = radius_out is not None and radius_out is not ... self.radius_out = radius_out if self.interpolate else self.radius_in self.marker = None if self.interpolate else "*" - def _transform_positions_to_radial(self, positions): - return np.linalg.norm( - positions["lattice_vectors"] @ positions["positions"].T, axis=0 - ) - def interpolate_bare_if_necessary(self, potential): return self._ohno_interpolation(potential) @@ -517,3 +514,15 @@ def _make_one_series(self, potential): # use only omega = 0 strength = potential.strength[0] return graph.Series(self.radius_out, strength, label=label, marker=self.marker) + + +def transform_positions_to_radial(positions, radius_max): + if not positions: + raise exception.DataMismatch("The output does not contain position data.") + radius = np.linalg.norm( + positions["lattice_vectors"] @ positions["positions"].T, axis=0 + ) + if radius_max is None: + return radius, slice(None) + mask = radius <= radius_max + return radius[mask], mask diff --git a/tests/calculation/test_effective_coulomb.py b/tests/calculation/test_effective_coulomb.py index 5b8ea1a1..2003dfbe 100644 --- a/tests/calculation/test_effective_coulomb.py +++ b/tests/calculation/test_effective_coulomb.py @@ -353,6 +353,7 @@ def check_radial_plot_raises_error(effective_coulomb): with pytest.raises(exception.DataMismatch): effective_coulomb.plot(radius=...) + def test_plot_radial_with_cutoff(nonpolarized_crpar, Assert): effective_coulomb = nonpolarized_crpar graph = effective_coulomb.plot("U", radius_max=10) @@ -360,7 +361,7 @@ def test_plot_radial_with_cutoff(nonpolarized_crpar, Assert): series = graph[0] mask = effective_coulomb.ref.radial_data["radius"] < 10 Assert.allclose(series.x, effective_coulomb.ref.radial_data["radius"][mask]) - Assert.allclose(series.y, effective_coulomb.ref.radial_data["screened U"][0,mask]) + Assert.allclose(series.y, effective_coulomb.ref.radial_data["screened U"][0, mask]) assert series.label == "screened U" @@ -467,6 +468,26 @@ def test_plot_radial_and_frequency(effective_coulomb, Assert): assert series.label == label +def test_plot_radial_and_frequency_with_cutoff(nonpolarized_crpar, Assert): + omega_data = nonpolarized_crpar.ref.omega_data + radial_data = nonpolarized_crpar.ref.radial_data + graph = nonpolarized_crpar.plot("U", omega=..., radius_max=10) + mask = radial_data["radius"] < 10 + assert len(graph) == sum(mask) + assert graph.xlabel == "Im(ω) (eV)" + assert graph.ylabel == "Coulomb potential (eV)" + expected_labels = [ + label + for label, is_within_cutoff in zip(radial_data["label for both"], mask) + if is_within_cutoff + ] + expected_lines = omega_data["U for both"][mask].real + for series, expected_line, label in zip(graph, expected_lines, expected_labels): + Assert.allclose(series.x, omega_data["frequencies"].imag) + Assert.allclose(series.y, expected_line) + assert series.label == label + + def test_plot_radial_and_frequency_nondefault_radius(nonpolarized_crpar, Assert): with pytest.raises(exception.NotImplemented): nonpolarized_crpar.plot(omega=..., radius=np.array([1.0, 2.0])) From 5a03766855d6ddc153e085cb7d67be8e5398c712 Mon Sep 17 00:00:00 2001 From: Martin Schlipf Date: Thu, 5 Feb 2026 12:43:05 +0100 Subject: [PATCH 3/3] Add documentation --- src/py4vasp/_calculation/effective_coulomb.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/py4vasp/_calculation/effective_coulomb.py b/src/py4vasp/_calculation/effective_coulomb.py index a0146c5c..90a306e6 100644 --- a/src/py4vasp/_calculation/effective_coulomb.py +++ b/src/py4vasp/_calculation/effective_coulomb.py @@ -262,6 +262,8 @@ def to_graph( interpolated to the selected radii. radius_max + Maximum radius for radial-dependent plots. If set, all data for radii + greater than this value will be ignored. config Configuration for the analytic continuation of the frequency-dependent data.