Feat: Coulomb radius max

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,10 @@ def to_graph(
             are used. You can also provide specific radii, then the data  will be
             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.
             Use this if you need to adjust the parameters of the analytic continuation.
@@ -271,25 +276,26 @@ 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:
-            if radius is not ...:
+    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 ... 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)
-        if radius is not None:
+            return _OmegaPlotter(omega_in, omega, config, positions, radius_max)
+        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)
@@ -415,7 +421,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:
@@ -425,7 +431,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)
@@ -442,51 +450,51 @@ 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)
 
 
 class _RadialPlotter:
     xlabel = "Radius (Å)"
 
-    def __init__(self, positions, radius_out):
-        if not positions:
-            raise exception.DataMismatch("The output does not contain position data.")
-        self.radius_in = self._transform_positions_to_radial(positions)
+    def __init__(self, positions, radius_out, radius_max):
+        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)
 
     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]
@@ -508,3 +516,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

tests/calculation/test_effective_coulomb.py

@@ -354,6 +354,17 @@ def check_radial_plot_raises_error(effective_coulomb):
         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):
     effective_coulomb = collinear_crpa
@@ -384,38 +395,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_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)
-        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_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
 
@@ -439,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]))