Progress on data_grid

DataModule/data_grid.py

@@ -8,12 +8,17 @@ Powered by xarray (xarray.pydata.org)
 """
 from .base import data_module_base
 import holoviews as hv
-hv.extension('bokeh')
 import holoviews.operation.datashader as hd
+import scipy.signal as sp_sig
+import matplotlib.pyplot as plt
 import numpy as np
 import xarray as xr
 from .data_table import data_table
 from IPython.display import display
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+
+hv.extension('bokeh')
+
 
 class data_grid(data_module_base):
     """Class for grid like data in multiple dimensions. Powered by the
@@ -32,13 +37,14 @@ class data_grid(data_module_base):
     def __init__(self, data_array, data_names=None):
         super().__init__()
         # Create default names
-        df_names = ['x{}'.format(i) for i in range(len(data_array))]
-        df_names[-1] = 'Value'
+        df_names = ['x', 'y']
+        df_names += ['x{}'.format(i) for i in range(2, len(data_array))]
+        df_names[-1] = 'Value'  # Last element is matrix/tensor
         # Replace default names by the names given
         if data_names:
             for idx in range(len(data_names)):
                 if data_names[idx]:
-                   df_names[idx] = data_names[idx]
+                    df_names[idx] = data_names[idx]
         # Create xarray
         self.df = xr.DataArray(data_array[-1],
                                dims=tuple(d_name for d_name in df_names[:-1]),
@@ -46,6 +52,16 @@ class data_grid(data_module_base):
                                        zip(df_names[:-1], data_array[:-1])},
                                name=df_names[-1])
 
+        # Idx variables for limiting data
+        self.x_min = 0
+        self.x_max = None
+        self.y_min = 0
+        self.y_max = None
+        # Easy access variables
+        self.name_x = df_names[0]
+        self.name_y = df_names[1]
+        self.name_v = df_names[-1]
+
     # Helpful data functions ###################################################
     def return_coord(self, coord_name):
         return np.array(self.df.coords[coord_name])
@@ -53,19 +69,89 @@ class data_grid(data_module_base):
     @property
     def x(self):
         """Default for two dim grid: Return first dimension"""
-        dim_name = self.df.dims[0]
-        return self.return_coord(dim_name)
+        return self.return_coord(self.name_x)[self.x_min:self.x_max]
+
+    @x.setter
+    def x(self, value):
+        self.df.coords[self.name_x][self.x_min:self.x_max] = value
 
     @property
     def y(self):
         """Default for two dim grid: Return second dimension"""
-        dim_name = self.df.dims[1]
-        return self.return_coord(dim_name)
+        return self.return_coord(self.name_y)[self.y_min:self.y_max]
+
+    @y.setter
+    def y(self, value):
+        self.df.coords[self.name_y][self.y_min:self.y_max] = value
 
     @property
     def z(self):
         """Default for two dim grid: Return values"""
-        return self.df.values
+        return np.array(self.df.values)[self.x_min:self.x_max,
+                                        self.y_min:self.y_max]
+
+    @z.setter
+    def z(self, values):
+        self.df.values[self.x_min:self.x_max, self.y_min:self.y_max] = values
+
+    @property
+    def values(self):
+        return np.array(self.df.values)[self.x_min:self.x_max,
+                                        self.y_min:self.y_max]
+
+    @values.setter
+    def values(self, values):
+        self.df.values[self.x_min:self.x_max, self.y_min:self.y_max] = values
+
+    def rename_x(self, new_name):
+        self.df = self.df.rename({self.name_x: new_name})
+        self.name_x = new_name
+
+    def rename_y(self, new_name):
+        self.df = self.df.rename({self.name_y: new_name})
+        self.name_y = new_name
+
+    def rename_values(self, new_name):
+        self.df = self.df.rename(new_name)
+        self.name_v = new_name
+
+    def rename_z(self, new_name):
+        return self.rename_values(new_name)  # alias
+
+    def select(self, xrng=None, yrng=None):
+        """Limit data between specified ranges.
+
+        This function will select the data in the specified range of the
+        x-axis and y-axis. If nothing is specified all the data will be
+        selected.
+
+        Note
+        -----
+        To undo a selection just run `select()` without an argument.
+
+        Parameters
+        -----------
+        xrng : list
+            Start and stop x value [x_start, x_stop]
+        yrng : list
+            Start and stop y value [y_start, y_stop]
+        """
+        if xrng is None:
+            x_idx = [0, None]
+        else:
+            x = self.df.coords[self.name_x]
+            x_idx = np.where((x >= xrng[0]) & (x <= xrng[1]))[0]
+
+        if yrng is None:
+            y_idx = [0, None]
+        else:
+            y = self.df.coords[self.name_y]
+            y_idx = np.where((y >= yrng[0]) & (y <= yrng[1]))[0]
+
+        self.y_min = y_idx[0]
+        self.y_max = y_idx[-1]
+        self.x_min = x_idx[0]
+        self.x_max = x_idx[-1]
 
     def extract_x(self, x0, plot=True):
         """Extract z values along  axis for specified x value x0.
@@ -139,6 +225,7 @@ class data_grid(data_module_base):
 
         return data
 
+    # Plotting #################################################################
     def plot(self, name_x=None, name_y=None, cmap='magma', height=400,
                 width=800):
         """Plot table with Holoviews
