Repository: wright-group/WrightTools

Issue:
`wt.close` method to close all open WrightTools objects

Discussion:
ksunden: Also cleans up a test which was closing files improperly, leading to them not being removed from tmp
ksunden: oops, commit message missed the w and said qt instead of wt, oh well


Patch:
diff --git a/WrightTools/__init__.py b/WrightTools/__init__.py
index e561ca0..667c127 100644
--- a/WrightTools/__init__.py
+++ b/WrightTools/__init__.py
@@ -19,6 +19,7 @@ from . import units
 from . import exceptions
 
 from ._open import *
+from ._close import *
 from .collection._collection import *
 from .data._data import *
 
diff --git a/WrightTools/_close.py b/WrightTools/_close.py
new file mode 100644
index 0000000..87ac5dc
--- /dev/null
+++ b/WrightTools/_close.py
@@ -0,0 +1,27 @@
+"""Function to close all open wt5 files."""
+
+
+# --- import -------------------------------------------------------------------------------------
+
+
+from . import _group as wt_group
+
+
+# --- define -------------------------------------------------------------------------------------
+
+
+__all__ = ["close"]
+
+
+# --- functions ----------------------------------------------------------------------------------
+
+
+def close():
+    """Close all open wt5 files.
+    
+    Warning
+    -------
+    This will make any open objects unusable and delete unsaved temporary files.
+    """
+    while len(wt_group.Group._instances) > 0:
+        wt_group.Group._instances.popitem()[1].close()
diff --git a/WrightTools/artists/_quick.py b/WrightTools/artists/_quick.py
index ffc6982..47078ac 100644
--- a/WrightTools/artists/_quick.py
+++ b/WrightTools/artists/_quick.py
@@ -156,6 +156,7 @@ def quick1D(
             if verbose:
                 print("image saved at", fpath)
             out.append(fpath)
+    chopped.close()
     return out
 
 
@@ -372,4 +373,5 @@ def quick2D(
             if verbose:
                 print("image saved at", fpath)
             out.append(fpath)
+    chopped.close()
     return out
diff --git a/WrightTools/data/_channel.py b/WrightTools/data/_channel.py
index 91a04ed..03d4973 100644
--- a/WrightTools/data/_channel.py
+++ b/WrightTools/data/_channel.py
@@ -152,14 +152,17 @@ class Channel(Dataset):
 
         factor : number (optional)
             Tolerance factor.  Default is 3.
-        replace : {'nan', 'mean', number} (optional)
+        replace : {'nan', 'mean', 'exclusive_mean', number} (optional)
             Behavior of outlier replacement. Default is nan.
 
             nan
                 Outliers are replaced by numpy nans.
 
             mean
-                Outliers are replaced by the mean of its neighborhood.
+                Outliers are replaced by the mean of its neighborhood, including itself.
+
+            exclusive_mean
+                Outilers are replaced by the mean of its neighborhood, not including itself.
 
             number
                 Array becomes given number.
@@ -177,6 +180,7 @@ class Channel(Dataset):
         warnings.warn("trim", category=wt_exceptions.EntireDatasetInMemoryWarning)
         outliers = []
         means = []
+        ex_means = []
         # find outliers
         for idx in np.ndindex(self.shape):
             slices = []
@@ -186,26 +190,33 @@ class Channel(Dataset):
                 slices.append(slice(start, stop, 1))
             neighbors = self[slices]
             mean = np.nanmean(neighbors)
+            sum_ = np.nansum(neighbors)
             limit = np.nanstd(neighbors) * factor
             if np.abs(self[idx] - mean) > limit:
                 outliers.append(idx)
                 means.append(mean)
+                # Note, "- 1" is to exclude the point itself, which is not nan, in order
+                # to enter this if block, as `np.abs(nan - mean)` is nan, which would
+                # evaluate to False
+                ex_means.append((sum_ - self[idx]) / (np.sum(~np.isnan(neighbors)) - 1))
+
         # replace outliers
         i = tuple(zip(*outliers))
-        if replace == "nan":
-            arr = self[:]
-            arr[i] = np.nan
-            self[:] = arr
-        elif replace == "mean":
-            arr = self[:]
-            arr[i] = means
-            self[:] = arr
-        elif isinstance(replace, numbers.Number):
-            arr = self[:]
-            arr[i] = replace
-            self[:] = arr
-        else:
-            raise KeyError("replace must be one of {nan, mean} or some number")
+
+        if len(i) == 0:
+            if verbose:
+                print("No outliers found")
+            return []
+
+        replace = {"nan": np.nan, "mean": means, "exclusive_mean": ex_means}.get(replace, replace)
+
+        # This may someday be available in h5py directly, but seems that day is not yet.
+        # This is annoying because it is the only reason we hold the whole set in memory.
+        # KFS 2019-03-21
+        arr = self[:]
+        arr[i] = replace
+        self[:] = arr
+
         # finish
         if verbose:
             print("%i outliers removed" % len(outliers))


Test patch:
diff --git a/tests/base.py b/tests/base.py
index e471faf..c58119a 100644
--- a/tests/base.py
+++ b/tests/base.py
@@ -5,6 +5,7 @@
 
 
 import os
+import pytest
 
 import WrightTools as wt
 from WrightTools import datasets
@@ -74,7 +75,8 @@ def test_nested():
     c = wt.Collection()
     cc = c.create_collection()
     assert c.fid.id == cc.fid.id
-    c.file.close()
+    c.close()
+    assert not os.path.isfile(c.filepath)
     assert c.id.valid == 0
     assert cc.id.valid == 0
 
@@ -87,6 +89,15 @@ def test_open_context():
     assert d.id.valid == 0
 
 
+def test_close():
+    d = wt.Data()
+    path = d.filepath
+    wt.close()
+    assert not os.path.isfile(path)
+    with pytest.raises(ValueError):
+        d.file
+
+
 if __name__ == "__main__":
     test_named_root_collection()
     test_named_root_data()
@@ -96,3 +107,4 @@ if __name__ == "__main__":
     test_tempfile_cleanup()
     test_nested()
     test_open_context()
+    test_close()
diff --git a/tests/data/trim.py b/tests/data/trim.py
index 2a94167..f096904 100644
--- a/tests/data/trim.py
+++ b/tests/data/trim.py
@@ -32,18 +32,23 @@ def test_trim_2Dgauss():
     d.create_channel("damaged1", arr2)
     d.create_channel("damaged2", arr2)
     d.create_channel("damaged3", arr2)
+    d.create_channel("damaged4", arr2)
     d.transform("x", "y")
     # trim
+    d.original.trim([2, 2], factor=2)
     d.damaged1.trim([2, 2], factor=2)
     d.damaged2.trim([2, 2], factor=2, replace="mean")
     d.damaged3.trim([2, 2], factor=2, replace=0.5)
+    d.damaged4.trim([2, 2], factor=2, replace="exclusive_mean")
     # now heal
     d.create_channel("healed_linear", d.damaged1[:])
     d.heal(channel="healed_linear", fill_value=0, method="linear")
     # check
-    assert np.allclose(d.original[:], d.healed_linear[:], rtol=1e-1, atol=1e-1)
-    assert np.allclose(d.original[:], d.damaged2[:], rtol=1e-1, atol=9e-1)
-    assert np.allclose(d.original[:], d.damaged3[:], rtol=1e-1, atol=5e-1)
+    np.testing.assert_allclose(d.original[:], d.original[:], rtol=1e-1, atol=1e-1)
+    np.testing.assert_allclose(d.original[:], d.healed_linear[:], rtol=1e-1, atol=1e-1)
+    np.testing.assert_allclose(d.original[:], d.damaged2[:], rtol=1e-1, atol=9e-1)
+    np.testing.assert_allclose(d.original[:], d.damaged3[:], rtol=1e-1, atol=5e-1)
+    np.testing.assert_allclose(d.original[:], d.damaged4[:], rtol=1e-1, atol=3e-1)
 
 
 def test_trim_3Dgauss():
@@ -67,7 +72,7 @@ def test_trim_3Dgauss():
     # trim
     d.damaged.trim([2, 2, 2], factor=2, replace="mean")
     # check
-    assert np.allclose(d.original[:], d.damaged[:], rtol=1e-1, atol=9e-1)
+    np.testing.assert_allclose(d.original[:], d.damaged[:], rtol=1e-1, atol=9e-1)
 
 
 if __name__ == "__main__":