More generic configuration filter

ipsuite/configuration_selection/__init__.py

@@ -1,7 +1,7 @@
 """Configuration Selection Nodes."""
 
 from ipsuite.configuration_selection.base import ConfigurationSelection
-from ipsuite.configuration_selection.filter import FilterOutlier
+from ipsuite.configuration_selection.filter import PropertyFilter
 from ipsuite.configuration_selection.index import IndexSelection
 from ipsuite.configuration_selection.kernel import KernelSelection
 from ipsuite.configuration_selection.random import RandomSelection
@@ -21,5 +21,5 @@
     "IndexSelection",
     "ThresholdSelection",
     "SplitSelection",
-    "FilterOutlier",
+    "PropertyFilter",
 ]

ipsuite/configuration_selection/filter.py

@@ -1,70 +1,111 @@
 import typing as t
 
+import ase
 import matplotlib.pyplot as plt
 import numpy as np
 import zntrack
 
-from ipsuite import base
+from ipsuite.configuration_selection import ConfigurationSelection
 
 
-class FilterOutlier(base.ProcessAtoms):
-    """Remove outliers from the data based on a given property.
+def direct_cutoff(values, threshold, cutoffs):
+    # Filtering the direct cutoff values
+    if cutoffs is None:
+        raise ValueError(
+            "cutoffs have to be specified for using the direct cutoff filter."
+        )
+    return (cutoffs[0], cutoffs[1])
+
+
+def cutoff_around_mean(values, threshold, cutoffs):
+    # Filtering in multiples of the standard deviation around the mean.
+    mean = np.mean(values)
+    std = np.std(values)
+
+    upper_cutoff = mean + threshold * std
+    lower_cutoff = mean - threshold * std
+    return (lower_cutoff, upper_cutoff)
+
+
+CUTOFF = {"direct": direct_cutoff, "around_mean": cutoff_around_mean}
+
+
+class PropertyFilter(ConfigurationSelection):
+    """Filter structures from the dataset based on a given property.
 
     Attributes
     ----------
     key : str, default="energy"
         The property to filter on.
-    threshold : float, default=3
-        The threshold for filtering in units of standard deviations.
+    cutoff_type : {"direct", "around_mean"}, default="around_mean"
+        Defines the cutoff type.
     direction : {"above", "below", "both"}, default="both"
         The direction to filter in.
