Source code for classes.utils_math

import sys
if 'omfit_tree' in sys.modules:
    print('Loading math utility functions...')

try:
    # framework is running
    from .startup_choice import *
except ImportError as _excp:
    # class is imported by itself
    if (
        'attempted relative import with no known parent package' in str(_excp)
        or 'No module named \'classes\'' in str(_excp)
        or "No module named '__main__.startup_choice'" in str(_excp)
    ):
        from startup_choice import *
    else:
        raise

from classes.utils_base import *
from classes.utils_base import _available_to_user_math

from six import string_types
import uncertainties
from uncertainties import unumpy
from uncertainties.unumpy import nominal_values, std_devs, uarray
import numpy as np
import scipy
from scipy import interpolate, optimize, integrate, ndimage
from scipy.interpolate.polyint import _Interpolator1D
from scipy.ndimage import _ni_support
import functools
import bisect
import itertools

try:
    from _ni_docstrings import docfiller  # v1.1.0+
except ImportError:
    docfiller = lambda x: x

#--------------------------------
# numpy error handling decorators
#--------------------------------
[docs]def np_errors(**kw): context = np.errstate(**kw) def decorator(f): @functools.wraps(f) def decorated(*args, **kw): with context: return f(*args, **kw) return decorated return decorator
np_ignored = np_errors(all='ignore') np_raised = np_errors(all='raise') np_printed = np_errors(all='print') np_warned = np_errors(all='warn') #--------------------- # Delete leading zeros #---------------------
[docs]def delete_leading_zeros(number_str): ''' Delete leading zeros from array indices :param number_str: A string containing comma separated dimension indices and colon separated ranges and strides, such as '008:010,0001:0002' :return: String with leading zeros stripped, such as '8:10,1:2' ''' return ','.join([':'.join([colon_split.lstrip('0') if colon_split.lstrip('0') else '0' for colon_split in comma_split.split(':')]) for comma_split in number_str.split(',')])
#------- # arrays #-------
[docs]@_available_to_user_math def get_array_hash(y): """ Get the hash for an array. A hash is an fixed sized integer that identifies a particular value, comparing this integer is faster than comparing the whole array (if you already have it stored). Example:: y1 = np.arange(3) y2 = np.arange(3) + 1 assert(get_array_hash(y1) == get_array_hash(y2)) # will raise an error :param y: np.ndarray. :return: hash. """ return bytes(y).__hash__()
[docs]@_available_to_user_math def ismember(A, B): ''' Mimics the Matlab ismember() function to look for occurrences of A into B :param A: number or list/array :param B: number or list/array :return: returns lia, locb lists. lia: returns 'True' where the data in A is found in B and 'False' elsewhere. locb: contains the lowest index in B for each value in A that is a member of B. while it contains 'None' elsewhere (where A is not a member of B) ''' bind = {} tmp = {} a=np.atleast_1d(A) b=np.atleast_1d(B) for i, elt in enumerate(a): if elt in b: tmp[elt] = True for i, elt in enumerate(b): if elt not in bind: bind[elt] = i lia = [tmp.get(itm, False) for itm in a] locb = [bind.get(itm, None) for itm in a] return lia, locb # two lists in output
[docs]@_available_to_user_math def map1d(vecfun, ndarr, axis=-1, args=[], **kwargs): r""" Map a vector operation to a single axis of any multi-dimensional array. input: :param vecfun: A function that takes a 1D vector input. :param ndarr: A np.ndarray. :param axis: Axis of the array on which vecfun will operate. :param \*args: Additional arguments passed to vecfun. :param \**kwargs: Additional key word arguments are also passed to vecfun. output: A new array of same dimensionality as ndarr. """ # don't modify original array nda = copy.copy(ndarr) if len(nda.shape)==1: nda = vecfun(nda,*args,**kwargs) else: # move axis of interest to the back nda = nda.swapaxes(axis,-1) newshape = list(nda.shape) # flatten all other dimensions nda = nda.reshape(-1,newshape[-1]) # apply vector operation for i,v in enumerate(nda): nda[i] = vecfun(v,*args,**kwargs) # reform nD array newshape[-1] = -1 nda = nda.reshape(newshape) # put axis back where it was nda = nda.swapaxes(-1,axis) return nda
[docs]@_available_to_user_math def uniform_resample_2D_line(X0, Y0, n): ''' Resampling of 2D line with uniform distribution of points along the line :param X0: input x coordinates of 2D path :param Y0: input y coordinates of 2D path :param n: number of points for output 2d path :return: tuple of x,y uniformly resampled path ''' X01=np.hstack((X0[-1],X0)) Y01=np.hstack((Y0[-1],Y0)) L0=np.cumsum(np.sqrt(np.diff(X01)**2+np.diff(Y01)**2)) l=np.linspace(min(L0),max(L0),n) x=interpolate.interp1d(L0,X0)(l) y=interpolate.interp1d(L0,Y0)(l) return x,y
[docs]@_available_to_user_math def stepwise_data_index(y,broaden=0): ''' This function returns the indices that one must use to reproduce the step-wise data with the minimum number of points. In the ascii-art below, it returns the indices of the crosses. The original data can then be reproduced by nearest-neighbor interpolation:: Y ^ | x.....x | x....x | x....x x......x | 0-----------------------------> X Hint: can also be used to compress linearly varying data:: i = stepwise_data_index(np.gradient(y)) The original data can then be reproduced by linear interpolation. :param y: input array :param broaden: return `broaden` points around the x's :return: indices for compression ''' y=np.abs(np.gradient(y)) if broaden>1: y=smooth(y,int(broaden//2)*2+1) return np.unique([0]+tolist(np.where(y)[0])+[len(y)-1])
[docs]def pack_points(n, x0, p): ''' Packed points distribution between -1 and 1 :param n: number of points :param x0: pack points around `x0`, a float between -1 and 1 :param p: packing proportional to `p` factor >0 :return: packed points distribution between -1 and 1 ''' x=np.linspace(-1,1,n) y=np.sinh((x-x0)*p) y=(y-min(y))/(max(y)-min(y))*(max(x)-min(x))+min(x) return y
[docs]def simplify_polygon(x, y, tolerance=None, preserve_topology=True): ''' Returns a simplified representation of a polygon :param x: array of x coordinates :param y: array of y coordinates :param tolerance: all points in the simplified object will be within the tolerance distance of the original geometry if tolerance is None, then a tolerance guess is returned :param preserve_topology: by default a slower algorithm is used that preserves topology :return: x and y coordinates of simplified polygon geometry if tolerance is None, then a tolerance guess is returned ''' if tolerance is None: dd = np.sqrt(np.gradient(x) ** 2 + np.gradient(y) ** 2) dd = dd[dd > 0] tolerance = np.median(dd) / 100.0 return tolerance with warnings.catch_warnings(): warnings.filterwarnings('ignore', category=ImportWarning, message="can't resolve package*") from shapely.geometry import Polygon polygon = Polygon(zip(x, y)) polygon = polygon.simplify(tolerance=tolerance, preserve_topology=preserve_topology) tmp = np.array(list(polygon.exterior.coords)) return tmp[:, 0], tmp[:, 1]
#------- # iterables #-------
[docs]def torecarray(data2D, names=None): ''' Converts a 2D list of lists, or a dictionary of dictionaries with 1d elements to a recarray :param data2D: 2D input data :param names: recarray columns names :return: recarray ''' if isinstance(data2D,dict): if isinstance(list(data2D[list(data2D.keys())[0]].keys())[0],dict): data2D=dictdict_2_dictlist(data2D) names=list(data2D.keys()) data2D=[data2D[k] for k in data2D] dtype=[] for k,name in enumerate(names): dtype.append((name,np.array(data2D[k]).dtype)) M=np.recarray(len(data2D[0]),dtype=dtype) for k,item in enumerate(names): M[item][:]=data2D[k][:] return M
[docs]def dictdict_2_dictlist(data): ''' dictionary of dictionaries with 1d elements, to dictionary of lists :param data: dictionary of dictionaries :return: dictionary of dictionaries ''' col=set() for k in data: item=data[k] col.update(list(map(str,list(item.keys())))) col=sorted(list(col)) data={} for c in col: data[c]=[] for k in data: item=data[k] if c in list(item.keys()): data[c].append(item[c]) else: data[c].append(np.nan) for c in col: data[c]=np.array(data[c]) return data
[docs]@_available_to_user_math def flatten_iterable(l): ''' flattens iterable of iterables structures :param l: input iterable structure :return: generator with items in the structure ''' import collections for el in l: if isinstance(el, collections.Iterable) and not isinstance(el, str): for sub in flatten_iterable(el): yield sub else: yield el
[docs]@_available_to_user_math def lists_loop_generator(*argv, **kwargs): r""" This function generates a list of items combining the entries from multiple lists :param \*args: multiple 1D lists (of equal length, or of 0 or 1 element length) :param permute: output all element permutations :return: list of items combining the entries from multiple lists Examples:: shots = 123 times = [1,2] runs = 'a01' shots, times, runs = lists_loop_generator(shots,times,runs) for shot, time, run in zip(shots,times,runs): myCode(shot,time,run) shots = [123,446] times = [1,2] runs = 'a01' shots, times, runs = lists_loop_generator(shots,times,runs) now permute will return lists of length 2 or length 4 """ # Define inputs and make sure that they are type list inList = list(argv) for i,arg in enumerate(inList): if not isinstance(arg,list): inList[i] = [arg] # Get permutation keyword argument permute = kwargs.pop('permute',False) # Check the lengths n = [] for el in inList: n.append(len(el)) # Check if length of all elements are equal lens = len(set(iter(n))) len_equal = (lens <= 1) myList = [] if len_equal and not permute: # Single values or all flat lists to zip myList = inList elif len_equal and permute: # All flat lists and permuteuting before zip') myTuples = list(itertools.product(*inList)) for i in range(len(inList)): myList.append([item[i] for item in myTuples]) elif not len_equal and not permute: # Only two allowable options: list has max(n) elements or list is 0 or 1D for i,el in enumerate(inList): if len(el) == np.max(n): # List element has appropriate number of entries continue elif len(el) == 1: # List element is single; duplicating inList[i] = el*np.max(n) else: printe('Error, list element is neither the same as others, nor 1D') raise OMFITexception('Inputs error') myList = inList elif not len_equal and permute: # Permuting before zip myTuples = list(itertools.product(*inList)) for i in range(len(inList)): myList.append([item[i] for item in myTuples]) return myList
# Uncertainties helpers
[docs]def is_uncertain(var): ''' :param var: Variable or array to test :return: True if input variable or array is uncertain ''' from builtins import any def _uncertain_check(x): return isinstance(x, uncertainties.core.AffineScalarFunc) if isinstance(var, str): return False elif np.iterable(var) or isinstance(var, np.ndarray): # isinstance needed for 0D arrays from squeeze tmp = np.array(var) if tmp.dtype not in ['O','object']: return False else: # the argument of any is a generator object (using a list slows things down) return any(_uncertain_check(x) for x in tmp.flat) else: return _uncertain_check(var)
# Modify uncertainties.unumpy.core.wrap_array_func to give unc_tag keyword holding # repr of uncertainty Variable being varied to the function being called
[docs]def wrap_array_func(func): from uncertainties.unumpy.core import array_derivative import uncertainties.core as uncert_core # !!! This function is not used in the code, and is also not # in the user guide: should it be promoted? # # !!! The implementation seems superficially similar to # !!! uncertainties.core.wrap(): is there code/logic duplication # !!! (which should be removed)? """ Return a version of the function func() that works even when func() is given a NumPy array that contains numbers with uncertainties, as first argument. This wrapper is similar to uncertainties.core.wrap(), except that it handles an array argument instead of float arguments, and that the result can be an array. However, the returned function is more restricted: the array argument cannot be given as a keyword argument with the name in the original function (it is not a drop-in replacement). func -- function whose first argument is a single NumPy array, and which returns a NumPy array. """ def wrapped_func(arr, *args, **kwargs): # Nominal value: arr_nominal_value = unumpy.nominal_values(arr) func_nominal_value = func(arr_nominal_value, *args, **kwargs) # The algorithm consists in numerically calculating the derivatives # of func: # Variables on which the array depends are collected: variables = set() for element in arr.flat: # floats, etc. might be present if isinstance(element, uncert_core.AffineScalarFunc): variables |= set(element.derivatives.keys()) # If the matrix has no variables, then the function value can be # directly returned: if not variables: return func_nominal_value # Calculation of the derivatives of each element with respect # to the variables. Each element must be independent of the # others. The derivatives have the same shape as the output # array (which might differ from the shape of the input array, # in the case of the pseudo-inverse). derivatives = np.vectorize(lambda _: {})(func_nominal_value) for var in variables: # A basic assumption of this package is that the user # guarantees that uncertainties cover a zone where # evaluated functions are linear enough. Thus, numerical # estimates of the derivative should be good over the # standard deviation interval. This is true for the # common case of a non-zero standard deviation of var. If # the standard deviation of var is zero, then var has no # impact on the uncertainty of the function func being # calculated: an incorrect derivative has no impact. One # scenario can give incorrect results, however, but it # should be extremely uncommon: the user defines a # variable x with 0 standard deviation, sets y = func(x) # through this routine, changes the standard deviation of # x, and prints y; in this case, the uncertainty on y # might be incorrect, because this program had no idea of # the scale on which func() is linear, when it calculated # the numerical derivative. # The standard deviation might be numerically too small # for the evaluation of the derivative, though: we set the # minimum variable shift. shift_var = max(var._std_dev, 1e-8*abs(var._nominal_value)) # An exceptional case is that of var being exactly zero. # In this case, an arbitrary shift is used for the # numerical calculation of the derivative. The resulting # derivative value might be quite incorrect, but this does # not matter as long as the uncertainty of var remains 0, # since it is, in this case, a constant. if not shift_var: shift_var = 1e-8 # Shift of all the elements of arr when var changes by shift_var: shift_arr = array_derivative(arr, var)*shift_var # Origin value of array arr when var is shifted by shift_var: shifted_arr_values = arr_nominal_value + shift_arr ### OMFIT addition kwargs['unc_tag'] = repr(var) ### End OMFIT addition func_shifted = func(shifted_arr_values, *args, **kwargs) numerical_deriv = (func_shifted-func_nominal_value)/shift_var # Update of the list of variables and associated # derivatives, for each element: for (derivative_dict, derivative_value) in ( list(zip(derivatives.flat, numerical_deriv.flat))): if derivative_value: derivative_dict[var] = derivative_value # numbers with uncertainties are built from the result: return np.vectorize(uncert_core.AffineScalarFunc)( func_nominal_value, np.vectorize(uncert_core.LinearCombination)(derivatives)) wrapped_func = uncert_core.set_doc("""\ Version of %s(...) that works even when its first argument is a NumPy array that contains numbers with uncertainties. Warning: elements of the first argument array that are not AffineScalarFunc objects must not depend on uncert_core.Variable objects in any way. Otherwise, the dependence of the result in uncert_core.Variable objects will be incorrect. Original documentation: %s""" % (func.__name__, func.__doc__))(wrapped_func) # It is easier to work with wrapped_func, which represents a # wrapped version of 'func', when it bears the same name as # 'func' (the name is used by repr(wrapped_func)). wrapped_func.__name__ = func.__name__ return wrapped_func
[docs]def chunks(l, n): """Yield successive n-sized chunks from l""" for i in range(0, len(l), n): yield l[i:i+n]
[docs]class sparray(object): """ Class for n-dimensional sparse array objects using Python's dictionary structure. based upon: http://www.janeriksolem.net/2010/02/sparray-sparse-n-dimensional-arrays-in.html """ def __init__(self, shape=None, default=np.nan, dtype=None, wrap_dim=0, offset=0, index_offset=False): if dtype is None: self.dtype = None elif callable(dtype): self.dtype = dtype = np.array(dtype(0)).dtype else: self.dtype = dtype = np.array(dtype).dtype self.default = np.array(default) if self.dtype is not None: self.default = np.array(default).astype(dtype) elif dtype is None: self.dtype = dtype = self.default.dtype if shape is None: self.shape = () else: self.shape = tuple(shape) self.ndim = len(self.shape) self.wrap_dim = wrap_dim self.__data = self.data = OrderedDict() self.offset = np.atleast_1d(offset) self.index_offset = index_offset def _index_offset(self, index, mult=-1): if not self.index_offset: return index index = np.atleast_1d(index) _index = [] offset = np.zeros(index.shape) + self.offset for i, o in zip(index, offset): if isinstance(i, slice): start, stop, step = i.start, i.stop, i.step if start is not None: start = int(start + o * mult) if stop is not None: stop = int(stop + o * mult) _index.append(slice(start, stop, step)) else: _index.append(int(i + o * mult)) return tuple(_index) def _setitem_base(self, index, value): """ set value to position given in index, where index is a tuple. """ # adjust number of dimensions index=tuple(map(int,index)) if len(index) > len(self.shape): tmp = list(np.array(index) + 1) tmp[:len(self.shape)] = self.shape self.shape = tuple(tmp) # adjust dimensions for k, dim in enumerate(np.array(index) + 1): if dim > self.shape[k]: tmp = list(self.shape) tmp[k] = dim self.shape = tuple(tmp) self.ndim = len(self.shape) if self.dtype is None: self.dtype = np.array(value).dtype self.default = np.array(self.default).astype(self.dtype) self.__data[index] = np.atleast_1d([value]).astype(self.dtype) def __setitem__(self, index, value): """ set value to position given in index, where index is a tuple. """ index = self._index_offset(index, -1) if isinstance(value, np.ndarray): _index = list(index) for k in range(len(value)): self._setitem_base(tuple(_index), value[k]) _index[self.wrap_dim] += 1 else: self._setitem_base(index, value) def __getitem__(self, index): """ get value at position given in index, where index is a tuple. """ index = self._index_offset(index, -1) try: return self.dense()[index] except IndexError: return index*0+self.default def __delitem__(self, index): """ index is tuples of element to be deleted. """ index = self._index_offset(index) if index in self.__data: del (self.__data[index]) def __binary_op__(self, other, op): if isinstance(other,sparray): if self.shape == other.shape: tmp = copy.deepcopy(self) for k in set.difference(set(tmp.__data.keys()), set(other.__data.keys())): tmp.__data[k] = eval('tmp.__data[k].' +op+ '(other.default)') tmp.default = eval('tmp.default.' + op+ '(other.default)') for k in list(other.__data.keys()): old_val = tmp.__data.setdefault(k, tmp.default) tmp.__data[k] = 'old_val.' +op+ '(other.__data[k])' return tmp else: raise ValueError('Array sizes do not match. ' + str(self.shape) + ' versus ' + str(other.shape)) else: return eval('self.dense().'+op+'(other)') def __unary_op__(self, op): tmp = copy.deepcopy(self) for k in tmp.__data.keys(): tmp.__data[k] = eval('tmp.__data[k].' +op+ '()') tmp.default = eval('tmp.default.' + op+ '()') return tmp def __str__(self): return str(self.dense())
[docs] def dense(self): """Convert to dense NumPy array""" out = (self.default * np.ones(self.shape)).astype(self.dtype) for ind in self.__data: out[ind] = self.__data[ind].astype(self.dtype) return out
[docs] def sum(self): """Sum of elements""" s = self.default * np.array(self.shape).prod() for ind in self.__data: s += (self.__data[ind] - self.default) return s
[docs] def fortran(self, index, value): # Drop leading zeros otherwise the index is evaluated as an octal. index = delete_leading_zeros(index) index = eval('(' + re.sub('([0-9]+):', r'\1', index.split('(')[-1].strip(')')) + ')') self.__setitem__(index, value) return self
[docs] def isnan(self,x): if np.iterable(x): x=np.array(x) i=np.where(x==-9223372036854775808) x=x.astype(float) x[i]=np.nan return np.isnan(np.array(x)) elif np.isnan(x): return True elif x==np.array(np.nan).astype(np.array(x).dtype): return x==np.array(np.nan).astype(np.array(x).dtype)
[docs] def fortran_repr(self): # print strips strips = OrderedDict() for item in list(self.__data.keys()): item = tuple(item) strips.setdefault(item[1:], item) strips[item[1:]] = tuple(map(int, np.min([strips[item[1:]], item], 0))) # prints outputs = OrderedDict() if not self.isnan(self.default): outputs['']=self.default for item in list(strips.values()): tmp=list(item) tmp[0]=slice(None) tmp=tuple(tmp) if self.isnan(self.default): ix_=np.where(self.isnan(self[tmp]))[0] else: ix_=np.where(self[tmp]==self.default)[0] #split where default value is found substrips_=np.split(self[tmp],ix_) substrips=[] ix=[] ix_=[0]+list(ix_) for ki,ks in zip(ix_,substrips_): if len(ks)>1: ix.append(ki) ix.append(ki+1) substrips.append(ks[:1]) substrips.append(ks[1:]) elif len(ks): ix.append(ki) substrips.append(ks) #print substrips #note: values that are equal to the default value or that are NaN will not be printed item=item+self.offset for k,v in enumerate(substrips): if not len(v): continue elif self.isnan(v[0]): continue elif v[0]==self.default: continue item[0]=ix[k]+self.offset[0] if len(item)==1: outputs['(%d)'%item.astype(int)] = v.astype(self.dtype) else: outputs[str(re.sub(' ', '', repr(tuple(item.astype(int)))))] = v.astype(self.dtype) return outputs
def __str__(self): tmp = self.fortran_repr() out = [] for item, value in zip(list(tmp.keys()), list(tmp.values())): out.append('%s: %s' % (str(item)[1:-1], str(value))) return '\n'.join(out) def __repr__(self): return self.__str__() def __tree_repr__(self): return self.dense(), [''] def __iter__(self): for i in self.dense(): yield i # binary operations
[docs] def lt(self, other): return self.__binary_op__(other, 'lt')
[docs] def le(self, other): return self.__binary_op__(other, 'le')
[docs] def eq(self, other): return self.__binary_op__(other, 'eq')
[docs] def ge(self, other): return self.__binary_op__(other, 'ge')
[docs] def gt(self, other): return self.__binary_op__(other, 'gt')
def __lt__(self, other): return self.__binary_op__(other, '__lt__') def __le__(self, other): return self.__binary_op__(other, '__le__') def __eq__(self, other): return self.__binary_op__(other, '__eq__') def __ne__(self, other): return self.__binary_op__(other, '__ne__') def __ge__(self, other): return self.__binary_op__(other, '__ge__') def __gt__(self, other): return self.__binary_op__(other, '__gt__')
[docs] def is_(self, other): return self.__binary_op__(other, 'is_')
[docs] def is_not(self, other): return self.__binary_op__(other, 'is_not')
[docs] def add(self, other): return self.__binary_op__(other, 'add')
def __add__(self, other): return self.__binary_op__(other, '__add__') def __radd__(self, other): return self.__binary_op__(other, '__radd__')
[docs] def and_(self, other): return self.__binary_op__(other, 'and_')
def __and__(self, other): return self.__binary_op__(other, '__and__') def __rand__(self, other): return self.__binary_op__(other, '__rand__')
[docs] def floordiv(self, other): return self.__binary_op__(other, 'floordiv')
def __floordiv__(self, other): return self.__binary_op__(other, '__floordiv__') def __rfloordiv__(self, other): return self.__binary_op__(other, '__rfloordiv__')
[docs] def index(self, other): return self.__binary_op__(other, 'index')
def __index__(self, other): return self.__binary_op__(other, '__index__')
[docs] def lshift(self, other): return self.__binary_op__(other, 'lshift')
def __lshift__(self, other): return self.__binary_op__(other, '__lshift__') def __rlshift__(self, other): return self.__binary_op__(other, '__rlshift__')
[docs] def mod(self, other): return self.__binary_op__(other, 'mod')
def __mod__(self, other): return self.__binary_op__(other, '__mod__') def __rmod__(self, other): return self.__binary_op__(other, '__rmod__')
[docs] def mul(self, other): return self.__binary_op__(other, 'mul')
def __mul__(self, other): return self.__binary_op__(other, '__mul__') def __rmul__(self, other): return self.__binary_op__(other, '__rmul__')
[docs] def matmul(self, other): return self.__binary_op__(other, 'matmul')
def __matmul__(self, other): return self.__binary_op__(other, '__matmul__')
[docs] def or_(self, other): return self.__binary_op__(other, 'or_')
def __or__(self, other): return self.__binary_op__(other, '__or__') def __ror__(self, other): return self.__binary_op__(other, '__ror__')
[docs] def pow(self, other): return self.__binary_op__(other, 'pow')
def __pow__(self, other): return self.__binary_op__(other, '__pow__') def __rpow__(self, other): return self.__binary_op__(other, '__rpow__')
[docs] def rshift(self, other): return self.__binary_op__(other, 'rshift')
def __rshift__(self, other): return self.__binary_op__(other, '__rshift__') def __rrshift__(self, other): return self.__binary_op__(other, '__rrshift__')
[docs] def sub(self, other): return self.__binary_op__(other, 'sub')
def __sub__(self, other): return self.__binary_op__(other, '__sub__') def __rsub__(self, other): return self.__binary_op__(other, '__rsub__')
[docs] def truediv(self, other): return self.__binary_op__(other, 'truediv')
def __truediv__(self, other): return self.__binary_op__(other, '__truediv__') def __rtruediv__(self, other): return self.__binary_op__(other, '__rtruediv__') def __divmod__(self, other): return self.__binary_op__(other, '__divmod__') def __rdivmod__(self, other): return self.__binary_op__(other, '__rdivmod__')
[docs] def xor(self, other): return self.__binary_op__(other, 'xor')
def __xor__(self, other): return self.__binary_op__(other, '__xor__') def __rxor__(self, other): return self.__binary_op__(other, '__rxor__') # unary operations
[docs] def bool(self): return self.__unary_op__('bool')
def __bool__(self): return self.__unary_op__('__bool__') def __nonzero__(self): return self.__unary_op__('__nonzero__')
[docs] def real(self): return self.__unary_op__('real')
[docs] def imag(self): return self.__unary_op__('imag')
[docs] def not_(self): return self.__unary_op__('not_')
def __not__(self): return self.__unary_op__('__not__')
[docs] def truth(self): return self.__unary_op__('truth')
[docs] def abs(self): return self.__unary_op__('abs')
def __abs__(self): return self.__unary_op__('__abs__')
[docs] def inv(self): return self.__unary_op__('inv')
[docs] def invert(self): return self.__unary_op__('invert')
def __inv__(self): return self.__unary_op__('__inv__') def __invert__(self): return self.__unary_op__('__invert__')
[docs] def neg(self): return self.__unary_op__('neg')
def __neg__(self): return self.__unary_op__('__neg__')
[docs] def pos(self): return self.__unary_op__('pos')
def __pos__(self): return self.__unary_op__('__pos__') def __float__(self): return self.__unary_op__('__float__') def __complex__(self): return self.__unary_op__('__complex__') def __oct__(self): return self.__unary_op__('__oct__') def __hex__(self): return self.__unary_op__('__hex__') def __trunc__(self): return self.__unary_op__('__trunc__') __hash__ = None
[docs]@_available_to_user_math def unsorted_unique(x): ''' Find the unique elements of an array The order of the elements is according to their first appearance :param x: input array or list :return: unique array or list (depending on whether the input was an array or a list) ''' tmp=[] for k in x: if k not in tmp: tmp.append(k) if isinstance(x,np.ndarray): return np.array(tmp) else: return tmp
[docs]@_available_to_user_math def closestIndex(my_list, my_number=0): """ Given a SORTED iterable (a numeric array or list of numbers) and a numeric scalar my_number, find the index of the number in the list that is closest to my_number :param my_list: Sorted iterable (list or array) to search for number closest to my_number :param my_number: Number to get close to in my_list :return: Index of my_list element closest to my_number :note: If two numbers are equally close, returns the index of the smallest number. """ if not hasattr(my_list, '__iter__'): raise TypeError("closestIndex() in utils_math.py requires an iterable as the first argument. Got " "instead: {:}".format(my_list)) if not is_numeric(my_number): if hasattr(my_number, '__iter__') and len(my_number) == 1 and is_numeric(my_number[0]): printw('Warning: closestIndex got a len()=1 iterable instead of a scalar for my_number. my_number[0] will ' 'be used, but please input a scalar next time.') # Before, the function would run without an error if given a one element array, but it would not return the # correct answer. my_number = my_number[0] printd('my_number is now {:}'.format(my_number)) else: raise TypeError("closestIndex() in utils_math.py requires a numeric scalar as the second argument. Got " "instead: {:}".format(my_number)) pos = bisect.bisect_left(my_list, my_number) if pos == 0: return 0 if pos == len(my_list): return pos-1 before = pos - 1 after = pos if my_list[after] - my_number < my_number - my_list[before]: return pos else: return pos-1
[docs]@_available_to_user_math def order_axes(M, order): ''' arbitrary sort axes of a multidimensional array :param M: input multidimensional array :param order: integers with list of axes to be ordered :return: M with reordered axes ''' source=''.join([chr(k+65) for k in sorted(order)]) dest=''.join([chr(k+65) for k in order]) return np.einsum(source+'->'+dest, M)
#----------- # binning #-----------
[docs]@_available_to_user_math def qgriddata(px, py, data, X, Y, func=np.sum, fill=np.nan): ''' :param px: :param py: :param data: :param X: :param Y: :param func: :return: ''' if len(X.shape)==2: x=X[1,:] else: x=X nx=len(x) if len(Y.shape)==2: y=Y[1,:] else: y=Y ny=len(y) indx=interpolate.interp1d(x,np.arange(nx)*1.0,kind='nearest',bounds_error=False)(px) indy=interpolate.interp1d(y,np.arange(ny)*1.0,kind='nearest',bounds_error=False)(py) res=np.zeros(nx*ny) res[:]=fill ind=np.array(indx+indy*ny,np.int64) tmp={} for kd,ki in enumerate(ind): if ki not in tmp: tmp[ki]=[data[kd]] else: tmp[ki].append(data[kd]) for ki in tmp: res[ki]=func(tmp[ki]) res=(res.reshape((nx,ny))) return res
#----------- # interpolation #-----------
[docs]class RectBivariateSplineNaN: def __init__(self, Z, R, Q, *args, **kw): tmp = Q.copy() bad = np.isnan(tmp) self.thereAreNaN = False if bad.any(): self.thereAreNaN = True tmp[bad] = 0 self.mask = interpolate.RectBivariateSpline(Z, R, bad, kx=1, ky=1) self.spline = interpolate.RectBivariateSpline(Z, R, tmp, *args, **kw) def __call__(self,*args): tmp=self.spline(*args) if self.thereAreNaN: mask=self.mask(*args) tmp[mask>0.01]=np.nan return tmp
[docs] def ev(self,*args): tmp=self.spline.ev(*args) if self.thereAreNaN: mask=self.mask.ev(*args) tmp[mask>0.01]=np.nan return tmp
if compare_version(scipy.__version__,'0.17.0')<0: def _check_broadcast_up_to(arr_from, shape_to, name): """Helper to check that arr_from broadcasts up to shape_to""" shape_from = arr_from.shape if len(shape_to) >= len(shape_from): for t, f in zip(shape_to[::-1], shape_from[::-1]): if f != 1 and f != t: break else: # all checks pass, do the upcasting that we need later if arr_from.size != 1 and arr_from.shape != shape_to: arr_from = np.ones(shape_to, arr_from.dtype) * arr_from return arr_from.ravel() # at least one check failed raise ValueError('%s argument must be able to broadcast up ' 'to shape %s but had shape %s' % (name, shape_to, shape_from)) def _do_extrapolate(fill_value): """Helper to check if fill_value == "extrapolate" without warnings""" return isinstance(fill_value, str) and fill_value == 'extrapolate' class interp1d(_Interpolator1D): """ Interpolate a 1-D function. `x` and `y` are arrays of values used to approximate some function f: ``y = f(x)``. This class returns a function whose call method uses interpolation to find the value of new points. Parameters ---------- x : (N,) array_like A 1-D array of real values. y : (...,N,...) array_like A N-D array of real values. The length of `y` along the interpolation axis must be equal to the length of `x`. kind : str or int, optional Specifies the kind of interpolation as a string ('linear', 'nearest', 'zero', 'slinear', 'quadratic, 'cubic' where 'slinear', 'quadratic' and 'cubic' refer to a spline interpolation of first, second or third order) or as an integer specifying the order of the spline interpolator to use. Default is 'linear'. axis : int, optional Specifies the axis of `y` along which to interpolate. Interpolation defaults to the last axis of `y`. copy : bool, optional If True, the class makes internal copies of x and y. If False, references to `x` and `y` are used. The default is to copy. bounds_error : bool, optional If True, a ValueError is raised any time interpolation is attempted on a value outside of the range of x (where extrapolation is necessary). If False, out of bounds values are assigned `fill_value`. By default, an error is raised unless `fill_value="extrapolate"`. fill_value : array-like or (array-like, array_like) or "extrapolate", optional - if a np.ndarray (or float), this value will be used to fill in for requested points outside of the data range. If not provided, then the default is NaN. The array-like must broadcast properly to the dimensions of the non-interpolation axes. - If a two-element tuple, then the first element is used as a fill value for ``x_new < x[0]`` and the second element is used for ``x_new > x[-1]``. Anything that is not a 2-element tuple (e.g., list or np.ndarray, regardless of shape) is taken to be a single array-like argument meant to be used for both bounds as ``below, above = fill_value, fill_value``. .. versionadded:: 0.17.0 - If "extrapolate", then points outside the data range will be extrapolated. ("nearest" and "linear" kinds only.) .. versionadded:: 0.17.0 assume_sorted : bool, optional If False, values of `x` can be in any order and they are sorted first. If True, `x` has to be an array of monotonically increasing values. Methods ------- __call__ See Also -------- splrep, splev Spline interpolation/smoothing based on FITPACK. UnivariateSpline : An object-oriented wrapper of the FITPACK routines. interp2d : 2-D interpolation Examples -------- >> import matplotlib.pyplot as plt >> from scipy import interpolate >> x = np.arange(0, 10) >> y = np.exp(-x/3.0) >> f = interpolate.interp1d(x, y) >> >> xnew = np.arange(0, 9, 0.1) >> ynew = f(xnew) # use interpolation function returned by `interp1d` >> plt.plot(x, y, 'o', xnew, ynew, '-') >> plt.show() """ def __init__(self, x, y, kind='linear', axis=-1, copy=True, bounds_error=None, fill_value=np.nan, assume_sorted=False): """ Initialize a 1D linear interpolation class.""" _Interpolator1D.__init__(self, x, y, axis=axis) self.bounds_error = bounds_error # used by fill_value setter self.copy = copy if kind in ['zero', 'slinear', 'quadratic', 'cubic']: order = {'nearest': 0, 'zero': 0, 'slinear': 1, 'quadratic': 2, 'cubic': 3}[kind] kind = 'spline' elif isinstance(kind, int): order = kind kind = 'spline' elif kind not in ('linear', 'nearest'): raise NotImplementedError("%s is unsupported: Use fitpack " "routines for other types." % kind) x = np.array(x, copy=self.copy) y = np.array(y, copy=self.copy) if not assume_sorted: ind = np.argsort(x) x = x[ind] y = np.take(y, ind, axis=axis) if x.ndim != 1: raise ValueError("the x array must have exactly one dimension.") if y.ndim == 0: raise ValueError("the y array must have at least one dimension.") # Force-cast y to a floating-point type, if it's not yet one if not issubclass(y.dtype.type, np.inexact): y = y.astype(np.float_) # Backward compatibility self.axis = axis % y.ndim # Interpolation goes internally along the first axis self.y = y self._y = self._reshape_yi(self.y) self.x = x del y, x # clean up namespace to prevent misuse; use attributes self._kind = kind self.fill_value = fill_value # calls the setter, can modify bounds_err # Adjust to interpolation kind; store reference to *unbound* # interpolation methods, in order to avoid circular references to self # stored in the bound instance methods, and therefore delayed garbage # collection. See: http://docs.python.org/2/reference/datamodel.html if kind in ('linear', 'nearest'): # Make a "view" of the y array that is rotated to the interpolation # axis. minval = 2 if kind == 'nearest': self.x_bds = (self.x[1:] + self.x[:-1]) / 2.0 self._call = self.__class__._call_nearest else: # Check if we can delegate to np.interp (2x-10x faster). cond = self.x.dtype == np.float_ and self.y.dtype == np.float_ cond = cond and self.y.ndim == 1 cond = cond and not _do_extrapolate(fill_value) if cond: self._call = self.__class__._call_linear_np else: self._call = self.__class__._call_linear else: minval = order + 1 self._spline = interpolate.splmake(self.x, self._y, order=order) self._call = self.__class__._call_spline if len(self.x) < minval: raise ValueError("x and y arrays must have at " "least %d entries" % minval) @property def fill_value(self): # backwards compat: mimic a public attribute return self._fill_value_orig @fill_value.setter def fill_value(self, fill_value): # extrapolation only works for nearest neighbor and linear methods if _do_extrapolate(fill_value): if self._kind not in ('nearest', 'linear'): raise ValueError("Extrapolation does not work with " "kind=%s" % self._kind) if self.bounds_error: raise ValueError("Cannot extrapolate and raise " "at the same time.") self.bounds_error = False self._extrapolate = True else: broadcast_shape = (self.y.shape[:self.axis] + self.y.shape[self.axis + 1:]) if len(broadcast_shape) == 0: broadcast_shape = (1,) # it's either a pair (_below_range, _above_range) or a single value # for both above and below range if isinstance(fill_value, tuple) and len(fill_value) == 2: below_above = [np.asarray(fill_value[0]), np.asarray(fill_value[1])] names = ('fill_value (below)', 'fill_value (above)') for ii in range(2): below_above[ii] = _check_broadcast_up_to( below_above[ii], broadcast_shape, names[ii]) else: fill_value = np.asarray(fill_value) below_above = [_check_broadcast_up_to( fill_value, broadcast_shape, 'fill_value')] * 2 self._fill_value_below, self._fill_value_above = below_above self._extrapolate = False if self.bounds_error is None: self.bounds_error = True # backwards compat: fill_value was a public attr; make it writeable self._fill_value_orig = fill_value def _call_linear_np(self, x_new): # Note that out-of-bounds values are taken care of in self._evaluate return np.interp(x_new, self.x, self.y) def _call_linear(self, x_new): # 2. Find where in the orignal data, the values to interpolate # would be inserted. # Note: If x_new[n] == x[m], then m is returned by searchsorted. x_new_indices = searchsorted(self.x, x_new) # 3. Clip x_new_indices so that they are within the range of # self.x indices and at least 1. Removes mis-interpolation # of x_new[n] = x[0] x_new_indices = x_new_indices.clip(1, len(self.x)-1).astype(int) # 4. Calculate the slope of regions that each x_new value falls in. lo = x_new_indices - 1 hi = x_new_indices x_lo = self.x[lo] x_hi = self.x[hi] y_lo = self._y[lo] y_hi = self._y[hi] # Note that the following two expressions rely on the specifics of the # broadcasting semantics. slope = (y_hi - y_lo) / (x_hi - x_lo)[:, None] # 5. Calculate the actual value for each entry in x_new. y_new = slope*(x_new - x_lo)[:, None] + y_lo return y_new def _call_nearest(self, x_new): """ Find nearest neighbour interpolated y_new = f(x_new).""" # 2. Find where in the averaged data the values to interpolate # would be inserted. # Note: use side='left' (right) to searchsorted() to define the # halfway point to be nearest to the left (right) neighbour x_new_indices = searchsorted(self.x_bds, x_new, side='left') # 3. Clip x_new_indices so that they are within the range of x indices. x_new_indices = x_new_indices.clip(0, len(self.x)-1).astype(intp) # 4. Calculate the actual value for each entry in x_new. y_new = self._y[x_new_indices] return y_new def _call_spline(self, x_new): return spleval(self._spline, x_new) def _evaluate(self, x_new): # 1. Handle values in x_new that are outside of x. Throw error, # or return a list of mask array indicating the outofbounds values. # The behavior is set by the bounds_error variable. x_new = asarray(x_new) y_new = self._call(self, x_new) if not self._extrapolate: below_bounds, above_bounds = self._check_bounds(x_new) if len(y_new) > 0: # Note fill_value must be broadcast up to the proper size # and flattened to work here y_new[below_bounds] = self._fill_value_below y_new[above_bounds] = self._fill_value_above return y_new def _check_bounds(self, x_new): """Check the inputs for being in the bounds of the interpolated data. Parameters ---------- x_new : array Returns ------- out_of_bounds : bool array The mask on x_new of values that are out of the bounds. """ # If self.bounds_error is True, we raise an error if any x_new values # fall outside the range of x. Otherwise, we return an array indicating # which values are outside the boundary region. below_bounds = x_new < self.x[0] above_bounds = x_new > self.x[-1] # !! Could provide more information about which values are out of bounds if self.bounds_error and below_bounds.any(): raise ValueError("A value in x_new is below the interpolation " "range.") if self.bounds_error and above_bounds.any(): raise ValueError("A value in x_new is above the interpolation " "range.") # !! Should we emit a warning if some values are out of bounds? # !! matlab does not. return below_bounds, above_bounds
[docs]@_available_to_user_math class interp1e(interpolate.interp1d): """ Shortcut for scipy.interpolate.interp1d with fill_value='extrapolate' and bounds_error=False as defaults """ __doc__ += interpolate.interp1d.__doc__.replace('\n ----------\n', ':\n').replace('\n -------\n', ':\n').replace('\n --------\n', ':\n') def __init__(self, x, y, *args, **kw): kw.setdefault('fill_value', 'extrapolate') kw.setdefault('bounds_error', False) interpolate.interp1d.__init__(self, x, y, *args, **kw)
[docs]@_available_to_user_math class interp1dPeriodic: """ 1D linear interpolation for periodic functions """ def __init__(self,t,y,*args,**kw): """ :param t: array Independent variable: Angle (rad). Doesn't have to be on any particular interval as it will be unwrapped. :param y: array matching length of t Dependent variable: values as a function of t """ self.t_other,ti=np.unique(np.unwrap(t),return_index=True) if self.t_other[0]>self.t_other[1]: self.t_other=-self.t_other self.tckp=interpolate.splrep(self.t_other, y[ti], k=1, per=1) def __call__(self,t_self,*args): """ :param t_self: float or array New basis for independent variable: Angle (rad). Doesn't have to be on any particular interval. :return: float or array matching length of t_self Interpolated values at new t values specified by t_self """ if np.iterable(t_self): t_self=np.unwrap(t_self) if t_self[0]<t_self[1]: t_self=-t_self tmp=interpolate.splev( (t_self-self.t_other[0])%max(self.t_other-self.t_other[0])+self.t_other[0] ,self.tckp, ext=2) if np.iterable(t_self): return tmp else: return tmp[0]
[docs]@_available_to_user_math class uinterp1d(interp1e): """ Adjusted scipy.interpolate.interp1d (documented below) to interpolate the nominal_values and std_devs of an uncertainty array. NOTE: uncertainty in the `x` data is neglected, only uncertainties in the `y` data are propagated. :param std_pow: float. Uncertainty is raised to this power, interpolated, then lowered. (Note std_pow=2 interpolates the variance, which is often used in fitting routines). Additional arguments and key word arguments are as in interp1e (documented below). :Examples: >> x = np.linspace(0,2*np.pi,30) >> u = unumpy.uarray(np.cos(x),np.random.rand(len(x))) >> >> fi = utils.uinterp1d(x,u,std_pow=2) >> xnew = np.linspace(x[0],x[-1],1e3) >> unew = fi(xnew) >> >> f = figure(num='uniterp example') >> f.clf() >> ax = f.add_subplot(111) >> uerrorbar(x,u) >> uband(xnew,unew) interp1e Documentation: """ __doc__+= interp1e.__doc__.replace('\n ----------\n',':\n').replace('\n -------\n',':\n').replace('\n --------\n',':\n') def __init__(self, x, y, kind='linear', axis=-1, copy=True, bounds_error=True, fill_value=np.nan, assume_sorted=False, std_pow=2, **kw): self._std_pow = std_pow self._y_is_uncertain = is_uncertain(y) # convert Variables to floats, shape to shape+[2] if self._y_is_uncertain: axis = axis % y.ndim yfloat=np.rollaxis(np.array([unumpy.nominal_values(y),unumpy.std_devs(y)**std_pow]),0,len(y.shape)+1) else: yfloat = y x=unumpy.nominal_values(x) interp1e.__init__(self, x, yfloat, kind=kind, axis=axis, copy=copy, bounds_error=bounds_error, fill_value=fill_value, assume_sorted=assume_sorted, **kw) # overwrite y attr (note _y used internally) self.y = y def _finish_y(self, y, x_shape): """Return y to uncertainty array before returning""" y = interp1e._finish_y(self, y, x_shape) if self._y_is_uncertain: y = unumpy.uarray(np.take(y,0,axis=-1),np.abs(np.take(y,1,axis=-1))**(1./self._std_pow)) return y
[docs]@_available_to_user_math class uinterp1e(uinterp1d): """ Adjusted unterp1d to extrapolate Arguments and key word arguments are as in uinterp1d (documented below). :uinterp1d Documentation: """ __doc__+= uinterp1d.__doc__.replace('\n ----------\n',':\n').replace('\n -------\n',':\n').replace('\n --------\n',':\n') def __init__(self, x, y, kind='linear', axis=-1, copy=True, assume_sorted=False, std_pow=2, **kw): kw.setdefault('bounds_error',False) kw.setdefault('fill_value','extrapolate') uinterp1d.__init__(self, x, y, kind=kind, axis=axis,copy=copy, assume_sorted=assume_sorted, **kw)
[docs]@_available_to_user_math class URegularGridInterpolator(interpolate.RegularGridInterpolator): """ Adjusted scipy.interpolate.RegularGridInterpolator (documented below) to interpolate the nominal_values and std_devs of an uncertainty array. :param std_pow: float. Uncertainty is raised to this power, interpolated, then lowered. (Note std_pow=2 interpolates the variance, which is often used in fitting routines). Additional arguments and key word arguments are as in RegularGridInterpolator. :Examples: Make some sample 2D data >> x = np.linspace(0,2*np.pi,30) >> y = np.linspace(0,2*np.pi,30) >> z = np.cos(x[:,np.newaxis]+y[np.newaxis,:]) >> u = unumpy.uarray(np.cos(x[:,np.newaxis]+y[np.newaxis,:]),np.random.rand(*z.shape)) Form interpolator >> fi = URegularGridInterpolator((x,y),u,std_pow=2) interpolate along the diagonal of (x,y) >> xnew = np.linspace(x[0],x[-1],1e3) >> unew = fi(zip(xnew,xnew)) Compare original and interpolated values. >> f = figure(num='URegularGridInterpolator example') >> f.clf() >> ax = f.add_subplot(111) >> uerrorbar(x,u.diagonal()) >> uband(xnew,unew) Note the interpolated uncertainty between points is curved by std_pow=2. The curve is affected by the uncertainty of nearby off diagonal points (not shown). :RegularGridInterpolator Documentation: """ __doc__+=interpolate.RegularGridInterpolator.__doc__.replace('\n ----------\n',':\n').replace('\n -------\n',':\n').replace('\n --------\n',':\n').replace('\n -----\n',':\n') def __init__(self, points, values, method="linear", bounds_error=True, fill_value=np.nan,std_pow=2,**kw): y = values self._std_pow = std_pow self._values_are_uncertain = is_uncertain(y) # convert Variables to floats, shape to shape+[2] if self._values_are_uncertain: yfloat=np.rollaxis(np.array([unumpy.nominal_values(y),unumpy.std_devs(y)**std_pow]),0,len(y.shape)+1) else: yfloat = y interpolate.RegularGridInterpolator.__init__(self, points, yfloat, method=method, bounds_error=bounds_error, fill_value=fill_value, **kw) # overwrite y attr (note _y used internally) self.uvalues = y def __call__(self, xi, method=None): """ Interpolate floats, then re-form uncertainty array. """ y = interpolate.RegularGridInterpolator.__call__(self, xi, method=method) if self._values_are_uncertain: y = unumpy.uarray(np.take(y,0,axis=-1),np.take(y,1,axis=-1)**(1./self._std_pow)) return y __call__.__doc__ = interpolate.RegularGridInterpolator.__call__.__doc__
#----------- # basic math #-----------
[docs]@_available_to_user_math def cumtrapz(y, x=None, dx=1.0, axis=-1, initial=0): """ This is a convenience function for scipy.integrate.cumtrapz. Notice that here `initial=0` which is what one most often wants, rather than the `initial=None`, which is the default for the scipy function. Cumulatively integrate y(x) using the composite trapezoidal rule. This is the right way to integrated derivatives quantities which were calculated with `gradient`. If a derivative was obtained with the `diff` command, then the `cumsum` command should be used for its integration. :param y: Values to integrate. :param x: The coordinate to integrate along. If None (default), use spacing dx between consecutive elements in y. :param dx: Spacing between elements of y. Only used if x is None :param axis : Specifies the axis to cumulate. Default is -1 (last axis). :param initial : If given, uses this value as the first value in the returned result. Typically this value should be 0. If None, then no value at x[0] is returned and the returned array has one element less than y along the axis of integration. :return: The result of cumulative integration of y along axis. If initial is None, the shape is such that the axis of integration has one less value than y. If initial is given, the shape is equal to that of y. """ return integrate.cumtrapz(y, x=x, dx=dx, axis=axis, initial=initial)
[docs]@_available_to_user_math def deriv(x, y): ''' This function returns the derivative of the 2nd order lagrange interpolating polynomial of y(x) :param x: x axis array :param y: y axis array :return: dy/dx ''' x = np.array(x) y = np.array(y) def dlip(ra,r,f): '''dlip - derivative of lagrange interpolating polynomial''' r1,r2,r3 = r f1,f2,f3 = f return ((ra-r1)+(ra-r2))/(r3-r1)/(r3-r2)*f3+((ra-r1)+(ra-r3))/(r2-r1)/(r2-r3)*f2+((ra-r2)+(ra-r3))/(r1-r2)/(r1-r3)*f1 return np.array([dlip(x[0],x[0:3],y[0:3])]+list(dlip(x[1:-1],[x[0:-2],x[1:-1],x[2:]],[y[0:-2],y[1:-1],y[2:]]))+[dlip(x[-1],x[-3:],y[-3:])])
[docs]@_available_to_user_math def reverse_enumerate(l): import itertools return zip(range(len(l)-1, -1, -1), reversed(l))
[docs]@_available_to_user_math def factors(n): """ get all the factors of a number > print(list(factors(100))) """ import math large_factors = [] for i in range(1, int(math.sqrt(n) + 1)): if n % i is 0: yield i if i is not n / i: large_factors.insert(0, n / i) for factor in large_factors: yield factor
[docs]def greatest_common_delta(input_int_array): ''' Given an array of integers it returns the greatest uniform delta step :param input_int_array: array of integers :return: greatest uniform delta step ''' input_int_array = np.atleast_1d(input_int_array).astype(int) if len(input_int_array) < 2: return 1 delta = np.unique(np.diff(np.unique(input_int_array))) if len(input_int_array) < 2 or len(delta) < 2: return delta[0] from math import gcd d = gcd(delta[0], delta[1]) for k in range(len(delta)): d = gcd(delta[k], d) return d
[docs]@_available_to_user_math def mad(x, axis=-1, std_norm=True): """ Median absolute deviation, defined as `1.4826 * np.median(np.abs(np.median(x)-x))` :param x: input data :param axis: axis along which to apply mad :param std_norm: if True (default) multiply output by 1.4826 to make MAD normalization agree with standard deviation for normal distribution. This normalization makes mad a robust estimator of the standard deviation. http://books.google.com/books?id=i2bD50PbIikC&lpg=PA118&dq=1.4826&pg=PA118#v=onepage&q=1.4826&f=false :return: Median absolute deviation of the input data """ norm = 1.0 if std_norm: norm = 1.4826 return norm * np.median( np.abs( np.expand_dims( np.median(x, axis=axis) ,axis=axis) - x) ,axis=axis)
[docs]@_available_to_user_math def mad_outliers(data, m = 3., outliers_or_valid='valid'): """ Function to identify outliers data based on median absolute deviation (mad) distance. Note: used median absolute deviation defined as `1.4826 * np.median(np.abs(np.median(x)-x))` :param data: input data array (if a dictionary of arrays, the `mad_outliers` function is applied to each of the values in the dictionary) :param m: mad distance multiplier from the median after which a point is considered an outlier :param outliers_or_valid: return valid/outlier points (`valid` is default) :return: boolean array indicating which data points are a within `m` mad from the median value (i.e. the valid points) """ if isinstance(data,dict): for k,item in enumerate(data.keys()): if k==0: out=mad_outliers(data[item], m, outliers_or_valid) else: out&=mad_outliers(data[item], m, outliers_or_valid) if outliers_or_valid == 'valid': return out else: return ~out d = 1.4826 * np.abs(data - np.median(data)) mdev = np.median(d) s = d/mdev if mdev else 0 if outliers_or_valid == 'valid': return s<m else: return s>=m
[docs]def bin_outliers(data, mincount, nbins, outliers_or_valid='valid'): """ Function to identify outliers data based on binning of data. The algorythm bins the data in `nbins` and then considers valid data only the data that falls within the bins that have at least `mincount` counts. :param data: input data array (if a dictionary of arrays, the `bin_outliers` function is applied to each of the values in the dictionary) :param mincount: minimum number of counts within a bin for data to be considered as valid :param nbins: number of bins for binning of data :param outliers_or_valid: return valid/outlier points (`valid` is default) :return: boolean array indicating which data points are a valid or not """ if isinstance(data,dict): for k,item in enumerate(data.keys()): if k==0: out=bin_outliers(data[item], mincount, nbins, outliers_or_valid) else: out&=bin_outliers(data[item], mincount, nbins, outliers_or_valid) if outliers_or_valid == 'valid': return out else: return ~out tmp=np.histogram(data.flat,nbins) ih=np.where(tmp[0]>mincount)[0] limits=np.nanmin(tmp[1][ih]),np.nanmax(tmp[1][ih+1]) if outliers_or_valid=='valid': return (data>=limits[0]) & (data<=limits[1]) else: return (data<limits[0]) | (data>limits[1])
[docs]@_available_to_user_math def powerlaw_fit(x, y): """ evaluates multidimensional power law for y based on inputs x :param x: 2D array of inputs [N,M] :param y: 1D array of output [M] :return: power law coefficient `p`, fitting function `f(p,x)` """ fitfunc = lambda p, x: 10**(np.sum(p[:-1][:,np.newaxis]*np.log10(x),0)+p[-1]) errfunc = lambda p, x, y: (np.log10(y) - np.log10(fitfunc(p, x))).astype(float) from scipy import optimize pinit = [1]*x.shape[0]+[0] out = optimize.leastsq(errfunc, pinit, args=(x, y), full_output=1) return out[0], fitfunc
[docs]@_available_to_user_math def exp_no_overflow(arg, factor=1.0, minpad=1e3, extra_pad=1e3, return_big_small=False): """ Performs exp(arg) but first limits the value of arg to prevent floating math errors. Checks sys.float_info to so it can avoid floating overflow or underflow. Can be informed of factors you plan on multiplying with the exponential result later in order to make the limits more restrictive (the limits are "padded") and avoid over/underflow later on in your code. :param arg: Argument of exponential function :param factor: Factor that might be multiplied in to exponential function later. Adjusts limits to make them more restrictive and prevent overflow later on. The adjustment to limits is referred to as padding. :param minpad: Force the padding to be at least a certain size. :param extra_pad: Extra padding beyond what's determined by minpad and factor :param return_big_small: T/F: flag to just return big and small numbers. You may be able to speed up execution in repeated calls by getting appropriate limits, doing your own cropping, and looping exp() instead of looping exp_no_overflow(). Even in this case, exp_no_overflow() can help you pick good limits. Or if you don't have time for any of that, you can probably use -70 and 70 as the limits on arg, which will get you to order 1e30. :return: exp(arg) with no math errors, or (big, small) if return_big_small is set """ # Sanitize: if minpad == 0: minpad = 1 extra_pad = max([extra_pad, 1.0]) # Define a pad. Padding gives a margin below overflow limit to allow for multiplying the exp() time some factor. # Make sure the big pad is >= 1, small pad is <= 1, and neither is 0. bigpad = max([np.atleast_1d(abs(factor)).max(), 1.0, minpad]) * extra_pad smallpad = max([min([np.atleast_1d(abs(factor)).min(), 1.0, 1.0/minpad]), np.sqrt(sys.float_info.min)]) / extra_pad # printd('bigpad = {:}, smallpad = {:}'.format(bigpad, smallpad)) # Take the operating system's floating max (min) and reduce (increase) it by pad, then take the log because this # thing will be compared to the argument of exp(). big = float(np.log(sys.float_info.max / bigpad)) small = float(np.log(sys.float_info.min / smallpad)) if return_big_small: return big, small # Apply the limits to get floating over/underflow protection scalar_input = not np.shape(arg) arg = np.atleast_1d(arg) arg2 = copy.copy(arg) arg2[np.isnan(arg2)] = 0 arg[arg2 > big] = big arg[arg2 < small] = small # printd('big = {:}, small = {:}'.format(big, small)) result = np.exp(arg) return result[0] if scalar_input else result
[docs]@_available_to_user_math def dimens(x): """ Returns the number of dimensions in a mixed list/array object. Handles mutli-level nested iterables. From Craig Burgler on stack exchange https://stackoverflow.com/a/39255678/6605826 This would be useful, for example on a case like: x = [np.array([1, 2, 3]), np.array([2, 5, 10, 20])] :param x: iterable (maybe nested iterable) :return: int Generalized dimension count, considering that some elements in a list might be iterable and others not. """ try: size = len(x) except TypeError: # Not an iterable return 0 else: if size: # Non-empty iterable return 1 + max(list(map(dimens, x))) else: # Empty iterable return 1
[docs]@_available_to_user_math def safe_divide(numerator, denominator, fill_value=0): """ Division function to safely compute the ratio of two lists/arrays. The fill_value input parameter specifies what value should be filled in for the result whenever the denominator is 0. :param numerator: numerator of the division :param denominator: denominator of the division :param fill_value: fill value when denominator is 0 :return: division with fill_value where nan or inf would have been instead """ division = np.zeros_like(numerator) + fill_value ind = denominator != 0 division[ind] = numerator[ind] / denominator[ind] return division
#----------- # inverse scale-length #-----------
[docs]def calcz(x, y, consistent_reconstruction=True): ''' Calculate Z: the inverse normalized scale-length The function is coded in such a way to avoid NaN and Inf where y==0 z = -dy/dx/y :param x: x axis array :param y: y axis array :param consistent_reconstruction: calculate z so that integration of z with integz exactly generates original profile >> z_cF =calcz(rho,ne,consistent_reconstruction=False) >> z_cT=calcz(rho,ne,consistent_reconstruction=True) >> plt.plot(rho,ne) >> plt.plot(rho,integz(rho,z_cT,rho[0],ne[0],rho),'.') >> plt.plot(rho,integz(rho,z_cF,rho[0],ne[0],rho),'x') :return: z calculated at x ''' if not consistent_reconstruction or np.any(y==0): tmp = np.nanmin(np.abs(y[np.where(y != 0)])) return -deriv(x, y) / (np.abs(y) + tmp * 1E-32) * np.sign(y) else: rat = np.log(y/y[-1]) z = np.zeros(y.size) for i in range(1,len(y)): z[i] = 2*(rat[i]-rat[i-1])/(x[i]-x[i-1])-z[i-1] return -z
[docs]def mergez(x0, z0, x1, z1, core, nml, edge, x): """ Perform merge of two Z profiles * `rho < core`: Z goes linearly to 0 on axis * `core < rho < nml`: profile 0 * `nml < rho < edge`: linearly connect Z at `nml` and `edge` * `rho > edge`: profile 1 :param x0: abscissa of first profile :param z0: values of first Z profile :param x1: abscissa of second profile :param z1: values of second Z profile :param core: abscissa of core boundary :param nml: abscissa of nml boundary :param edge: abscissa of edge boundary :param: returns tuple of merged x and z points """ if edge < min(x1): raise Exception('edge must be > min(x1)') elif core < 0: raise Exception('core must be >= 0') elif core < nml and nml > max(x0): raise Exception('nml must be < max(x0)') # capture all of the important points x_ = [] if core < nml: x_.extend(x0[np.where(((x0 > 0) & (x0 > core) & (x0 <= nml)))[0]]) x_.extend(x1[np.where(x1 >= edge)[0]]) x_.append(nml) x_.append(edge) x_.extend(x) x_ = np.unique(x_) # profile outside of edge (included) x1_ = x_[np.where((x_ >= edge) & (x_ > core))[0]] z1_ = interp1e(x1, z1)(x1_) # profile between core (included) and nml (included) if core < nml: x0_ = x_[np.where((x_ > 0) & (x_ > core) & (x_ <= nml))[0]] z0_ = interp1e(x0, z0)(x0_) elif core != 0: x0_ = [core] z0_ = [z1_[0]] else: x0_ = [] z0_ = [] x_ = np.hstack((0, x0_, x1_)) z_ = np.hstack((0, z0_, z1_)) x_, i = np.unique(x_, return_index=True) z_ = z_[i] return x, interp1e(x_, z_)(x)
[docs]def integz(x0, z0, xbc, ybc, x, clipz=True): ''' Integrate inverse scale-length Z to get profile :param x0: abscissa of the Z profile :param z0: Z profile :param xbc: abscissa of the boundary condition :param ybc: value of the boundary condition :param x: integrate the Z profile at these points :param clipz: use clipped z for extrapolation (alternatively linear) :return: integrated Z profile at points defined by x ''' # make sure that the boundary condtion point is part of the itegration x_ = x x = np.unique(np.hstack((x_, xbc))) # evaluate z at requested points and allow for clipping z = uinterp1e(x0, z0)(x) if clipz: inside = np.where((x >= min(x0)) & (x <= max(x0)))[0].astype(int) low = np.where(x < min(x0))[0].astype(int) high = np.where(x > max(x0))[0].astype(int) z[low] = z[inside][0] z[high] = z[inside][-1] # backward integration <<< i0 = np.where(x <= xbc)[0] x0 = x[i0] z0 = z[i0] y0 = [] if len(x0): t0 = -(z0[:-1] + z0[1:]) * np.diff(x0) / 2. y0 = np.cumprod([ybc] + (unumpy.exp(-t0[::-1]).tolist()))[::-1] # forward integration >>> i1 = list(range(len(x)))[max(np.hstack((0, i0))):] x1 = x[i1] z1 = z[i1] y1 = [] if len(x1): t1 = (z1[:-1] + z1[1:]) * np.diff(x1) / 2. y1 = np.cumprod([ybc] + (unumpy.exp(-t1).tolist()))[1:] return uinterp1e(x, np.hstack((y0, y1)))(x_)
#----------- # peak detection #-----------
[docs]@_available_to_user_math def detect_peaks(x, mph=None, mpd=1, threshold=0, edge='rising', kpsh=False, valley=False, show=False, ax=None): """Detect peaks in data based on their amplitude and other features. Parameters ---------- x : 1D array_like data. mph : {None, number}, optional (default = None) detect peaks that are greater than minimum peak height. mpd : positive integer, optional (default = 1) detect peaks that are at least separated by minimum peak distance (in number of data). threshold : positive number, optional (default = 0) detect peaks (valleys) that are greater (smaller) than `threshold` in relation to their immediate neighbors. edge : {None, 'rising', 'falling', 'both'}, optional (default = 'rising') for a flat peak, keep only the rising edge ('rising'), only the falling edge ('falling'), both edges ('both'), or don't detect a flat peak (None). kpsh : bool, optional (default = False) keep peaks with same height even if they are closer than `mpd`. valley : bool, optional (default = False) if True (1), detect valleys (local minima) instead of peaks. show : bool, optional (default = False) if True (1), plot data in matplotlib figure. ax : a matplotlib.axes.Axes instance, optional (default = None). Returns ------- ind : 1D array_like indeces of the peaks in `x`. Notes ----- The detection of valleys instead of peaks is performed internally by simply negating the data: `ind_valleys = detect_peaks(-x)` The function can handle NaN's See this IPython Notebook [1]_. References ---------- .. [1] http://nbviewer.ipython.org/github/demotu/BMC/blob/master/notebooks/DetectPeaks.ipynb Examples -------- >> from detect_peaks import detect_peaks >> x = np.random.randn(100) >> x[60:81] = np.nan >> # detect all peaks and plot data >> ind = detect_peaks(x, show=True) >> print(ind) >> x = np.sin(2*np.pi*5*np.linspace(0, 1, 200)) + np.random.randn(200)/5. >> # set minimum peak height = 0 and minimum peak distance = 20 >> detect_peaks(x, mph=0, mpd=20, show=True) >> x = [0, 1, 0, 2, 0, 3, 0, 2, 0, 1, 0] >> # set minimum peak distance = 2 >> detect_peaks(x, mpd=2, show=True) >> x = np.sin(2*np.pi*5*np.linspace(0, 1, 200)) + np.random.randn(200)/5. >> # detection of valleys instead of peaks >> detect_peaks(x, mph=0, mpd=20, valley=True, show=True) >> x = [0, 1, 1, 0, 1, 1, 0] >> # detect both edges >> detect_peaks(x, edge='both', show=True) >> x = [-2, 1, -2, 2, 1, 1, 3, 0] >> # set threshold = 2 >> detect_peaks(x, threshold = 2, show=True) """ def _plot(x, mph, mpd, threshold, edge, valley, ax, ind): """Plot results of the detect_peaks function, see its help.""" if ax is None: from matplotlib import pyplot as plt _, ax = plt.subplots(1, 1, figsize=(8, 4)) ax.plot(x, 'b', lw=1) if ind.size: label = 'valley' if valley else 'peak' label = label + 's' if ind.size > 1 else label ax.plot(ind, x[ind], '+', mfc=None, mec='r', mew=2, ms=8, label='%d %s' % (ind.size, label)) ax.legend(loc='best', framealpha=.5, numpoints=1) ax.set_xlim(-.02*x.size, x.size*1.02-1) ymin, ymax = x[np.isfinite(x)].min(), x[np.isfinite(x)].max() yrange = ymax - ymin if ymax > ymin else 1 ax.set_ylim(ymin - 0.1*yrange, ymax + 0.1*yrange) ax.set_xlabel('Data #', fontsize=14) ax.set_ylabel('Amplitude', fontsize=14) mode = 'Valley detection' if valley else 'Peak detection' ax.set_title("%s (mph=%s, mpd=%d, threshold=%s, edge='%s')" % (mode, str(mph), mpd, str(threshold), edge)) x = np.atleast_1d(x).astype('float64') if x.size < 3: return np.array([], dtype=int) if valley: x = -x # find indices of all peaks dx = x[1:] - x[:-1] # handle NaN's indnan = np.where(np.isnan(x))[0] if indnan.size: x[indnan] = np.inf dx[np.where(np.isnan(dx))[0]] = np.inf ine, ire, ife = np.array([[], [], []], dtype=int) if not edge: ine = np.where((np.hstack((dx, 0)) < 0) & (np.hstack((0, dx)) > 0))[0] else: if edge.lower() in ['rising', 'both']: ire = np.where((np.hstack((dx, 0)) <= 0) & (np.hstack((0, dx)) > 0))[0] if edge.lower() in ['falling', 'both']: ife = np.where((np.hstack((dx, 0)) < 0) & (np.hstack((0, dx)) >= 0))[0] ind = np.unique(np.hstack((ine, ire, ife))) # handle NaN's if ind.size and indnan.size: # NaN's and values close to NaN's cannot be peaks ind = ind[np.in1d(ind, np.unique(np.hstack((indnan, indnan-1, indnan+1))), invert=True)] # first and last values of x cannot be peaks if ind.size and ind[0] == 0: ind = ind[1:] if ind.size and ind[-1] == x.size-1: ind = ind[:-1] # remove peaks < minimum peak height if ind.size and mph is not None: ind = ind[x[ind] >= mph] # remove peaks - neighbors < threshold if ind.size and threshold > 0: dx = np.min(np.vstack([x[ind]-x[ind-1], x[ind]-x[ind+1]]), axis=0) ind = np.delete(ind, np.where(dx < threshold)[0]) # detect small peaks closer than minimum peak distance if ind.size and mpd > 1: ind = ind[np.argsort(x[ind])][::-1] # sort ind by peak height idel = np.zeros(ind.size, dtype=bool) for i in range(ind.size): if not idel[i]: # keep peaks with the same height if kpsh is True idel = idel | (ind >= ind[i] - mpd) & (ind <= ind[i] + mpd) \ & (x[ind[i]] > x[ind] if kpsh else True) idel[i] = 0 # Keep current peak # remove the small peaks and sort back the indices by their occurrence ind = np.sort(ind[~idel]) if show: if indnan.size: x[indnan] = np.nan if valley: x = -x _plot(x, mph, mpd, threshold, edge, valley, ax, ind) return ind
[docs]@_available_to_user_math def find_feature(y, x=None, M=0.01, k=5.0, xmin=None, xmax=None, retall=False): """ Identify edges and center of a sharp feature (pedestal, peak, or well) in a otherwise smooth profile. :param y: np.ndarray. Parameter of interest, such as T_e :param x: np.ndarray. Position basis, such as psi_N (Default is np.arange(len(y)) :param M: float. Gaussian smoothing sigma in units of position basis :param k: float. Difference of gaussians factor (second smoothing sigma is M*k) :param xmin: float. Lower limit on range of X values to consider in search :param xmax: float. Upper limit on range of X values to consider in search :param retall: bool. Return gaussian smoothed functions :returns : ( Inner edge, outer edge, center, value at inner edge, value at outer edge ) :Examples: Here, is a simple 1D example using find_feature to identify the steep gradient region of a tanh. >> x = np.linspace(0,1,1000) >> center = 0.8 >> halfwidth = 0.05 >> y = np.tanh((center-x)/halfwidth)+1 >> dydx = deriv(x,y) >> indx = find_feature(dydx,x=x) >> xpt = x[indx] >> ypt = dydx[indx] >> foundwidth = x[indx[2]]-x[indx[0]] >> print("Width within {:.2f}%".format((1-2*halfwidth/foundwidth)*100)) Note we chose M to set the appropriate scale of interest for the above problem. We can see the sensitivity to this by scanning the tanh width in a 2D example. >> xs = x.reshape(1,-1) >> halfwidths = np.linspace(0.01,0.1,50).reshape(-1,1) >> ys = np.tanh((center-xs)/halfwidths)+1 >> dydxs = np.gradient(ys)[1]/np.gradient(x) M sets the scale for the steepness of interest using one M for a range of bump sizes isolates only the scale of interest >> indxs = np.apply_along_axis(find_feature,1,dydxs,x=x,M=0.01,k=5) Tracking the scale of the bump with M approximates tanh widths >> #indxs = [find_feature(dy,x=x,M=2*hw/10.,k=5) for dy,hw in zip(dydxs,halfwidths)] >> # found peak and edge points of steep gradient region >> foundwidths = map(lambda indx: x[indx[2]]-x[indx[0]], indxs) >> xpts = [x[i] for i in indxs] >> ypts = [yy[i] for yy,i in zip(ys,indxs)] >> dypts= [yy[i] for yy,i in zip(dydxs,indxs)] >> # tanh half width points >> # Note np.tanh(1) = 0.76, and d/dxtanh = sech^2 = 0.42 >> ihws = np.array([(np.abs(x-center-hw).argmin(), >> np.abs(x-center+hw).argmin()) for hw in halfwidths[:,0]]) >> xhws = [x[i] for i in ihws] >> yhws = [yy[i] for yy,i in zip(ys,ihws)] >> dyhws= [yy[i] for yy,i in zip(dydxs,ihws)] Visualize the comparison between tanh widths and the identified region of steep gradient >> close('all') >> f,ax = plt.subplots(3,1) >> ax[0].set_ylabel('y = np.tanh((c-x)2/w)+1') >> ax[0].set_xlabel('x') >> ax[1].set_ylabel('dy/dx') >> ax[1].set_xlabel('x') >> for i in [0,24,49]: ... l, = ax[0].plot(x,ys[i]) ... w1, = ax[0].plot(xhws[i],yhws[i],marker='o',ls='',fillstyle='none',color=l.get_color()) ... w2, = ax[0].plot(xpts[i],ypts[i],marker='x',ls='',fillstyle='none',color=l.get_color()) ... l, = ax[1].plot(x,dydxs[i]) ... w1, = ax[1].plot(xhws[i],dyhws[i],marker='o',ls='',fillstyle='none',color=l.get_color()) ... w2, = ax[1].plot(xpts[i],dypts[i],marker='x',ls='',fillstyle='none',color=l.get_color()) >> ax[-1].set_ylabel('Edge Width') >> ax[-1].set_xlabel('Tanh Width') >> ax[-1].plot([0,2*halfwidths[-1,0]],[0,2*halfwidths[-1,0]]) >> ax[-1].plot(2*halfwidths[:,0],foundwidths,marker='o',lw=0) """ # must be 1D yy = np.atleast_1d(y).ravel() # isolate x range if x is None: x = np.arange(y.shape[0]) if xmin is None: xmin = np.NINF if xmax is None: xmax = np.inf xx = np.atleast_1d(x).ravel() w = np.where((xx >= xmin) & (xx <= xmax))[0] xx, yy = xx[w], yy[w] # difference of gaussian smooths dx = np.mean(np.diff(xx)) s_1 = ndimage.filters.gaussian_filter(yy, M/dx) s_2 = ndimage.filters.gaussian_filter(yy, M*k/dx) sdiff = s_1 - s_2 # feature location peak_i = np.argmax(np.abs(sdiff)) # find width from gaussian smooths crossing on either side of peak left_i = 0 left_all = np.where(np.diff(np.signbit(sdiff[:peak_i])))[0] if len(left_all)>0: left_i += left_all[-1] right_i = peak_i right_all = np.where(np.diff(np.signbit(sdiff[peak_i:])))[0] if len(right_all)>0: right_i += right_all[0] result = np.array((left_i, peak_i, right_i)) if retall: result = (result, s_1, s_2) return result
#----------- # fit parabolas #-----------
[docs]def parabola(x, y): """ y = a*x^2 + b*x + c :return: a,b,c """ if np.any(np.isnan(x)) or np.any(np.isnan(y)): raise np.linalg.LinAlgError('parabola could not be fit with x=%s y=%s'%(x,y)) # A=np.matrix([[x[0]**2,x[0],1],[x[1]**2,x[1],1],[x[2]**2,x[2],1]]) # a,b,c=np.array(A.I*np.matrix(y[0:3]).T).flatten().tolist() #polyfit is equally fast but more robust when x values are similar, and matrix A becomes singular, eg. # x,y=[-1E-16,1,1+1E-16],[1,1,1] with warnings.catch_warnings(): warnings.simplefilter('ignore', np.RankWarning) a,b,c=np.polyfit(x,y, 2) return a,b,c
[docs]def parabolaMax(x, y, bounded=False): ''' Calculate a parabola through x,y, then return the extremum point of the parabola :param x: At least three abcissae points :param y: The corresponding ordinate points :param bounded: False, 'max', or 'min' - False: The extremum is returned regardless of location relative to x (default) - 'max' ('min'): The extremum location must be within the bounds of x, and if not return the location and value of max(y) (min(y)) :return: x_max,y_max - The location and value of the extremum ''' a,b,c=parabola(x,y) x_max=-b/(2*a) y_max=a*x_max**2+b*x_max+c if bounded and (x_max > max(x) or x_max < min(x)): printe('ParabolaMax found the maximum outside the input abcissae; using arg%s(y) instead'%bounded) iy = eval('arg%s(y)'%bounded) x_max = x[iy] y_max = y[iy] return x_max,y_max
[docs]def parabolaCycle(X, Y, ix): ix=np.array([-1,0,1])+ix if (X[0]-X[-1]) < 1E-15: if ix[0]<0: ix[0]=len(X)-1-1 if ix[-1]>(len(X)-1): ix[-1]=0+1 else: if ix[0]<0: ix[0]=len(X)-1 if ix[-1]>(len(X)-1): ix[-1]=0 return parabola(X[ix],Y[ix])
[docs]def parabolaMaxCycle(X, Y, ix, bounded=False): ''' Calculate a parabola through X[ix-1:ix+2],Y[ix-1:ix+2], with proper wrapping of indices, then return the extremum point of the parabola :param X: The abcissae points: an iterable to be treated as periodic :param y: The corresponding ordinate points :param ix: The index of X about which to find the extremum :param bounded: False, 'max', or 'min' - False: The extremum is returned regardless of location relative to x (default) - 'max' ('min'): The extremum location must be within the bounds of x, and if not return the location and value of max(y) (min(y)) :return: x_max,y_max - The location and value of the extremum ''' ix=np.array([-1,0,1])+ix if (X[0]-X[-1]) < 1E-15: if ix[0]<0: ix[0]=len(X)-1-1 if ix[-1]>(len(X)-1): ix[-1]=0+1 else: if ix[0]<0: ix[0]=len(X)-1 if ix[-1]>(len(X)-1): ix[-1]=0 try: return parabolaMax(X[ix],Y[ix],bounded=bounded) except Exception: printe([ix,X[ix],Y[ix]]) return X[ix[1]],Y[ix[1]]
[docs]def paraboloid(x, y, z): ''' z = ax*x^2 + bx*x + ay*y^2 + by*y + c NOTE: This function uses only the first 5 points of the x, y, z arrays to evaluate the paraboloid coefficients :return: ax,bx,ay,by,c ''' if np.any(np.isnan(x.flatten())) or np.any(np.isnan(y.flatten())) or np.any(np.isnan(z.flatten())): raise np.linalg.LinAlgError('paraboloid could not be fit with x=%s y=%s z=%s'%(x,y,z)) A=[] for k in range(5): A.append([x[k]**2,x[k],y[k]**2,y[k],1]) A=np.array(A) ax,bx,ay,by,c=np.dot(np.linalg.inv(A),np.array(z[:5])) return ax,bx,ay,by,c
#----------- # filtering #-----------
[docs]@_available_to_user_math def smooth(x, window_len=11, window='hanning', axis=0): """smooth the data using a window with requested size. This method is based on the convolution of a scaled window with the signal. The signal is prepared by introducing reflected copies of the signal (with the window size) in both ends so that transient parts are minimized in the beginning and end part of the output signal. input: :param x: the input signal :param window_len: the dimension of the smoothing window; should be an odd integer; is ignored if `window` is an array :param window: the window function to use, see scipy.signal.get_window documentation for list of available windows 'flat' or 'boxcar' will produce a moving average smoothing Can also be an array, in which case it is used as the window function and `window_len` is ignored :param axis: The axis that is smoothed output: the smoothed signal :Examples: >> t=np.linspace(-2,2,100) >> x=np.sin(t)+randn(len(t))*0.1 >> y=smooth(x,11) see also: scipy.signal.get_window, np.convolve, scipy.signal.lfilter """ def smooth1d(x, win): window_len = win.size s = np.r_[x[window_len-1:0:-1],x,x[-1:-window_len:-1]] y = np.convolve(win/win.sum(), s, mode='valid') start = int(np.ceil(window_len/2.)-1) stop = start+len(x) return y[start:stop] window_len = int(window_len) if window_len < 3: return x if window == 'flat': # backwards compatibility window = 'boxcar' if isinstance(window, string_types) or isinstance(window, tuple): if window_len%2 == 0: # Use odd-length window to avoid shifting data window_len += 1 win = scipy.signal.get_window(window, window_len, fftbins=False) else: win = np.asarray(window) if len(win.shape) != 1: raise ValueError('window must be 1-D') if x.shape[axis] < win.size: raise ValueError("Input vector needs to be bigger than window length " "of {}".format(win.size)) return np.apply_along_axis(smooth1d, axis, x, win)
[docs]@_available_to_user_math def smoothIDL(data, n_smooth=3, timebase=None, causal=False, indices=False, keep=False): """ An efficient top hat smoother based on the `smooth` IDL routine. The use of cumsum-shift(cumsum) means that execution time is 2xN flops compared to 2 x n_smooth x N for a convolution. If supplied with a timebase, the shortened timebase is returned as the first of a tuple. :param data: (timebase, data) is a shorthand way to pass timebase :param n_smooth: smooth bin size :param timebase: if passed, a tuple with (timebase,smoothed_data) gets returned :param causal: If True, the smoothed signal never preceded the input, otherwise, the smoothed signal is "centred" on the input (for n_smooth odd) and close (1//2 timestep off) for even :param indices: if True, return the timebase indices instead of the times :param keep: Better to throw the partially cooked ends away, but if you want to keep them use keep=True. This is useful for quick filtering applications so that original and filtered signals are easily compared without worrying about timebase >> smooth([1,2,3,4],3) np.array([ 2., 3.]) >> smooth([1.,2.,3.,4.,5.],3) np.array([ 2., 3., 4.]) >> smooth([1,2,3,4,5],timebase=np.array([1,2,3,4,5]),n_smooth=3, causal=False) (np.array([2, 3, 4]), np.array([ 2., 3., 4.])) >> smooth([0,0,0,3,0,0,0],timebase=[1,2,3,4,5,6,7],n_smooth=3, causal=True) ([3, 4, 5, 6, 7], np.array([ 0., 1., 1., 1., 0.])) >> smooth([0,0,0,3,0,0,0],timebase=[1,2,3,4,5,6,7],n_smooth=3, causal=True, indices=True) ([2, 3, 4, 5, 6], np.array([ 0., 1., 1., 1., 0.])) >> smooth([0, 0, 0, 0, 5, 0, 0, 0, 0, 0, 0], 5, keep=1) np.array([ 0., 0., 1., 1., 1., 1., 1., 0., 0., -1., -1.]) """ if len(data) == 2: # a tuple (timebase, data) timebase=data[0] data = data[1] csum = np.cumsum(data) if not(keep): sm_sig = csum[n_smooth-1:] # just the valid bit #sm_sig[1:] -= csum[0:-n_smooth] # should work? - bug in python? sm_sig[1:] = sm_sig[1:] - csum[0:-n_smooth] sm_sig=sm_sig/float(n_smooth) else: # try to retain full length for quick tests - only allow causal=False if causal: raise ValueError("can't be causal and keep ends") else: sm_sig=np.array(data)*0 # get a vector of the right length half_offs = int(n_smooth//2) # use integer round down, only perfect for odd n # See test code [0, 0, 0, 1, 0, 0, 0] # for the n=3 case, want sm[3] to be cs[4]-cs[1] so load cs[1:] into sm, # then subtract cs[0] from sm[2] etc # for n=5, want sm[5] to be cs[7]-cs[2] so load cs[2:] into sm # and subtr cs[0] from sm[3] sm_sig[0:-half_offs] = csum[half_offs:] # get the first bit left shifted sm_sig[half_offs+1:] = sm_sig[half_offs+1:] - csum[0:-(half_offs+1)] sm_sig=sm_sig/float(n_smooth) if timebase is None: return(sm_sig) else: if causal: offset = n_smooth-1 elif keep: offset = 0 else: offset = int(n_smooth//2) if indices: sm_timebase = list(range(offset,offset+len(sm_sig))) else: sm_timebase = timebase[offset:offset+len(sm_sig)] return((sm_timebase, sm_sig))
[docs]@_available_to_user_math def pcs_rc_smooth(xx, yy, s): """ RC lowpass smoothing filter using same formula as in the PCS. Can be used to reproduce PCS lowpass filtering. :param xx: array Independent variable (like time). Does not have to be evenly spaced. :param yy: array matching xx Depending variable :param s: float Smoothing time constant in units matching xx :return: array Smoothed/filtered version of yy """ printd(' start rc_smooth w/ s = {:}'.format(s), topic='pcs_rc_smooth') rc = s # Renaming it makes its meaning clearer dx = np.diff(xx) dx = np.append(dx[0], dx) a = dx / (rc + dx) printd(' mean(dx) = {:}, mean(a) ={:}'.format(np.mean(dx), np.mean(a)), topic='pcs_rc_smooth') ys = yy * 0 ys[0] = yy[0] printd(' min(a) = {:}, max(a) = {:}'.format(min(a), max(a)), topic='pcs_rc_smooth') with np.errstate(under='ignore'): # Ignore harmless underflow (might happen if yy -> 0 & ys asymptotes to 0) for ii in range(1, len(yy)): ys[ii] = ys[ii - 1] * (1 - a[ii]) + yy[ii] * a[ii] printd(' ys[-1] = {:}'.format(ys[-1]), topic='pcs_rc_smooth') return ys
[docs]@_available_to_user_math def smooth_by_convolution(yi, xi=None, xo=None, window_size=None, window_function='gaussian', axis=0, causal=False, interpolate=False, std_dev=2): """ Convolution of a non-uniformly discretized array with window function. The output values are np.nan where no points are found in finite windows (weight is zero). The gaussian window is infinite in extent, and thus returns values for all xo. Supports uncertainties arrays. If the input --does not-- have associated uncertainties, then the output will --not-- have associated uncertainties. :param yi: array_like (...,N,...). Values of input array :param xi: array_like (N,). Original grid points of input array (default y indicies) :param xo: array_like (M,). Output grid points of convolution array (default xi) :param window_size: float. Width of passed to window function (default maximum xi step). For the Gaussian, sigma=window_size/4. and the convolution is integrated across +/-4.*sigma. :param window_function: str/function. Accepted strings are 'hanning','bartlett','blackman','gaussian', or 'boxcar'. Function should accept x and window_size as arguments and return a corresponding weight. :param axis: int. Axis of y along which convolution is performed :param causal: int. Forces f(x>0) = 0. :param interpolate: False or integer number > 0 Paramter indicating to interpolate data so that there are`interpolate` number of data points within a time window. This is useful in presence of sparse data, which would result in stair-case output if not interpolated. The integer value sets the # of points per window size. :param std_dev: str/int Accepted strings are 'propagate', 'expand', 'none', 'deviation', 'variance'. Setting to an integer will convolve the error uncertainties to the std_dev power before taking the std_dev root. std_dev = 'propagate' is true propagation of errors (slow if not interpolating) std_dev = 'expand' is propagation of errors weighted by w~1/window_function std_dev = 'deviation' is equivalent to std_dev=1 std_dev = 'variance' is equivalent to std_dev=2 :return: convolved array on xo >> M=300 >> ninterp = 10 >> window=['hanning','gaussian','boxcar'][1] >> width=0.05 >> f = figure(num='nu_conv example') >> f.clf() >> ax = f.add_subplot(111) >> >> xo=np.linspace(0,1,1000) >> >> x0=xo >> y0=(x0>0.25)&(x0<0.75) >> plt.plot(x0,y0,'b-',label='function') >> >> x1=np.cumsum(rand(M)) >> x1=(x1-min(x1))/(max(x1)-min(x1)) >> y1=(x1>0.25)&(x1<0.75) >> plt.plot(x1,y1,'b.',ms=10,label='subsampled function') >> if window=='hanning': >> ys=smooth(interp1e(x0,y0)(xo),int(len(xo)*width)) >> plt.plot(xo,ys,'g-',label='smoothed function') >> yc=smooth(interp1e(x1,y1)(xo),int(len(xo)*width)) >> plt.plot(xo,yc,'m--',lw=2,label='interp subsampled then convolve') >> >> y1=unumpy.uarray(y1,y1*0.1) >> a=time.time() >> yo=nu_conv(y1,xi=x1,xo=xo,window_size=width,window_function=window,std_dev='propagate',interpolate=ninterp) >> print('nu_conv time: {:}'.format(time.time()-a)) >> ye=nu_conv(y1,xi=x1,xo=xo,window_size=width,window_function=window,std_dev='expand',interpolate=ninterp) >> >> uband(x1,y1,color='b',marker='.',ms=10) >> uband(xo,yo,color='r',lw=3,label='convolve subsampled') >> uband(xo,ye,color='c',lw=3,label='convolve with expanded-error propagation') >> >> legend(loc=0) >> plt.ylim([-0.1,1.1]) >> plt.title('%d points , %s window, %3.3f width, interpolate %s'%(M,window,width, ninterp)) """ if len(yi.shape)>1: yo = np.apply_along_axis(smooth_by_convolution, axis, yi, xi=xi, xo=xo, window_size=window_size, window_function=window_function, causal=causal, interpolate=interpolate, std_dev=std_dev) return yo from scipy import interpolate as scipy_interpolate # named functions if isinstance(window_function,str): win_mult=1.0 if window_function=='gaussian': f=lambda x,a:unumpy.exp(-0.5*(4.0*x/a)**2)/(a/4.0)/np.sqrt(2*np.pi) win_mult=2.0 elif window_function=='hanning': f=lambda x,a:(np.abs(x)<=a/2.)*(1-unumpy.cos(2*np.pi*(x-a/2.)/a))/a elif window_function=='bartlett': f=lambda x,a:(np.abs(x)<=a/2.)*(1-np.abs(x)*2./a) elif window_function=='blackman': f=lambda x,a:(np.abs(x)<=a/2.)*(0.42-0.5*unumpy.cos(2.0*np.pi*(x+a/2.)/a) + 0.08*unumpy.cos(4.0*np.pi*(x+a/2.)/a)) elif window_function in ['boxcar','median']: f=lambda x,a:1.0*((x>=-a/2.)*(x<=a/2.))/a elif window_function in ['triangle']: f = lambda x_, a: (a >= abs(x_)) * (1 - abs(x_ / a)) win_mult = 2.0 else: raise Exception('Valid window functions are: gaussian, hanning, bartlett, blackman, boxcar') else: f=window_function # handle std_dev if isinstance(std_dev,str): std_dev = std_dev.lower() if std_dev == 'variance': std_dev = 2 elif std_dev == 'deviation': std_dev = 1 elif std_dev is None: std_dev='none' # defaults if xi is None: xi = np.arange(yi.shape[0]) if xo is None: xo = xi if window_size is None: window_size=max(np.diff(xi))/2. # initialize input/output arrays yo=np.zeros(xo.shape) if is_uncertain(yi): ei = unumpy.std_devs(yi) yi = unumpy.nominal_values(yi) eo = np.zeros(xo.shape) else: ei = 0.0*yi eo = np.zeros(xo.shape) std_dev='none' # remove nans i=np.where( (~np.isnan(xi)) & (~np.isnan(yi)) ) xi=xi[i] yi=yi[i] ei=ei[i] # If there is no non-Nan input data, no point continuing if yi.shape[0] == 0: printw(' Smooth by convolution: Zero length input, returning Nan output') if std_dev == 'none': return xo * np.nan else: return unumpy.uarray(xo * np.nan, xo * np.nan) if np.any(np.isnan(yi)): printe(yi) raise RuntimeError('nans detected') if not xi.size: if std_dev == 'none': return xo * np.nan else: return unumpy.uarray(xo * np.nan, xo * np.nan) # interpolation (# of points per window size) if interpolate and len(xi)>1: if not (interpolate is False or is_int(interpolate) or interpolate>0): raise ValueError('interpolate argument should be False or a number') xi1=np.linspace(min(xi), max(xi), int(np.ceil((max(xi)-min(xi))/(window_size*win_mult)*interpolate)) ) yi=scipy_interpolate.interp1d(xi,yi)(xi1) if ei is not None: ei=scipy_interpolate.interp1d(xi,ei)(xi1) xi=xi1 # grid-center values for trapezoidal integration dx = np.ediff1d(xi, to_end=0) y_center = np.hstack([(yi[:-1] + yi[1:])*0.5,0]) # convolution for k,xo_k in enumerate(xo): # find distances to requested x point x=xo_k-xi # select points within time window if not causal: i=np.abs(x)<=window_size*win_mult/2. else: i=np.abs(x-window_size*win_mult/4.)<=window_size*win_mult/4. npts = np.count_nonzero(i) # set output to nan if points are outside of convolution window if npts == 0: yo[k]=np.nan eo[k]=np.nan # explicitly use single data point if its the only thing available elif npts == 1: yo[k] = yi[i][0] eo[k] = ei[i][0] # median value and median std elif window_function == 'median': yo[k] = np.median(yi[i]) eo[k] = np.median(ei[i]) # use weighted convolution if multiple points in window else: # generate window and it's normalization ff = f(x[i], window_size) if causal: ff *= (x[i] >= 0) ff_center = ff[:-1] + np.diff(ff) * 0.5 norm = np.sum(ff_center * dx[i][:-1]) # trapezoidal integration # set output to nan if the norm is zero if norm == 0: yo[k] = np.nan eo[k] = np.nan else: # normalize window ff = ff / norm # treat value and error together if std_dev in ['propagate', 'expand']: # propagate uncertainties (slow) # grid-center points of weighted values for trapezoidal integration if std_dev == 'propagate': # simply propogate the uncertainties through the convolution integral fy = unumpy.uarray(ff * yi[i], ff * ei[i]) else: # expand the errorbar of points in the integral bounds by # an amount proportional to 1/window_function before propagating them, # increasing the error associated with the convolution. # Note integrals with only far off points will have larger errors then any of the given points fy = unumpy.uarray(ff * yi[i], (f(0, window_size) / norm) * ei[i]) fy_center = (fy[:-1] + fy[1:]) * 0.5 # trapezoidal integration (uncertainties propagate themselves) # considerable speedup is possible if input grid is uniform if interpolate or np.all(dx[i][:-1]==dx[i][0]): yk = np.sum(fy_center * dx[i][0]) else: yk = np.sum(fy_center * dx[i][:-1]) # re-separate values and uncertainties yo[k] = unumpy.nominal_values(yk) eo[k] = unumpy.std_devs(yk) # convolve value and estimate error separately else: # grid-center points of weighted values for trapezoidal integration fy = ff * yi[i] fy_center = (fy[:-1] + fy[1:]) * 0.5 # trapezoidal integration yo[k] = np.sum(fy_center * dx[i][:-1]) # associated error estimate if std_dev == 'population': # this is the "weighted standard deviation" # it would be strictly correct if yo[k] was a weighted np.mean like np.sum(ff_center*dx*y_center) / np.sum(ff_center*dx) # (http://stats.stackexchange.com/questions/6534/how-do-i-calculate-a-weighted-standard-deviation-in-excel) w = ff_center*dx[i][:-1] eo[k] = np.sqrt(np.sum(w * (y_center[i][:-1] - yo[k]) ** 2) / (np.sum(w) * (npts - 1) / npts)) #eo[k] = np.sqrt(np.sum(w * (y_center[i][:-1] - yo[k]) ** 2)) elif type(std_dev) in [float, int]: # convolve error/covariance/etc. (gives error envelope) # grid-center points of nth power of weighted values for trapezoidal integration fe = np.power(ff * ei[i], std_dev) fe_center = (fe[:-1] + fe[1:]) * 0.5 # trapezoidal integration and nth root eo[k] = np.power(np.sum(fe_center * dx[i][:-1]), 1. / std_dev) elif std_dev == 'none': pass elif std_dev != 'propagate': raise ValueError('std_dev=%s is an invalid optional key word argument' % repr(std_dev)) if std_dev == 'none': return yo with warnings.catch_warnings(): warnings.filterwarnings('ignore', category=RuntimeWarning) uo = unumpy.uarray(yo,eo) return uo
[docs]@_available_to_user_math def trismooth(*args, **kw): """ Smooth with triangle kernel, designed to mimic smooth in reviewplus :param args: y: float array, y, window_size: float array, float y, x, window_size: float array, float array, float OR y, window_size[optional] as OMFITmdsValue, float If x or window_size are not defined, smooth_by_convolution will pick values automatically. window_size is the half width fo the kernel. The default window_size set by smooth_by_convolution will probably be too small, unless you are providing a much higher resolution output grid with xo. You should probably set window_size by providing at least two arguments. :param kw: Keywords passed to smooth_by_convolution. :return: float array or uarray or tuple of arrays ysmo = result from smooth_by_convolution; array (default) or uarray (std_dev is modified) if OMFITmdsValue: (x, ysmo) else: ysmo """ y = args[0] from classes.omfit_mds import OMFITmdsValue if isinstance(y, OMFITmdsValue): m = y x = y.dim_of(0) y = y.data() if len(args) > 1: window_size = args[1] else: window_size = None else: m = None if len(args) == 2: window_size = args[1] x = None elif len(args) >= 3: x = args[1] window_size = args[2] else: window_size = x = None kw.setdefault('std_dev', 'none') kw.setdefault('xi', x) kw.setdefault('xo', x) kw.setdefault('window_size', window_size) kw['window_function'] = 'triangle' ys = smooth_by_convolution(y, **kw) if m is not None: return x, ys # MDS mode; got an object with x,y, so return an object with x,y. Allows plt.plot(*trismooth()). else: return ys
[docs]@_available_to_user_math def cicle_fourier_smooth(R, Z, keep_M_harmonics, inPlace=False, equalAngle=False, doPlot=False): ''' smooth closed periodic boundary by zeroing out high harmonics components :param R: x coordinates of the boundary :param Z: y coordinates of the boundary :param keep_M_harmonics: how many harmonics to keep :param equalAngle: use equal angle interpolation, and if so, how many points to use :param doPlot: plot plasma boundary before and after :return: smoothed R and Z coordinates ''' closes = False if R[0] == R[-1] and Z[0] == Z[-1]: closes = True Rc = R Zc = Z R = R[:-1] Z = Z[:-1] else: Rc = np.hstack((R, R[0])) Zc = np.hstack((Z, Z[0])) R1 = R Z1 = Z if equalAngle: R0 = np.mean(R) Z0 = np.mean(Z) T = np.unwrap(np.arctan2(Z - Z0, R - R0)) T1 = np.linspace(0, 2 * np.pi, equalAngle + 1)[:-1] R1 = interp1dPeriodic(T, R - R0)(T1) + R0 Z1 = interp1dPeriodic(T, Z - Z0)(T1) + Z0 n = int(len(R1) // 2 - keep_M_harmonics) # fourier transform Rf = np.fft.fftshift(np.fft.fft(R1)) Zf = np.fft.fftshift(np.fft.fft(Z1)) # zero out unwanted harmonics Rf[:n] *= 0 Rf[-n:] *= 0 Zf[:n] *= 0 Zf[-n:] *= 0 # inverse fourier transform RS = np.real(np.fft.ifft(np.fft.ifftshift(Rf))) ZS = np.real(np.fft.ifft(np.fft.ifftshift(Zf))) RSc = np.hstack((RS, RS[0])) ZSc = np.hstack((ZS, ZS[0])) # plotting if doPlot: from matplotlib import pyplot pyplot.gca().plot(Rc, Zc, 'k', label='original') pyplot.gca().plot(RSc, ZSc, 'r', label='smoothed M=%d' % keep_M_harmonics) pyplot.gca().set_aspect('equal') if closes: return RS, ZS else: return RSc, ZSc
# for backwards compatability nu_conv = smooth_by_convolution
[docs]class FilteredNDInterpolator: """ A class for smoothing ND data. Useful for down-sampling and gridding. """ def __init__(self, points, values, size_scales=None, check_finite=False, std_dev=True, filter_type='median'): """ Mean or median filter interpolate in N dimensions. Calculates the mean or median of all values within a size_scale "sphere" of the interpolation point. Unlike linear interpolation, you can incorporate information from all your data when down-sampling to a regular grid. Of course, it would be better to do a weight function (gaussian, hanning, etc) convolution, but the "volume" elements required for integration get very computationally expensive in multiple dimensions. In that case, try linear interpolation followed by regular-grid convolution (ndimage processing). :param points: np.ndarray of floats shape (M, D), where D is the number of dimensions :param values: np.ndarray of floats shape (M,) :param size_scales: tuple (D,) scales of interest in each dimension :param check_finite: bool. Cleans non-finite points and/or values prior to interpolation :param std_dev: bool. Estimate uncertainty of mean or median interpolation :param filter_type: str. Accepts 'mean' or 'median' :Examples: This example shows the interpolator is reasonably fast, but note it is nowhere near as scalable to 100k+ points as linear interpolation + ndimage processing. >> ngrid = 100 >> width = (0.05, 0.05) >> noise_factor = 0.5 >> def func(x, y): >> return x * (1 - x) * np.cos(4 * np.pi * x) * np.sin(4 * np.pi * y ** 2) ** 2 >> x = np.linspace(0, 1, ngrid) >> y = np.linspace(0, 1, ngrid) >> grid_x, grid_y = np.meshgrid(x, y) >> xy = np.array(np.meshgrid(x, y)).T.reshape(-1, 2) >> fig, axs = plt.subplots(3, 2, sharex=True, sharey=True, figsize=(8, 9)) >> axs[0, 0].set_title('Original') >> data = func(grid_x, grid_y) >> im = axs[0, 0].imshow(data, extent=(0, 1, 0, 1), origin='lower', aspect='auto') >> kw = dict(extent=(0, 1, 0, 1), origin='lower', aspect='auto', clim=im.get_clim()) >> axs[0, 1].set_title('Noisy') >> data = data + (np.random.random(data.shape) - 0.5) * np.ptp(data) * noise_factor >> im = axs[0, 1].imshow(data, **kw ) >> ns = [400, 1600, 6400] >> times, nums = [], [] >> for npts, axi in zip(ns, axs[1:]): >> print('npts = {:}'.format(npts)) >> points = np.random.rand(npts, 2) >> values = func(points[:, 0], points[:, 1]) >> values = values + (np.random.random((npts,)) - 0.5) * np.ptp(values) * noise_factor >> strt = datetime.datetime.now() >> # ci = ConvolutionNDInterpolator(points, values, xi=xy, window_size=width, window_function='boxcar') >> ci = FilteredNDInterpolator(points, values, size_scales=width, filter_type='mean')(grid_x, grid_y) >> ctime = datetime.datetime.now() - strt >> print(' > Convolution took {:}'.format(ctime)) >> times.append(ctime.total_seconds()) >> nums.append(npts) >> ax = axi[1] >> ax.set_title('Interpolated {:} pts'.format(npts)) >> yi = LinearNDInterpolator(points, values)(grid_x, grid_y) >> im = ax.imshow(nominal_values(yi), **kw) >> l, = ax.plot(points[:, 0], points[:, 1], 'k.', ms=1) >> ax = axi[0] >> ax.set_title('Convolved {:} pts'.format(npts)) >> im = ax.imshow(nominal_values(ci), **kw) >> l, = ax.plot(points[:, 0], points[:, 1], 'k.', ms=1) >>fig.tight_layout() >>axs[0, 0].set_xlim(0, 1) >>axs[0, 0].set_ylim(0, 1) >>fig, axs = plt.subplots() >>axs.plot([0] + nums, [0] + times, marker='o') >>axs.set_xlabel('# of points') >>axs.set_ylabel('time (s)') >>print('sec/k ~ {:.3e}'.format(1e3 * np.mean(np.array(times) / ns))) This example shows that the median filter is "edge preserving" and contrasts both filters with a more sophisticated convolution. Note the boxcar convolution is not identical to the mean filter for low sampling because the integral weights isolated points more than very closely spaced points. >> npts = 50 >> window = ['hanning', 'gaussian', 'boxcar'][1] >> width = 0.05 >> ped = 0.1 >> pedwidth = 0.03 >> fun = lambda x: (0.9-x)*2*(x < 0.9) + 1.0 - np.tanh((x-(1.0-ped))/(0.25*pedwidth)) >> fig, axs = plt.subplots(3, sharex=True, sharey=True) >> axs[0].set_title('Pedestal width {:}, Filter Length Scale {:}'.format(pedwidth, width)) >> for npts, ax in zip([20, 50, 200], axs): >> x0 = np.linspace(0, 1, 1000) >> x1 = np.cumsum(rand(npts)) >> x1 = (x1 - min(x1)) / (max(x1) - min(x1)) >> y1 = fun(x1) + np.random.random(x1.shape) * 0.4 * fun(x1) >> y1 = uarray(y1, y1 * 0.1) >> t0 = datetime.datetime.now() >> yc = nu_conv(y1, xi=x1, xo=x0, window_size=width, window_function=window, std_dev='propagate') >> tc = datetime.datetime.now() - t0 >> print('nu_conv time: {:}'.format(tc)) >> t0 = datetime.datetime.now() >> yd = FilteredNDInterpolator(x1.reshape(-1, 1), y1, size_scales=width / 2.)(x0) >> td = datetime.datetime.now() - t0 >> print('median filter time: {:}'.format(td)) >> t0 = datetime.datetime.now() >> ym = FilteredNDInterpolator(x1.reshape(-1, 1), y1, size_scales=width / 2., filter_type='mean')(x0) >> tm = datetime.datetime.now() - t0 >> print('median filter time: {:}'.format(td)) >> #ax.plot(x0, fun(x0), label='function') >> uband(x0, yc, ax=ax, lw=3, label='{:} convolution, {:.2f}s'.format(window.title(), tc.total_seconds())) >> uband(x0, ym, ax=ax, lw=3, label='Mean Filter, {:.2f}s'.format(tm.total_seconds())) >> uband(x0, yd, ax=ax, lw=3, label='Median Filter, {:.2f}s'.format(td.total_seconds())) >> uerrorbar(x1, y1, ax=ax, markersize=4, ls='-', lw=0.5, label='{:} data points'.format(npts), color='k') >> ax.legend(loc=0, frameon=False) >> ax.set_ylim([-0.1, 5]) >> fig.tight_layout() """ # clean up the input sigmas = unumpy.std_devs(values) values = unumpy.nominal_values(values) points = np.atleast_2d(points) nd = points.shape[1] if nd >= 10: raise RuntimeError("I don't want to do a 10+ dimension convolution. Do you?") if is_int(size_scales): size_scales = [size_scales] * nd # remove nans if check_finite: valid = np.isfinite(values) * np.product(np.isfinite(points), axis=-1) points = points[valid] values = values[valid] sigmas = sigmas[valid] if not values.size: return unumpy.uarray(np.arange(xi.shape[0]) * np.nan, np.arange(xi.shape[0]) * np.nan) # normalize to relevant length scales points_norm = points / size_scales offsets = np.mean(points_norm, axis=0) points_norm -= offsets[np.newaxis, :] # don't bother with errors if there are none if np.sum(sigmas) == 0: std_dev = False self._points_norm = points_norm.astype(np.float16) # normalized to order 1-100ish by size_scales self._size_scales = size_scales self._offsets = offsets self._std_dev = std_dev self._values = values self._sigmas = sigmas self._ndim = nd self._filter_type = filter_type.lower() self._force_fast = False def __call__(self, *args): r""" :param \*args: Tuple of D np.ndarrays representing the coordinates of the points to interpolate to. :return: Interpolated values at each point (same dimensions as the arrays in args). """ from scipy.spatial import distance_matrix t0 = datetime.datetime.now() xx = np.array(args) # inform if bad inputs if xx.shape[0] != self._ndim: raise ValueError("MeanNDInterpolator must be called with the correct number of dimensions") # normalize the requested points xi = xx.T.reshape(-1, self._ndim) # (M, D) xi_norm = (xi / self._size_scales - self._offsets).astype(np.float16) printd('x cleaning took {:}'.format(datetime.datetime.now() - t0)) # trim points outside region of interest valid = (self._points_norm > xi_norm.min(axis=0) - 1) & ( self._points_norm < xi_norm.max(axis=0) + 1) valid = np.where(np.product(valid, axis=-1))[0] points_norm = self._points_norm[valid] values = self._values[valid] sigmas = self._sigmas[valid] # this performs better than anything I did myself t0 = datetime.datetime.now() distance = distance_matrix(points_norm, xi_norm, p=2, threshold=1000000) printd('Distance took {:}'.format(datetime.datetime.now() - t0)) tn = datetime.datetime.now() near = np.less(distance, 1).astype(np.float16) num_near = np.sum(near, axis=0) near[near == 0] = np.nan printd('Sphere nearness took {:}'.format(datetime.datetime.now() - tn)) tv = datetime.datetime.now() if self._filter_type == 'mean': vi = np.nanmean(values[:, np.newaxis] * near, axis=0) else: vi = np.nanmedian(values[:, np.newaxis] * near, axis=0) printd('Values took {:}'.format(datetime.datetime.now() - tv)) ts = datetime.datetime.now() if self._std_dev: si = np.sqrt(np.nansum(sigmas[:, np.newaxis] ** 2 * near, axis=0)) / num_near if self._filter_type != 'mean': si *= 1.253 else: si = np.zeros_like(vi) printd('Sigmas took {:}'.format(datetime.datetime.now() - ts)) result = unumpy.uarray(vi, si).reshape(*xx.shape[1:][::-1]).T return result
[docs]@_available_to_user_math def weighted_avg(x, err=None, minerr=None, return_stddev_mean=False, return_stddev_pop=False): """ Weighted average using uncertainty Propagates errors and calculates weighted standard deviation While nu_conv does these for a sliding window vs. time, this function is simpler and does calculations for a single mean of an array. :param x: The input data :param err: Uncertainty in x. Should have the same units as x. Should have the same length as x. Special case: a single value of err will be used to propagate errors for the standard deviation of the mean, but will produce uniform (boring) weights. If no uncertainty is provided (err==None), then uniform weighting is used. :param minerr: Put a floor on the uncertainties before calculating weight. This prevents a few points with unreasonbly low error from stealing the calculation. Should be a scalar with same units as x. :param return_stddev_mean: True/False flag for whether the standard deviation of the mean (propagated uncertainty) should be returned with the mean. :param return_stddev_pop: Flag for whether the standard deviation of the population (weighted standard deviation) should be returned with the mean. :return: mean OR mean, xpop, xstdm OR mean, xpop OR mean, xstdm depending on the return_stddev flags. xstdm and xpop are the propagated error and the weighted standard deviation. """ x=np.atleast_1d(x) #make sure we have an array and not a list nx=len(x) ne=len(np.atleast_1d(err)) invalid_err = err is None or np.all(err==0) or (ne!=nx and ne!=1) if invalid_err: #no valid uncertainties provided weight=x*0+1 printd('assigned all weights = 1 because no errors, zero errors, or mismatched lengths',topic='weighted_avg') else: err=np.atleast_1d(err) #make sure we have an array and not a list if ne==1: #if the user provided only one uncertainty value, they must want us to propagate errors for them. The weighting is trivial in this case, as is the standard deviation of the population. The stddev of the mean is easy, too, but we can calculate it weight=x*0+err[0] printd('Received only one uncertainty measurement. This will not do anything interesting to the weights, but it can be used in calculating the error in the mean, so we will take it.',topic='weighted_avg') else: weight=copy.copy(err) printd('Got array of errors',topic='weighted_avg') eeq0=err==0 #Err EQuals 0 weight[eeq0]=1.0 #prevent division by 0 because we hate math errors if minerr is not None: printd('enforcing minimum uncertainty of %f'%minerr,topic='weighted_avg') weight[weight<minerr]=minerr #weight is still holding onto uncertainty weight=1.0/weight**2 #weight = uncertainty^-2 weight[eeq0]=0.0 #put the zeros back in printd('Formed weights',topic='weighted_avg') p=weight/np.sum(weight) #normalized weight printd('weights have been normalized',topic='weighted_avg') xmean=np.sum(x*p) #weighted average of x printd('the weighted mean is %f'%xmean,topic='weighted_avg') if return_stddev_mean or return_stddev_pop: if invalid_err: xstdm=None xpop=None else: printd('estimating errors in the mean...',topic='weighted_avg') #calculate standard deviations xstdm=np.sqrt(np.sum((p*err)**2)) #standard deviation of the mean / propagated error in the weighted average (this will be smaller than the errors in individual x measurements) xpop=np.sqrt(np.sum(p*(x-xmean)**2)) #standard deviation of the population / scatter in data while accounting for weight printd('standard deviation of the mean (propagated error) = {:} ({:} %)'.format(xstdm,xstdm/xmean*100),topic='weighted_avg') printd('standard deviation of the population = {:} ({:} %)'.format(xpop,xpop/xmean*100),topic='weighted_avg') if return_stddev_mean: if return_stddev_pop: printd('return xmean,xpop,xstdm',topic='weighted_avg') return xmean,xpop,xstdm else: printd('return xmean,xstdm',topic='weighted_avg') return xmean,xstdm else: if return_stddev_pop: printd('return xmean,xpop',topic='weighted_avg') return xmean,xpop else: printd('return xmean',topic='weighted_avg') return xmean
[docs]@_available_to_user_math def firFilter(time_s, data, cutoff_hz, sharpness=0.5, full_output=False): """ Filter signal data with FIR filter (lowpass, highpass) :param time_s: time basis of the data signal in seconds :param data: data signal :param cutoff_hz: a list of two elements with low and high cutoff frequencies [lowFreq,highFreq] :param sharpness: sharpness of the filter (taps of the FIR filter = sample_rate/2. * sharpness) :param full_output: return filter window and frequency in addition to filtered data :returns: filtered data or tuple with filtered data and filter window and frequency as a tuple """ if len(np.atleast_1d(cutoff_hz)) < 2: raise Exception('argumen `cutoff_hz` must be a list of two elements [lowFreq,highFreq]') from scipy import signal sample_rate = 1. / np.abs(time_s[1] - time_s[0]) nyq_rate = sample_rate / 2. numtaps = int(nyq_rate * sharpness / 2) * 2 cutoff_hz = np.array(list(map(float, np.atleast_1d(cutoff_hz)))) cutoff_hz[np.where(cutoff_hz <= 0)[0]] = 0 cutoff_hz[np.where(cutoff_hz >= nyq_rate)[0]] = nyq_rate if cutoff_hz[0] == 0 and cutoff_hz[-1] == nyq_rate: return data if cutoff_hz[0] == 0: fir_coeff = signal.firwin(numtaps, cutoff_hz[-1] / nyq_rate) elif cutoff_hz[-1] == nyq_rate: fir_coeff = signal.firwin(numtaps, cutoff_hz[0] / nyq_rate, pass_zero=False) else: fir_coeff = signal.firwin(numtaps, cutoff_hz / nyq_rate, pass_zero=False) data = np.hstack(([data[0]] * numtaps, data, [data[-1]] * numtaps)) out = signal.lfilter(fir_coeff, 1.0, data) output = out[int(numtaps * 3 // 2):-int(numtaps // 2)] if not full_output: return output else: f = fftfreq(numtaps, 1 / (2 * nyq_rate)) return output, (f, np.abs(np.fft(fir_coeff)))
[docs]@_available_to_user_math def butter_smooth(xx, yy, timescale=None, cutoff=None, laggy=False, order=1, nan_screen=True, btype='low'): """ Butterworth smoothing lowpass filter. Similar to firFilter, but with a notable difference in edge effects (butter_smooth may behave better at the edges of an array in some cases). :param xx: 1D array Independent variable. Should be evenly spaced for best results. :param yy: array matching dimension of xx :param timescale: float [specifiy either timescale or cutoff] Smoothing timescale. Units should match xx. :param cutoff: float [specify either timescale or cutoff] Cutoff frequency. Units should be inverse of xx. (xx in seconds, cutoff in Hz; xx in ms, cutoff in kHz, etc.) :param laggy: bool True: causal filter: smoothed output lags behind input False: acausal filter: uses information from the future so that smoothed output doesn't lag input :param order: int Order of butterworth filter. Lower order filters seem to have a longer tail after the ELM which is helpful for detecting the tail of the ELM. :param nan_screen: bool Perform smoothing on only the non-NaN part of the array, then pad the result out with NaNs to maintain length :param btype: string low or high. For smoothing, always choose low. You can do a highpass filter instead of lowpass by choosing high, though. :return: array matching dimension of xx Smoothed version of yy """ if nan_screen: ok = ~np.isnan(xx) & ~np.isnan(yy) yout = np.empty(np.shape(yy)) yout[:] = np.NaN xx = xx[ok] yy = yy[ok] else: yout = ok = None cutoff = 1./timescale if cutoff is None else cutoff dx = np.mean(np.diff(xx)) nyquist = 0.5 / dx cutoff_norm = cutoff / nyquist cutoff_norm = 0.9 if cutoff_norm > 0.9 else cutoff_norm printd(' butter_smooth: nyquist = {}, timescale = {}, cutoff = {}, cutoff_norm = {}'.format( nyquist, timescale, cutoff, cutoff_norm), topic='butter_smooth') b, a = scipy.signal.butter(order, cutoff_norm, btype=btype, analog=False) with warnings.catch_warnings(): warnings.filterwarnings( "ignore", message="Using a non-tuple sequence for multidimensional indexing is deprecated") warnings.filterwarnings("ignore", message="underflow encountered in multiply") if laggy: zi = scipy.signal.lfilter_zi(b, a) ys, _ = scipy.signal.lfilter(b, a, yy, zi=zi * yy[0]) # This one lags (the lag actually can help!) else: ys = scipy.signal.filtfilt(b, a, yy) # This one doesn't lag if nan_screen: yout[ok] = ys return yout else: return ys
[docs]@_available_to_user_math def fourier(y,t=None): """ Calculate fourier transform :param y: signal :param t: time basis of the signal (uniformly spaced) :return: tuple with Fourier transform of the signal and frequency basis >> t=r_[0:1:0.1] >> y=np.cos(2*np.pi*t) >> >> Y,W=fourier(y,t) >> >> ax=plt.subplot(2,1,1) >> ax.plot(t,y) >> ax.set_xlabel('t [s]') >> ax.set_ylabel('$y$') >> >> ax=plt.subplot(2,1,2) >> ax.plot(W/2./np.pi,Y) >> ax.set_xlabel('f [Hz]') >> ax.set_ylabel('$Y$') """ N = len(y) if np.all(np.isreal(y)): tmp = np.fft.rfft(y) Y = np.hstack((np.conj(np.flipud(tmp[1:])), tmp[:-1])) / N else: tmp = np.fft.fft(y) Y = np.fft.fftshift(tmp) / N if t is None: return Y else: DT = max(t) - min(t) Wmax = (2.0 * np.pi * (N - 1)) / DT / 2.0 if np.mod(N, 2) == 0: w = np.linspace(-Wmax, Wmax, N + 1)[:-1] else: w = np.linspace(-Wmax, Wmax, N) return Y, w
[docs]@_available_to_user_math def windowed_FFT(t, y, nfft='auto', overlap=0.95, hanning_window=True, subtract_mean=True, real_input=None): """ Bin data into windows and compute FFT for each. Gives amplitude vs. time and frequency. Useful for making spectrograms. input: :param t: 1D time vector in ms :param y: 1D parameter as a function of time :param nfft: Number of points in each FFT bin. More points = higher frequency resolution but lower time resolution. :param overlap: Fraction of window that overlaps previous/next window. :param hanning_window: Apply a Hanning window to y(t) before each FFT. :param subtract_mean: Subtract the average y value (of the entire data set, not individually per window) before calculating FFT :param real_input: T/F: Use rfft instead of fft because all inputs are real numbers. Set to None to decide automatically. output: :return: spectral density (time,frequency) :return: array of times at the center of each FFT window :return: frequency :return: amplitude(time, positive frequency) :return: power(time, positive frequency) :return: positive-only frequency array :return: nfft (helpful if you let it be set automatically) more on overlap for windows 0, 1, 2: overlap = 0.0 : no overlap, window 0 ends where window 1 begins overlap = 0.5 : half overlap, window 0 ends where window 2 begins overlap = 0.99: this is probably as high as you should go. It will look nice and smooth, but will take longer. overlap = 1.0 : FAIL, all the windows would stack on top of each other and infinite windows would be required. more on nfft: Set nfft='auto' and the function will pick a power of two that should give a reasonable view of the dataset. It won't choose nfft < 16. """ printd('Running windowed_FFT()...') nt = len(t) # Length of input array time base dt = np.mean(np.diff(t)) # Time step (the time step should be nearly uniform or else FFT doesn't make sense. # A little bit of noise is probably fine.) printd(' nt = {:}, dt = {:}'.format(nt, dt)) # First round of data sanitization and error checking if len(y) != nt: # Check for input errors printe('Windowed FFT Error! Data dimensions do not match! The longer one will be cropped. This could give junk ' 'results! BE CAREFUL!') nt = min([nt, len(y)]) # Crop to shorter of the two and try to keep going, but will probably fail if nt == 0: # Check for input errors raise OMFITexception('Windowed FFT FATAL ERROR! NO DATA SUPPLIED') # Auto detect complex data if real_input is None: printd('Checking for complex numbers in input...') real_input = ~np.any(np.atleast_1d(np.iscomplex(y))) printd('Chose real_input = {:}'.format(real_input)) # Basic setup calculations if nfft == 'auto': # Option to pick nfft automatically based on dataset length. It will always pick a power of 2. nfpow = (np.round(np.log10(nt / 50.0) / np.log10(2))) if nfpow < 4: nfpow = 4 nfft = 2 ** nfpow printd('auto selected nfft = {:}'.format(nfft)) nfft = int(nfft) # Make sure it's an int to prevent problems later binSpacingInPoints = int(np.round(nfft * (1.0 - overlap))) # Space between start of an interval and start of next interval binSpacingInTime = binSpacingInPoints * dt # Second round of data sanitization and error checking if nt < nfft: # Check for input errors raise OMFITexception('Windowed FFT Error! len(t < nfft! Your data set is smaller than the requested FFT window.') if overlap > (nfft - 1.0) / nfft: # Check for input errors raise OMFITexception('Windowed FFT Error! overlap too big! overlap>(1-nfft)/nfft') if binSpacingInPoints < 1: # Check for input errors # This is almost the same check as the overlap check, which is okay. We really don't want a problem here. raise OMFITexception('Windowed FFT Error! FFT bin-bin spacing too small (try reducing overlap)') # Setup nBins = 1 + int((nt - nfft) // binSpacingInPoints) # Number of bins or windows printd('number of FFT windows =', nBins) printd('number of points in input =', nt) printd('number of points in each FFT window =', nfft) printd('number of points between the starts of successive bins =', binSpacingInPoints) if nBins < 1: nBins = 1 # It's nonsensical not to have at least one bin if real_input: if nfft % 2: nfft_out = int((nfft + 1) // 2) else: nfft_out = int((nfft // 2)) + 1 else: nfft_out = nfft f = np.zeros([nBins, nfft_out], dtype='complex') # Make the array that will receive the data # Get a window function if hanning_window: window = np.hanning(nfft) else: window = 1.0 # No window if subtract_mean: # This is useful if your entire data set has a large DC offset. This tries to suppress the zero frequency # component and make the plot easier to read. However, the average of each window isn't the same as the average # across the entire data set, so the zero frequency component isn't perfectly suppressed by this method. offset = -np.mean(y) else: offset = 0.0 # Do the FFT on each bin for i in range(nBins): yy = (y[i * binSpacingInPoints:i * binSpacingInPoints + nfft] + offset) * window if real_input: ff = np.fft.rfft(yy) f[i, :] = ff else: ff = np.fft.fft(yy) f[i, :] = ff # Get the frequency freq = np.abs(np.fft.fftfreq(nfft, d=dt)[0: nfft_out]) # Fold in the spectral density from negative frequencies to get amplitude vs. positive frequency only if real_input: amp = np.abs(f) if nfft % 2: # nfft is odd, so avoid doubling the zero frequency component amp[:, 1:] *= 2 else: # nfft is even, so avoid doubling the zero frequency component and also the nyquist frequency component amp[:, 1:-1] *= 2 freqpos = freq else: # Get the positive frequency array: # If nfft is even, then 0:nfft/2+1 includes the most negative frequency and np.abs() makes it positive. # If nfft is odd, then the frequency array is symmetric and 0:nfft/2+1 gets all of the non-negative frequencies. freqpos = np.abs(freq[0:int(nfft // 2) + 1]) if nfft % 2: # Get the amplitude of all positive frequency fluctuations (including 0) amppos = np.abs(f[:, 0:int(nfft // 2) + 1]) # Goes from 0 to nyquist # Get the amplitude of all negative freq fluctuations ampneg = np.fliplr(np.abs(f[:, int(nfft // 2) + 1:])) # Goes from -nyquist to -df else: # Get the amplitude of all positive frequency fluctuations (including 0) amppos = np.abs(f[:, 0:int(nfft // 2)]) # Goes from 0 to nyquist - df, where df = freq[1] - freq[0] # Get the amplitude of all negative freq fluctuations ampneg = np.fliplr(np.abs(f[:, int(nfft // 2):])) # Goes from -df to -nyquist (after flipping) # Add a slot for the highest absolute frequency to the positive freq array amppos = np.concatenate([amppos, np.zeros([nBins, 1])], axis=1) # Now this goes from 0 to nyquist # Add a slot for zero frequency to the negative array ampneg = np.concatenate([np.zeros([nBins, 1]), ampneg], axis=1) # Now this goes from 0 to -nyquist # The amplitudes for positive and negative frequencies have the same array dimensions and can now be added. # If you just double the positive frequencies, you count zero twice, which is naughty. If nfft is even, the # nyquist frequency would be counted twice as well. amp = amppos + ampneg power = amp ** 2 # Power is just amplitude squared timebins = t[np.arange(int(nfft // 2), nt - int(nfft // 2) + 1, binSpacingInPoints)] # These are times at the center of each FFT window return f.T, timebins, freq, amp.T, power.T, freqpos, nfft
# ----------- # Signal processing # -----------
[docs]@_available_to_user_math def noise_estimator( t, y, cutoff_timescale=None, cutoff_omega=None, window_dt=0, dt_var_thresh=1e-9, avoid_jumps=False, debug=False, debug_plot=None, restrict_cutoff=False, ): """ Estimates uncertainty in a signal by assuming that high frequency (above cutoff frequency) variation is random noise :param t: 1D float array with regular spacing (ms) Time base for the input signal :param y: 1D float array with length matching t (arbitrary) Input signal value as a function of t :param cutoff_timescale: float (ms) For a basic RC filter, this would be tau = R*C. Define either this or cutoff_omega. :param cutoff_omega: float (krad/s) The cutoff angular frequency, above which variation is assumed to be noise. Overrides cutoff_timescale if specified. cutoff_timescale = 1.0/cutoff_omega :param window_dt: float or None (ms) Window half width for windowed_fft, used in some strategies for time varying noise. If <= 0, one scalar noise estimate is made for the entire range of t, using FFT methods. Set to None to choose window size automatically based on cutoff_timescale. This option does not seem to be as good as the standard method. :param debug: bool Flag for activating debugging tests, plots (unless debug_plot is explicitly disabled), and reports. Also returns all four estimates instead of just the best one. :param debug_plot: bool [optional] By default, debug_plot copies the value of debug, but you can set it to False to disable plots and still get other debugging features. Setting debug_plot without debug is not supported and will be ignored. :param dt_var_thresh: float (fraction) Threshold for variability in dt. t must be evenly spaced, so nominally dt will be a constant. Because of various real world issues, there could be some slight variation. As long as this is small, everything should be fine. You can adjust the threshold manually, but be careful: large variability in the spacing of the time base breaks assumptions in the signal processing techniques used by this method. If std(dt)/mean(dt) exceeds this threshold, an exception will be raised. :param avoid_jumps: bool Large jumps in signal will have high frequency components which shouldn't be counted with high frequency noise. You SHOULD pick a time range to avoid these jumps while estimating noise, but if you don't want to, you can try using this option instead. If this flag is true, an attempt will be made to identify time periods when jumps are happening. The time derivative, dy/dt, is evaluated at times where there are no detected jumps and interpolated back onto the full time base before being integrated to give a new signal. The new signal should have a continuous derivative with no spikes, such that its high frequency component should now be just noise. This will prevent the high frequency components of a jump in y from bleeding into the noise estimate near the jump. The noise estimate during the jump may not be accurate, but you were never going to get a good fix on that, anyway. The big problem that is solved by this option is that the jump causes spurious spectral noise which extends well before and after the jump itself. Cutting the jump out somehow confines the problem to the relatively narrow time range when the jump itself is happening. :param restrict_cutoff: bool Some versions of scipy throw an error if cutoff_frequency > nyquist frequency, and others do not. If your version hates high frequency cutoffs, set this to True and cutoff will be reduced to nyqusit - df/2.0, where df is the frequency increment of the FFT, if cutoff >= nyquist. :return: 1D uncertain float array with length matching t, or set of four such arrays with different estimates. Lowpass smoothed y with uncertainty (dimensions and units match input y) The standard estimate is a hilbert envelope of the high frequency part of the signal, times a constant for correct normalization:: ylf = butterworth_lowpass(y, cutoff_frequency) yhf = butterworth_highpass(y, cutoff_frequency) uncertainty = smooth(hilbert_env(yhf)) * norm_factor return = unumpy.uarray(ylf, uncertainty) where smoothing of the envelope is accomplished by a butterworth lowpass filter using cutoff_frequency. ``norm_factor = np.sqrt(0.5) = std_dev(sin(x))`` There are other estimates (accessible by setting the debug flag) based on the fluctuation amplitude in the windowed FFT above the cutoff frequency. """ from scipy import interpolate as scipy_interpolate printd('noise_estimator() called with len(t) = {}, len(y) = {}, cutoff_timescale = {}, cutoff_omega = {}'.format( len(t), len(y), cutoff_timescale, cutoff_omega)) # Setup ----------------------- # Check inputs if debug_plot is None: debug_plot = debug if len(t) != len(y): raise ValueError('Dimensions of t and y do not match: {} vs {}'.format(len(t), len(y))) if (cutoff_omega is None) and (cutoff_timescale is None): raise ValueError('Must specify either cutoff_timescale or cutoff_omega!') if (cutoff_omega is not None) and (cutoff_timescale is not None): printw('Warning: both cutoff_timescale and cutoff_omega have been specified! ' 'cutoff_omega will override cutoff_timescale!') dt = np.diff(t) dt_var = np.std(dt)/np.mean(dt) if dt_var > dt_var_thresh: raise ValueError('t is not evenly spaced: std(dt)/mean(dt) = {} > {}'.format(dt_var, dt_var_thresh)) dt = np.mean(dt) # (ms) sample_frequency = 1.0/dt # (kHz) freq = np.fft.fftfreq(len(t), dt) df = freq[1] - freq[0] if cutoff_omega is None: cutoff_omega = 1.0/cutoff_timescale # (krad/s) cutoff_timescale = 1.0/cutoff_omega # (ms) cutoff_frequency = cutoff_omega/(2*np.pi) # (kHz) if window_dt is None: window_dt = cutoff_timescale * 5 nt = len(t) nyquist_frequency = sample_frequency * 0.5 old_cut = None if (cutoff_frequency >= nyquist_frequency) and restrict_cutoff: # Correct illegal settings printw('NOISE_ESTIMATOR: CUTOFF FREQUENCY MUST BE < NYQUIST FREQUENCY.') old_cut = cutoff_frequency cutoff_frequency = nyquist_frequency - df * 0.5 elif (cutoff_frequency <= 0) and restrict_cutoff: old_cut = cutoff_frequency printw('CUTOFF FREQUENCY MUST BE > 0') cutoff_frequency = df * 0.5 if old_cut is not None: print('Cutoff frequency modified to obey the rules; it is now {} (was {})'.format(cutoff_frequency, old_cut)) cutoff_omega = cutoff_frequency * 2 * np.pi cutoff_timescale = 1.0 / cutoff_omega norm_factor = np.sqrt(0.5) # std(sin(x)) = np.sqrt(0.5), hilbert_envelope(sin(x)) = 1 printd('noise_estimator(): cutoff_timescale = {} ms, cutoff_omega = {} krad/s, cutoff_frequency = {} kHz; ' 'dt = {} ms, sample_frequency = {} kHz, window_dt = {} ms'.format( cutoff_timescale, cutoff_omega, cutoff_frequency, dt, sample_frequency, window_dt)) # Calculate ----------------------- def butter_filter(data, cutoff, fs, order=5, btype='high'): nyq = 0.5 * fs # kHz normal_cutoff = cutoff / nyq b, a = scipy.signal.butter(order, normal_cutoff, btype=btype, analog=False) return scipy.signal.filtfilt(b, a, data) def est_from_fft(fft_out, df_): return abs(np.sum(np.real(fft_out) * df_, 0)) if avoid_jumps: printd(' Avoiding jumps in signal') # Detect jumps by checking for unusually large spikes in the derivative dydt = np.gradient(y)/dt dydts = abs(butter_filter(dydt, cutoff_frequency, sample_frequency, btype='low')) thresh = np.percentile(dydts, 90) * 15 # Interpolate the non-jumping dydt back onto the full timebase to cut jumps out smoothly dydt2 = scipy_interpolate.interp1d( t[dydts < thresh], dydt[dydts < thresh], bounds_error=False, fill_value='extrapolate')(t) # Integrate y_use = integrate.cumtrapz(dydt2, t, initial=0) + y[0] else: y_use = y import warnings with warnings.catch_warnings(): warnings.filterwarnings('ignore', category=FutureWarning) yhf = butter_filter(y_use, cutoff_frequency, sample_frequency) fyh = np.fft.fft(yhf) est1 = est_from_fft(fyh, df) est4 = abs(scipy.signal.hilbert(yhf)) * norm_factor est4f = abs(butter_filter(est4, cutoff_frequency, sample_frequency, btype='low')) if window_dt > 0: overlap = 0.95 n_windows_0 = (t.max() - t.min())/(window_dt*2) n_windows = n_windows_0*(1.0/(1-overlap)) nfft = int(np.ceil(float(nt/n_windows_0))) fyht, tw, freqw, amp, power, freqpos, nfft_out = windowed_FFT( t, yhf, nfft=nfft, overlap=overlap, hanning_window=False) printd('noise_estimator(): nt = {}, overlap = {}, n_windows = {}, nfft = {}'.format( nt, overlap, n_windows, nfft)) printd('noise_estimator(): np.shape(fyht) = {}, np.shape(tw) = {}, np.shape(freqw) = {}'.format( np.shape(fyht), np.shape(tw), np.shape(freqw))) dfw = freqw[1] - freqw[0] # When doing windowed_fft, the phase of these things vs. time can change, giving different +/- signs at # different times. Compensate by multiplying the whole thing by the sign of the highest absolute value at each # time. fsign = np.array([np.sign(np.real(fyht)[i, abs(fyht[i, :]).argmax()]) for i in range(np.shape(fyht)[0])]) fyht2 = fyht * fsign[:, np.newaxis] est3_a = est_from_fft(fyht2, dfw) * norm_factor est3 = abs(scipy_interpolate.interp1d(tw, est3_a, bounds_error=False, fill_value=np.NaN)(t)) est3f = butter_filter(est3[~np.isnan(est3)], cutoff_frequency / 2., sample_frequency, btype='low') est3 = abs(scipy_interpolate.interp1d( t[~np.isnan(est3)], est3f, bounds_error=False, fill_value=(est3f[0], est3f[-1]))(t)) if debug: est2_a = est_from_fft(fyht, dfw) * norm_factor est2 = scipy_interpolate.interp1d(tw, est2_a, bounds_error=False, fill_value=np.NaN)(t) est2f = butter_filter(est2[~np.isnan(est2)], cutoff_frequency / 2., sample_frequency, btype='low') est2 = abs(scipy_interpolate.interp1d( t[~np.isnan(est2)], est2f, bounds_error=False, fill_value=(est2f[0], est2f[-1]))(t)) else: est2 = 0 printd("noise_estimator(): np.shape(est3_a) = {}, np.shape(t) = {}, np.shape(est3) = {}".format( np.shape(est3_a), np.shape(t), np.shape(est3))) else: est3 = est2 = freqw = fyht2 = dfw = tw = None ylf = butter_filter(y_use, cutoff_frequency, sample_frequency, btype='low') if debug_plot and debug: from matplotlib import pyplot as plt from utils_plot import uband fig, axs = plt.subplots(2, 2, sharex='col') # Plot 1: y vs. t with low frequency components shown ax = axs[0, 0] ax.plot(t, y, label='y (input)') ax.plot(t, y_use, label='y used') ax.plot(t, ylf, label='y lowpass') # ax.plot(t, yhf+ylf, label='y lowpass + y highpass') # uband(t, unumpy.uarray(ylf, est1), ax=ax, label='lowpass w/ simple est uncert') ax.plot(t, yhf, label='y highpass') ax.set_ylabel('y') uband(t, unumpy.uarray(ylf, est4f), ax=ax, label='lowpass w/ est uncert') # Plot 2: y vs. t with zero frequency component not included ax = axs[1, 0] ax.plot(t, yhf, label='y highpass') ax.axhline(est1, linestyle='--', label='noise level estimate 1 = {}'.format(est1)) ax.axhline(-est1, linestyle='--') ax.set_ylabel('y') ax.set_xlabel('Time (ms)') if window_dt > 0: est2line = ax.plot(t, est2, label='noise level estimate 2') ax.plot(t, -est2, color=est2line[0].get_color()) est3line = ax.plot(t, est3, label='noise level estimate 3') ax.plot(t, -est3, color=est3line[0].get_color()) est4line = ax.plot(t, est4f, label='noise level estimate 4; avg = {}'.format(np.mean(est4f))) ax.plot(t, -est4f, color=est4line[0].get_color()) ax.axhline(np.std(yhf), linestyle=':', color='k', label='std(yhf) = {}'.format(np.std(yhf))) # Plot 3: FFTs ax = axs[1, 1] ax.set_xlabel('f (kHz)') ax.set_ylabel('|FFT|*df') fy = np.fft.fft(y) fyl = np.fft.fft(ylf) # fyhl = np.fft.fft(ylf + yhf) fyline = ax.plot(freq[freq > 0], abs(fy[freq > 0])*df, label='FFT(y)') ax.plot(freq[freq > 0], abs(fyl[freq > 0])*df, label='FFT(y lowpass)') # ax.set_xscale("log", nonposx='clip') ax.set_yscale("log", nonposy='clip') # ax.plot(freq[freq > 0], abs(fyhl[freq > 0])*df, label='FFT(y lowpass + y highpass)') ax.plot(freq[freq > 0], abs(fyh[freq > 0])*df, label='FFT(y highpass)') ax.axhline(abs(fy[freq == 0])*df, color=fyline[0].get_color(), linestyle='--', label='FFT(y) @ f=0 ({})'.format( abs(fy[freq == 0])*df)) ax.axhline(est1*df, color='r', label='noise est 1', linestyle='--') ax.axvline(cutoff_frequency, color='k', linestyle='--', label='cutoff = {} kHz'.format(cutoff_frequency)) if window_dt > 0: ax.plot(freqw, abs(fyht2[:, 0])*dfw, label='first window in windowed_FFT') printd('np.shape(tw) = {}, np.shape(fyht2) = {}'.format(np.shape(tw), np.shape(fyht2))) # Plot 4: windowed FFT ax = axs[0, 1] if window_dt > 0: ax.imshow(np.real(fyht2.T), extent=(freqw.min(), freqw.max(), tw.min(), tw.max()), aspect='auto') ax.axvline(cutoff_frequency, color='k', linestyle='--', label='cutoff = {} kHz'.format(cutoff_frequency)) else: # Not needed ax.axis('off') for ax in axs.flatten(): ax.legend(loc=0) if debug: printd( "np.shape(ylf) = {}, np.shape(est1) = {}, np.shape(est2) = {}, np.shape(est3), np.shape(est4f) = {}".format( np.shape(ylf), np.shape(est1), np.shape(est2), np.shape(est3), np.shape(est4f)) ) u1 = unumpy.uarray(ylf, est1) u2 = unumpy.uarray(ylf, est2) if est2 is not None else None u3 = unumpy.uarray(ylf, est3) if est3 is not None else None u4 = unumpy.uarray(ylf, est4f) return u1, u2, u3, u4 else: return unumpy.uarray(ylf, est4f)
#----------- # inverse hyperbolic sine transform #-----------
[docs]@_available_to_user_math def ihs_opti_transform(x, shift=False, lmbda=None, off=None): """ Calculate the optimal coefficient for the inverse hyperbolic sine transformation of the data :param x: input data :param shift: calculate and return offset coefficients :param lmbda: initial guess for lmbda parameter :param off: initial guess for offset parameter :return: lmbda and offset coefficients (offset=0 if shift==False) """ N = len(x) # compute uniform median statistics Ui = np.zeros(N)*1.0 Ui[-1] = 0.5**(1.0/N) Ui[0] = 1-Ui[-1] i = np.arange(2,N) Ui[1:-1] = (i-0.3175)/(N+0.365) # this function computes the x-axis values of the probability plot # and computes a linear regression (including the correlation) # and returns 1-r so that a minimization function maximizes the # correlation from scipy import stats xvals = stats.distributions.norm.ppf(Ui) m=np.median(np.abs(xvals)) off_guess=off lmbda_guess=lmbda def tempfunc1_lmbda(lmbda, off, xvals, samps): y = ihs(samps,lmbda,off) yvals = np.sort(y) r, prob = stats.pearsonr(xvals, yvals) c=(1-r)*(1+1E-3*np.log10(1+np.abs(lmbda))) return c def tempfunc1_off(off, lmbda, xvals, samps): y = ihs(samps,lmbda,off) yvals = np.sort(y) r, prob = stats.pearsonr(xvals, yvals) c=(1-r)*(1+1E-3*np.log10(1+np.abs(off/m))) return c def tempfunc2(Xopt, xvals, samps): lmbda,off=Xopt y = ihs(samps,lmbda,off) yvals = np.sort(y) r, prob = stats.pearsonr(xvals, yvals) c=(1-r)*( 1 + 1E-3*np.log10(1+np.abs(lmbda)) + 1E-3*np.log10(1+np.abs(off/m) ) ) return c if not shift: return [np.abs(optimize.brent(tempfunc1_lmbda, brack=(0.1,1.0), args=(0.,xvals, x),tol=1E-2)),0.] else: off=off_guess lmbda=lmbda_guess if lmbda_guess is None: if off_guess is None: off=min(x) lmbda=np.abs(optimize.brent(tempfunc1_lmbda, brack=(0.1,1.0), args=(off,xvals, x), tol=1E-0)) if off_guess is None: off=optimize.brent(tempfunc1_off, brack=(min(x),np.mean(x)), args=(lmbda,xvals, x), tol=1E-0) return optimize.fmin_l_bfgs_b(tempfunc2, [lmbda,off], bounds=((1E-6,None),(None,None)), args=(xvals, x), disp=0, approx_grad=True)[0]
[docs]@_available_to_user_math def ihs(x, lmbda=1., off=0.): """ Inverse hyperbolic sine transform: y=np.arcsinh(lmbda*(x-off))/lmbda :param x: input data :param lmbda: transformation coefficient :param off: offset coefficient :return: transformed data """ x=np.atleast_1d(x) if lmbda==0: return x.copy()-off else: return np.arcsinh(lmbda*(x-off))/lmbda
[docs]@_available_to_user_math def iihs(y, lmbda=1., off=0.): """ Inverse of Inverse hyperbolic sine transform: x=np.sinh(y*lmbda)/lmbda+off :param y: transformed data :param lmbda: transformation coefficient :param off: offset coefficient :return: original input data """ y=np.atleast_1d(y) if lmbda==0: return y.copy()+off else: return np.sinh(lmbda*y)/lmbda+off
[docs]@_available_to_user_math def array_snippet(a): """ Gives a string summarizing array or list contents: either elements [0, 1, 2, ..., -3, -2, -1] or all elements :param a: Array, list, or other compatible iterable to summarize :return: String with summary of array, or just str(a) for short arrays """ return '[{}, {}, {}, ..., {}, {}, {}]'.format( a[0], a[1], a[2], a[-3], a[-2], a[-1]) if len(np.atleast_1d(a)) > 6 else str(a)
[docs]def ordinal(n, fmt="%d%s"): ''' Return ordinal of an integer '1st', '2nd', '3rd', '4th', '5th', ... Based on: https://stackoverflow.com/a/20007730/6605826 :param n: integer :param fmt: string format to use ("%d%s" is the default) :return: string with ordinal number ''' return fmt % (n, "tsnrhtdd"[(np.floor(n // 10) % 10 != 1) * (n % 10 < 4) * n % 10::4])
#----------- # roman numerals: : http://diveintopython.org/ #----------- # Define exceptions
[docs]class RomanError(Exception): pass
[docs]class RomanOutOfRangeError(RomanError): pass
[docs]class RomanNotIntegerError(RomanError): pass
[docs]class RomanInvalidNumeralError(RomanError): pass
# Define digit mapping _romanNumeralMap= (('M', 1000), ('CM', 900), ('D', 500), ('CD', 400), ('C', 100), ('XC', 90), ('L', 50), ('XL', 40), ('X', 10), ('IX', 9), ('V', 5), ('IV', 4), ('I', 1))
[docs]@_available_to_user_math def toRoman(n): """convert integer to Roman numeral""" if not (0 < n < 5000): raise RomanOutOfRangeError("number out of range (must be 1..4999)") if int(n) != n: raise RomanNotIntegerError("decimals can not be converted") result = "" for numeral, integer in _romanNumeralMap: while n >= integer: result += numeral n -= integer return result
# Define pattern to detect valid Roman numerals _romanNumeralPattern = re.compile(""" ^ # beginning of string M{0,4} # thousands - 0 to 4 M's (CM|CD|D?C{0,3}) # hundreds - 900 (CM), 400 (CD), 0-300 (0 to 3 C's), # or 500-800 (D, followed by 0 to 3 C's) (XC|XL|L?X{0,3}) # tens - 90 (XC), 40 (XL), 0-30 (0 to 3 X's), # or 50-80 (L, followed by 0 to 3 X's) (IX|IV|V?I{0,3}) # ones - 9 (IX), 4 (IV), 0-3 (0 to 3 I's), # or 5-8 (V, followed by 0 to 3 I's) $ # end of string """ ,re.VERBOSE)
[docs]@_available_to_user_math def fromRoman(s): """convert Roman numeral to integer""" if not s: raise RomanInvalidNumeralError('Input can not be blank') if not _romanNumeralPattern.search(s): raise RomanInvalidNumeralError('Invalid Roman numeral: %s' % s) result = 0 index = 0 for numeral, integer in _romanNumeralMap: while s[index:index+len(numeral)] == numeral: result += integer index += len(numeral) return result
#----------- # elements #----------- _atomic_table=[]
[docs]@_available_to_user_math def atomic_element(symbol_A=None, symbol=None, name=None, Z=None, Z_ion=None, mass=None, A=None, abundance=None, use_D_T=True, return_most_abundant=True, ): """ Returns dictionary with name, symbol, symbol_A, Z, mass, A, abundance information of all the elements that match a a given query. Most of the information was gathered from: http://www.sisweb.com/referenc/source/exactmas.htm returns dictionary with name, symbol, symbol_A, Z, mass, A, abundance information of all the elements that match a a given query :param symbol_A: string Atomic symbol followed by the mass number in parenthesis eg. H(2) for Deuterium :param symbol: string Atomic symbol * can be followed by the mass number eg. H2 for Deuterium * can be prepreceded the mass number and followed by the ion charge number eg. 2H1 for Deuterium :param name: string Long name of the atomic element :param Z: int Atomic number (proton count in nucleus) :param Z_ion: int Charge number of the ion (if not specified, it is assumed `Z_ion = Z`) :param mass: float Mass of the atomic element in AMU For matching, it will be easier to use A :param A: int Mass of the atomic element rounded to the closest integer :param abundance: float Abundance of the atomic element as a fraction 0 < abundance <= 1 :param use_D_T: bool Whether to use deuterium and tritium for isotopes of hydrogen :param return_most_abundant: bool Whether only the most abundant element should be returned for a query that matches multiple isotopes :return: dictionary with all the elements that match a query """ # Load atomic data table if not done already if not len(_atomic_table): import json with open(os.path.split(__file__)[0] + os.sep + 'atomic_elements.json') as f: _atomic_table.append(json.load(f)) if symbol: match = re.match(r'(\d*)([a-zA-Z]+)(\d*)$', symbol) if not match: raise ValueError('Wrong form of symbol, expected H, H2, or 2H1 or similar, but got %s' % symbol) symbol = match.groups()[1] if match.groups()[0] == '' and match.groups()[2] != '': A = int(match.groups()[2]) elif match.groups()[0] != '' and match.groups()[2] != '': A = int(match.groups()[0]) Z_ion = int(match.groups()[2]) # Query dictionary kw = {'name': name, 'symbol': symbol, 'symbol_A': symbol_A, 'Z': Z, 'mass': mass, 'A': A, 'abundance': abundance} kw = {item: kw[item] for item in kw if kw[item] is not None} # Find matches matches = {} for item in _atomic_table[0]: match = True for query in kw: if kw[query] and _atomic_table[0][item][query] != kw[query]: match = False break if match: matches[item] = copy.deepcopy(_atomic_table[0][item]) if Z_ion is None: matches[item][u'Z_ion'] = int(matches[item]['Z']) else: matches[item][u'Z_ion'] = int(Z_ion) # Return most abundant if return_most_abundant: tmp = {} for item in matches: if matches[item]['symbol'] not in tmp or \ tmp[matches[item]['symbol']]['abundance'] < matches[item]['abundance']: tmp[matches[item]['symbol']] = matches[item] matches.clear() for item in tmp: matches[tmp[item]['symbol_A']] = tmp[item] # Remove D(2) and T(3) if they are not allowed (use_D_T=False) if len(matches) > 1: for item in list(matches.keys()): if not use_D_T and (matches[item]['symbol_A'] == 'D(2)' or matches[item]['symbol_A'] == 'T(3)'): del matches[item] elif use_D_T and (matches[item]['symbol_A'] == 'H(2)' or matches[item]['symbol_A'] == 'H(3)'): del matches[item] # Handle no matches if not len(matches): raise ValueError('No atomic element satisfies %s' % ', '.join(['%s=%s' % (k[0], repr(k[1])) for k in list(kw.items())])) return matches
[docs]@_available_to_user_math def element_symbol(z, n=None): """ Return the element symbol given the charge (and possibly the number of neutrons) :param z: Ion's charge (number of protons) :param n: Ion's number of nucleons :return: Element symbol (1 or 2 character string) :note: If z==-1, then 'elec' is returned """ if z == -1: return 'elec' else: return list(atomic_element(Z=z, A=n).values())[0]['symbol']
[docs]@_available_to_user_math def element_label(**kw): """ Return a LaTeX formatted string with the element's symbol, charge state as superscript, and optionally mass number as subscript. :param z_n or zn: int Nuclear charge / atomic number :param z_a or za: int [optional] Ionic charge, including any electrons which may be bound. Defaults to displaying fully-stripped ion if not specified (z_a = z_n). :param zamin: int Minimum for a range of Zs in a bundled charge state (like in SOLPS-ITER) :param zamax: int Minimum for a range of Zs in a bundled charge state (like in SOLPS-ITER) :param a or am: int [optional] Mass number. Provides additional filter on results. :return: string Prettily formatted symbol """ z_n = kw.pop('z_n', kw.pop('zn', None)) assert z_n is not None, 'Must specify z_n or zn' z_a = kw.pop('z_a', kw.pop('za', None)) zamin = kw.pop('zamin', None) zamax = kw.pop('zamax', None) a = kw.pop('a', kw.pop('a_m', None)) if z_n == -1: return '$e^-$' else: if z_a is not None: za = '%0d' % int(z_a) elif z_a is None and zamin is None and zamax is None: za = '%0d' % int(z_n) elif z_a is None and zamax > zamin: za = '{:0d}-{:0d}'.format(int(zamin), int(zamax)) elif z_a is None and zamax == zamin and zamax is not None: za = '%0d' % int(zamin) else: za = '?' printw('WARNING: This code should not be reached. You may have weird z_a/zamin/zamax settings.') data = list(atomic_element(Z=z_n, A=a, return_most_abundant=True, use_D_T=a is not None).values())[0] return '${symbol:}^{{+{z_a:}}}$'.format(symbol=data['symbol'], z_a=za)
[docs]def z_from_symbol(symbol): ''' Given an atomic element symbol it returns it charge Z :param symbol: element symbol :return: element charge Z ''' return list(atomic_element(symbol=symbol).values())[0]['Z']
#----------- # tension splines #----------- def _coefsplinet(t,y,tau): ''' Spline coefficients By VICTOR AGUIAR NUMERICAL ANALYSIS OF KINCAID :param t: nodes location :param y: value at the nodes :param x: return values at :param tau: tension ''' ## step 1: ## order t and y n=len(t) ## Create the arrays that will be filled up. delta=h=list(range(0,n-1)) ## Order the table from less to high. ind = t.argsort() t=t[ind] y=y[ind] ## step 2: ## compute h, a,b,o as described in Kincaid book h=t[1:n]-t[0:n-1] a=(1/h)-(tau/np.sinh(tau*h)) b=(tau*np.cosh(tau*h)/np.sinh(tau*h))-(1/h) g=(tau**2)*(y[1:n]-y[0:n-1])/h ## step 3: ## initialize Zi or coefficients of spline z=np.zeros(n-1) ## step 4: ## solve tri-diagonal system for Zi 1<=i<=n-1 ## output nz=len(z) v=g[1:nz]-g[0:nz-1] ## matrix of coefficients diagel=b[0:nz-1]+b[1:nz] A=np.diag(diagel) ## complete the matrix na=len(a) a2=a[1:na-1] a3=a2=np.diag(a2) dum1=np.zeros((1,na-2)) dum2=np.zeros((na-1,1)) a2=np.vstack([a2,dum1]) a2=np.hstack([dum2,a2]) a3=np.vstack([dum1,a3]) a3=np.hstack([a3,dum2]) ## solve system using scipy library!! ## get zi A=A+a2+a3 z=np.linalg.solve(A,v) z=np.hstack([0,z,0]) return z,h
[docs]@_available_to_user_math def splinet(t,y,x,tau): ''' Tension spline evaluator By VICTOR AGUIAR NUMERICAL ANALYSIS OF KINCAID :param t: nodes location :param y: value at the nodes :param x: return values at :param tau: tension ''' ind = t.argsort() t=t[ind] y=y[ind] x=np.sort(x) n=np.size(t) nx=np.size(x) # x can be a vector ## INPUTS x=np.array(x) # Evaluates t,y in x sy=list(range(nx)) #sy is the image of x under the spline function. z,h=_coefsplinet(t,y,tau) m=np.size(z) ## evaluation for i in range(nx): dif=list((x[i]-t)>=0) if dif[n-1]!=True: index=dif.index(False) if (index==0 or index==1): index=0 sum1=z[index]*np.sinh(tau*(t[index+1]-x[i])) sum2=z[index+1]*np.sinh(tau*(x[i]-t[index])) div1=(tau**2)*np.sinh(tau*h[index]) sum3=(y[index]-(z[index]/tau**2))*(t[index+1]-x[i])/h[index] sum4=(y[index+1]-(z[index+1]/(tau**2)))*(x[i]-t[index])/h[index] sy[i]=((sum1+sum2)/div1)+sum3+sum4 else: index=index-1 sum1=z[index]*np.sinh(tau*(t[index+1]-x[i])) sum2=z[index+1]*np.sinh(tau*(x[i]-t[index])) div1=(tau**2)*np.sinh(tau*h[index]) sum3=(y[index]-(z[index]/tau**2))*(t[index+1]-x[i])/h[index] sum4=(y[index+1]-(z[index+1]/tau**2))*(x[i]-t[index])/h[index] sy[i]=(sum1+sum2)/div1+sum3+sum4 else: index=n-2 sum1=z[index]*np.sinh(tau*(t[index+1]-x[i])) sum2=z[index+1]*np.sinh(tau*(x[i]-t[index])) div1=(tau**2)*np.sinh(tau*h[index]) sum3=(y[index]-(z[index]/tau**2))*(t[index+1]-x[i])/h[index] sum4=(y[index+1]-(z[index+1]/tau**2))*(x[i]-t[index])/h[index] sy[i]=(sum1+sum2)/div1+sum3+sum4 ##return the image return sy
[docs]class CLSQTensionSpline(object): """ Constrained least square tension spline. :param x: np.ndarray. Data grid. :param y: np.ndarray. Data values. :param t: int or np.ndarray. Knot number or locations. :param tau: float. Tension (higher is smoother). :param w: np.ndarray. Data weights (usually ~1/std_devs) :param xbounds: tuple. Minimum and maximum x of knot locations. :param min_separation: float. Minumum separation between knot locations. :param xy_constraints: list of tuples. Spline is constrained to have these values at these locations. :param xyprime_constraints: Spline is constrained to have these derivatives as these locations. :param optimize_knot_type: str. choose 'xy' to simultaneously optimize knot (x,y) values, 'x' to optimize x and y separately, and 'y' to simply use the prescribed knot locations. :Examples: Using the same data from the LSQUnivariateSpline examples, >> x = np.linspace(-3, 3, 50) >> y = np.exp(-x**2) + 0.1 * np.random.randn(50) We can fit a tension spline. We can even set some boundary constraints. >> t = [-1, 0, 1] >> spl = CLSQTensionSpline(x, y, t, tau=0.1, xy_constraints=[(-3,0)]) >> xs = np.linspace(-3, 3, 1000) >> plt.subplots() >> plt.plot(x, y, marker='o', lw=0) >> plt.plot(xs, spl(xs)) Note the xknots are optimized by default. We can compare to the un-optimized knot locations, but (for historical reasons) we wont be able to set constraints. >> sply = CLSQTensionSpline(x, y, t, tau=0.1, xy_constraints=[(-3,0)], optimize_knot_type='y') >> plt.plot(xs, sply(xs)) """ def __init__(self, x, y, t, tau=1, w=None, xbounds=(None, None), min_separation=0, xy_constraints=(), xyprime_constraints=(), optimize_knot_type='xy'): """ :param x: np.ndarray. Data grid. :param y: np.ndarray. Data values. :param t: int or np.ndarray. Knot number or locations. :param tau: float. Tension (higher is smoother). :param w: np.ndarray. Data weights (usually ~1/std_devs) :param xbounds: tuple. Minimum and maximum x of knot locations. :param min_separation: float. Minumum separation between knot locations. :param xy_constraints: list of tuples. Spline is constrained to have these values at these locations. :param xyprime_constraints: Spline is constrained to have these derivatives as these locations. :param optimize_knot_type: str. choose 'xy' to simultaneously optimize knot (x,y) values, 'x' to optimize x and y separately, and 'y' to simply use the prescribed knot locations. """ # store variables if is_int(t): self._xknots = np.linspace(min(x), max(x), t) else: self._xknots = np.atleast_1d(t) if w is None: w = np.ones_like(x) # be robust to typical shorthand if xy_constraints is None: xy_constraints = () if np.shape(xy_constraints) == (2,): xy_constraints = (xy_constraints,) if xyprime_constraints is None: xyprime_constraints = () if np.shape(xyprime_constraints) == (2,): xyprime_constraints = (xyprime_constraints,) self._data = np.array([x, y]) self._weights = w self._tau = tau self._min_sep = min_separation self._xyc = xy_constraints self._xyp = xyprime_constraints if xbounds is None or np.shape(xbounds) != (2,): xbounds = (None, None) self._xbounds = np.array(xbounds) # optimize knot values optimize_knot_type = optimize_knot_type.lower() if optimize_knot_type == 'xy': self._set_xyknots() # optimize knot x,y pairs elif optimize_knot_type == 'x': self._set_xyknots() # optimize the x of knots, optimizing y at each step elif optimize_knot_type == 'y': if len(xy_constraints) or len(xyprime_constraints): printe("WARNING: Constraints are not considered if only optimizing knot y values.") self._yknots = self._get_yknots() # just get the best y for each given knot x def __call__(self, x): """ :param x: np.ndarray. Desired x grid. :return: Interpolated y values. """ return splinet(self._xknots, self._yknots, np.atleast_1d(x), self._tau) def _set_xyknots(self, xknots=None): """Simultaneous optimization of the x and y of each knot.""" if xknots is None: xknots = self._xknots * 1.0 # fit knots in a 1 by 1 box xnorm = np.ptp(self._data[0]) x = self._data[0] / xnorm xbounds = self._xbounds ibounds = [i for i, v in enumerate(xbounds) if v is not None] xbounds[ibounds] /= xnorm xknots = xknots / xnorm min_sep = self._min_sep / xnorm ynorm = np.ptp(self._data[1]) y = self._data[1] / ynorm w = self._weights * ynorm xyc = [(xc / xnorm, yc / ynorm) for xc, yc in self._xyc] xyp = [(xc / xnorm, yc / ynorm) for xc, yc in self._xyp] tau = self._tau nk = len(xknots) def cost(xyknots): xknots = xyknots[:nk] yknots = xyknots[nk:] try: f = splinet(xknots, yknots, x, tau) except Exception as error: print(error) f = 0 * y return np.sum(((y - f) * w) ** 2) # initial guess of knots args0 = np.ravel([xknots, interp1e(x, y)(xknots)]) # global bounds bounds = [xbounds] * nk + [(min(y) - np.ptp(y), max(y) + np.ptp(y))] * nk # relational constraints cons = [] if self._min_sep > 0: for i in range(nk - 1): cons.append({'type': 'ineq', 'fun': lambda x, i=i: (x[i + 1] - x[i]) - min_sep}) for xc, yc in xyc: cons.append({'type': 'eq', 'fun': lambda x, xc=xc, yc=yc: splinet(x[:nk], x[nk:], [xc], tau)[0] - yc}) eps = 1e-8 # consistent with default minimize steps for xp, yp in xyp: cons.append({'type': 'eq', 'fun': lambda x, xp=xp, yp=yp: (splinet(x[:nk], x[nk:], [xp + eps / 2.], tau)[0] - splinet(x[:nk], x[nk:], [xp - eps / 2.], tau)[0]) / eps - yp}) options = {"eps": 1e-8} self.opt = optimize.minimize(cost, args0, bounds=bounds, constraints=cons, options=options) self._xknots = self.opt.x[:nk] * xnorm self._yknots = self.opt.x[nk:] * ynorm if not self.opt['success']: printe("WARNING: Optimizer failed due to\n{:}".format(self.opt['message'])) def _get_yknots(self, xknots=None): """Original optimization of yknots given fixed xknots""" x, y = self._data w = self._weights tau = self._tau if xknots is None: xknots = self._xknots def cost(y0): f = splinet(xknots, y0, x, tau) return np.sum(((y - f) * w) ** 2) y0 = interp1e(x, y)(xknots) bounds = [(min(y) - np.ptp(y), max(y) + np.ptp(y))] * len(xknots) y0 = optimize.fmin_l_bfgs_b(cost, y0, args=(), bounds=bounds, approx_grad=True)[0] return y0 def _set_xknots(self, xknots=None): """ Optimize the xknots, using the original LSQUnivariateSplineT method for finding the yknots given xknots. """ if xknots is None: xknots = self._xknots * 1.0 # fit knots in a unit box xnorm = np.ptp(self._data[0]) x = self._data[0] / xnorm xbounds = self._xbounds ibounds = [i for i, v in enumerate(xbounds) if v is not None] xbounds[ibounds] /= xnorm xknots = xknots / xnorm min_sep = self._min_sep / xnorm ynorm = np.ptp(self._data[1]) y = self._data[1] / ynorm w = self._weights * ynorm xyc = [(xc / xnorm, yc / ynorm) for xc, yc in self._xyc] xyp = [(xc / xnorm, yc / ynorm) for xc, yc in self._xyp] tau = self._tau nk = len(xknots) def temp_spline(xknots, xnew): try: yknots = _get_yknots(x, y, xknots, tau=tau, w=w) s = splinet(xknots, yknots, xnew, tau) except Exception as error: print(error) s = 0 * xnew return s def cost(xknots): s = temp_spline(xknots, x) return np.sum(((y - s) * w) ** 2) # initial guess of knots args0 = xknots # global bounds bounds = [xbounds] * nk # relational constraints cons = [] if self._min_sep > 0: for i in range(nk - 1): cons.append({'type': 'ineq', 'fun': lambda x, i=i: (x[i + 1] - x[i]) - min_sep}) for xc, yc in self._xyc: cons.append({'type': 'eq', 'fun': lambda x: temp_spline(x, [xc])[0] - yc}) eps = 1e-8 # consistent with default minimize steps for xc, yp in self._xyp: cons.append({'type': 'eq', 'fun': lambda x: (temp_spline(x, [xc + eps / 2.])[0] - temp_spline(x, [xc - eps / 2.])[0]) / eps - yp}) options = {"eps": 1e-8} self.opt = optimize.minimize(cost, args0, bounds=bounds, constraints=cons, options=options) self._xknots = opt.x * xnorm self._yknots = _set_yknots(x, y, self._xknots, tau=tau, w=w) * ynorm if not self.opt['success']: printe("WARNING: Optimizer failed due to\n{:}".format(self.opt['message']))
[docs] def get_knots(self): """ Get the x locations of knots :return: tuple. list of (x, y) tuples of knot location and value. """ knots = list(zip(self._xknots * 1, self._yknots * 1)) return knots
[docs] def get_knot_locations(self): """ Get the x locations of knots :return: tuple. (x, y) knot location and value. """ return self._xknots * 1
[docs] def get_knot_values(self): """ Get the x locations of knots :return: tuple. (x, y) knot location and value. """ return self._yknots * 1
[docs]class AutoUnivariateSpline(interpolate.UnivariateSpline): """ Class of univariate spline with smoothing optimized for reduced chi squared assuming w=1/std_devs. If no weight is supplied, it is taken to be 1/std_devs(y) and these should all be finite for this to make sense. :Examples: If we have some curve of data with meaningful deviations beyond the errorbars, >> nx = 50 >> x = np.linspace(-3, 3, nx) >> y = np.exp(-x**2) + 0.3 * (np.random(nx) - 0.5) >> u = unumpy.uarray(y, 0.1 * np.ones(nx)) >> fig, ax = plt.subplots() >> uerrorbar(x, u, marker='o', label='data', ax=ax) The scipy default spline uses smoothing s = len(x), >> xs = np.linspace(min(x), max(x), 1000) >> spl = UnivariateSpline(x, y, w = 1/std_devs(u)) >> ax.plot(xs, spl(xs), label='default spline') but the unceratinties spline optimizes the smoothing for reduced chi squared ~1, >> uspl = AutoUnivariateSpline(x, u) >> ax.plot(xs, uspl(xs), label='auto spline') >> ax.legend() The figure shows the uncertainties spline more accurately captures the meaningful deviations beyond the errorbars. In numbers, >> default_smooth = spl._data[6] >> default_rchisq = spl.get_residual() / (nx - (len(spl.get_coeffs()) + len(spl.get_knots()) - 2)) >> print('Default smoothing is {:.1f} results in reduced chi squared {:.1f}'.format(default_smooth, default_rchisq)) >> print('Optimal smoothing of {:.1f} results in reduced chi squared {:.1f}'.format(uspl.get_smoothing_factor(), ... uspl.get_reduced_chisqr())) If the difference is not large, try running again (the deviations are random!). To see how the optimizer arrived at the result, you can get the full evolution. Remember, it is targeting a reduced chi squared of unity. >> s, f = uspl.get_evolution(norm=False) >> fig, ax = plt.subplots() >> ax.plot(s, f, marker='o', ls='') # all points tested >> ax.plot([uspl.get_smoothing_factor()], [uspl.get_reduced_chisqr() - 1], marker='s') # final value >> ax.set_xscale('log') >> ax.set_xlabel('s') >> ax.set_ylabel('Reduced chi squared - 1') """ def __init__(self, x, y, w=None, bbox=(None, None), k=3, ext=0, check_finite=False, max_interior_knots=None, verbose=False): """ :param x: np.ndarray. Must be increasing :param y: unumpy.uarray. Uncertainties array from uncertainties.unumpy. :param w: np.ndarray. Optionally overrides uncertainties from y. Assumed to be 1/std_devs of gaussian errors. :param bbox: (2,) array_like. 2-sequence specifying the boundary of the approximation interval. :param k: int. Degree of the smoothing spline. Must be 1 <= k <= 5. Default is k=3, a cubic spline. :param ext: int or str. Controls the extrapolation mode for elements not in the knot interval. Default 0. if ext=0 or 'extrapolate', return the extrapolated value. if ext=1 or 'zeros', return 0 if ext=2 or 'raise', raise a ValueError if ext=3 of 'const', return the boundary value. :param check_finite: bool. Whether to check that the input arrays contain only finite numbers. :param max_interior_knots: int. Maximum number of interior knots in a successful optimization. Use this to enforce over smoothing. :param verbose: bool. Print warnings """ printd('Initializing AutoUnivariateSpline') # split the uncertainties array into its components yn = unumpy.nominal_values(y) ye = unumpy.std_devs(y) if w is None: if is_uncertain(y): w = 1. / ye else: w = np.ones_like(y) / (0.01 * np.mean(np.abs(y))) # assume ~1% error # Sanitize x, y, e data if check_finite: w[w == np.inf] = max(w[w != np.inf]) * 1e6 # inf weight makes all residuals inf - can't optimize that valid = np.isfinite(x) * np.isfinite(yn) * np.isfinite(w) xv = x[valid] yv = yn[valid] wv = w[valid] printd(' - removed {:} of {:} values checking finite status'.format(len(valid) - len(xv), len(valid))) check_finite = False else: xv, yv, wv = x, yn, w # defaults nx = len(xv) # key word arguments always passed to smoothed splines kwargs = dict(w=wv, bbox=bbox, k=k, ext=ext, check_finite=check_finite) self._fev = [] self._sev = [] def cost(s): """Reduced chi squared of a smoothed spline""" fit = interpolate.UnivariateSpline(xv, yv, s=max(s, 0.0) * nx, **kwargs) res = fit.get_residual() nfree = len(fit.get_coeffs()) + len(fit.get_knots()) - 2 nob = fit._data[0].shape[0] if (nob - nfree) <= 0: # plateau when vastly under smoothing printd(' ** forcing the over-smooth plateau on vastly under-smoothed spline') fit = interpolate.LSQUnivariateSpline(xv, yv, t=(), **kwargs) res = fit.get_residual() nfree = len(fit.get_coeffs()) + len(fit.get_knots()) - 2 nob = fit._data[0].shape[0] fun = np.abs(res / (nob - nfree) - 1) printd('s={:}, res = {:}, cost = {:}'.format(s, res, fun)) self._sev.append(s[0]) self._fev.append(fun) return fun # optimize the scaling factor on the default s=nx, so the optimizer should be dealing in order 1 # Multiplied the sum((wv - yv)**2) to be robust to crazy weights like 1 when the data is ~1e19 shigh = np.sum((wv * yv) ** 2) / nx # very smooth (ok with a all 0 fit) slow = 1e-4 * shigh if max_interior_knots is not None: def check_max_knots(s): if isinstance(s, list): s = s[0] check = (max_interior_knots + 2.0 - len(interpolate.UnivariateSpline(xv, yv, s=max(s, 0.0) * nx, **kwargs).get_knots())) printd('s = {:}, check = {:}'.format(s, check)) return check # make sure we start somewhere valid for f1, f2 in zip(np.logspace(-7, 1, 9)[:-1], np.logspace(-7, 1, 9)[1:]): if check_max_knots(shigh * f2) >= 0: # fine tune between higher and lower logscale factors for l1, l2 in zip(np.linspace(f1, f2, 5, True)[:-1], np.linspace(f1, f2, 5, True)[1:]): if check_max_knots(shigh * l2) >= 0: slow = shigh * l1 # lower bound shigh *= l2 break break cons = [{'type': 'ineq', 'fun': check_max_knots}] bounds = [(slow, 10 * shigh)] # I don't think COBYLA actually uses this, but SLSQP doesn't do ineq's method = 'COBYLA' options = dict(rhobeg=1e-3 * (bounds[0][1] - bounds[0][0])) else: check_max_knots = lambda s: 1 cons = () bounds = [(slow, shigh)] method = None options = {} # optimize the knots/coeffs self.opt = optimize.minimize(cost, (slow + shigh) * 0.5, bounds=bounds, constraints=cons, method=method, options=options) if not self.opt['success'] or not check_max_knots(self.opt.x) >= 0: # try over smoothing? # s1 = s0 * 100 # self.opt = optimize.minimize(cost, s1, bounds=[(0.1 * s1, 100 * s1)], constraints=cons, method=method) # just bound it from our brute force (robust but might miss a smaller s0 optimum for fear of breaking cons) bounds = [(slow, shigh)] printd(' ** Falling back to unconstrained optimization between {:}'.format(bounds)) self.opt = optimize.minimize(cost, shigh, bounds=bounds, method='L-BFGS-B') if not self.opt['success']: printe("WARNING: Optimizer failed due to\n{:}".format(self.opt['message'])) if max_interior_knots is not None: self.opt.x = shigh # go with our brute search lowest smoothing that met the constraint optimum if not check_max_knots(self.opt.x) >= 0: printe("WARNING: Optimizer failed to minimize red-chi2 with only {:} knots".format(max_interior_knots)) # for some reason it sometimes sticks on a boundary point despite having found an obvious internal minimum if self.opt.fun > np.min(self._fev): i = np.argmin(self._fev) check = check_max_knots(self._sev[i]) if check < 0: if verbose: printw('WARNING: Autoknotting found {:} (>{:}) knots lowers the reduced chi squared'.format( max_interior_knots - int(check), max_interior_knots)) else: self.opt.x, self.opt.fun = self._sev[i], self._fev[i] interpolate.UnivariateSpline.__init__(self, xv, yv, s=self.opt.x * nx, **kwargs) printd("Final number of knots is {:} of maximum {:}".format(len(self.get_knots()) - 2, max_interior_knots))
[docs] def get_reduced_chisqr(self): """ :return: float. Reduced chi squared of the spline. """ res = self.get_residual() nfree = len(self.get_coeffs()) + (len(self.get_knots()) - 2) nob = self._data[0].shape[0] return res / (nob - nfree)
[docs] def get_smoothing_factor(self): """ :return: float. Smoothing factor using in UnivariateSpline. """ return self._data[6] * 1.0
[docs] def get_evolution(self, norm=True): """ :param norm: bool. Normalize s to the UnivariateSpline default npts. :return: tuple. sev, fev where sev is a record of all smoothing values tried in the optimization and fev is the corresponding reduced chi squared costs. """ s, f = np.array(self._sev), np.array(self._fev) if not norm: s *= len(self._data[0]) return s, f
[docs]class CLSQUnivariateSpline(interpolate.LSQUnivariateSpline): """ Constrained least square univariate spline. This class sacrifices the generality of UnivariateSpline's smoothing but enables the ability to constrain values and/or derivatives of the spline. The constraints are used in an optimization of the knot locations, not fundamentally enforced in the underlying equations. Thus, constraints far from the natural spline will cause errors. :Examples: Using the same data from the LSQUnivariateSpline examples, >> x = np.linspace(-2, 2, 50) >> nominal_values = np.exp(-x**2) >> std_devs = 0.1 * np.ones_like(nominal_values) >> y = nominal_values + np.random.normal(loc=0, scale=std_devs) If we simply use this to fit a spline it will use AutoUnivariateSpline to optimize the knot locations. >> spl = CLSQUnivariateSpline(x, y, w=1./std_devs) >> xs = np.linspace(-2, 2, 1000) >> plt.subplots() >> plt.errorbar(x, y, std_devs, marker='o', ls='') >> plt.plot(xs, spl(xs), label='unconstrained') But the new part of this class is that is enables additional constraints on the spline. For example, we can request the spline have zero derivative at the left boundary. >> splc = CLSQUnivariateSpline(x, y, w=1/std_devs, min_separation=0.1, xyprime_constraints=[(-2,0)]) >> plt.plot(xs, splc(xs), label='constrained') >> plt.legend() """ def __init__(self, x, y, w=None, bbox=(None, None), k=3, ext=0, check_finite=False, t=None, optimize_knots=True, min_separation=0, xy_constraints=(), xyprime_constraints=(), maxiter=100, verbose=False): """ Initialize a instance of a constrained least square univariate spline. :param x: (N,) array_like. Input dimension of data points. Must be increasing :param y: (N,) array_like. Input dimension of data points :param w: (N,) array_like. Weights for spline fitting. Must be positive. Default is equal weighting. :param bbox: (2,) array_like. 2-sequence specifying the boundary of the approximation interval. :param k: int. Degree of the smoothing spline. Must be 1 <= k <= 5. Default is k=3, a cubic spline. :param ext: int or str. Controls the extrapolation mode for elements not in the knot interval. Default 0. if ext=0 or 'extrapolate', return the extrapolated value. if ext=1 or 'zeros', return 0 if ext=2 or 'raise', raise a ValueError if ext=3 of 'const', return the boundary value. :param check_finite: bool. Whether to check that the input arrays contain only finite numbers. :param t: (M,) array_like or integer. Interior knots of the spline in ascending order (t in LSQUnivariateSplien) or maximum number of interior knots (max_interior_knots in AutoUnivariateSpline). :param optimize_knots: bool. Allow optimizer to change knot locations after initial guess from t or AutoUnivariateSpline. :param min_separation: float. Minimum separation between knot locations if not explicitely specified by t. :param xy_constraints: list of tuples. Spline is constrained to have these values at these locations. :param xyprime_constraints: Spline is constrained to have these derivatives as these locations. x and y separately, and 'y' to simply use the prescribed knot locations. :param maxiter: int. Maximum number of iterations for spline coeff optimization under constraints. :param verbose: bool. Print warnings """ debug = False # set defaults if w is None: if is_uncertain(y): w = 1. / unumpy.std_devs(y) else: w = np.ones_like(y) y = unumpy.nominal_values(y) # Sanitize x, y, e data if check_finite: valid = np.isfinite(x) * np.isfinite(y) * np.isfinite(w) x = x[valid] y = y[valid] w = w[valid] printd(' - removed {:} of {:} values checking finite status'.format(len(valid) - len(x), len(valid))) check_finite = False # be robust to typical shorthand if xy_constraints is None: xy_constraints = () if np.shape(xy_constraints) == (2,): xy_constraints = (xy_constraints,) if xyprime_constraints is None: xyprime_constraints = () if np.shape(xyprime_constraints) == (2,): xyprime_constraints = (xyprime_constraints,) # store unique variables self._min_sep = min_separation self._xyc = xy_constraints self._xyp = xyprime_constraints # optimize knot locations on a unit scale xnorm = np.ptp(x) x0 = x / xnorm s0 = max(1e-8, min_separation / xnorm) ynorm = np.ptp(y) y0 = y / ynorm w0 = w * ynorm xyc = [(xc / xnorm, yc / ynorm) for xc, yc in xy_constraints] xyp = [(xc / xnorm, yc / ynorm) for xc, yc in xyprime_constraints] # helpful shorthand kwargs = dict(bbox=bbox, k=k, ext=ext, check_finite=check_finite) eps = 1e-8 # consistent with default minimize steps # initial guess if t is None or is_int(t): sp0 = AutoUnivariateSpline(x0, y0, w=w0, max_interior_knots=t, verbose=verbose, **kwargs) t0 = sp0.get_knots()[1:-1] auto_t = True else: printd(" CLSQUnivariateSpline starting with given knots") t0 = t / xnorm auto_t = False sp0 = interpolate.LSQUnivariateSpline(x0, y0, t0, w=w0, **kwargs) args0 = np.hstack((sp0.get_knots()[1:-1], sp0.get_coeffs())) # helpful constants nx = len(x0) nt = len(t0) nc = len(sp0.get_coeffs()) nfree = len(args0[nt * int(optimize_knots):]) - len(xyc) - len(xyp) nu = nx - nfree if nu <= 0: raise ValueError("Under constrained fit. Reduce number of knots.") t_front = [sp0._eval_args[0][0]] * (k + 1) t_back = [sp0._eval_args[0][-1]] * (k + 1) c_back = np.zeros(k + 1) if debug: print('eval args = {:}'.format(sp0._eval_args)) print('guess tc = {:}'.format(args0)) def form_tck0(tc_args): """Form a full (t,c,k) tuple from only the varying knots and coeffs.""" t = np.hstack((t_front, tc_args[:nt], t_back)) c = np.hstack((tc_args[nt:], c_back)) return (t, c, k) def tc_cost(tc_args): """ Reduced chi squared given interior knots and all coeffs, and assuming w = 1/std_devs. """ try: yf = interpolate.splev(x0, form_tck0(tc_args)) reduced_chisqr = np.sum((w0 * (y0 - yf)) ** 2) / nu c = np.abs(reduced_chisqr - 1) except ValueError as error: # protect against "Interior knots t must satisfy Schoenberg-Whitney conditions" c = np.sum((y0 * w0) ** 2) return c def c_cost(c_args): """Reduced chi squared if only changing the coeffs.""" return tc_cost(np.hstack((t0, c_args))) if (auto_t or not optimize_knots) and (len(xy_constraints) + len(xyprime_constraints) == 0): # just a normal scipy coefficient optimization with the given or automatic knots printd(" CLSQUnivariateSpline defaulting to plain LSQUnivariateSpline") printd(" with knots at {}".format(t0 * xnorm)) interpolate.LSQUnivariateSpline.__init__(self, x, y, t0 * xnorm, w=w, **kwargs) else: if optimize_knots: printd('CLSQUnivariateSpline optimizing knot locations') # global bounds bounds = [(min(x0) + max(eps, s0), max(x0) - max(eps, s0))] * nt bounds += [(None, None)] * nc # relational constraints cons = [] if s0 > 0: eps = 1e-3 * s0 for xc, yc in xyc: cons.append({'type': 'eq', 'fun': lambda args, xc=xc, yc=yc: interpolate.splev(xc, form_tck0(args)) - yc}) for xp, yp in xyp: cons.append({'type': 'eq', 'fun': lambda args, xp=xp, yp=yp: interpolate.splev(xp, form_tck0(args), der=1) - yp}) # additional constraints for knot separation for i in range(1, nt): cons.append({'type': 'ineq', 'fun': lambda t, i=i: (t[i] - t[i - 1]) - s0}) # optimize the knot locations under given constraints self.opt = optimize.minimize(tc_cost, args0, bounds=bounds, constraints=cons, options={"eps": eps, "maxiter": maxiter, "ftol":1e-3}) if not self.opt['success']: printe("WARNING: Optimizer failed due to\n{:}".format(self.opt['message'])) tck = form_tck0(self.opt.x) # initialize self with all the data (unconstrained coeffs) interpolate.LSQUnivariateSpline.__init__(self, x, y, self.opt.x[:nt] * xnorm, w=w, **kwargs) else: printd('CLSQUnivariateSpline not optimizing knot locations') # relational constraints cons = [] for xc, yc in xyc: cons.append({'type': 'eq', 'fun': lambda args, xc=xc, yc=yc: interpolate.splev(xc, form_tck0(np.hstack((t0, args)))) - yc}) for xp, yp in xyp: cons.append({'type': 'eq', 'fun': lambda args, xp=xp, yp=yp: interpolate.splev(xp, form_tck0(np.hstack((t0, args))), der=1) - yp}) # optimize the coefficients to satisfy the constraints self.opt = optimize.minimize(c_cost, args0[nt:], constraints=cons, options={"eps": eps, "maxiter": maxiter, "ftol":1e-3}) if not self.opt['success']: printe("WARNING: Optimizer failed due to\n{:}".format(self.opt['message'])) tck = form_tck0(np.hstack((t0, self.opt.x))) # initialize self with all the data (unconstrained coeffs) interpolate.LSQUnivariateSpline.__init__(self, x, y, t, w=w, **kwargs) if debug: plt.subplots() errorbar(x0, y0, 1 / w) scatter(x0, interpolate.splev(x0, tck)) # renormalize the optimal knots and coeffs tck = (tck[0] * xnorm, tck[1] * ynorm, tck[2]) # set private t, c to the optimized t & c tmp = list(self._data) tmp[8:10] = tck[0:2] # reset the residual for the changed t & c tmp[10] = np.sum((w * (y - interpolate.splev(x, tck))) ** 2) self._data = tuple(tmp) # make sure they take effect for callable methods self._reset_class()
[docs] def get_reduced_chisqr(self): """ :return: float. Reduced chi squared of the spline. """ res = self.get_residual() nfree = len(self.get_coeffs()) + len(self.get_knots()) - 2 nob = self._data[0].shape[0] return res / (nob - nfree)
[docs]@_available_to_user_math class MonteCarloSpline(CLSQUnivariateSpline): """ Monte Carlo Uncertainty propagation through python spline fits. The concept follows https://gist.github.com/thriveth/4680e3d3cd2cfe561a57 by Th/oger Rivera-Thorsen (thriveth), and essentially forms n_trials unique spline instances with randomly perturbed data assuming w=1/std_devs of gaussian noise. Note, calling instances of this class returns unumpy.uarrays of Variable objects using the uncertainties package. :Examples: Using the same data from the LSQUnivariateSpline examples, >> x = np.linspace(-2, 2, 50) >> nominal_values = np.exp(-x**2) >> std_devs = 0.1 * np.ones_like(y) >> y = nominal_values + np.random.normal(loc=0, scale=std_devs) We can fit a monte carlo spline to get an errorbar on the interpolation. >> spl = MonteCarloSpline(x, y, w=1/std_devs) >> xs = np.linspace(-2, 2, 1000) >> fig, ax = plt.subplots() >> eb = ax.errorbar(x, y, std_devs, marker='o', ls='') >> ub = uband(xs, spl(xs), label='unconstrained') The individual Monte Carlo splines are CLSQUnivariateSpline instances, so we can set hard constraints as well. >> splc = MonteCarloSpline(x, y, w=1/std_devs, min_separation=0.1, xyprime_constraints=[(-2,0)]) >> ub = uband(xs, splc(xs), label='constrained') >> ax.legend() Note, this class is a child of the scipy.interpolate.LSQUnivariateSpline class, and has all of the standard spline class methods. Where appropriate, these methods dig into the montecarlo trials to return uncertainties. For example, >> print('knots are fixed at {}'.format(splc.get_knots())) >> print('coeffs vary around {}'.format(splc.get_coeffs())) """ def __init__(self, x, y, w=None, bbox=(None, None), k=3, ext=0, check_finite=False, t=None, optimize_knots=True, min_separation=0, xy_constraints=(), xyprime_constraints=(), maxiter=200, n_trials=100): """ Initialize a instance of a MonteCarlo constrained least square univariate spline. :param x: (N,) array_like. Input dimension of data points. Must be increasing :param y: (N,) array_like. Input dimension of data points :param w: (N,) array_like. Weights for spline fitting. Must be positive. Default is equal weighting. :param bbox: (2,) array_like. 2-sequence specifying the boundary of the approximation interval. :param k: int. Degree of the smoothing spline. Must be 1 <= k <= 5. Default is k=3, a cubic spline. :param ext: int or str. Controls the extrapolation mode for elements not in the knot interval. Default 0. if ext=0 or 'extrapolate', return the extrapolated value. if ext=1 or 'zeros', return 0 if ext=2 or 'raise', raise a ValueError if ext=3 of 'const', return the boundary value. :param check_finite: bool. Whether to check that the input arrays contain only finite numbers. :param t: (M,) array_like or integer. Interior knots of the spline in ascending order or maximum number of interior knots. :param optimize_knots: bool. Allow optimizer to change knot locations after initial guess from t or AutoUnivariateSpline. :param min_separation: float. Minimum separation between knot locations if not explicitely specified by t. :param xy_constraints: list of tuples. Spline is constrained to have these values at these locations. :param xyprime_constraints: Spline is constrained to have these derivatives as these locations. x and y separately, and 'y' to simply use the prescribed knot locations. :param maxiter: int. Maximum number of iterations for spline coeff optimization under constraints. :param n_trials: int. Number of Monte Carlo spline iterations used to form errorbars. """ printd('MonteCarloSpline') t_0 = time.time() # initial spline sets the base attributes, methods, etc. printd(' - Initializing first CLSQUnivariateSpline') kwargs = dict(w=w, bbox=bbox, k=k, ext=ext, check_finite=check_finite, t=t, optimize_knots=optimize_knots, min_separation=min_separation, xy_constraints=xy_constraints, xyprime_constraints=xyprime_constraints, maxiter=maxiter) CLSQUnivariateSpline.__init__(self, x, y, verbose=True, **kwargs) printd(' > Initialization first CLSQUnivariateSpline took {:.3e} seconds', (time.time() - t_0)) # repeat n times for statistics printd(' - Initialization of all MC splines') printd(' > Initializing with knots = {:}'.format(t)) self.montecarlo = [CLSQUnivariateSpline(x, y, **kwargs)] kwargs['optimize_knots'] = False kwargs['t'] = self.get_knots()[1:-1] # fix the knots for speed & consistency printd(' > Fixing knots = {:}'.format(kwargs['t'])) w = self._data[2] # changed if check_finite removed nans kwargs['w'] = self._data[2] # keep weighting by error? # kwargs['w'] = np.ones_like(self._data[2]) * np.mean(self._data[2]) # Monte Carlo determines UQ (don't double count!) kwargs['check_finite'] = False # data already cleaned by first call for i in range(max(n_trials, 0)): xi = self._data[0] # _data == x,y,w,xb,xe,k,s,n,t,c,fp,fpint,nrdata,ier yi = self._data[1] + np.random.normal(loc=0, scale=1 / w) spli = CLSQUnivariateSpline(xi, yi, **kwargs) self.montecarlo.append(spli) printd(' > Initialization of all MC splines took {:.3e} seconds', (time.time() - t_0)) # reset some data that relies on (patched) methods tmp = list(self._data) tmp[10] = np.sum((tmp[2] * (tmp[1] - unumpy.nominal_values(self(tmp[0])))) ** 2, axis=0) self._data = tuple(tmp) # make sure they take effect for callable methods self._reset_class() # at this point the parent class get_knots and get_residual methods are valid def __call__(self, x, nu=0, ext=None): """ Evaluate spline (or its nu-th derivative) at positions x. :param x: array_like. A 1-D array of points at which to return the value of the smoothed spline or its derivatives. Note: x can be unordered but the evaluation is more efficient if x is (partially) ordered. :param nu : int. The order of derivative of the spline to compute. :param ext: int. Controls the value returned for elements of ``x`` not in the interval defined by the knot sequence. * if ext=0 or 'extrapolate', return the extrapolated value. * if ext=1 or 'zeros', return 0 * if ext=2 or 'raise', raise a ValueError * if ext=3 or 'const', return the boundary value. The default value is 0, passed from the initialization of the class instance. """ t_0 = time.time() printd('Evaluating MonteCarloSpline fit results') ss = np.vstack([spl(x, nu=nu, ext=ext) for spl in self.montecarlo]) result = unumpy.uarray(np.mean(ss, axis=0), np.std(ss, axis=0)) printd(' - Evaluation of Monte Carlo splines took {:.3e} seconds'.format(time.time() - t_0)) return result
[docs] def antiderivative(self, n=1): """ Construct a new spline representing the antiderivative of this spline. :param n: int. Order of antiderivative to evaluate. """ base = CLSQUnivariateSpline.antiderivative(self, n=n) result = copy.deepcopy(self) for i, spl in enumerate(self.montecarlo): result.montecarlo[i] = spl.antiderivative(n=n) result._data = base._data return result
[docs] def derivative(self, n=1): """ Construct a new spline representing the derivative of this spline. :param n: int. Order of antiderivative to evaluate. """ base = CLSQUnivariateSpline.antiderivative(self, n=n) result = copy.deepcopy(self) for i, spl in enumerate(self.montecarlo): result.montecarlo[i] = spl.derivative(n=n) result._data = base._data return result
[docs] def derivatives(self, x): """Return all derivatives of the spline at the point x.""" ss = np.vstack([spl.derivatives(x) for spl in self.montecarlo]) result = unumpy.uarray(np.mean(ss, axis=0), np.std(ss, axis=0)) return result
[docs] def get_coeffs(self): """Return spline coefficients.""" ss = np.vstack([spl.get_coeffs() for spl in self.montecarlo]) result = unumpy.uarray(np.mean(ss, axis=0), np.std(ss, axis=0)) return result
[docs] def integral(self, a, b): """Return definite integral of the spline between two given points.""" ss = np.vstack([spl.integral(a, b) for spl in self.montecarlo]) result = unumpy.uarray(np.mean(ss, axis=0), np.std(ss, axis=0)) return result
[docs] def roots(self): """Return the zeros of the spline.""" # todo: make this robust to some instances having a different number of roots ss = np.vstack([spl.roots() for spl in self.montecarlo]) result = unumpy.uarray(np.mean(ss, axis=0), np.std(ss, axis=0)) return result
[docs] def set_smoothing_factor(self, s): """Continue spline computation with the given smoothing factor s and with the knots found at the last call.""" self.set_smoothing_factor(s) for spl in self.montecarlo: spl.set_smoothing_factor(s)
#----------- # paths #-----------
[docs]@_available_to_user_math def contourPaths(x, y, Z, levels, remove_boundary_points=False, smooth_factor=1): ''' :param x: 1D x coordinate :param y: 1D y coordinate :param Z: 2D data :param levels: levels to trace :param remove_boundary_points: remove traces at the boundary :param smooth_factor: smooth contours by cranking up grid resolution :return: list of segments ''' import matplotlib if compare_version(matplotlib.__version__,'2.1')>=0: import matplotlib._contour as _contour else: from matplotlib import _cntr sf = int(round(smooth_factor)) if sf > 1: x = scipy.ndimage.zoom(x, sf) y = scipy.ndimage.zoom(y, sf) Z = scipy.ndimage.zoom(Z, sf) [X,Y]=np.meshgrid(x,y) if compare_version(matplotlib.__version__,'2.1')>=0: contour_generator = _contour.QuadContourGenerator(X, Y, Z, None, True, 0) else: Cntr = matplotlib._cntr.Cntr(X,Y,Z) mx=min(x) Mx=max(x) my=min(y) My=max(y) allsegs = [] for level in levels: if compare_version(matplotlib.__version__,'2.1')>=0: segs=contour_generator.create_contour(level) else: nlist = Cntr.trace(level) nseg = len(nlist)//2 segs = nlist[:nseg] if not remove_boundary_points: segs_ = segs else: segs_ = [] for segarray in segs: x_ = segarray[:,0] y_ = segarray[:,1] valid = [] for i in range(len(x_)-1): if np.isclose(x_[i],x_[i+1]) and (np.isclose(x_[i],Mx) or np.isclose(x_[i],mx)): continue if np.isclose(y_[i],y_[i+1]) and (np.isclose(y_[i],My) or np.isclose(y_[i],my)): continue valid.append((x_[i],y_[i])) if i==len(x_): valid.append(x_[i+1],y_[i+1]) if len(valid): segs_.append(np.array(valid)) segs=list(map(matplotlib.path.Path,segs_)) allsegs.append(segs) return allsegs
[docs]@_available_to_user_math def streamPaths(xm, ym, u, v, start_points, minlength=0, maxlength=1E10, bounds_error=True): ''' Integrate vector field and returns stream line :params xm: uniformly spaced x grid :params xm: uniformly spaced y grid :params u: vector field in the x direction :params v: vector field in the y direction :params start_points: 2D array of seed x,y coordinate points used to start integration :params minlength: reject trajectories shorter than this length :params maxlength: stop tracing trajectory when this length is reached :param bounds_error: raise error if trajectory starts outside of bounds ''' import matplotlib def _get_integrator(u, v, dmap, minlength, maxlength): # rescale velocity onto grid-coordinates for integrations. u, v = dmap.data2grid(u, v) # speed (path length) will be in axes-coordinates u_ax = u / dmap.grid.nx v_ax = v / dmap.grid.ny speed = np.ma.sqrt(u_ax ** 2 + v_ax ** 2) def forward_time(xi, yi): ds_dt = matplotlib.streamplot.interpgrid(speed, xi, yi) if ds_dt == 0: raise TerminateTrajectory() dt_ds = 1. / ds_dt ui = matplotlib.streamplot.interpgrid(u, xi, yi) vi = matplotlib.streamplot.interpgrid(v, xi, yi) return ui * dt_ds, vi * dt_ds def backward_time(xi, yi): dxi, dyi = forward_time(xi, yi) return -dxi, -dyi def integrate(x0, y0): """Return x, y grid-coordinates of trajectory based on starting point. Integrate both forward and backward in time from starting point in grid coordinates. Integration is terminated when a trajectory reaches a domain boundary or when it crosses into an already occupied cell in the StreamMask. The resulting trajectory is None if it is shorter than `minlength`. """ try: dmap.start_trajectory(x0, y0) except Exception: return None sf, xf_traj, yf_traj = matplotlib.streamplot._integrate_rk12(x0, y0, dmap, forward_time, maxlength) dmap.reset_start_point(x0, y0) sb, xb_traj, yb_traj = matplotlib.streamplot._integrate_rk12(x0, y0, dmap, backward_time, maxlength) # combine forward and backward trajectories stotal = sf + sb x_traj = xb_traj[::-1] + xf_traj[1:] y_traj = yb_traj[::-1] + yf_traj[1:] if stotal > minlength: return x_traj, y_traj else: # reject short trajectories dmap.undo_trajectory() return None return integrate x=np.linspace(-1,1,len(xm)) y=np.linspace(-1,1,len(ym)) grid = matplotlib.streamplot.Grid(x, y) mask = matplotlib.streamplot.StreamMask(1) dmap = matplotlib.streamplot.DomainMap(grid, mask) start_points=copy.deepcopy(np.atleast_2d(start_points)).T start_points[:,0]=((start_points[:,0]-min(xm))/(max(xm)-min(xm))-0.5)*2 start_points[:,1]=((start_points[:,1]-min(ym))/(max(ym)-min(ym))-0.5)*2 ## Sanity checks. if (u.shape != grid.shape) or (v.shape != grid.shape): msg = "'u' and 'v' must be of shape 'Grid(x,y)'" raise ValueError(msg) u = np.ma.masked_invalid(u) v = np.ma.masked_invalid(v) integrate = _get_integrator(u, v, dmap, minlength, maxlength) sp2 = np.asanyarray(start_points, dtype=np.float).copy() if not bounds_error: index=((sp2[:,0]<grid.x_origin) + (sp2[:,0]>grid.x_origin+grid.width) + (sp2[:,1]<grid.y_origin) + (sp2[:,1]>grid.y_origin+grid.height)) sp2[index,:]=np.nan # Check if start_points are outside the data boundaries for xs, ys in sp2: if (xs < grid.x_origin or xs > grid.x_origin + grid.width or ys < grid.y_origin or ys > grid.y_origin + grid.height): raise ValueError("Starting point ({}, {}) outside of" " data boundaries".format(xs, ys)) sp2[:, 0] -= grid.x_origin sp2[:, 1] -= grid.y_origin allsegs=[] for xs, ys in sp2: dmap.mask._mask*=0 xg, yg = dmap.data2grid(xs, ys) t = integrate(xg, yg) if t is not None: X= (((np.array(t[0])/(len(xm)+0)-0.5)*2)/2.+0.5)*(max(xm)-min(xm))+min(xm) Y= (((np.array(t[1])/(len(ym)+0)-0.5)*2)/2.+0.5)*(max(ym)-min(ym))+min(ym) allsegs.append(np.array([X,Y])) return allsegs
def _ccw(A,B,C): return (C[1]-A[1])*(B[0]-A[0]) >= (B[1]-A[1])*(C[0]-A[0]) def _intersect(A,B,C,D): return _ccw(A,C,D) != _ccw(B,C,D) and _ccw(A,B,C) != _ccw(A,B,D) def _perp( a ) : b = np.empty_like(a) b[0] = -a[1] b[1] = a[0] return b def _seg_intersect(a1,a2,b1,b2) : if not _intersect(a1,a2,b1,b2): return None da = a2-a1 db = b2-b1 dp = a1-b1 dap = _perp(da) denom = np.dot( dap, db) num = np.dot( dap, dp ) return (num / denom)*db + b1
[docs]@_available_to_user_math def line_intersect(path1,path2,return_indices=False): """ intersection of two 2D paths :param path1: array of (x,y) coordinates defining the first path :param path2: array of (x,y) coordinates defining the second path :param return_indices: return indices of segments where intersection occurred :return: array of intersection points (x,y) """ warnings.filterwarnings('ignore', category=RuntimeWarning, message='invalid value.*') ret=[] index=[] path1=np.array(path1).astype(float) path2=np.array(path2).astype(float) switch=False if len(path1[:,0])>len(path2[:,0]): path1,path2=path2,path1 switch=True m1=(path1[:-1,:]+path1[1:,:])/2. m2=(path2[:-1,:]+path2[1:,:])/2. l1=np.sqrt(np.diff(path1[:,0])**2+np.diff(path1[:,1])**2) l2=np.sqrt(np.diff(path2[:,0])**2+np.diff(path2[:,1])**2) for k1 in range(len(path1)-1): d=np.sqrt((m2[:,0]-m1[k1,0])**2+(m2[:,1]-m1[k1,1])**2) for k2 in np.where(d<=(l2+l1[k1])/2.)[0]: tmp=_seg_intersect(path1[k1],path1[k1+1],path2[k2],path2[k2+1]) if tmp is not None: index.append([k1,k2]) ret.append(tmp) ret=np.array(ret) if not len(ret): return [] if switch: index=np.array(index) i=np.argsort(index[:,1]) index=[index[i,1],index[i,0]] index=np.array(index).T.astype(int).tolist() ret=ret[i] if return_indices: return ret,index return ret
[docs]@_available_to_user_math def point_to_line(px,py,x1,y1,x2,y2): ''' Calc minimum distance from a point and a line segment (i.e. consecutive vertices in a polyline) ''' x1=np.atleast_1d(x1) x2=np.atleast_1d(x2) y1=np.atleast_1d(y1) y2=np.atleast_1d(y2) LineMag = np.hypot(x2-x1,y2-y1) u1 = (((px - x1) * (x2 - x1)) + ((py - y1) * (y2 - y1))) u = u1 / (LineMag * LineMag) e1 = np.hypot(x1-px,y1-py) e2 = np.hypot(x2-px,y2-py) #handle extremes ix = x1 + u * (x2 - x1) iy = y1 + u * (y2 - y1) d = np.hypot(ix-px,iy-py) #handle extremes i = np.where( ((u < 0.00001) | (u > 1)) )[0] d[i] = e1[i] ii = np.where( e1[i]>e2[i] )[0] d[i[ii]]=e2[ii] return d
[docs]@_available_to_user_math def point_in_poly(x, y, poly): """ Determine if a point is inside a given polygon or not. Polygon is a list of (x,y) pairs. This function returns True or False. The algorithm is called the "Ray Casting Method". Source: http://geospatialpython.com/2011/01/point-in-polygon.html , retrieved 20160105 18:39 :param x, y: floats Coordinates of the point to test :param poly: List of (x,y) pairs defining a polygon. :return: bool Flag indicating whether or not the point is within the polygon. To test: polygon = [(0,10),(10,10),(10,0),(0,0)] point_x = 5 point_y = 5 inside = point_in_poly(point_x, point_y, polygon) print(inside) """ n = len(poly) inside_ = False p1x, p1y = poly[0] for i in range(n+1): p2x, p2y = poly[i % n] if y > min(p1y, p2y): if y <= max(p1y, p2y): if x <= max(p1x, p2x): if p1y != p2y: xints = (y-p1y)*(p2x-p1x)/(p2y-p1y)+p1x if p1x == p2x or x <= xints: inside_ = not inside_ p1x, p1y = p2x, p2y return inside_
[docs]@_available_to_user_math def centroid(x, y): ''' Calculate centroid of polygon :param x: x coordinates of the polygon :param y: y coordinates of the polygon :return: tuple with x and y coordinates of centroid ''' x = np.array(x) y = np.array(y) dy = np.diff(y) dx = np.diff(x) x0 = (x[1:] + x[:-1]) * 0.5 y0 = (y[1:] + y[:-1]) * 0.5 A = np.sum(dy * x0) x_c = -np.sum(dx * y0 * x0) / A y_c = np.sum(dy * x0 * y0) / A return x_c, y_c
#----------- # patch np #----------- if compare_version(np.__version__,'1.11')<0: def gradient(f, *varargs, **kwargs): r""" Return the gradient of an N-dimensional array. The gradient is computed using second order accurate central differences in the interior and either first differences or second order accurate one-sides (forward or backwards) differences at the boundaries. The returned gradient hence has the same shape as the input array. Parameters ---------- f : array_like An N-dimensional array containing samples of a scalar function. varargs : scalar or list of scalar, optional N scalars specifying the sample distances for each dimension, i.e. `dx`, `dy`, `dz`, ... Default distance: 1. single scalar specifies sample distance for all dimensions. if `axis` is given, the number of varargs must equal the number of axes. edge_order : {1, 2}, optional Gradient is calculated using N\ :sup:`th` order accurate differences at the boundaries. Default: 1. .. versionadded:: 1.9.1 axis : None or int or tuple of ints, optional Gradient is calculated only along the given axis or axes The default (axis = None) is to calculate the gradient for all the axes of the input array. axis may be negative, in which case it counts from the last to the first axis. .. versionadded:: 1.11.0 Returns ------- gradient : list of np.ndarray Each element of `list` has the same shape as `f` giving the derivative of `f` with respect to each dimension. Examples -------- >> x = np.array([1, 2, 4, 7, 11, 16], dtype=np.float) >> np.gradient(x) np.array([ 1. , 1.5, 2.5, 3.5, 4.5, 5. ]) >> np.gradient(x, 2) np.array([ 0.5 , 0.75, 1.25, 1.75, 2.25, 2.5 ]) For two dimensional arrays, the return will be two arrays ordered by axis. In this example the first array stands for the gradient in rows and the second one in columns direction: >> np.gradient(np.array([[1, 2, 6], [3, 4, 5]], dtype=np.float)) [np.array([[ 2., 2., -1.], [ 2., 2., -1.]]), np.array([[ 1. , 2.5, 4. ], [ 1. , 1. , 1. ]])] >> x = np.array([0, 1, 2, 3, 4]) >> dx = np.gradient(x) >> y = x**2 >> np.gradient(y, dx, edge_order=2) np.array([-0., 2., 4., 6., 8.]) The axis keyword can be used to specify a subset of axes of which the gradient is calculated >> np.gradient(np.array([[1, 2, 6], [3, 4, 5]], dtype=np.float), axis=0) np.array([[ 2., 2., -1.], [ 2., 2., -1.]]) """ f = np.asanyarray(f) N = len(f.shape) # number of dimensions axes = kwargs.pop('axis', None) if axes is None: axes = tuple(range(N)) # check axes to have correct type and no duplicate entries if isinstance(axes, int): axes = (axes,) if not isinstance(axes, tuple): raise TypeError("A tuple of integers or a single integer is required") # normalize axis values: axes = tuple(x + N if x < 0 else x for x in axes) if max(axes) >= N or min(axes) < 0: raise ValueError("'axis' entry is out of bounds") if len(set(axes)) != len(axes): raise ValueError("duplicate value in 'axis'") n = len(varargs) if n == 0: dx = [1.0]*N elif n == 1: dx = [varargs[0]]*N elif n == len(axes): dx = list(varargs) else: raise SyntaxError( "invalid number of arguments") edge_order = kwargs.pop('edge_order', 1) if kwargs: raise TypeError('"{}" are not valid keyword arguments.'.format( '", "'.join(list(kwargs.keys())))) if edge_order > 2: raise ValueError("'edge_order' greater than 2 not supported") # use central differences on interior and one-sided differences on the # endpoints. This preserves second order-accuracy over the full domain. outvals = [] # create slice objects --- initially all are [:, :, ..., :] slice1 = [slice(None)]*N slice2 = [slice(None)]*N slice3 = [slice(None)]*N slice4 = [slice(None)]*N otype = f.dtype.char if otype not in ['f', 'd', 'F', 'D', 'm', 'M']: otype = 'd' # Difference of datetime64 elements results in timedelta64 if otype == 'M': # Need to use the full dtype name because it contains unit information otype = f.dtype.name.replace('datetime', 'timedelta') elif otype == 'm': # Needs to keep the specific units, can't be a general unit otype = f.dtype # Convert datetime64 data into ints. Make dummy variable `y` # that is a view of ints if the data is datetime64, otherwise # just set y equal to the array `f`. if f.dtype.char in ["M", "m"]: y = f.view('int64') else: y = f for i, axis in enumerate(axes): if y.shape[axis] < 2: raise ValueError( "Shape of array too small to calculate a numerical gradient, " "at least two elements are required.") # Numerical differentiation: 1st order edges, 2nd order interior if y.shape[axis] == 2 or edge_order == 1: # Use first order differences for time data out = np.empty_like(y, dtype=otype) slice1[axis] = slice(1, -1) slice2[axis] = slice(2, None) slice3[axis] = slice(None, -2) # 1D equivalent -- out[1:-1] = (y[2:] - y[:-2])/2.0 out[slice1] = (y[slice2] - y[slice3])/2.0 slice1[axis] = 0 slice2[axis] = 1 slice3[axis] = 0 # 1D equivalent -- out[0] = (y[1] - y[0]) out[slice1] = (y[slice2] - y[slice3]) slice1[axis] = -1 slice2[axis] = -1 slice3[axis] = -2 # 1D equivalent -- out[-1] = (y[-1] - y[-2]) out[slice1] = (y[slice2] - y[slice3]) # Numerical differentiation: 2st order edges, 2nd order interior else: # Use second order differences where possible out = np.empty_like(y, dtype=otype) slice1[axis] = slice(1, -1) slice2[axis] = slice(2, None) slice3[axis] = slice(None, -2) # 1D equivalent -- out[1:-1] = (y[2:] - y[:-2])/2.0 out[slice1] = (y[slice2] - y[slice3])/2.0 slice1[axis] = 0 slice2[axis] = 0 slice3[axis] = 1 slice4[axis] = 2 # 1D equivalent -- out[0] = -(3*y[0] - 4*y[1] + y[2]) / 2.0 out[slice1] = -(3.0*y[slice2] - 4.0*y[slice3] + y[slice4])/2.0 slice1[axis] = -1 slice2[axis] = -1 slice3[axis] = -2 slice4[axis] = -3 # 1D equivalent -- out[-1] = (3*y[-1] - 4*y[-2] + y[-3]) out[slice1] = (3.0*y[slice2] - 4.0*y[slice3] + y[slice4])/2.0 # divide by step size out /= dx[i] outvals.append(out) # reset the slice object in this dimension to ":" slice1[axis] = slice(None) slice2[axis] = slice(None) slice3[axis] = slice(None) slice4[axis] = slice(None) if len(axes) == 1: return outvals[0] else: return outvals np.gradient=gradient # helpful generalization to handle uncertainties or not
[docs]def usqrt(y): """Handle uncertainties if needed""" if is_uncertain(y): result = unumpy.sqrt(y) else: result = np.sqrt(nominal_values(y)) return result
# helpful generalization to handle uncertainties or not
[docs]def ulog(y): """Handle uncertainties if needed""" if is_uncertain(y): result = unumpy.log(y) else: result = np.log(nominal_values(y)) return result
# copy scipy.ndimage filters (v1.1.0) and add causal key word argument def _gaussian_kernel1d(sigma, order, radius, causal=False): """ Computes a 1D Gaussian convolution kernel. """ if order < 0: raise ValueError('order must be non-negative') p = np.polynomial.Polynomial([0, 0, -0.5 / (sigma * sigma)]) x = np.arange(-radius, radius + 1) phi_x = np.exp(p(x), dtype=np.double) if causal: phi_x[radius + 1:] = 0.0 if order > 0: raise AssertionError('Cannot calculate causal gaussian kernel for order > 0') phi_x /= phi_x.sum() if order > 0: q = np.polynomial.Polynomial([1]) p_deriv = p.deriv() for _ in range(order): # f(x) = q(x) * phi(x) = q(x) * exp(p(x)) # f'(x) = (q'(x) + q(x) * p'(x)) * phi(x) q = q.deriv() + q * p_deriv phi_x *= q(x) return phi_x
[docs]@_available_to_user_math @docfiller def gaussian_filter1d(input, sigma, axis=-1, order=0, output=None, mode="reflect", cval=0.0, truncate=4.0, causal=False): """One-dimensional Gaussian filter. Parameters ---------- %(input)s sigma : scalar standard deviation for Gaussian kernel %(axis)s order : int, optional An order of 0 corresponds to convolution with a Gaussian kernel. A positive order corresponds to convolution with that derivative of a Gaussian. %(output)s %(mode)s %(cval)s truncate : float, optional Truncate the filter at this many standard deviations. Default is 4.0. causal : bool or sequence of bools Remove all forward weightings. Returns ------- gaussian_filter1d : np.ndarray Examples -------- >> from scipy.ndimage import gaussian_filter1d >> gaussian_filter1d([1.0, 2.0, 3.0, 4.0, 5.0], 1) np.array([ 1.42704095, 2.06782203, 3. , 3.93217797, 4.57295905]) >> gaussian_filter1d([1.0, 2.0, 3.0, 4.0, 5.0], 4) np.array([ 2.91948343, 2.95023502, 3. , 3.04976498, 3.08051657]) >> import matplotlib.pyplot as plt >> np.random.seed(280490) >> x = np.random.randn(101).cumsum() >> y3 = gaussian_filter1d(x, 3) >> y6 = gaussian_filter1d(x, 6) >> plt.plot(x, 'k', label='original data') >> plt.plot(y3, '--', label='filtered, sigma=3') >> plt.plot(y6, ':', label='filtered, sigma=6') >> plt.legend() >> plt.grid() >> plt.show() Causality example, >> x = np.arange(0,100) >> y = 1.* (x > 40) * (x < 60) >> fig, ax = plt.subplots() >> ax.plot(x, y, 'x-') >> ax.plot(x, gaussian_filter1d(y,3.)) >> ax.plot(x, gaussian_filter1d(y,3., causal=True)) >> ax.set_ylim(0,1.5) """ sd = float(sigma) # make the radius of the filter equal to truncate standard deviations lw = int(truncate * sd + 0.5) if order == 0: # symmetric, use forward order for consistency with causal kernel weights = _gaussian_kernel1d(sigma, order, lw, causal) else: # Since we are calling correlate, not convolve, revert the kernel weights = _gaussian_kernel1d(sigma, order, lw, causal)[::-1] return ndimage.correlate1d(input, weights, axis, output, mode, cval, 0)
[docs]@_available_to_user_math @docfiller def gaussian_filter(input, sigma, order=0, output=None, mode="reflect", cval=0.0, truncate=4.0, causal=False): """Multidimensional Gaussian filter. Parameters ---------- %(input)s sigma : scalar or sequence of scalars Standard deviation for Gaussian kernel. The standard deviations of the Gaussian filter are given for each axis as a sequence, or as a single number, in which case it is equal for all axes. order : int or sequence of ints, optional The order of the filter along each axis is given as a sequence of integers, or as a single number. An order of 0 corresponds to convolution with a Gaussian kernel. A positive order corresponds to convolution with that derivative of a Gaussian. %(output)s %(mode_multiple)s %(cval)s truncate : float Truncate the filter at this many standard deviations. Default is 4.0. causal : bool or sequence of bools Remove all forward weightings. Returns ------- gaussian_filter : np.ndarray Returned array of same shape as `input`. Notes ----- The multidimensional filter is implemented as a sequence of one-dimensional convolution filters. The intermediate arrays are stored in the same data type as the output. Therefore, for output types with a limited precision, the results may be imprecise because intermediate results may be stored with insufficient precision. Examples -------- >> from scipy.ndimage import gaussian_filter >> a = np.arange(50, step=2).reshape((5,5)) >> a np.array([[ 0, 2, 4, 6, 8], [10, 12, 14, 16, 18], [20, 22, 24, 26, 28], [30, 32, 34, 36, 38], [40, 42, 44, 46, 48]]) >> gaussian_filter(a, sigma=1) np.array([[ 4, 6, 8, 9, 11], [10, 12, 14, 15, 17], [20, 22, 24, 25, 27], [29, 31, 33, 34, 36], [35, 37, 39, 40, 42]]) >> from scipy import misc >> import matplotlib.pyplot as plt >> fig = plt.figure() >> plt.gray() # show the filtered result in grayscale >> ax1 = fig.add_subplot(121) # left side >> ax2 = fig.add_subplot(122) # right side >> ascent = misc.ascent() >> result = gaussian_filter(ascent, sigma=5) >> ax1.imshow(ascent) >> ax2.imshow(result) >> plt.show() Here is a nice little demo of the added OMFIT causal feature, >> nn = 24 >> a = np.zeros((nn, nn)) >> a[int(nn//2),int(nn//2)] = 1 >> fig, axs = plt.subplots(2, 2) >> ax = axs[0, 0] >> ax.imshow(scipy.ndimage.gaussian_filter(a, 3, mode='nearest'), origin='lower') >> ax.set_title('scipy version') >> ax = axs[0, 1] >> ax.imshow(gaussian_filter(a, 3, mode='nearest'), origin='lower') >> ax.set_title('OMFIT version') >> ax = axs[1, 0] >> ax.imshow(gaussian_filter(a, 3, causal=True, mode='nearest'), origin='lower') >> ax.set_title('causal True') >> ax = axs[1, 1] >> ax.imshow(gaussian_filter(a,3, causal=(True, False), mode='nearest'), origin='lower') >> ax.set_title('causal True, False') """ input = np.asarray(input) # output = _ni_support._get_output(output, input) # this function's behavior is different for earlier scipys, # just hardcore v1.1.0 _get_output below shape = input.shape if output is None: output = np.zeros(shape, dtype=input.dtype.name) elif type(output) in [type(type), type(np.zeros((4,)).dtype)]: output = np.zeros(shape, dtype=output) elif isinstance(output, str): output = np.typeDict[output] output = np.zeros(shape, dtype=output) elif output.shape != shape: raise RuntimeError("output shape not correct") orders = _ni_support._normalize_sequence(order, input.ndim) sigmas = _ni_support._normalize_sequence(sigma, input.ndim) modes = _ni_support._normalize_sequence(mode, input.ndim) causals = _ni_support._normalize_sequence(causal, input.ndim) axes = list(range(input.ndim)) axes = [(axes[ii], sigmas[ii], orders[ii], modes[ii], causals[ii]) for ii in range(len(axes)) if sigmas[ii] > 1e-15] if len(axes) > 0: for axis, sigma, order, mode, causal in axes: gaussian_filter1d(input, sigma, axis, order, output, mode, cval, truncate, causal) input = output else: output[...] = input[...] return output