Commit dee44cc4 authored by Christian Schneider's avatar Christian Schneider
Browse files

Progress on data_grid

parent 9574b8cf
......@@ -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
......@@ -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:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment