Merge pull request #1203 from wright-group/norm_for_each

norm_for_each

Author: rpm4

CHANGELOG.md

@@ -5,9 +5,17 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/).
 
 ## [Unreleased]
 
+### Added
+- `Data.norm_for_each`: easier norm-by-axis syntax
+- `kit.from_list_of_objects`: convenience method for grabbing an channel/variable/axis/etc. using multiple specifiers
+- `Data.get_var`: get variable object by index, name, or variable itself
+- `Data.get_axis`: get axis object by index, name, or variable itself
+- `Data.get_channel`: get channel object by index, name, or variable itself
+
 ### Fixed
 - quick2D works with contours on
 - quick2D works with `pixelated=False` (i.e. contourf)
+- `data.from_LabRAM`:  small tweaks to deal with newly emerged test failure in python 3.11 tests
 
 ### Changed
 - default colorbar is sampled at 4096 points (was 256)

WrightTools/data/_data.py

@@ -738,6 +738,68 @@ def gradient(self, axis, *, channel=0):
         else:
             new[:] = np.gradient(channel[:], self[axis].points, axis=axis_index)
 
+    def get_var(self, hint: str | int | Variable) -> Variable:
+        return wt_kit.from_list_of_objects(self.variables, self.variable_names, hint)
+
+    def get_channel(self, hint: str | int | Channel) -> Channel:
+        return wt_kit.from_list_of_objects(self.channels, self.channel_names, hint)
+
+    def get_axis(self, hint: str | int | Axis) -> Axis:
+        return wt_kit.from_list_of_objects(self.axes, self.axis_expressions, hint)
+
+    def norm_for_each(
+        self,
+        var: str | Variable | int,
+        channel: str | Channel | int = 0,
+        new_channel: dict = {},
+    ):
+        """normalize the data for each var slice
+        var array must at least one trivial dimension (or else norm will return an array of ones)
+
+        Parameters
+        ----------
+        var : str, int, or WrightTools.data.Variable
+            variable to apply normalization at each unique point.
+        channel : str, int or WrightTools.data.Channel (default 0)
+            channel to apply normalization.  Channel should have more non-trivial dimensions than variable
+        new_channel : dict
+            Default is empty, and channel is overwriten with norm values.
+            If not empty, a new channel will be created.
+            Fields (e.g. name) can be supplied by supplying a dictionary (consult `Data.create_channel`).
+
+
+        Examples
+        --------
+        import WrightTools.datasets as ds
+        import WrightTools as wt
+
+        d = wt.open(ds.wt5.v1p0p1_MoS2_TrEE_movie)
+        d.norm_for_each("d2", "ai0")  # equivalent to d.ai0[:] /= d.ai0[:].max(axis=(0,1))[None, None, :]
+
+        """
+        variable = self.get_var(var)
+        channel = self.get_channel(channel)
+        trivial = {i for i, si in enumerate(variable.shape) if si == 1}
+        if not trivial:
+            raise wt_exceptions.WrightToolsWarning(
+                f"Variable {variable.natural_name} and Channel {channel.natural_name} have the same shape {variable.shape}. "
+                + "Produces a ones array channel."
+            )
+        # nontrivial = tuple({i for i in range(self.ndim)} - trivial)
+        trivial = tuple(trivial)
+        norm_vals = np.expand_dims(np.nanmax(channel[:], axis=trivial), trivial)
+        new = (channel[:] - channel.null) / (norm_vals - channel.null)
+        if new_channel:
+            self.create_channel(
+                new_channel.pop("name", f"{channel.natural_name}_{variable.natural_name}_norm"),
+                values=new,
+                **new_channel,
+            )
+        else:
+            channel[:] = new
+            channel.null = 0
+        return
+
     def moment(self, axis, channel=0, moment=1, *, resultant=None):
         """Take the nth moment the dataset along one axis, adding lower rank channels.
 

WrightTools/data/_labram.py

@@ -1,7 +1,6 @@
 # --- import --------------------------------------------------------------------------------------
 
 
-import os
 import pathlib
 import warnings
 import time
@@ -46,7 +45,6 @@ def from_LabRAM(filepath, name=None, parent=None, verbose=True) -> Data:
         New data object(s).
     """
     # parse filepath
-    filestr = os.fspath(filepath)
     filepath = pathlib.Path(filepath)
 
     if not ".txt" in filepath.suffixes:
@@ -55,7 +53,7 @@ def from_LabRAM(filepath, name=None, parent=None, verbose=True) -> Data:
     if not name:
         name = filepath.name.split(".")[0]
 
-    kwargs = {"name": name, "kind": "Horiba", "source": filestr}
+    kwargs = {"name": name, "kind": "Horiba", "source": str(filepath)}
 
     # create data
     if parent is None:
