add FilterOutlier (#229)

Author: PythonFZ

Diff for: ipsuite/configuration_selection/__init__.py

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

Diff for: ipsuite/configuration_selection/filter.py

@@ -0,0 +1,70 @@
+import typing as t
+
+import matplotlib.pyplot as plt
+import numpy as np
+import zntrack
+
+from ipsuite import base
+
+
+class FilterOutlier(base.ProcessAtoms):
+    """Remove outliers from the data 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.
+    direction : {"above", "below", "both"}, default="both"
+        The direction to filter in.
+    """
+
+    key: str = zntrack.params("energy")
+    threshold: float = zntrack.params(3)
+    direction: t.Literal["above", "below", "both"] = zntrack.params("both")
+
+    filtered_indices: list = zntrack.outs()
+    histogram: str = zntrack.outs_path(zntrack.nwd / "histogram.png")
+
+    def run(self):
+        values = [x.calc.results[self.key] for x in self.data]
+        mean = np.mean(values)
+        std = np.std(values)
+
+        if self.direction == "above":
+            self.filtered_indices = [
+                i for i, x in enumerate(values) if x > mean + self.threshold * std
+            ]
+        elif self.direction == "below":
+            self.filtered_indices = [
+                i for i, x in enumerate(values) if x < mean - self.threshold * std
+            ]
+        else:
+            self.filtered_indices = [
+                i
+                for i, x in enumerate(values)
+                if x > mean + self.threshold * std or x < mean - self.threshold * std
+            ]
+
+        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],
+            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]

Diff for: ipsuite/nodes.py

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

Diff for: tests/integration/configuration_selection/test_index.py

@@ -137,3 +137,16 @@ def test_exclude_configurations_list(proj_path, traj_file):
     assert train_data[0].selected_configurations == {"AddData": [5, 6, 7, 8, 9]}
     assert test_data[0].selected_configurations == {"AddData": [0, 1, 2, 3, 4]}
     assert validation_data.selected_configurations == {"AddData": [10, 11, 12, 13, 14]}
+
+
+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"
+        )
+
+    project.run()
+
+    filtered_data.load()
+    assert len(filtered_data.atoms) == 13