@@ -167,6 +254,109 @@ class data_grid(data_module_base):
             y_vals = name_y
         hv.opts({'Image': {'plot': {'width': width, 'height': height},
                            'style': {'cmap': cmap}}})
-        ds = hv.Dataset(self.df)
+        ds = hv.Dataset(self.df[self.x_min:self.x_max, self.y_min:self.y_max])
         return hd.regrid(ds.to(hv.Image, [x_vals, y_vals]))
 
+    def pcolormesh(self):
+        """Simple color plot without options. Quick and light"""
+        df = self.df[self.x_min:self.x_max, self.y_min:self.y_max]
+        df.plot.pcolormesh(self.name_x, self.name_y)
+
+    def imshow(self, colormap='magma', levels=256, zlabel_pos="right",
+               labelpad=0, cbar_y=0.5, **args):
+        """Color plot using matplotlib
+
+        Note
+        -----
+        Plots just the .select() data.
+
+        Parameters
+        -----------
+        colormap : str
+            Choose colormap from 'Magma' (Def), 'Inferno', 'Plasma', 'Viridis'
+        levels : int
+            Color levels. Default is 256
+        data_type : str
+            Choose if data is linear 'Lin', logarithmic 'Log' or 'Amplitude'
+        xlabel : str
+            Label for x axis
+        ylabel : str
+            Label for y axis
+        zlabel : str
+            Label for colorbar
+        zlabel_pos : "top", "right", "right_invert"
+            Position and orientation of zlabel next to colorbar. If location is
+            wrong, play with labelpad and cbar_y
+        """
+        # Get colormap from matplotlib
+        cmap = plt.cm.get_cmap(colormap, levels)
+
+        # Just get select range
+        df = self.df[self.x_min:self.x_max, self.y_min:self.y_max]
+
+        # Plot
+        ax = plt.subplot(111)
+        im = df.plot.imshow(self.name_x, self.name_y, cmap=cmap,
+                            add_colorbar=False, ax=ax)
+
+        # Customizing colorbar
+        divider = make_axes_locatable(ax)
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        cbar = plt.colorbar(im, cax=cax)
+        if zlabel_pos == "top":
+            cbar.set_label(self.name_v, labelpad=labelpad-32, y=cbar_y+0.58,
+                           rotation=0)
+        elif zlabel_pos == "right":
+            cbar.set_label(self.name_v, labelpad=labelpad+5, y=cbar_y,
+                           rotation=90)
+        else:
+            cbar.set_label(self.name_v, labelpad=labelpad+15, y=cbar_y,
+                           rotation=-90)
+        return ax
+
+    def contourf(self):
+        df = self.df[self.x_min:self.x_max, self.y_min:self.y_max]
+        df.plot.contourf(self.name_x, self.name_y)
+
+    # Data processing functions ################################################
+    def smoothx(self, nnb=21, polyorder=2):
+        """Smooth data along x axis using the Savitzky-Golay filter
+
+        Currently just for xyz data.
+
+        Parameters
+        -----------
+        nnb : int
+            Window length of filter. Must be an odd integer
+        polyorder : int
+            Polynomial order for the fit
+
+        Returns
+        --------
+        np.array
+            Filtered z matrix.
+        """
+        zsel = self.values
+        for idx, row in enumerate(zsel):
+            zsel[idx, :] = sp_sig.savgol_filter(row, nnb, polyorder)
+        return zsel
+
+    def smoothy(self, nnb=11, polyorder=2):
+        """"Smooth data along y axis using the Savitzky-Golay filter
+
+        Parameters
+        -----------
+        nnb : int
+            Window length of filter. Must be an odd integer
+        polyorder : int
+            Polynomial order for the fit
+
+        Returns
+        --------
+        np.array
+            Filtered z matrix.
+        """
+        zsel = self.values
+        for idx, row in enumerate(zsel.T):
+            zsel[:, idx] = sp_sig.savgol_filter(row, nnb, polyorder)
+        return zsel
\ No newline at end of file

DataModule/data_table.py

@@ -88,7 +88,7 @@ class data_table(data_module_base):
 
     @x.setter
     def x(self, value):
-        self.df[self.name_x] = value
+        self.df[self.name_x][self.idx_min:self.idx_max] = value
 
     @property
     def y(self):
@@ -96,7 +96,7 @@ class data_table(data_module_base):
 
     @y.setter
     def y(self, value):
-        self.df[self.name_y] = value
+        self.df[self.name_y][self.idx_min:self.idx_max] = value
 
     def rename_x(self, new_name):
         """Rename x variable"""
@@ -299,7 +299,7 @@ class data_table(data_module_base):
                     fig.line(x*xscale, y*yscale, line_color=c,
                              line_width=linewidth, legend=legend, **kwargs)
                 elif kw == 'o':
-                    fig.circle(x, y, fill_color=c, line_color='#000000',
+                    fig.circle(x, y, fill_color=c, line_color=c,
                                size=markersize, legend=legend, **kwargs)
             # Plot Fit
             if plot_fit and self._fit_executed: