Source code for nutils.numeric

# Copyright (c) 2014 Evalf
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.

The numeric module provides methods that are lacking from the numpy module.

from . import warnings
import numpy, numbers, builtins,

_abc = 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ' # indices for einsum

def round(arr):
  return numpy.round(arr).astype(int)

def sign(arr):
  return numpy.sign(arr).astype(int)

def floor(arr):
  return numpy.floor(arr).astype(int)

def ceil(arr):
  return numpy.ceil(arr).astype(int)

[docs]def overlapping(arr, axis=-1, n=2): 'reinterpret data with overlaps' arr = numpy.asarray(arr) if axis < 0: axis += arr.ndim assert 0 <= axis < arr.ndim shape = arr.shape[:axis] + (arr.shape[axis]-n+1,n) + arr.shape[axis+1:] strides = arr.strides[:axis] + (arr.strides[axis],arr.strides[axis]) + arr.strides[axis+1:] overlapping = numpy.lib.stride_tricks.as_strided(arr, shape, strides) overlapping.flags.writeable = False return overlapping
[docs]def normdim(ndim, n): 'check bounds and make positive' assert isint(ndim) and ndim >= 0, 'ndim must be positive integer, got {}'.format(ndim) if n < 0: n += ndim assert 0 <= n < ndim, 'argument out of bounds: {} not in [0,{})'.format(n, ndim) return n
[docs]def get(arr, axis, item): 'take single item from array axis' arr = numpy.asarray(arr) axis = normdim(arr.ndim, axis) return arr[(slice(None),) * axis + (item,)]
[docs]def contract(A, B, axis=-1): 'contract' A = numpy.asarray(A) B = numpy.asarray(B) maxdim = max(A.ndim, B.ndim) m = _abc[maxdim-A.ndim:maxdim] n = _abc[maxdim-B.ndim:maxdim] axes = sorted([normdim(maxdim,axis)] if isinstance(axis,int) else [normdim(maxdim,ax) for ax in axis]) o = _abc[:maxdim-len(axes)] if axes == range(maxdim-len(axes), maxdim) \ else ''.join(_abc[a+1:b] for a, b in zip([-1]+axes, axes+[maxdim]) if a+1 != b) return numpy.einsum('{},{}->{}'.format(m,n,o), A, B, optimize=False)
[docs]def dot(A, B, axis=-1): '''Transform axis of A by contraction with first axis of B and inserting remaining axes. Note: with default axis=-1 this leads to multiplication of vectors and matrices following linear algebra conventions.''' A = numpy.asarray(A) B = numpy.asarray(B) m = _abc[:A.ndim] x = _abc[A.ndim:A.ndim+B.ndim-1] n = m[axis] + x o = m[:axis] + x if axis != -1: o += m[axis+1:] return numpy.einsum('{},{}->{}'.format(m,n,o), A, B, optimize=False)
[docs]def meshgrid(*args): 'multi-dimensional meshgrid generalisation' args = [numpy.asarray(arg) for arg in args] shape = [len(args)] + [arg.size for arg in args if arg.ndim] dtype = int if all(isintarray(a) for a in args) else float grid = numpy.empty(shape, dtype=dtype) n = len(shape)-1 for i, arg in enumerate(args): if arg.ndim: n -= 1 grid[i] = arg[(slice(None),)+(numpy.newaxis,)*n] else: grid[i] = arg assert n == 0 return grid
def takediag(A, axis=-2, rmaxis=-1): axis = normdim(A.ndim, axis) rmaxis = normdim(A.ndim, rmaxis) assert axis < rmaxis fmt = _abc[:rmaxis] + _abc[axis] + _abc[rmaxis:A.ndim-1] + '->' + _abc[:A.ndim-1] return numpy.einsum(fmt, A, optimize=False)
[docs]def normalize(A, axis=-1): 'devide by normal' s = [slice(None)] * A.ndim s[axis] = numpy.newaxis return A / numpy.linalg.norm(A, axis=axis)[tuple(s)]
[docs]def diagonalize(arg, axis=-1, newaxis=-1): 'insert newaxis, place axis on diagonal of axis and newaxis' axis = normdim(arg.ndim, axis) newaxis = normdim(arg.ndim+1, newaxis) assert 0 <= axis < newaxis <= arg.ndim diagonalized = numpy.zeros(arg.shape[:newaxis]+(arg.shape[axis],)+arg.shape[newaxis:], arg.dtype) diag = takediag(diagonalized, axis, newaxis) assert diag.base is diagonalized diag.flags.writeable = True diag[:] = arg return diagonalized
def eig(A): warnings.deprecation('numeric.eig is deprecated; use numpy.linalg.eig instead') return numpy.linalg.eig(A) isarray = lambda a: isinstance(a, (numpy.ndarray, const)) isboolarray = lambda a: isarray(a) and a.dtype == bool isbool = lambda a: isboolarray(a) and a.ndim == 0 or type(a) == bool isint = lambda a: isinstance(a, (numbers.Integral,numpy.integer)) isnumber = lambda a: isinstance(a, (numbers.Number,numpy.generic)) isintarray = lambda a: isarray(a) and numpy.issubdtype(a.dtype, numpy.integer) asobjvector = lambda v: numpy.array((None,)+tuple(v), dtype=object)[1:] # 'None' prevents interpretation of objects as axes def blockdiag(args): args = [numpy.asarray(arg) for arg in args] args = [arg[numpy.newaxis,numpy.newaxis] if arg.ndim == 0 else arg for arg in args] assert all(arg.ndim == 2 for arg in args) shapes = numpy.array([arg.shape for arg in args]) blockdiag = numpy.zeros(shapes.sum(0)) for arg, (i,j) in zip(args, shapes.cumsum(0)): blockdiag[i-arg.shape[0]:i, j-arg.shape[1]:j] = arg return blockdiag def nanjoin(args, axis=0): args = [numpy.asarray(arg) for arg in args] assert args assert axis >= 0 shape = list(args[0].shape) shape[axis] = sum(arg.shape[axis] for arg in args) + len(args) - 1 concat = numpy.empty(shape, dtype=float) concat[:] = numpy.nan i = 0 for arg in args: j = i + arg.shape[axis] concat[(slice(None),)*axis+(slice(i,j),)] = arg i = j + 1 return concat
[docs]def ix(args): 'version of :func:`numpy.ix_` that allows for scalars' args = tuple(numpy.asarray(arg) for arg in args) assert all(0 <= arg.ndim <= 1 for arg in args) idims = numpy.cumsum([0] + [arg.ndim for arg in args]) ndims = idims[-1] return [arg.reshape((1,)*idim+(arg.size,)+(1,)*(ndims-idim-1)) for idim, arg in zip(idims, args)]
def kronecker(arr, axis, length, pos): axis = normdim(arr.ndim+1, axis) kron = numpy.zeros(arr.shape[:axis]+(length,)+arr.shape[axis:], arr.dtype) kron[(slice(None),)*axis + (pos,)] = arr return kron class Broadcast1D: def __init__(self, arg): self.arg = numpy.asarray(arg) self.shape = self.arg.shape self.size = self.arg.size def __iter__(self): return ((item,) for item in self.arg.flat) broadcast = lambda *args: numpy.broadcast(*args) if len(args) > 1 else Broadcast1D(args[0]) def det_exact(A): # for some reason, numpy.linalg.det suffers from rounding errors A = numpy.asarray(A) assert A.ndim == 2 and A.shape[0] == A.shape[1] if len(A) == 0: det = 1. elif len(A) == 1: det = A[0,0] elif len(A) == 2: ((a,b),(c,d)) = A det = a*d - b*c elif len(A) == 3: ((a,b,c),(d,e,f),(g,h,i)) = A det = a*e*i + b*f*g + c*d*h - c*e*g - b*d*i - a*f*h else: raise NotImplementedError('shape=' + str(A.shape)) return det
[docs]def ext(A): """Exterior For array of shape (n,n-1) return n-vector ex such that ex.array = 0 and det(arr;ex) = ex.ex""" A = numpy.asarray(A) assert A.ndim == 2 and A.shape[0] == A.shape[1]+1 if len(A) == 1: ext = numpy.ones(1) elif len(A) == 2: ((a,),(b,)) = A ext = numpy.array((b,-a)) elif len(A) == 3: ((a,b),(c,d),(e,f)) = A ext = numpy.array((c*f-e*d,e*b-a*f,a*d-c*b)) else: raise NotImplementedError('shape={}'.format(A.shape)) return ext
def power(a, b): a = numpy.asarray(a) b = numpy.asarray(b) if a.dtype == int and b.dtype == int: b = b.astype(float) return numpy.power(a, b) def serialized(array, nsig, ndec): if array.ndim > 0: return '[{}]'.format(','.join(serialized(a, nsig, ndec) for a in array)) if not numpy.isfinite(array): # nan, inf return str(array) i = builtins.round(float(array) * 10**ndec) while abs(i) >= 10**nsig: ndec -= 1 i = builtins.round(float(array) * 10**ndec) return ('{}e{}' if i and ndec else '{}').format(i, -ndec) def encode64(array, nsig, ndec): import zlib, binascii assert isinstance(array, numpy.ndarray) and array.dtype == float binary = zlib.compress('{},{},{}'.format(nsig, ndec, serialized(array, nsig, ndec)).encode(), 9) data = binascii.b2a_base64(binary).decode().rstrip() assert_allclose64(array, data) return data def decode64(data): import zlib, binascii serialized = zlib.decompress(binascii.a2b_base64(data)) nsig, ndec, array = eval(serialized, numpy.__dict__) return nsig, ndec, numpy.array(array, dtype=float) def assert_allclose64(actual, data=None): try: nsig, ndec, desired = decode64(data) except Exception as e: status = str(e) nsig = 4 ndec = 15 else: try: numpy.testing.assert_allclose(actual, desired, atol=1.5*10**-ndec, rtol=10**(1-nsig)) except Exception as e: status = str(e) else: return status += '\n\nIf this is expected, use the following base64 string to test up to nsig={}, ndec={}:'.format(nsig, ndec) data = encode64(actual, nsig=nsig, ndec=ndec) while data: status += '\n{!r}'.format(data[:80]) data = data[80:] raise Exception(status)
[docs]class const( __slots__ = '__base', '__hash' @staticmethod def full(shape, fill_value): return const(numpy.lib.stride_tricks.as_strided(fill_value, shape, [0]*len(shape)), copy=False) def __new__(cls, base, copy=True, dtype=None): if isinstance(base, const): return base self = object.__new__(cls) self.__base = numpy.array(base, dtype=dtype) if copy or not isinstance(base, numpy.ndarray) or dtype and dtype != base.dtype else base self.__base.flags.writeable = False self.__hash = hash((self.__base.shape, self.__base.dtype, tuple(self.__base.flat[::self.__base.size//32+1]) if self.__base.size else ())) # NOTE special case self.__base.size == 0 necessary for numpy<1.12 return self @property def __array_struct__(self): return self.__base.__array_struct__ def __reduce__(self): return const, (self.__base, False) def __eq__(self, other): if self is other: return True if not isinstance(other, const): return False if self.__base is other.__base: return True if self.__hash != other.__hash or self.__base.dtype != other.__base.dtype or self.__base.shape != other.__base.shape or numpy.not_equal(self.__base, other.__base).any(): return False # deduplicate self.__base = other.__base return True def __lt__(self, other): if not isinstance(other, const): return NotImplemented return self != other and (self.dtype < other.dtype or self.dtype == other.dtype and (self.shape < other.shape or self.shape == other.shape and self.__base.tolist() < other.__base.tolist())) def __le__(self, other): if not isinstance(other, const): return NotImplemented return self == other or (self.dtype < other.dtype or self.dtype == other.dtype and (self.shape < other.shape or self.shape == other.shape and self.__base.tolist() < other.__base.tolist())) def __gt__(self, other): if not isinstance(other, const): return NotImplemented return self != other and (self.dtype > other.dtype or self.dtype == other.dtype and (self.shape > other.shape or self.shape == other.shape and self.__base.tolist() > other.__base.tolist())) def __ge__(self, other): if not isinstance(other, const): return NotImplemented return self == other or (self.dtype > other.dtype or self.dtype == other.dtype and (self.shape > other.shape or self.shape == other.shape and self.__base.tolist() > other.__base.tolist())) def __getitem__(self, item): retval = self.__base.__getitem__(item) return const(retval, copy=False) if isinstance(retval, numpy.ndarray) else retval dtype = property(lambda self: self.__base.dtype) shape = property(lambda self: self.__base.shape) size = property(lambda self: self.__base.size) ndim = property(lambda self: self.__base.ndim) flat = property(lambda self: self.__base.flat) T = property(lambda self: const(self.__base.T, copy=False)) __len__ = lambda self: self.__base.__len__() __repr__ = lambda self: 'const'+self.__base.__repr__()[5:] __str__ = lambda self: self.__base.__str__() __add__ = lambda self, other: self.__base.__add__(other) __radd__ = lambda self, other: self.__base.__radd__(other) __sub__ = lambda self, other: self.__base.__sub__(other) __rsub__ = lambda self, other: self.__base.__rsub__(other) __mul__ = lambda self, other: self.__base.__mul__(other) __rmul__ = lambda self, other: self.__base.__rmul__(other) __truediv__ = lambda self, other: self.__base.__truediv__(other) __rtruediv__ = lambda self, other: self.__base.__rtruediv__(other) __floordiv__ = lambda self, other: self.__base.__floordiv__(other) __rfloordiv__ = lambda self, other: self.__base.__rfloordiv__(other) __pow__ = lambda self, other: self.__base.__pow__(other) __hash__ = lambda self: self.__hash __int__ = lambda self: self.__base.__int__() __float__ = lambda self: self.__base.__float__() __abs__ = lambda self: self.__base.__abs__() __neg__ = lambda self: self.__base.__neg__() tolist = lambda self, *args, **kwargs: self.__base.tolist(*args, **kwargs) copy = lambda self, *args, **kwargs: self.__base.copy(*args, **kwargs) astype = lambda self, *args, **kwargs: self.__base.astype(*args, **kwargs) take = lambda self, *args, **kwargs: self.__base.take(*args, **kwargs) any = lambda self, *args, **kwargs: self.__base.any(*args, **kwargs) all = lambda self, *args, **kwargs: self.__base.all(*args, **kwargs) sum = lambda self, *args, **kwargs: self.__base.sum(*args, **kwargs) min = lambda self, *args, **kwargs: self.__base.min(*args, **kwargs) max = lambda self, *args, **kwargs: self.__base.max(*args, **kwargs) prod = lambda self, *args, **kwargs:*args, **kwargs) dot = lambda self, *args, **kwargs:*args, **kwargs) swapaxes = lambda self, *args, **kwargs: const(self.__base.swapaxes(*args, **kwargs), copy=False) ravel = lambda self, *args, **kwargs: const(self.__base.ravel(*args, **kwargs), copy=False) reshape = lambda self, *args, **kwargs: const(self.__base.reshape(*args, **kwargs), copy=False) transpose = lambda self, *args, **kwargs: const(self.__base.transpose(*args, **kwargs), copy=False) cumsum = lambda self, *args, **kwargs: const(self.__base.cumsum(*args, **kwargs), copy=False) nonzero = lambda self, *args, **kwargs: const(self.__base.nonzero(*args, **kwargs), copy=False) def insertaxis(self, axis, length): base = self.__base return const(numpy.lib.stride_tricks.as_strided(base, shape=base.shape[:axis]+(length,)+base.shape[axis:], strides=base.strides[:axis]+(0,)+base.strides[axis:]))
def binom(n, k): a = b = 1 for i in range(1, k+1): a *= n+1-i b *= i return a // b def poly_outer_product(left, right): left, right = numpy.asarray(left), numpy.asarray(right) nleft, nright = left.ndim-1, right.ndim-1 pshape = left.shape[1:] if not nright else right.shape[1:] if not nleft else (max(left.shape[1:])+max(right.shape[1:])-1,) * (nleft + nright) outer = numpy.zeros((left.shape[0], right.shape[0], *pshape), dtype=numpy.common_type(left, right)) a = slice(None) outer[(a,a,*(map(slice, left.shape[1:]+right.shape[1:])))] = left[(a,None)+(a,)*nleft+(None,)*nright]*right[(None,a)+(None,)*nleft+(a,)*nright] return const(outer.reshape(left.shape[0] * right.shape[0], *pshape), copy=False) def poly_stack(coeffs): coeffs = tuple(coeffs) n = max(icoeffs.shape[0] for icoeffs in coeffs) ndim = coeffs[0].ndim dest = numpy.zeros((len(coeffs),)+(n,)*ndim, dtype=float) for i, j in enumerate(coeffs): dest[(i,*map(slice, j.shape))] = j return const(dest, copy=False) def poly_grad(coeffs, ndim): I = range(ndim) dcoeffs = [coeffs[(...,*(slice(1,None) if i==j else slice(0,-1) for j in I))] for i in I] if coeffs.shape[-1] > 2: a = numpy.arange(1, coeffs.shape[-1]) dcoeffs = [a[tuple(slice(None) if i==j else numpy.newaxis for j in I)] * c for i, c in enumerate(dcoeffs)] dcoeffs = numpy.stack(dcoeffs, axis=coeffs.ndim-ndim) return const(dcoeffs, copy=False) def poly_eval(coeffs, points): assert points.ndim == 2 if coeffs.shape[-1] == 0: return const.full((points.shape[0],)+coeffs.shape[1:coeffs.ndim-points.shape[-1]], 0.) for dim in reversed(range(points.shape[-1])): result = numpy.empty((points.shape[0], *coeffs.shape[1:-1]), dtype=float) result[:] = coeffs[...,-1] points_dim = points[(slice(None),dim,*(numpy.newaxis,)*(result.ndim-1))] for j in reversed(range(coeffs.shape[-1]-1)): result *= points_dim result += coeffs[...,j] coeffs = result return const(coeffs, copy=False) def poly_mul(p, q): assert p.ndim == q.ndim pq = numpy.zeros([n+m-1 for n, m in zip(p.shape, q.shape)]) if q.size < p.size: p, q = q, p # loop over the smallest of the two arrays for i, pi in numpy.ndenumerate(p): if pi: pq[tuple(slice(o, o+m) for o, m in zip(i, q.shape))] += pi * q return pq def poly_pow(p, n): assert isint(n) and n >= 0 if n == 0: return numpy.ones((1,)*p.ndim) if n == 1: return p q = poly_pow(poly_mul(p, p), n//2) if n%2: return poly_mul(q, p) return q # vim:shiftwidth=2:softtabstop=2:expandtab:foldmethod=indent:foldnestmax=2