@@ -64,7 +62,7 @@ def from_LabRAM(filepath, name=None, parent=None, verbose=True) -> Data:
         data = parent.create_data(**kwargs)
 
     ds = DataSource(None)
-    f = ds.open(filestr, "rt", encoding="ISO-8859-1")
+    f = ds.open(str(filepath), "rt", encoding="ISO-8859-1")
 
     # header
     header = {}
@@ -110,7 +108,7 @@ def from_LabRAM(filepath, name=None, parent=None, verbose=True) -> Data:
     # dimensionality
     extra_dims = np.isnan(wm).sum()
 
-    if extra_dims == 0:  # single spectrum; we extracted wm wrong, so go back in file
+    if not extra_dims:  # single spectrum; we extracted wm wrong, so go back in file
         f.seek(0)
         wm, arr = np.genfromtxt(f, delimiter="\t", unpack=True)
         f.close()
@@ -121,13 +119,13 @@ def from_LabRAM(filepath, name=None, parent=None, verbose=True) -> Data:
         arr = np.genfromtxt(f, delimiter="\t")
         f.close()
         wm = wm[extra_dims:]
+        x = arr[:, 0]
 
         if extra_dims == 1:  # spectrum vs (x or survey)
             data.create_variable("wm", values=wm[:, None], units=spectral_units)
             data.create_channel(
                 "signal", values=arr[:, 1:].T / acquisition_time, units=channel_units
             )
-            x = arr[:, 0]
             if np.all(x == np.arange(x.size) + 1):  # survey
                 data.create_variable("index", values=x[None, :])
                 data.transform("wm", "index")
@@ -136,9 +134,7 @@ def from_LabRAM(filepath, name=None, parent=None, verbose=True) -> Data:
                 data.transform("wm", "x")
         elif extra_dims == 2:  # spectrum vs x vs y
             # fold to 3D
-            x = sorted(
-                set(arr[:, 0]), reverse=arr[0, 0] > arr[-1, 0]
-            )  # 0th column is stepped always (?)
+            x = sorted(set(x), reverse=bool(x[0] > x[-1]))  # 0th column is stepped always (?)
             x = np.array(list(x))
             x = x.reshape(1, -1, 1)
             y = arr[:, 1].reshape(1, x.size, -1)

WrightTools/kit/_list.py

@@ -7,7 +7,7 @@
 # --- define --------------------------------------------------------------------------------------
 
 
-__all__ = ["flatten_list", "intersperse", "get_index", "pairwise"]
+__all__ = ["flatten_list", "intersperse", "get_index", "from_list_of_objects", "pairwise"]
 
 
 # --- functions -----------------------------------------------------------------------------------
@@ -98,6 +98,43 @@ def get_index(lis, argument):
         return lis.index(argument)
 
 
+def from_list_of_objects(values: list, attrs: list, hint):
+    """
+    convenience wrapper of `get_index`.
+    Retrieve an object from a list of objects using either:
+    * its index,
+    * matching to an attr, or
+    * passing the value itself
+    the method normalizes a variety of input types to get the object itself.
+
+    arguments
+    ---------
+    values: list
+        list of objects to be searched
+    attrs: list
+        list of attrs that correspond to objects in the list `values`
+    hint: str, list or object
+        if hint is a member of values, return the hint itself
+
+    returns
+    -------
+    match: desired_type
+        the member of `values` that matches the given hint.
+        Raises ValueError if a match is not found.
+
+    see also
+    --------
+    WrightTools.kit.get_index
+    """
+
+    if isinstance(hint, int) or isinstance(hint, str):
+        return values[get_index(attrs, hint)]
+    elif hint in values:
+        return hint
+    else:
+        raise ValueError(f"could not find a member in values based on hint {hint}")
+
+
 def pairwise(iterable):
     """s -> (s0,s1), (s1,s2), (s2, s3), ...
 

tests/data/norm_for_each.py

@@ -0,0 +1,21 @@
+import numpy as np
+import WrightTools as wt
+from WrightTools import datasets
+
+
+def test_3D():
+    data = wt.open(datasets.wt5.v1p0p1_MoS2_TrEE_movie)
+
+    data.norm_for_each("w1", 0)
+    assert np.all(data.channels[0][:].max(axis=(0, 2)) == 1)
+
+    data.norm_for_each("d2", 0, new_channel={"name": "ai0_d2_norm"})
+    assert np.all(data.channels[-1][:].max(axis=(0, 1)) == 1)
+
+    data.norm_for_each("d1", 0, new_channel={"name": "ai0_d1_norm"})
+    data.channels[0].normalize()
+    assert np.all(np.isclose(data.ai0_d1_norm[:], data.channels[0][:]))
+
+
+if __name__ == "__main__":
+    test_3D()