Update _data.py

Author: ddkohler

WrightTools/data/_data.py

@@ -717,22 +717,20 @@ def gradient(self, axis, *, channel=0):
         else:
             raise wt_exceptions.TypeError("axis: expected {int, str}, got %s" % type(axis))
 
-        channel_index = wt_kit.get_index(self.channel_names, channel)
-        channel = self.channel_names[channel_index]
+        channel = self.get_channel(channel)
 
-        if self[channel].shape[axis_index] == 1:
+        if channel.shape[axis_index] == 1:
             raise wt_exceptions.ValueError(
                 "Channel '{}' has a single point along Axis '{}', cannot compute gradient".format(
                     channel, axis
                 )
             )
-        rtype = np.result_type(self[channel].dtype, float)
+        rtype = np.result_type(channel.dtype, float)
         new = self.create_channel(
             "{}_{}_gradient".format(channel, axis),
-            values=np.empty(self[channel].shape, dtype=rtype),
+            values=np.empty(channel.shape, dtype=rtype),
         )
 
-        channel = self[channel]
         if axis == axis_index:
             new[:] = np.gradient(channel[:], axis=axis_index)
         else:
@@ -1268,13 +1266,7 @@ def get_nadir(self, channel=0) -> tuple:
             Coordinates in units for each axis.
         """
         # get channel
-        if isinstance(channel, int):
-            channel_index = channel
-        elif isinstance(channel, str):
-            channel_index = self.channel_names.index(channel)
-        else:
-            raise TypeError("channel: expected {int, str}, got %s" % type(channel))
-        channel = self.channels[channel_index]
+        channel = self.get_channel(channel)
         # get indicies
         idx = channel.argmin()
         # finish
@@ -1294,13 +1286,7 @@ def get_zenith(self, channel=0) -> tuple:
             Coordinates in units for each axis.
         """
         # get channel
-        if isinstance(channel, int):
-            channel_index = channel
-        elif isinstance(channel, str):
-            channel_index = self.channel_names.index(channel)
-        else:
-            raise TypeError("channel: expected {int, str}, got %s" % type(channel))
-        channel = self.channels[channel_index]
+        channel:Channel = self.get_channel(channel)
         # get indicies
         idx = channel.argmax()
         # finish
@@ -1336,15 +1322,8 @@ def heal(self, channel=0, method="linear", fill_value=np.nan, verbose=True):
         warnings.warn("heal", category=wt_exceptions.EntireDatasetInMemoryWarning)
         timer = wt_kit.Timer(verbose=False)
         with timer:
-            # channel
-            if isinstance(channel, int):
-                channel_index = channel
-            elif isinstance(channel, str):
-                channel_index = self.channel_names.index(channel)
-            else:
-                raise TypeError("channel: expected {int, str}, got %s" % type(channel))
-            channel = self.channels[channel_index]
-            values = self.channels[channel_index][:]
+            channel = self.get_channel(channel)
+            values = channel[:]
             points = [axis[:] for axis in self._axes]
             xi = tuple(np.meshgrid(*points, indexing="ij"))
             # 'undo' gridding
@@ -1358,7 +1337,7 @@ def heal(self, channel=0, method="linear", fill_value=np.nan, verbose=True):
             tup = tuple([arr[i] for i in range(len(arr) - 1)])
             # grid data
             out = griddata(tup, arr[-1], xi, method=method, fill_value=fill_value)
-            self.channels[channel_index][:] = out
+            channel[:] = out
         # print
         if verbose:
             print(
@@ -1384,8 +1363,7 @@ def level(self, channel, axis, npts, *, verbose=True):
             Toggle talkback. Default is True.
         """
         warnings.warn("level", category=wt_exceptions.EntireDatasetInMemoryWarning)
-        channel_index = wt_kit.get_index(self.channel_names, channel)
-        channel = self.channels[channel_index]
+        channel = self.get_channel(channel)
         # verify npts not zero
         npts = int(npts)
         if npts == 0:
@@ -1443,8 +1421,7 @@ def map_variable(
             New data object.
         """
         # get variable index
-        variable_index = wt_kit.get_index(self.variable_names, variable)
-        variable = self.variables[variable_index]
+        variable = self.get_var(variable)
         # get points
         if isinstance(points, int):
             points = np.linspace(variable.min(), variable.max(), points)