internlm/SWE-Fixer-Train-110K huggingface dataset:
Repository: wright-group/WrightTools

Issue:
data.split:  mixes up unit conversions
This works:
```
>>> tune_test = wt.data.from_PyCMDS(path)
>>> tune_test.transform('wm', 'w3_Grating')
>>> tune_test = tune_test.split('w3_Grating', [30, 40])
```
But if the axes are transposed, it thinks 'wm' has units of `None`:
```
>>> tune_test.transform('w3_Grating', 'wm')
>>> tune_test = tune_test.split('w3_Grating', [30, 40])


  File "...\WrightTools\data\_data.py", line 1555, in split
    ax.convert(u)

  File "...\WrightTools\data\_axis.py", line 189, in convert
    raise wt_exceptions.UnitsError(valid, destination_units)
UnitsError(...), "expected units of ('wn', 'eV', 'meV', 'Hz', 'THz', 'GHz') (recieved None)
```
My guess is that it mixes up the units of 'w3_Grating' and 'wm'?



Discussion:
you say tune_test, but these look like autotune/motortune from the axes, correct?
I am having trouble reproducing this bug, trying with the `autotune.data` dataset which is in the tests directory


Resolving patch:
diff --git a/WrightTools/data/_data.py b/WrightTools/data/_data.py
--- a/WrightTools/data/_data.py
+++ b/WrightTools/data/_data.py
@@ -1544,13 +1544,15 @@ def split(
             try:
                 d.transform(*old_expr)
                 keep = []
+                keep_units = []
                 for ax in d.axes:
                     if ax.size > 1:
                         keep.append(ax.expression)
+                        keep_units.append(ax.units)
                     else:
                         d.create_constant(ax.expression, verbose=False)
                 d.transform(*keep)
-                for ax, u in zip(d.axes, old_units):
+                for ax, u in zip(d.axes, keep_units):
                     ax.convert(u)
             except IndexError:
                 continue


Test patch:
diff --git a/tests/data/split.py b/tests/data/split.py
--- a/tests/data/split.py
+++ b/tests/data/split.py
@@ -5,15 +5,17 @@
 # --- import --------------------------------------------------------------------------------------
 
 
-import pytest
 import numpy as np
 
 import WrightTools as wt
 from WrightTools import datasets
 
+import pathlib
 
 # --- test ----------------------------------------------------------------------------------------
 
+here = pathlib.Path(__file__).parent
+
 
 def test_split():
     p = datasets.PyCMDS.wm_w2_w1_000
@@ -155,8 +157,19 @@ def test_split_constants():
     split.close()
 
 
-if __name__ == "__main__":
+def test_autotune():
+    p = here / "test_data" / "autotune.data"
+    data = wt.data.from_PyCMDS(str(p))
+    data.transform("w2", "w2_BBO")
+    split = data.split("w2_BBO", [42])
+    split.close()
+    data.transform("w2_BBO", "w2")
+    split = data.split("w2_BBO", [42, 41.5])
+    split.close()
+    data.close()
+
 
+if __name__ == "__main__":
     test_split()
     test_split_edge()
     test_split_multiple()
@@ -168,3 +181,4 @@ def test_split_constants():
     test_split_expression()
     test_split_hole()
     test_split_constants()
+    test_autotune()
Repository: wright-group/WrightTools

Issue:
Document default_aspect in create_figure
missing from docstring


Discussion:

Resolving patch:
diff --git a/WrightTools/artists/_helpers.py b/WrightTools/artists/_helpers.py
--- a/WrightTools/artists/_helpers.py
+++ b/WrightTools/artists/_helpers.py
@@ -283,7 +283,11 @@ def create_figure(
         aspects larger than 1 will be taller than wide and vice-versa for
         aspects smaller than 1. You may only define the aspect for one subplot
         in each row. If no aspect is defined for a particular row, the leftmost
-        subplot will have an aspect of ``default_aspect``. Default is [].
+        subplot will have an aspect of ``default_aspect``. Default is given by
+        default_aspect kwarg.
+    default_aspect : number (optional)
+        Default aspect of left-most column, if no aspect is defined for a
+        given row.
 
     Returns
     -------


Test patch:
Repository: wright-group/WrightTools

Issue:
additional constants documentation
- [ ] label
- [ ] format spec
- [ ] round spec


Discussion:

Resolving patch:
diff --git a/examples/custom_fig.py b/examples/custom_fig.py
new file mode 100644
--- /dev/null
+++ b/examples/custom_fig.py
@@ -0,0 +1,83 @@
+# -*- coding: utf-8 -*-
+"""
+Custom Figure
+=============
+
+Example of custom figure layout, beautification, and saving.
+
+"""
+
+import matplotlib.pyplot as plt
+
+import numpy as np
+
+import WrightTools as wt
+from WrightTools import datasets
+
+
+# obtain and process data
+p = datasets.wt5.v1p0p1_MoS2_TrEE_movie
+data = wt.open(p)
+data.level(0, 2, -3)
+data.convert("eV", convert_variables=True, verbose=False)
+data.smooth([2, 2, 2])
+data.ai0.symmetric_root(2)
+data.ai0.normalize()
+data.ai0.clip(min=0, replace="value")
+# chop out data of interest
+d2_vals = [-50, -500]
+w2_vals = [1.7, 1.8, 1.9, 2.0]
+wigners = [data.chop("w1=wm", "d2", at={"w2": [w2, "eV"]})[0] for w2 in w2_vals]
+traces1 = [
+    data.chop("w1=wm", at={"w2": [w2, "eV"], "d2": [d2_vals[0], "fs"]})[0] for w2 in w2_vals
+]
+traces2 = [
+    data.chop("w1=wm", at={"w2": [w2, "eV"], "d2": [d2_vals[1], "fs"]})[0] for w2 in w2_vals
+]
+tracess = [traces1, traces2]
+# prepare spine colors
+wigner_colors = ["C0", "C1", "C2", "C3"]
+trace_colors = ["#FE4EDA", "#00B7EB"]
+# prepare figure gridspec
+cols = [1, 1, "cbar"]
+aspects = [[[0, 0], .3]]
+fig, gs = wt.artists.create_figure(
+    width="double", cols=cols, nrows=3, aspects=aspects, wspace=.35, hspace=.35
+)
+# plot wigners
+indxs = [(row, col) for row in range(1, 3) for col in range(2)]
+for indx, wigner, color in zip(indxs, wigners, wigner_colors):
+    ax = plt.subplot(gs[indx])
+    ax.pcolor(wigner, vmin=0, vmax=1)  # global colormpa
+    ax.contour(wigner)  # local contours
+    ax.grid()
+    wt.artists.set_ax_spines(ax=ax, c=color)
+    # set title as value of w2
+    wigner.constants[0].format_spec = ".2f"
+    wigner.round_spec = -1
+    wt.artists.corner_text(wigner.constants[0].label, ax=ax)
+    # plot overlines
+    for d2, t_color in zip(d2_vals, trace_colors):
+        ax.axhline(d2, color=t_color, alpha=.5, linewidth=6)
+    # plot w2 placement
+    ax.axvline(wigner.w2.points, color="grey", alpha=.75, linewidth=6)
+# plot traces
+indxs = [(0, col) for col in range(2)]
+for indx, color, traces in zip(indxs, trace_colors, tracess):
+    ax = plt.subplot(gs[indx])
+    for trace, w_color in zip(traces, wigner_colors):
+        ax.plot(trace, color=w_color, linewidth=1.5)
+    ax.grid()
+    ax.set_xlim(trace.axes[0].min(), trace.axes[0].max())
+    wt.artists.set_ax_spines(ax=ax, c=color)
+# plot colormap
+cax = plt.subplot(gs[1:3, -1])
+ticks = np.linspace(data.ai0.min(), data.ai0.max(), 11)
+wt.artists.plot_colorbar(cax=cax, label="amplitude", cmap="default", ticks=ticks)
+# set axis labels
+wt.artists.set_fig_labels(xlabel=data.w1__e__wm.label, ylabel=data.d2.label, col=slice(0, 1))
+# ylabel of zeroth row
+ax = plt.subplot(gs[0, 0])
+ax.set_ylabel("amplitude")
+# saving the figure as a png
+wt.artists.savefig("custom_fig.png", fig=fig, close=False)


Test patch:
Repository: wright-group/WrightTools

Issue:
clip and heal example
in docs gallery
Test heal method



Discussion:

Resolving patch:
diff --git a/examples/heal.py b/examples/heal.py
new file mode 100644
--- /dev/null
+++ b/examples/heal.py
@@ -0,0 +1,39 @@
+"""
+Heal
+====
+
+An example of how heal works.
+"""
+
+import numpy as np
+from matplotlib import pyplot as plt
+import WrightTools as wt
+
+# create original arrays
+x = np.linspace(-3, 3, 31)[:, None]
+y = np.linspace(-3, 3, 31)[None, :]
+arr = np.exp(-1 * (x ** 2 + y ** 2))
+# create damaged array
+arr2 = arr.copy()
+np.random.seed(11)  # set seed for reproducibility
+arr2[np.random.random(arr2.shape) < .2] = np.nan
+# create data object
+d = wt.data.Data()
+d.create_variable("x", values=x)
+d.create_variable("y", values=y)
+d.create_channel("original", arr)
+d.create_channel("damaged", arr2)
+d.create_channel("healed", arr2)
+d.transform("x", "y")
+# heal
+d.heal(channel="healed")
+# create figure
+fig, gs = wt.artists.create_figure(cols=[1, 1, 1])
+for i in range(3):
+    ax = plt.subplot(gs[i])
+    ax.pcolor(d, channel=i)
+    ax.set_title(d.channel_names[i])
+ticks = [-2, 0, 2]
+wt.artists.set_fig_labels(
+    xlabel=d.axes[0].label, ylabel=d.axes[1].label, xticks=ticks, yticks=ticks
+)


Test patch:
diff --git a/tests/data/heal.py b/tests/data/heal.py
new file mode 100644
--- /dev/null
+++ b/tests/data/heal.py
@@ -0,0 +1,48 @@
+#! /usr/bin/env python3
+"""Test data.heal."""
+
+
+# --- import --------------------------------------------------------------------------------------
+
+
+import numpy as np
+
+import WrightTools as wt
+
+import matplotlib.pyplot as plt
+
+
+# --- test ----------------------------------------------------------------------------------------
+
+
+def test_heal_gauss():
+    # create original arrays
+    x = np.linspace(-3, 3, 31)[:, None]
+    y = np.linspace(-3, 3, 31)[None, :]
+    arr = np.exp(-1 * (x ** 2 + y ** 2))
+    # create damaged array
+    arr2 = arr.copy()
+    np.random.seed(11)  # set seed for reproducibility
+    arr2[np.random.random(arr2.shape) < .2] = np.nan
+    # create data object
+    d = wt.data.Data()
+    d.create_variable("x", values=x)
+    d.create_variable("y", values=y)
+    d.create_channel("original", arr)
+    d.create_channel("damaged", arr2)
+    d.create_channel("healed_linear", arr2)
+    d.create_channel("healed_nearest", arr2)
+    d.create_channel("healed_cubic", arr2)
+    d.transform("x", "y")
+    # heal
+    d.heal(channel="healed_linear", fill_value=0, method="linear")
+    d.heal(channel="healed_nearest", fill_value=0, method="nearest")
+    d.heal(channel="healed_cubic", fill_value=0, method="cubic")
+    # check
+    assert np.allclose(d.original[:], d.healed_linear[:], rtol=1e-1, atol=1e-1)
+    assert np.allclose(d.original[:], d.healed_nearest[:], rtol=5e-1, atol=5e-1)
+    assert np.allclose(d.original[:], d.healed_cubic[:], rtol=1e-1, atol=1e-1)
+
+
+if __name__ == "__main__":
+    test_heal_gauss()
Repository: wright-group/WrightTools

Issue:
map_variable example in gallery



Discussion:

Resolving patch:
diff --git a/examples/map-variable.py b/examples/map-variable.py
new file mode 100644
--- /dev/null
+++ b/examples/map-variable.py
@@ -0,0 +1,40 @@
+"""
+Map-Variable
+============
+
+An example of how map-variable works.
+"""
+
+import numpy as np
+from matplotlib import pyplot as plt
+import WrightTools as wt
+
+# create original arrays
+x = np.linspace(1, 5, 11)[:, None]
+y = np.linspace(1, 5, 11)[None, :]
+arr = np.exp(-1 * (((x - 3) / .6) ** 2 + ((y - 3) / .6) ** 2))
+# create data object
+d = wt.data.Data(name="original")
+d.create_variable("x", values=x)
+d.create_variable("y", values=y)
+d.create_channel("z", arr)
+d.transform("x", "y")
+# create new data objects
+pointsx = np.linspace(x.min(), x.max(), 31)  # linear spacing
+pointsy = points = np.logspace(0, .7, 11)  # log spacing
+dx = d.map_variable("x", points=pointsx)  # just linear along x
+dy = d.map_variable("y", points=pointsy)  # just log along y
+dxy = dx.map_variable("y", points=pointsy)  # linear along x and log along y
+ds = [d, dx, dy, dxy]
+# create figure
+fig, gs = wt.artists.create_figure(width="double", cols=[1, 1, 1, 1])
+for i, d in enumerate(ds):
+    ax = plt.subplot(gs[i])
+    ax.pcolor(d)
+    ax.set_title(d.natural_name)
+    ax.set_xlim(1, 5)
+    ax.set_ylim(1, 5)
+ticks = [1, 3, 5]
+wt.artists.set_fig_labels(
+    xlabel=d.axes[0].label, ylabel=d.axes[1].label, xticks=ticks, yticks=ticks
+)


Test patch:
Repository: wright-group/WrightTools

Issue:
BUG: Add support for lists of pathlib objects
https://github.com/wright-group/WrightTools/blob/671b66ff8cfdc66e56fd5bd225ab6fefa8f14bc9/WrightTools/_group.py#L70-L85

See #639, Just extending to lists


Discussion:
Should be a matter of changing, the isinstance check to os.PathLike from str. `np.array(..., dtype="S")` works with pathlikes (well, at least pathlib objects, probably just calls str, not os.fspath)


Resolving patch:
diff --git a/WrightTools/_group.py b/WrightTools/_group.py
--- a/WrightTools/_group.py
+++ b/WrightTools/_group.py
@@ -70,7 +70,7 @@ def __init__(self, file=None, parent=None, name=None, **kwargs):
             try:
                 if isinstance(value, pathlib.Path):
                     value = str(value)
-                elif isinstance(value, list) and len(value) > 0 and isinstance(value[0], str):
+                elif isinstance(value, list) and len(value) > 0 and isinstance(value[0], os.PathLike):
                     value = np.array(value, dtype="S")
                 elif sys.version_info > (3, 6):
                     try:


Test patch:
Repository: wright-group/WrightTools

Issue:
Documentation for create_figure should reflect subclassed returns
https://github.com/wright-group/WrightTools/blob/cecddd9025678f390b55df0498bd27530faa37c4/WrightTools/artists/_helpers.py#L291

Currently we say we return matplotlib.figure.Figure/matplotlib.gridspec.GridSpec

We actually return WrightTools.artists.Figure and WrightTools.artists.GridSpec objects, which are subclasses of the matplotlib versions which have knowledge of WrightTools.artists.Axes, which know how to plot data objects. (Granted, WrightTools.artists.GridSpec currently does exactly nothing other than change the namespace)


Discussion:

Resolving patch:
diff --git a/WrightTools/artists/_helpers.py b/WrightTools/artists/_helpers.py
--- a/WrightTools/artists/_helpers.py
+++ b/WrightTools/artists/_helpers.py
@@ -292,7 +292,7 @@ def create_figure(
     Returns
     -------
     tuple
-        (matplotlib.figure.Figure, matplotlib.gridspec.GridSpec). GridSpec
+        (WrightTools.artists.Figure, WrightTools.artists.GridSpec). GridSpec
         contains SubplotSpec objects that can have axes placed into them.
         The SubplotSpec objects can be accessed through indexing: [row, col].
         Slicing works, for example ``cax = plt.subplot(gs[:, -1])``. See


Test patch:
Repository: wright-group/WrightTools

Issue:
syntax for opening wt5 files in a context manager
useful for quickly opening within a larger project

```python
with wt.open('data.wt5') as data:
    do_something
```


Discussion:

Resolving patch:
diff --git a/WrightTools/_group.py b/WrightTools/_group.py
--- a/WrightTools/_group.py
+++ b/WrightTools/_group.py
@@ -99,6 +99,12 @@ def __init__(self, file=None, parent=None, name=None, **kwargs):
                 parent.attrs["item_names"], self.natural_name.encode()
             )
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args, **kwargs):
+        self.close()
+
     def __getattr__(self, key):
         """Gets called if attribute not in self.__dict__.
 


Test patch:
diff --git a/tests/base.py b/tests/base.py
--- a/tests/base.py
+++ b/tests/base.py
@@ -7,6 +7,7 @@
 import os
 
 import WrightTools as wt
+from WrightTools import datasets
 
 
 # --- test ----------------------------------------------------------------------------------------
@@ -78,6 +79,14 @@ def test_nested():
     assert cc.id.valid == 0
 
 
+def test_open_context():
+    p = datasets.wt5.v1p0p1_MoS2_TrEE_movie
+    with wt.open(p) as d:
+        assert os.path.isfile(d.filepath)
+    assert not os.path.isfile(d.filepath)
+    assert d.id.valid == 0
+
+
 if __name__ == "__main__":
     test_named_root_collection()
     test_named_root_data()
@@ -86,3 +95,4 @@ def test_nested():
     test_single_instance_data()
     test_tempfile_cleanup()
     test_nested()
+    test_open_context()
Repository: wright-group/WrightTools

Issue:
data.prune() does not work correctly when passing a bare string
When I set the kwarg `keep_channels='mean'` the method does not have the desired result. But when I pass in `('mean')` the method has the desired result.


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
@@ -1365,6 +1365,8 @@ def prune(self, keep_channels=True, *, verbose=True):
                 self.remove_variable(v.natural_name, implied=False, verbose=verbose)
         if keep_channels is not True:
             try:
+                if isinstance(keep_channels, str):
+                    raise TypeError
                 indexes = tuple(keep_channels)
             except TypeError:
                 indexes = (keep_channels,)


Test patch:
diff --git a/tests/data/prune.py b/tests/data/prune.py
--- a/tests/data/prune.py
+++ b/tests/data/prune.py
@@ -47,6 +47,18 @@ def test_prune_number():
     data.close()
 
 
+def test_prune_string():
+    p = datasets.PyCMDS.wm_w2_w1_000
+    data = wt.data.from_PyCMDS(p)
+    data.create_constant("d1")
+    data.prune("pyro1")
+    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)
@@ -63,4 +75,5 @@ def test_prune_tuple():
     test_prune_default()
     test_prune_false()
     test_prune_number()
+    test_prune_string()
     test_prune_tuple()
Repository: wright-group/WrightTools

Issue:
Recover contour for Quick2D
https://github.com/wright-group/WrightTools/blob/86ebafa5910fb7614552feec13279f53299e13b6/WrightTools/artists/_quick.py#L312


Discussion:
Consider looking at old implementation:
https://github.com/wright-group/WrightTools/blob/806c92da0bea2148a9e7953e4864962c6c451de5/WrightTools/artists.py#L2416

Resolving patch:
diff --git a/WrightTools/artists/_quick.py b/WrightTools/artists/_quick.py
--- a/WrightTools/artists/_quick.py
+++ b/WrightTools/artists/_quick.py
@@ -309,7 +309,23 @@ def quick2D(
             ax.contourf(d, channel=channel_index, cmap=cmap, levels=levels)
         # contour lines ---------------------------------------------------------------------------
         if contours:
-            raise NotImplementedError
+            # get contour levels
+            # force top and bottom contour to be data range then clip them out
+            if channel.signed:
+                if contours_local:
+                    limit = channel.mag()
+                else:
+                    limit = data_channel.mag()
+                contour_levels = np.linspace(
+                    -limit + channel.null, limit + channel.null, contours + 2
+                )[1:-1]
+            else:
+                if contours_local:
+                    limit = channel.max()
+                else:
+                    limit = data_channel.max()
+                contour_levels = np.linspace(channel.null, limit, contours + 2)[1:-1]
+            ax.contour(d, channel=channel_index, levels=contour_levels)
         # decoration ------------------------------------------------------------------------------
         plt.xticks(rotation=45, fontsize=14)
         plt.yticks(fontsize=14)


Test patch:
Repository: wright-group/WrightTools

Issue:
close data objects in quick1D and quick2D after use



Discussion:
see also #727 
May also apply to interact, but since the object is desired longer, may need to rely on GC in that case, still (though maybe we can catch some signal from mpl that the figure was closed)

Resolving patch:
diff --git a/WrightTools/__init__.py b/WrightTools/__init__.py
--- a/WrightTools/__init__.py
+++ b/WrightTools/__init__.py
@@ -19,6 +19,7 @@
 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
--- /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
--- 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


Test patch:
diff --git a/tests/base.py b/tests/base.py
--- 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 @@ def test_open_context():
     test_tempfile_cleanup()
     test_nested()
     test_open_context()
+    test_close()