Repository: wright-group/WrightTools

Issue:
syntax to prune data objects
remove all unused variables, remove all-but-first channel (as kwarg, off by default)


Discussion:

Resolving patch:
diff --git a/WrightTools/data/_data.py b/WrightTools/data/_data.py
--- a/WrightTools/data/_data.py
+++ b/WrightTools/data/_data.py
@@ -1220,6 +1220,38 @@ def print_tree(self, *, verbose=True):
         print("{0} ({1})".format(self.natural_name, self.filepath))
         self._print_branch("", depth=0, verbose=verbose)
 
+    def prune(self, keep_channels=True, *, verbose=True):
+        """Remove unused variables and (optionally) channels from the Data object.
+            
+        Unused variables are those that are not included in either axes or constants.
+        Unused channels are those not specified in keep_channels, or the first channel.
+
+        Parameters
+        ----------
+        keep_channels : boolean or int or str or tuple
+            If False, removes all but the first channel.
+            If int or str, removes all but that index/name channel.
+            If tuple, removes all channels except those in the tuple by index or name.
+            Default is True: do not delete channels
+        verbose : boolean
+            Toggle talkback. Default is True.
+        """
+        for v in self.variables:
+            for var in wt_kit.flatten_list([ax.variables for ax in self._axes + self._constants]):
+                if v == var:
+                    break
+            else:
+                self.remove_variable(v.natural_name, implied=False, verbose=verbose)
+        if keep_channels is not True:
+            try:
+                indexes = tuple(keep_channels)
+            except TypeError:
+                indexes = (keep_channels,)
+
+            for i, ch in enumerate(self.channels):
+                if i not in indexes and not ch.natural_name in indexes:
+                    self.remove_channel(ch.natural_name, verbose=verbose)
+
     def remove_channel(self, channel, *, verbose=True):
         """Remove channel from data.
 


Test patch:
diff --git a/tests/data/prune.py b/tests/data/prune.py
new file mode 100644
--- /dev/null
+++ b/tests/data/prune.py
@@ -0,0 +1,66 @@
+"""Test remove channel."""
+
+
+# --- import --------------------------------------------------------------------------------------
+
+
+import WrightTools as wt
+from WrightTools import datasets
+
+
+# --- test ----------------------------------------------------------------------------------------
+
+
+def test_prune_default():
+    p = datasets.PyCMDS.wm_w2_w1_000
+    data = wt.data.from_PyCMDS(p)
+    data.create_constant("d1")
+    num_channels = len(data.channels)
+    data.prune()
+    assert len(data.variables) == 4
+    assert set(data.variable_names) == {"wm", "w2", "w1", "d1"}
+    assert len(data.channels) == num_channels
+    data.close()
+
+
+def test_prune_false():
+    p = datasets.PyCMDS.wm_w2_w1_000
+    data = wt.data.from_PyCMDS(p)
+    data.create_constant("d1")
+    data.prune(False)
+    assert len(data.variables) == 4
+    assert set(data.variable_names) == {"wm", "w2", "w1", "d1"}
+    assert len(data.channels) == 1
+    assert data.channel_names[0] == "signal_diff"
+    data.close()
+
+
+def test_prune_number():
+    p = datasets.PyCMDS.wm_w2_w1_000
+    data = wt.data.from_PyCMDS(p)
+    data.create_constant("d1")
+    data.prune(2)
+    assert len(data.variables) == 4
+    assert set(data.variable_names) == {"wm", "w2", "w1", "d1"}
+    assert len(data.channels) == 1
+    assert data.channel_names[0] == "pyro1"
+    data.close()
+
+
+def test_prune_tuple():
+    p = datasets.PyCMDS.wm_w2_w1_000
+    data = wt.data.from_PyCMDS(p)
+    num_channels = len(data.channels)
+    data.prune(("pyro1", 3, 4), verbose=False)
+    assert len(data.variables) == 3
+    assert set(data.variable_names) == {"wm", "w2", "w1"}
+    assert len(data.channels) == 3
+    assert set(data.channel_names) == {"pyro1", "pyro2", "pyro3"}
+    data.close()
+
+
+if __name__ == "__main__":
+    test_prune_default()
+    test_prune_false()
+    test_prune_number()
+    test_prune_tuple()