+    threshold : float, default=3
+        The threshold for filtering in units of standard deviations.
+    cutoffs : list(float), default=None
+        Lower and upper cutoff.
     """
 
     key: str = zntrack.params("energy")
-    threshold: float = zntrack.params(3)
+    cutoff_type: t.Literal["direct", "around_mean"] = zntrack.params("around_mean")
     direction: t.Literal["above", "below", "both"] = zntrack.params("both")
+    threshold: float = zntrack.params(3)
+    cutoffs: t.Union[t.List[float], None] = zntrack.params(None)
+
+    def select_atoms(self, atoms_lst: t.List[ase.Atoms]) -> t.List[int]:
+        values = [atoms.calc.results[self.key] for atoms in atoms_lst]
 
-    filtered_indices: list = zntrack.outs()
-    histogram: str = zntrack.outs_path(zntrack.nwd / "histogram.png")
+        # get maximal atomic value per struckture
+        if isinstance(values[0], np.ndarray):
+            if values[0].ndim == 2:
+                # calculates the maximal magnetude of atomic cartesian property
+                values = [
+                    np.max(np.linalg.norm(value, axis=1), axis=0) for value in values
+                ]
+            elif values[0].ndim == 1:
+                # calculates the maximal atomic property
+                values = [np.max(value, axis=0) for value in values]
 
-    def run(self):
-        values = [x.calc.results[self.key] for x in self.data]
-        mean = np.mean(values)
-        std = np.std(values)
+        lower_limit, upper_limit = CUTOFF[self.cutoff_type](
+            values,
+            self.threshold,
+            self.cutoffs,
+        )
 
         if self.direction == "above":
-            self.filtered_indices = [
-                i for i, x in enumerate(values) if x > mean + self.threshold * std
-            ]
+            selection = [i for i, x in enumerate(values) if x > upper_limit]
         elif self.direction == "below":
-            self.filtered_indices = [
-                i for i, x in enumerate(values) if x < mean - self.threshold * std
-            ]
+            selection = [i for i, x in enumerate(values) if x < lower_limit]
         else:
-            self.filtered_indices = [
-                i
-                for i, x in enumerate(values)
-                if x > mean + self.threshold * std or x < mean - self.threshold * std
+            selection = [
+                i for i, x in enumerate(values) if x > lower_limit and x < upper_limit
             ]
 
+        return selection
+
+    def _get_plot(self, atoms_lst: t.List[ase.Atoms], indices: t.List[int]):
+        values = [atoms.calc.results[self.key] for atoms in atoms_lst]
+
+        # get maximal atomic value per struckture
+        if isinstance(values[0], np.ndarray):
+            if values[0].ndim == 2:
+                # calculates the maximal magnetude of atomic cartesian property
+                values = [
+                    np.max(np.linalg.norm(value, axis=1), axis=0) for value in values
+                ]
+            elif values[0].ndim == 1:
+                # calculates the maximal atomic property
+                values = [np.max(value, axis=0) for value in values]
+
         fig, ax = plt.subplots(3, figsize=(10, 10))
         ax[0].hist(values, bins=100)
         ax[0].set_title("All")
-        ax[1].hist(
-            [values[i] for i in range(len(values)) if i not in self.filtered_indices],
+        ax[1].hist([values[i] for i in indices], bins=100)
+        ax[1].set_title("Selected")
+        ax[2].hist(
+            [values[i] for i in range(len(values)) if i not in indices],
             bins=100,
         )
-        ax[1].set_title("Filtered")
-        ax[2].hist([values[i] for i in self.filtered_indices], bins=100)
         ax[2].set_title("Excluded")
-        fig.savefig(self.histogram, bbox_inches="tight")
-
-    @property
-    def atoms(self):
-        return [
-            self.data[i] for i in range(len(self.data)) if i not in self.filtered_indices
-        ]
-
-    @property
-    def excluded_atoms(self):
-        return [self.data[i] for i in self.filtered_indices]
+        fig.savefig(self.img_selection, bbox_inches="tight")

ipsuite/nodes.py

@@ -26,7 +26,7 @@ class _Nodes:
     )
     UniformTemporalSelection = "ipsuite.configuration_selection.UniformTemporalSelection"
     ThresholdSelection = "ipsuite.configuration_selection.ThresholdSelection"
-    FilterOutlier = "ipsuite.configuration_selection.FilterOutlier"
+    PropertyFilter = "ipsuite.configuration_selection.PropertyFilter"
     BatchKernelSelection = "ipsuite.models.apax.BatchKernelSelection"
 
     # Configuration Comparison

tests/integration/configuration_selection/test_index.py

@@ -142,8 +142,12 @@ def test_exclude_configurations_list(proj_path, traj_file):
 def test_filter_outlier(proj_path, traj_file):
     with ips.Project() as project:
         data = ips.AddData(file=traj_file)
-        filtered_data = ips.configuration_selection.FilterOutlier(
-            data=data.atoms, key="energy", threshold=1, direction="both"
+        filtered_data = ips.configuration_selection.PropertyFilter(
+            data=data.atoms,
+            key="energy",
+            cutoff_type="around_mean",
+            threshold=1,
+            direction="both",
         )
 
     project.run()

tests/unit_tests/configuration_selection/test_filter.py

@@ -0,0 +1,36 @@
+import numpy as np
+import pytest
+
+from ipsuite.configuration_selection import PropertyFilter
+
+
+@pytest.mark.parametrize(
+    "key, cutoff_type, direction, cutoffs",
+    [
+        ("forces", "direct", "both", [7, 13]),
+        ("forces", "direct", "both", None),
+        ("forces", "around_mean", "both", None),
+    ],
+)
+def test_get_selected_atoms(atoms_list, key, cutoff_type, direction, cutoffs):
+    for idx, atoms in enumerate(atoms_list):
+        atoms.calc.results[key] = np.array([[idx, 0, 0], [0, 0, 0]])
+
+    filter = PropertyFilter(
+        key=key,
+        cutoff_type=cutoff_type,
+        direction=direction,
+        data=None,
+        cutoffs=cutoffs,
+        threshold=0.4,
+    )
+
+    if "direct" in cutoff_type and cutoffs is None:
+        with pytest.raises(ValueError):
+            selected_atoms = filter.select_atoms(atoms_list)
+    else:
+        test_selection = [8, 9, 10, 11, 12]
+        selected_atoms = filter.select_atoms(atoms_list)
+        assert isinstance(selected_atoms, list)
+        assert len(set(selected_atoms)) == 5
+        assert selected_atoms == test_selection