from typing import Callable, Optional
import scipy as scp
import numpy as np
from devctrl.data.data import Columns, NDArray, Transform
import functools
import logging
log = logging.getLogger(__name__)
"""
REGISTERING TRANSFORMS
Transforms can be registered using the decorator `_register_transform`.
Registered transforms can then be redefined with different default parameters,
for example to allow specialized transform modules to set default columns to skip.
"""
_REGISTERED_TRANSFORMS: dict[str, Transform] = {}
def _register_transform(func: Callable) -> Callable:
"""
Decorator to register a transform
"""
_REGISTERED_TRANSFORMS[func.__name__] = func
return func
[docs]
def export_functions(prefix='', **override_kwargs) -> dict[str, Callable]:
"""
Create new versions of all registered functions with new default arguments `override_kwargs`
"""
global _REGISTERED_TRANSFORMS
# print("reg trans:", _REGISTERED_TRANSFORMS)
new_funcs = {}
for name, func in _REGISTERED_TRANSFORMS.items():
new_name = f"{prefix}{name}"
# have to use factory function here because otherwise all wrappers are the same object and will always point the function that was last registered
def make_wrapper(f):
@functools.wraps(f)
def wrapper(*args, **kwargs):
return f(*args, **(override_kwargs | kwargs))
return wrapper
new_funcs[new_name] = make_wrapper(func)
return new_funcs
[docs]
def apply_to_columns(f: Callable, sdata: NDArray, columns: Columns, skipcols: Optional[Columns]=None, docols: Optional[Columns]=None, *args, **kwargs):
"""
Apply a transform to all selected columns.
:param f: Transform taking a numpy array (data column) and *args and **kwargs and returns new numpy array of same length
:param sdata: 2D-data
:param columns: List of columns
:param skipcols: Optional list of columns to not apply the transform to (takes precedence over docols)
:param docols: Optional list of columns to apply the transform to
:return: data, columns
"""
for i, col in enumerate(columns):
if skipcols is not None and col in skipcols: continue
if docols is not None and col not in docols: continue
sdata[:,i] = f(sdata[:,i], *args, **kwargs)
return sdata, columns
[docs]
@_register_transform
def remove_nans(sdata: NDArray, columns: Columns):
"""Delete all rows where any number is nan"""
rows_no_nan = np.isfinite(sdata).all(axis=1)
return sdata[rows_no_nan,:], columns
def _shift(data: NDArray, value):
return data + value
[docs]
@_register_transform
def shift(sdata, columns, value: float, **apply_kw):
"""
Add `shift` to data
"""
return apply_to_columns(_shift, sdata, columns, **apply_kw, value=value)
def _shift_to_zero(data: NDArray):
_min = np.min(data)
_max = np.max(data)
if _max < 0:
return data - _max
elif _min > 0:
return data - _min
return data
[docs]
@_register_transform
def shift_to_zero(sdata, columns, **apply_kw):
"""
Shift data so that either the min or max value is zero
"""
return apply_to_columns(_shift_to_zero, sdata, columns, **apply_kw)
def _scale(data: NDArray, factor):
return data * factor
[docs]
@_register_transform
def scale(sdata, columns, factor: float, **apply_kw):
"""
Multiply data by `factor`
"""
return apply_to_columns(_scale, sdata, columns, **apply_kw, factor=factor)
[docs]
def get_data_normalize_scale(data: np.ndarray):
"""
Return the scale factor that data must be multiplied with to be entirely within the interval [-1, 1]
"""
min_ = np.abs(np.nanmin(data))
max_ = np.abs(np.nanmax(data))
max_abs = max(min_, max_)
scale = 1/max_abs
return scale
[docs]
def normalize_data(data: np.ndarray):
"""
Return data scaled to be within [-1, 1]
"""
return data * get_data_normalize_scale(data)
[docs]
@_register_transform
def normalize(sdata, columns, **apply_kw):
"""
Normalize all columns except skipcols so that the largest absolute value becomes either 1 or -1.
"""
return apply_to_columns(normalize_data, sdata, columns, **apply_kw)
def _normalize_to_index(data, idx):
val = np.abs(data[idx])
return data / val
def _normalize_to_index_range(data, idx_lo, idx_hi):
val = np.abs(np.mean(data[idx_lo:idx_hi]))
return data / val
[docs]
@_register_transform
def normalize_to(sdata, columns, xkey: str, x: float, dx: float|None=None, **apply_kw):
"""
Normalize the data so that for each column, the value at `x` or the mean value of the range `x-dx,x+dx` is set to (-)1.
:param xkey: Column key of the x-data
:param x: x value
:param dx: Optional delta x
"""
if not xkey in columns:
raise KeyError(f"'{xkey}' is not in columns '{columns}'")
xdata = sdata[:,columns.index(xkey)]
if dx is not None:
idx1 = np.argmin(np.abs(xdata - (x - dx)))
idx2 = np.argmin(np.abs(xdata - (x + dx)))
idx_lo = min(idx1, idx2)
idx_hi = max(idx1, idx2)
# interval so small it falls on a single value
if idx_lo != idx_hi:
return apply_to_columns(_normalize_to_index_range, sdata, columns, idx_lo=idx_lo, idx_hi=idx_hi, **apply_kw)
log.warning(f"Range dx={dx} too small, using single value")
idx = np.argmin(np.abs(xdata - x))
return apply_to_columns(_normalize_to_index, sdata, columns, idx=idx, **apply_kw)
"""
CALCULUS
"""
def _integrate(data, n=1, pre_func=None):
for i in range(n):
if pre_func:
data = pre_func(data)
data = scp.integrate.cumulative_trapezoid(data, initial=0)
return data
[docs]
@_register_transform
def integrate(sdata: NDArray, columns: Columns, n=1, pre_func=None, **apply_kw):
"""Numerically integrate the data of column(s) <key> n times"""
return apply_to_columns(_integrate, sdata, columns, n=n, pre_func=pre_func, **apply_kw)
def _derivative(data, n=1, pre_func=None):
for i in range(n):
if pre_func:
data = pre_func(data)
data = np.gradient(data)
return data
[docs]
@_register_transform
def derivative(sdata: NDArray, columns: Columns, n=1, pre_func=None, **apply_kw):
"""
Calculate the n-th derivative of column(s) <key> numerically.
:param n: The derivative to calculate
:param pre_func: A function taking and returning a 1d numpy array. If given, it is applied before every derivative.
This may, for example, be useful if the data should be smoothened before taking the derivatives.
"""
return apply_to_columns(_derivative, sdata, columns, n=n, pre_func=pre_func, **apply_kw)
"""
FILTERING
"""
def _notch(data, w0:float, Q: float, **iirnotch_kw):
b, a = scp.signal.iirnotch(w0, Q, **iirnotch_kw)
filtered = scp.signal.lfilter(b, a, data)
return filtered
[docs]
@_register_transform
def notch(sdata, columns, w0: float, Q: float, **kw):
"""
Notch filter at frequency w0 with quality factor Q.
See scipy.signal.iirnotch
"""
return apply_to_columns(_notch, sdata, columns, w0=w0, Q=Q, **kw)
def _smoothen(data, **savgol_kw):
data = scp.signal.savgol_filter(data, **(dict(window_length=100, polyorder=4, mode="nearest") | (savgol_kw or {})))
return data
[docs]
@_register_transform
def smoothen(sdata, columns, **kw):
return apply_to_columns(_smoothen, sdata, columns, **kw)
"""
DOWNSAMPLING
"""
try:
from mqutil.calc.downsample import get_gradient_downsample_indices, get_log_downsample_indices
[docs]
def downsample_log(data, columns, N, **kw):
indices = get_log_downsample_indices(data.shape[0], N, **kw)
data = data[indices,:]
return data, columns
except ImportError as e:
log.warning(f"Downsampling transforms not available: {e}")
