Source code for nutils.function

# -*- coding: utf8 -*-
#
# Module FUNCTION
#
# Part of Nutils: open source numerical utilities for Python. Jointly developed
# by HvZ Computational Engineering, TU/e Multiscale Engineering Fluid Dynamics,
# and others. More info at http://nutils.org <info@nutils.org>. (c) 2014

"""
The function module defines the :class:`Evaluable` class and derived objects,
commonly referred to as nutils functions. They represent mappings from a
:mod:`nutils.topology` onto Python space. The notabe class of :class:`ArrayFunc`
objects map onto the space of Numpy arrays of predefined dimension and shape.
Most functions used in nutils applicatons are of this latter type, including the
geometry and function bases for analysis.

Nutils functions are essentially postponed python functions, stored in a tree
structure of input/output dependencies. Many :class:`ArrayFunc` objects have
directly recognizable numpy equivalents, such as :class:`Sin` or
:class:`Inverse`. By not evaluating directly but merely stacking operations,
coplex operations can be defined prior to entering a quadrature loop, allowing
for a higher lever style programming. It also allows for automatic
differentiation and code optimization.

It is important to realize that nutils functions do not map for a physical
xy-domain but from a topology, where a point is characterized by the combination
of an element and its local coordinate. This is a natural fit for typical finite
element operations such as quadrature. Evaluation from physical coordinates is
possible only via inverting of the geometry function, which is a fundamentally
expensive and currently unsupported operation.
"""

from . import util, numpy, numeric, log, prop, core, _
import sys, warnings

ELEM    = object()
POINTS  = object()
WEIGHTS = object()

[docs]class Evaluable( object ): 'Base class' __slots__ = 'operations', 'data', '__args', '__evalf' def __init__( self, args, evalf ): 'constructor' self.__args = tuple(args) self.__evalf = evalf self.operations = None
[docs] def recurse_index( self, data, operations, cbuild ): 'compile' indices = numpy.empty( len(self.__args), dtype=int ) for iarg, arg in enumerate( self.__args ): if isinstance(arg,Evaluable): for idx, (op,evalf,idcs) in enumerate( operations ): if op == arg: break else: idx = arg.recurse_index(data,operations,cbuild) elif arg is ELEM: idx = -3 elif arg is POINTS: idx = -2 elif arg is WEIGHTS: idx = -1 else: data.insert( 0, arg ) idx = -len(data)-3 indices[iarg] = idx if cbuild and self.__cdef: evalf = cbuild[ self.__cdef() ] if prop.use_c_funcs == 'debug': evalf = numeric.check_equal_wrapper( evalf, self.__evalf ) else: evalf = self.__evalf operations.append( (self,evalf,indices) ) return len(operations)-1
[docs] def compile( self ): 'compile' log.context( 'compiling' ) if self.operations is None: cbuild = getattr( prop, 'use_c_funcs', False ) and CBuilder() self.data = [] operations = [] self.recurse_index( self.data, operations, cbuild ) # compile expressions self.operations = [ (evalf,idcs) for (op,evalf,idcs) in operations ] # break cycles! if getattr( prop, 'dot', False ): self.graphviz() if cbuild: cbuild.compile()
def __call__( self, elem, ischeme ): 'evaluate' if isinstance( ischeme, dict ): ischeme = ischeme[elem] if isinstance( ischeme, str ): points, weights = elem.eval( ischeme ) elif isinstance( ischeme, tuple ): points, weights = ischeme assert points.shape[-1] == elem.ndims assert points.shape[:-1] == weights.shape, 'non matching shapes: points.shape=%s, weights.shape=%s' % ( points.shape, weights.shape ) elif isinstance( ischeme, numpy.ndarray ): points = ischeme weights = None assert points.shape[-1] == elem.ndims elif ischeme is None: points = weights = None else: raise Exception, 'invalid integration scheme of type %r' % type(ischeme) self.compile() N = len(self.data) + 3 values = self.data + [ elem, points, weights ] for evalf, indices in self.operations: args = [ values[N+i] for i in indices ] try: retval = evalf( *args ) except KeyboardInterrupt: raise except: etype, evalue, traceback = sys.exc_info() raise EvaluationError, ( etype, evalue, self, values ), traceback values.append( retval ) return values[-1]
[docs] def graphviz( self, title='graphviz' ): 'create function graph' log.context( title ) self.compile() import os, subprocess dotpath = getattr( prop, 'dot', True ) if not isinstance( dotpath, str ): dotpath = 'dot' imgtype = getattr( prop, 'imagetype', 'png' ) imgpath = util.getpath( 'dot{0:03x}.' + imgtype ) try: dot = subprocess.Popen( [dotpath,'-Tjpg'], stdin=subprocess.PIPE, stdout=open(imgpath,'w') ) except OSError: log.error( 'error: failed to execute', dotpath ) return False print >> dot.stdin, 'digraph {' print >> dot.stdin, 'graph [ dpi=72 ];' data = self.data + [ '<elem>', '<points>', '<weights>' ] for i, (evalf,indices) in enumerate( self.operations ): args = [ '%%%d=%s' % ( iarg, _obj2str( data[idx] ) ) for iarg, idx in enumerate( indices ) if idx < 0 ] vertex = 'label="%s"' % r'\n'.join( [ '%d. %s' % ( i, evalf.__name__ ) ] + args ) #if evalf == Inflate.inflate: # vertex += ', style=filled, fillcolor=black, fontcolor=white' print >> dot.stdin, '%d [%s]' % ( i, vertex ) for iarg, idx in enumerate( indices ): if idx >= 0: print >> dot.stdin, '%d -> %d [label="%%%d"];' % ( idx, i, iarg ); print >> dot.stdin, '}' dot.stdin.close() log.path( os.path.basename(imgpath) )
[docs] def stackstr( self, values=None ): 'print stack' self.compile() if values is None: values = self.data + [ '<elem>', '<points>', '<weights>' ] N = len(self.data) + 3 lines = [] for i, (evalf,indices) in enumerate( self.operations ): line = ' %%%d =' % i args = [ '%%%d' % idx if idx >= 0 else _obj2str(values[N+idx]) for idx in indices ] try: code = evalf.func_code names = code.co_varnames[ :code.co_argcount ] names += tuple( '%s[%d]' % ( code.co_varnames[ code.co_argcount ], n ) for n in range( len(indices) - len(names) ) ) args = [ '%s=%s' % item for item in zip( names, args ) ] except: pass line += ' %s( %s )' % ( evalf.__name__, ', '.join( args ) ) lines.append( line ) if N+i == len(values): break return '\n'.join( lines )
def __eq__( self, other ): 'compare' return self is other or ( self.__class__ == other.__class__ and self.__evalf == other.__evalf and len( self.__args ) == len( other.__args ) and all( _equal(arg1,arg2) for arg1, arg2 in zip( self.__args, other.__args ) ) ) def __ne__( self, other ): 'not equal' return not self == other
[docs] def asciitree( self ): 'string representation' key = self.__evalf.__name__ lines = [] indent = '\n' + ' ' + ' ' * len(key) for it in reversed( self.__args ): s = it.asciitree() if isinstance(it,Evaluable) else _obj2str(it) lines.append( indent.join( s.splitlines() ) ) indent = '\n' + '|' + ' ' * len(key) indent = '\n' + '+' + '-' * (len(key)-1) + ' ' return key + ' ' + indent.join( reversed( lines ) )
[docs]class EvaluationError( Exception ): 'evaluation error' def __init__( self, etype, evalue, evaluable, values ): 'constructor' self.etype = etype self.evalue = evalue self.evaluable = evaluable self.values = values def __repr__( self ): return 'EvaluationError%s' % self def __str__( self ): 'string representation' return '\n%s --> %s: %s' % ( self.evaluable.stackstr( self.values ), self.etype.__name__, self.evalue )
[docs]class Tuple( Evaluable ): 'combine' __slots__ = 'items', def __init__( self, items ): 'constructor' self.items = tuple( items ) Evaluable.__init__( self, args=self.items, evalf=self.vartuple ) def __iter__( self ): 'iterate' return iter(self.items) def __len__( self ): 'length' return len(self.items) def __getitem__( self, item ): 'get item' return self.items[item] def __add__( self, other ): 'add' return Tuple( self.items + tuple(other) ) def __radd__( self, other ): 'add' return Tuple( tuple(other) + self.items ) @staticmethod
[docs] def vartuple( *f ): 'evaluate' return f
[docs]class PointShape( Evaluable ): 'shape of integration points' __slots__ = () def __init__( self ): 'constructor' return Evaluable.__init__( self, args=[POINTS], evalf=self.pointshape ) @staticmethod
[docs] def pointshape( points ): 'evaluate' return points.shape[:-1]
[docs]class Cascade( Evaluable ): 'point cascade: list of (elem,points) tuples' __slots__ = 'ndims', 'side' def __init__( self, ndims, side=0 ): 'constructor' self.ndims = ndims self.side = side Evaluable.__init__( self, args=[ELEM,POINTS,ndims,side], evalf=self.cascade ) def transform( self, ndims ): if ndims == self.ndims: return eye( ndims ) assert self.ndims > ndims return Transform( self, Cascade(ndims,self.side) ) @staticmethod
[docs] def cascade( elem, points, ndims, side ): 'evaluate' while elem.ndims != ndims \ or elem.interface and elem.interface[side][0].ndims < elem.ndims: # TODO make less dirty elem, transform = elem.interface[side] if elem.interface \ else elem.context or elem.parent points = transform.eval( points ) cascade = [ (elem,points) ] while elem.parent: elem, transform = elem.parent points = transform.eval( points ) cascade.append( (elem,points) ) return cascade
@property def inv( self ): return Cascade( self.ndims, 1-self.side ) # ARRAYFUNC # # The main evaluable. Closely mimics a numpy array.
[docs]class ArrayFunc( Evaluable ): 'array function' __array_priority__ = 1. # http://stackoverflow.com/questions/7042496/numpy-coercion-problem-for-left-sided-binary-operator/7057530#7057530 __slots__ = 'shape', 'ndim', 'dtype' def __init__( self, evalf, args, shape, dtype=float ): 'constructor' self.shape = tuple(shape) self.ndim = len(self.shape) assert dtype is int or dtype is float self.dtype = dtype Evaluable.__init__( self, evalf=evalf, args=args ) # mathematical operators def __mul__( self, other ): return multiply( self, other ) def __rmul__( self, other ): return multiply( other, self ) def __div__( self, other ): return divide( self, other ) def __rdiv__( self, other ): return divide( other, self ) def __add__( self, other ): return add( self, other ) def __radd__( self, other ): return add( other, self ) def __sub__( self, other ): return subtract( self, other ) def __rsub__( self, other ): return subtract( other, self ) def __neg__( self ): return negative( self ) def __pow__( self, n ): return power( self, n ) def sum( self, axes=-1 ): return sum( self, axes ) # standalone methods
[docs] def vector( self, ndims ): 'vectorize' return vectorize( [self] * ndims )
[docs] def dot( self, weights, axis=0 ): 'array contraction' weights = numpy.asarray( weights, dtype=float ) assert weights.ndim == 1 s = [ numpy.newaxis ] * self.ndim s[axis] = slice(None) return dot( self, weights[tuple(s)], axes=axis )
def __getitem__( self, item ): 'get item, general function which can eliminate, add or modify axes.' myitem = list( item if isinstance( item, tuple ) else [item] ) n = 0 arr = self while myitem: it = myitem.pop(0) if isinstance(it,int): # retrieve one item from axis arr = get( arr, n, it ) elif it == _: # insert a singleton axis arr = insert( arr, n ) n += 1 elif it == slice(None): # select entire axis n += 1 elif it == Ellipsis: # skip to end remaining_items = len(myitem) - myitem.count(_) skip = arr.ndim - n - remaining_items assert skip >= 0, 'shape=%s, item=%s' % ( self.shape, _obj2str(item) ) n += skip elif isinstance(it,slice) and it.step in (1,None) and it.stop == ( it.start or 0 ) + 1: # special case: unit length slice arr = insert( get( arr, n, it.start or 0 ), n ) n += 1 elif isinstance(it,(slice,list,tuple,numpy.ndarray)): # modify axis (shorten, extend or renumber one axis) arr = take( arr, it, n ) n += 1 else: raise NotImplementedError assert n <= arr.ndim return arr def __iter__( self ): 'split first axis' if not self.shape: raise TypeError, 'scalar function is not iterable' return ( self[i,...] for i in range(self.shape[0]) )
[docs] def find( self, elem, C ):#target, start, tol=1e-10, maxiter=999 ): 'iteratively find x for f(x) = target, starting at x=start' raise NotImplementedError assert self.ndim == 1 points = start Jinv = inverse( localgradient( self, elem.ndims ) ) r = target - self( elem, points ) niter = 0 while numpy.any( numeric.contract( r, r, axis=-1 ) > tol ): niter += 1 if niter >= maxiter: raise Exception, 'failed to converge in %d iterations' % maxiter points = points.offset( numeric.contract( Jinv( elem, points ), r[:,_,:], axis=-1 ) ) r = target - self( elem, points ) return points
[docs] def normalized( self ): 'normalize last axis' return self / norm2( self, axis=-1 )
[docs] def normal( self, ndims=-1 ): 'normal' assert len(self.shape) == 1 if ndims <= 0: ndims += self.shape[0] grad = localgradient( self, ndims ) if grad.shape == (2,1): normal = concatenate([ grad[1,:], -grad[0,:] ]).normalized() elif grad.shape == (3,2): normal = cross( grad[:,0], grad[:,1], axis=0 ).normalized() elif grad.shape == (3,1): normal = cross( grad[:,0], self.normal(), axis=0 ).normalized() elif grad.shape == (1,0): normal = OrientationHack()[_] else: raise NotImplementedError, 'cannot compute normal for %dx%d jacobian' % ( self.shape[0], ndims ) return normal
[docs] def curvature( self, ndims=-1 ): 'curvature' return self.normal().div( self, ndims=ndims ) #if ndims <= 0: # ndims += self.shape[0] #assert ndims == 1 and self.shape == (2,) #J = localgradient( self, ndims ) #H = localgradient( J, ndims ) #dx, dy = J[:,0] #ddx, ddy = H[:,0,0] #return ( dx * ddy - dy * ddx ) / norm2( J[:,0], axis=0 )**3
[docs] def swapaxes( self, n1, n2 ): 'swap axes' return swapaxes( self, (n1,n2) )
[docs] def transpose( self, trans=None ): 'transpose' return transpose( self, trans )
[docs] def grad( self, coords, ndims=0 ): 'gradient' assert coords.ndim == 1 if ndims <= 0: ndims += coords.shape[0] J = localgradient( coords, ndims ) if J.shape[0] == J.shape[1]: Jinv = inverse( J ) elif J.shape[0] == J.shape[1] + 1: # gamma gradient G = ( J[:,:,_] * J[:,_,:] ).sum( 0 ) Ginv = inverse( G ) Jinv = ( J[_,:,:] * Ginv[:,_,:] ).sum() else: raise Exception, 'cannot invert %sx%s jacobian' % J.shape return sum( localgradient( self, ndims )[...,_] * Jinv, axes=-2 )
[docs] def laplace( self, coords, ndims=0 ): 'laplacian' return self.grad(coords,ndims).div(coords,ndims)
[docs] def add_T( self, axes=(-2,-1) ): 'add transposed' return add_T( self, axes )
[docs] def symgrad( self, coords, ndims=0 ): 'gradient' return .5 * add_T( self.grad( coords, ndims ) )
[docs] def div( self, coords, ndims=0 ): 'gradient' return trace( self.grad( coords, ndims ), -1, -2 )
[docs] def dotnorm( self, coords, ndims=0 ): 'normal component' return sum( self * coords.normal( ndims-1 ) )
[docs] def ngrad( self, coords, ndims=0 ): 'normal gradient' return dotnorm( self.grad( coords, ndims ), coords, ndims )
[docs] def nsymgrad( self, coords, ndims=0 ): 'normal gradient' return dotnorm( self.symgrad( coords, ndims ), coords, ndims )
@property
[docs] def T( self ): 'transpose' return transpose( self )
def __str__( self ): 'string representation' return '%s<%s>' % ( self.__class__.__name__, ','.join(map(str,self.shape)) ) __repr__ = __str__
[docs]class ElemArea( ArrayFunc ): 'element area' __slots__ = () def __init__( self, weights ): 'constructor' assert weights.ndim == 0 ArrayFunc.__init__( self, args=[weights], evalf=self.elemarea, shape=weights.shape ) @staticmethod
[docs] def elemarea( weights ): 'evaluate' return numpy.sum( weights )
[docs]class ElemInt( ArrayFunc ): 'elementwise integration' __slots__ = () def __init__( self, func, weights ): 'constructor' assert _isfunc( func ) and _isfunc( weights ) assert weights.ndim == 0 ArrayFunc.__init__( self, args=[weights,func,func.ndim], evalf=self.elemint, shape=func.shape ) @staticmethod
[docs] def elemint( w, f, ndim ): 'evaluate' if f.ndim == ndim: # the missing point axis problem return f * w.sum() return numeric.dot( w, f ) if w.size else numpy.zeros( f.shape[1:] )
[docs]class Align( ArrayFunc ): 'align axes' __slots__ = 'func', 'axes' def __init__( self, func, axes, ndim ): 'constructor' assert func.ndim == len(axes) self.func = func assert all( 0 <= ax < ndim for ax in axes ) self.axes = tuple(axes) shape = [ 1 ] * ndim for ax, sh in zip( self.axes, func.shape ): shape[ax] = sh negaxes = [ ax-ndim for ax in self.axes ] ArrayFunc.__init__( self, args=[func,negaxes,ndim], evalf=self.align, shape=shape ) @staticmethod
[docs] def align( arr, trans, ndim ): 'align' extra = arr.ndim - len(trans) return numeric.align( arr, range(extra)+trans, ndim+extra )
def _elemint( self, weights ): return align( elemint( self.func, weights ), self.axes, self.ndim ) def _align( self, axes, ndim ): newaxes = [ axes[i] for i in self.axes ] return align( self.func, newaxes, ndim ) def _takediag( self ): if self.ndim-1 not in self.axes: return align( self.func, self.axes, self.ndim-1 ) if self.ndim-2 not in self.axes: axes = [ ax if ax != self.ndim-1 else self.ndim-2 for ax in self.axes ] return align( self.func, axes, self.ndim-1 ) if self.axes[-2:] in [ (self.ndim-2,self.ndim-1), (self.ndim-1,self.ndim-2) ]: axes = self.axes[:-2] + (self.ndim-2,) return align( takediag( self.func ), axes, self.ndim-1 ) def _get( self, i, item ): axes = [ ax - (ax>i) for ax in self.axes if ax != i ] if len(axes) == len(self.axes): return align( self.func, axes, self.ndim-1 ) n = self.axes.index( i ) return align( get( self.func, n, item ), axes, self.ndim-1 ) def _sum( self, axis ): if axis in self.axes: idx = self.axes.index( axis ) func = sum( self.func, idx ) else: func = self.func trans = [ ax - (ax>axis) for ax in self.axes if ax != axis ] return align( func, trans, self.ndim-1 ) def _localgradient( self, ndims ): return align( localgradient( self.func, ndims ), self.axes+(self.ndim,), self.ndim+1 ) def _multiply( self, other ): if not _isfunc(other) and len(self.axes) == other.ndim: return align( self.func * transpose( other, self.axes ), self.axes, self.ndim ) if isinstance( other, Align ) and self.axes == other.axes: return align( self.func * other.func, self.axes, self.ndim ) def _add( self, other ): if not _isfunc(other) and len(self.axes) == self.ndim: return align( self.func + transpose( other, self.axes ), self.axes, self.ndim ) if isinstance( other, Align ) and self.axes == other.axes: return align( self.func + other.func, self.axes, self.ndim ) def _take( self, indices, axis ): n = self.axes.index( axis ) return align( take( self.func, indices, n ), self.axes, self.ndim ) def _opposite( self ): return align( opposite(self.func), self.axes, self.ndim )
[docs]class Get( ArrayFunc ): 'get' __slots__ = 'func', 'axis', 'item' def __init__( self, func, axis, item ): 'constructor' self.func = func self.axis = axis self.item = item assert 0 <= axis < func.ndim, 'axis is out of bounds' assert 0 <= item < func.shape[axis], 'item is out of bounds' s = (Ellipsis,item) + (slice(None),)*(func.ndim-axis-1) shape = func.shape[:axis] + func.shape[axis+1:] ArrayFunc.__init__( self, args=(func,s), evalf=numpy.ndarray.__getitem__, shape=shape ) def _localgradient( self, ndims ): f = localgradient( self.func, ndims ) return get( f, self.axis, self.item ) def _get( self, i, item ): tryget = get( self.func, i+(i>=self.axis), item ) if not isinstance( tryget, Get ): # avoid inf recursion return get( tryget, self.axis, self.item ) def _take( self, indices, axis ): return get( take( self.func, indices, axis+(axis>=self.axis) ), self.axis, self.item ) def _opposite( self ): return get( opposite(self.func), self.axis, self.item )
[docs]class Product( ArrayFunc ): 'product' __slots__ = 'func', 'axis' def __init__( self, func, axis ): 'constructor' self.func = func self.axis = axis assert 0 <= axis < func.ndim shape = func.shape[:axis] + func.shape[axis+1:] ArrayFunc.__init__( self, args=[func,axis-func.ndim], evalf=numpy.product, shape=shape ) def _localgradient( self, ndims ): return self[...,_] * ( localgradient(self.func,ndims) / self.func[...,_] ).sum( self.axis ) def _get( self, i, item ): func = get( self.func, i+(i>=self.axis), item ) return product( func, self.axis-(i<self.axis) ) def _opposite( self ): return product( opposite(self.func), self.axis )
[docs]class IWeights( ArrayFunc ): 'integration weights' __slots__ = () def __init__( self ): 'constructor' ArrayFunc.__init__( self, args=[ELEM,WEIGHTS], evalf=self.iweights, shape=() ) @staticmethod
[docs] def iweights( elem, weights ): 'evaluate' return elem.root_det * weights
[docs]class OrientationHack( ArrayFunc ): 'orientation hack for 1d elements; VERY dirty' __slots__ = 'side', def __init__( self, side=0 ): 'constructor' self.side = side ArrayFunc.__init__( self, args=[ELEM,side], evalf=self.orientation, shape=[] ) @staticmethod
[docs] def orientation( elem, side ): 'evaluate' pelem, trans = elem.interface[side] if elem.interface else elem.context offset, = trans.offset return numpy.sign( offset - .5 )
def _opposite( self ): return OrientationHack( 1-self.side )
[docs]class Transform( ArrayFunc ): 'transform' __slots__ = 'fromcascade', 'tocascade', 'side' def __init__( self, fromcascade, tocascade, side=0 ): 'constructor' assert fromcascade.ndims > tocascade.ndims self.fromcascade = fromcascade self.tocascade = tocascade self.side = side ArrayFunc.__init__( self, args=[fromcascade,tocascade,side], evalf=self.transform, shape=(fromcascade.ndims,tocascade.ndims) ) @staticmethod
[docs] def transform( fromcascade, tocascade, side ): 'transform' fromelem = fromcascade[0][0] toelem = tocascade[0][0] elem = toelem T = elem.inv_root_transform while elem is not fromelem: elem, transform = elem.interface[side] if elem.interface \ else elem.context or elem.parent T = numpy.dot( transform.transform, T ) T = numpy.dot( elem.root_transform, T ) return T
def _localgradient( self, ndims ): return _zeros( self.shape + (ndims,) ) def _opposite( self ): return Transform( self.fromcascade, self.tocascade, 1-self.side )
[docs]class Function( ArrayFunc ): 'function' __slots__ = 'cascade', 'stdmap', 'igrad' def __init__( self, cascade, stdmap, igrad, axis ): 'constructor' self.cascade = cascade self.stdmap = stdmap self.igrad = igrad ArrayFunc.__init__( self, args=(cascade,stdmap,igrad), evalf=self.function, shape=(axis,)+(cascade.ndims,)*igrad ) @staticmethod
[docs] def function( cascade, stdmap, igrad ): 'evaluate' fvals = [] for elem, points in cascade: std = stdmap.get(elem) if not std: continue if isinstance( std, tuple ): std, keep = std F = std.eval(points,grad=igrad)[(Ellipsis,keep)+(slice(None),)*igrad] else: F = std.eval(points,grad=igrad) for axis in range(-igrad,0): F = numeric.dot( F, elem.inv_root_transform, axis ) fvals.append( F ) assert fvals, 'no function values encountered' return fvals[0] if len(fvals) == 1 else numpy.concatenate( fvals, axis=-1-igrad )
def _opposite( self ): return Function( self.cascade.inv, self.stdmap, self.igrad, self.shape[0] ) def _localgradient( self, ndims ): assert ndims <= self.cascade.ndims grad = Function( self.cascade, self.stdmap, self.igrad+1, self.shape[0] ) return grad if ndims == self.cascade.ndims \ else dot( grad[...,_], self.cascade.transform( ndims ), axes=-2 )
[docs]class Choose( ArrayFunc ): 'piecewise function' __slots__ = 'level', 'choices' def __init__( self, level, choices ): 'constructor' self.level = level self.choices = tuple(choices) shape = _jointshape( *[ choice.shape for choice in choices ] ) level = level[ (_,)*(len(shape)-level.ndim) ] assert level.ndim == len( shape ) ArrayFunc.__init__( self, args=(level,)+self.choices, evalf=self.choose, shape=shape ) @staticmethod
[docs] def choose( level, *choices ): 'choose' return numpy.choose( level, choices )
def _localgradient( self, ndims ): grads = [ localgradient( choice, ndims ) for choice in self.choices ] if not any( grads ): # all-zero special case; better would be allow merging of intervals return _zeros( self.shape + (ndims,) ) return Choose( self.level[...,_], grads ) def _opposite( self ): return choose( opposite(self.level), tuple(opposite(c) for c in self.choices) )
[docs]class Choose2D( ArrayFunc ): 'piecewise function' __slots__ = () def __init__( self, coords, contour, fin, fout ): 'constructor' shape = _jointshape( fin.shape, fout.shape ) ArrayFunc.__init__( self, args=(coords,contour,fin,fout), evalf=self.choose2d, shape=shape ) @staticmethod
[docs] def choose2d( xy, contour, fin, fout ): 'evaluate' from matplotlib import nxutils mask = nxutils.points_inside_poly( xy.T, contour ) out = numpy.empty( fin.shape or fout.shape ) out[...,mask] = fin[...,mask] if fin.shape else fin out[...,~mask] = fout[...,~mask] if fout.shape else fout return out
[docs]class Inverse( ArrayFunc ): 'inverse' __slots__ = 'func', def __init__( self, func ): 'constructor' assert func.shape[-1] == func.shape[-2] self.func = func ArrayFunc.__init__( self, args=[func], evalf=numeric.inverse, shape=func.shape ) def _localgradient( self, ndims ): G = localgradient( self.func, ndims ) H = sum( self[...,_,:,:,_] * G[...,:,:,_,:], -3 ) I = sum( self[...,:,:,_,_] * H[...,_,:,:,:], -3 ) return -I def _opposite( self ): return Inverse( opposite(self.func) )
[docs]class DofMap( ArrayFunc ): 'dof axis' __slots__ = 'cascade', 'dofmap' def __init__( self, cascade, dofmap, axis ): 'new' self.cascade = cascade self.dofmap = dofmap ArrayFunc.__init__( self, args=(cascade,dofmap), evalf=self.evalmap, shape=[axis], dtype=int ) @staticmethod
[docs] def evalmap( cascade, dofmap ): 'evaluate' alldofs = [] for elem, points in cascade: dofs = dofmap.get( elem ) if dofs is not None: alldofs.append( dofs ) assert alldofs, 'no dofs encountered' return alldofs[0] if len(alldofs) == 1 else numpy.concatenate( alldofs )
def _opposite( self ): return DofMap( self.cascade.inv, self.dofmap, self.shape[0] )
[docs]class Concatenate( ArrayFunc ): 'concatenate' __slots__ = 'funcs', 'axis' def __init__( self, funcs, axis=0 ): 'constructor' funcs = [ asarray(func) for func in funcs ] ndim = funcs[0].ndim assert all( func.ndim == ndim for func in funcs ) axis = numeric.normdim( ndim, axis ) lengths = [ func.shape[axis] for func in funcs ] if any( n == None for n in lengths ): assert all( n == None for n in lengths ) sh = None else: sh = sum( lengths ) shape = _jointshape( *[ func.shape[:axis] + (sh,) + func.shape[axis+1:] for func in funcs ] ) self.funcs = tuple(funcs) self.axis = axis ArrayFunc.__init__( self, args=(axis-ndim,)+self.funcs, evalf=self.concatenate, shape=shape ) @staticmethod
[docs] def concatenate( iax, *arrays ): 'evaluate' ndim = numpy.max([ array.ndim for array in arrays ]) axlen = util.sum( array.shape[iax] for array in arrays ) shape = _jointshape( *[ (1,)*(ndim-array.ndim) + array.shape[:iax] + (axlen,) + ( array.shape[iax+1:] if iax != -1 else () ) for array in arrays ] ) dtype = float if any( array.dtype == float for array in arrays ) else int retval = numpy.empty( shape, dtype=dtype ) n0 = 0 for array in arrays: n1 = n0 + array.shape[iax] retval[(slice(None),)*( iax if iax >= 0 else iax + ndim )+(slice(n0,n1),)] = array n0 = n1 assert n0 == axlen return retval
@property def blocks( self ): n = 0 for func in self.funcs: for f, ind in blocks( func ): yield f, ind[:self.axis] + (ind[self.axis]+n,) + ind[self.axis+1:] n += func.shape[self.axis] def _get( self, i, item ): if i == self.axis: for f in self.funcs: if item < f.shape[i]: fexp = expand( f, self.shape[:self.axis] + (f.shape[self.axis],) + self.shape[self.axis+1:] ) return get( fexp, i, item ) item -= f.shape[i] axis = self.axis - (self.axis > i) return concatenate( [ get( f, i, item ) for f in self.funcs ], axis=axis ) def _localgradient( self, ndims ): funcs = [ localgradient( func, ndims ) for func in self.funcs ] return concatenate( funcs, axis=self.axis ) def _multiply( self, other ): funcs = [] n0 = 0 for func in self.funcs: n1 = n0 + func.shape[ self.axis ] funcs.append( func * take( other, slice(n0,n1), self.axis ) ) n0 = n1 assert n0 == self.shape[ self.axis ] return concatenate( funcs, self.axis ) def _cross( self, other, axis ): if axis == self.axis: n = 1, 2, 0 m = 2, 0, 1 return take(self,n,axis) * take(other,m,axis) - take(self,m,axis) * take(other,n,axis) funcs = [] n0 = 0 for func in self.funcs: n1 = n0 + func.shape[ self.axis ] funcs.append( cross( func, take( other, slice(n0,n1), self.axis ), axis ) ) n0 = n1 assert n0 == self.shape[ self.axis ] return concatenate( funcs, self.axis ) def _add( self, other ): if isinstance( other, Concatenate ) and self.axis == other.axis: i = 0 N1 = numpy.cumsum( [0] + [f1.shape[self.axis] for f1 in self.funcs] ) N2 = numpy.cumsum( [0] + [f2.shape[self.axis] for f2 in other.funcs] ) ifun1 = ifun2 = 0 funcs = [] while i < self.shape[self.axis]: j = min( N1[ifun1+1], N2[ifun2+1] ) funcs.append( take( self.funcs[ifun1], slice(i-N1[ifun1],j-N1[ifun1]), self.axis ) + take( other.funcs[ifun2], slice(i-N2[ifun2],j-N2[ifun2]), self.axis )) i = j ifun1 += i >= N1[ifun1+1] ifun2 += i >= N2[ifun2+1] assert ifun1 == len(self.funcs) assert ifun2 == len(other.funcs) return concatenate( funcs, axis=self.axis ) funcs = [] n0 = 0 for func in self.funcs: n1 = n0 + func.shape[ self.axis ] funcs.append( func + take( other, slice(n0,n1), self.axis ) ) n0 = n1 assert n0 == self.shape[ self.axis ] return concatenate( funcs, self.axis ) def _sum( self, axis ): if axis == self.axis: return util.sum( sum( func, axis ) for func in self.funcs ) funcs = [ sum( func, axis ) for func in self.funcs ] axis = self.axis - (axis<self.axis) return concatenate( funcs, axis ) def _align( self, axes, ndim ): funcs = [ align( func, axes, ndim ) for func in self.funcs ] axis = axes[ self.axis ] return concatenate( funcs, axis ) def _takediag( self ): if self.axis < self.ndim-2: return concatenate( [ takediag(f) for f in self.funcs ], axis=self.axis ) axis = self.ndim-self.axis-3 # -1=>-2, -2=>-1 n0 = 0 funcs = [] for func in self.funcs: n1 = n0 + func.shape[self.axis] funcs.append( takediag( take( func, slice(n0,n1), axis ) ) ) n0 = n1 assert n0 == self.shape[self.axis] return concatenate( funcs, axis=-1 ) def _inflate( self, dofmap, length, axis ): assert not isinstance( self.shape[axis], int ) return concatenate( [ inflate(func,dofmap,length,axis) for func in self.funcs ], self.axis ) def _take( self, indices, axis ): if axis != self.axis: return concatenate( [ take(func,indices,axis) for func in self.funcs ], self.axis ) funcs = [] while len(indices): n = 0 for func in self.funcs: if n <= indices[0] < n + func.shape[axis]: break n += func.shape[axis] else: raise Exception, 'index out of bounds' length = 1 while length < len(indices) and n <= indices[length] < n + func.shape[axis]: length += 1 funcs.append( take( func, indices[:length]-n, axis ) ) indices = indices[length:] assert funcs, 'empty slice' if len( funcs ) == 1: return funcs[0] return concatenate( funcs, axis=axis ) def _dot( self, other, naxes ): axes = range( self.ndim-naxes, self.ndim ) n0 = 0 funcs = [] for f in self.funcs: n1 = n0 + f.shape[self.axis] funcs.append( dot( f, take( other, slice(n0,n1), self.axis ), axes ) ) n0 = n1 if self.axis >= self.ndim - naxes: return util.sum( funcs ) return concatenate( funcs, self.axis ) def _negative( self ): return concatenate( [ -func for func in self.funcs ], self.axis ) def _opposite( self ): return concatenate( [ opposite(func) for func in self.funcs ], self.axis ) def _power( self, n ): return concatenate( [ power( func, n ) for func in self.funcs ], self.axis )
[docs]class Interpolate( ArrayFunc ): 'interpolate uniformly spaced data; stepwise for now' __slots__ = () def __init__( self, x, xp, fp, left=None, right=None ): 'constructor' xp = numpy.array( xp ) fp = numpy.array( fp ) assert xp.ndim == fp.ndim == 1 if not numpy.all( numpy.diff(xp) > 0 ): warnings.warn( 'supplied x-values are non-increasing' ) assert x.ndim == 0 ArrayFunc.__init__( self, args=[x,xp,fp,left,right], evalf=numpy.interp, shape=() )
[docs]class Cross( ArrayFunc ): 'cross product' __slots__ = 'func1', 'func2', 'axis' def __init__( self, func1, func2, axis ): 'contructor' self.func1 = func1 self.func2 = func2 self.axis = axis shape = _jointshape( func1.shape, func2.shape ) assert 0 <= axis < len(shape), 'axis out of bounds: axis={0}, len(shape)={1}'.format( axis, len(shape) ) ArrayFunc.__init__( self, args=(func1,func2,axis-len(shape)), evalf=numeric.cross, shape=shape ) def _localgradient( self, ndims ): return cross( self.func1[...,_], localgradient(self.func2,ndims), axis=self.axis ) \ - cross( self.func2[...,_], localgradient(self.func1,ndims), axis=self.axis ) def _take( self, index, axis ): if axis != self.axis: return cross( take(self.func1,index,axis), take(self.func2,index,axis), self.axis ) def _opposite( self ): return cross( opposite(self.func1), opposite(self.func2), self.axis )
[docs]class Determinant( ArrayFunc ): 'normal' __slots__ = 'func', def __init__( self, func ): 'contructor' self.func = func ArrayFunc.__init__( self, args=[func], evalf=numeric.determinant, shape=func.shape[:-2] ) def _localgradient( self, ndims ): Finv = swapaxes( inverse( self.func ) ) G = localgradient( self.func, ndims ) return self[...,_] * sum( Finv[...,_] * G, axes=[-3,-2] ) def _opposite( self ): return determinant( opposite(self.func) )
[docs]class DofIndex( ArrayFunc ): 'element-based indexing' __slots__ = 'array', 'iax', 'index' def __init__( self, array, iax, index ): 'constructor' assert index.ndim >= 1 assert isinstance( array, numpy.ndarray ) self.array = array assert 0 <= iax < self.array.ndim self.iax = iax self.index = index shape = self.array.shape[:iax] + index.shape + self.array.shape[iax+1:] item = [ slice(None) ] * self.array.ndim item[iax] = index ArrayFunc.__init__( self, args=[self.array]+item, evalf=self.dofindex, shape=shape ) @staticmethod
[docs] def dofindex( arr, *item ): 'evaluate' return arr[item]
def _get( self, i, item ): if self.iax <= i < self.iax + self.index.ndim: index = get( self.index, i - self.iax, item ) return take( self.array, index, self.iax ) return take( get( self.array, i, item ), self.index, self.iax if i > self.iax else self.iax-1 ) def _add( self, other ): if isinstance( other, DofIndex ) and self.iax == other.iax and self.index == other.index: return take( self.array + other.array, self.index, self.iax ) def _multiply( self, other ): if not _isfunc(other) and other.ndim == 0: return take( self.array * other, self.index, self.iax ) def _localgradient( self, ndims ): return _zeros( self.shape + (ndims,) ) def _concatenate( self, other, axis ): if isinstance( other, DofIndex ) and self.iax == other.iax and self.index == other.index: array = numpy.concatenate( [ self.array, other.array ], axis ) return take( array, self.index, self.iax ) def _opposite( self ): return take( self.array, opposite(self.index), self.iax ) def _negative( self ): return take( -self.array, self.index, self.iax )
[docs]class Multiply( ArrayFunc ): 'multiply' __slots__ = 'funcs', def __init__( self, func1, func2 ): 'constructor' shape = _jointshape( func1.shape, func2.shape ) self.funcs = func1, func2 ArrayFunc.__init__( self, args=self.funcs, evalf=numpy.multiply, shape=shape )
[docs] def cdef( self ): 'generate C code' arg1 = 'arg1[%s]' % _cindex( self.funcs[0].shape, self.shape ) arg2 = 'arg2[%s]' % _cindex( self.funcs[1].shape, self.shape ) body = 'retval[0] = %s * %s;\nretval++;' % ( arg1, arg2 ) shape = _cshape( self.shape ) for i in reversed( range(self.ndim) ): body = _cfor( varname='i%d'%i, count=shape[i], do=body ) if self.ndim: body = 'int %s;\n' % ', '.join( 'i%d' % idim for idim in range(self.ndim) ) + body args = [ '%s* retval' % _dtypestr(self), '%s* arg1' % _dtypestr(self.funcs[0]), '%s* arg2' % _dtypestr(self.funcs[1]) ] \ + [ 'int %s' % s for s in shape if isinstance(s,str) ] return _cfunc( args, body ), init_args_for( self, *self.funcs )
def __eq__( self, other ): 'compare' return self is other or ( isinstance( other, Multiply ) and _matchpairs( self.funcs, other.funcs ) ) def _sum( self, axis ): func1, func2 = self.funcs return dot( func1, func2, [axis] ) def _get( self, i, item ): func1, func2 = self.funcs return get( func1, i, item ) * get( func2, i, item ) def _add( self, other ): func1, func2 = self.funcs if _equal( other, func1 ): return func1 * (func2+1) if _equal( other, func2 ): return func2 * (func1+1) if isinstance( other, Multiply ): common = _findcommon( self.funcs, other.funcs ) if common: f, (g1,g2) = common return f * add( g1, g2 ) def _determinant( self ): if self.funcs[0].shape[-2:] == (1,1): return determinant( self.funcs[1] ) * self.funcs[0][...,0,0] if self.funcs[1].shape[-2:] == (1,1): return determinant( self.funcs[0] ) * self.funcs[1][...,0,0] def _product( self, axis ): func1, func2 = self.funcs return product( func1, axis ) * product( func2, axis ) def _multiply( self, other ): func1, func2 = self.funcs func1_other = multiply( func1, other ) if Multiply( func1, other ) != func1_other: return multiply( func1_other, func2 ) func2_other = multiply( func2, other ) if Multiply( func2, other ) != func2_other: return multiply( func1, func2_other ) def _localgradient( self, ndims ): func1, func2 = self.funcs return func1[...,_] * localgradient( func2, ndims ) \ + func2[...,_] * localgradient( func1, ndims ) def _takediag( self ): func1, func2 = self.funcs return takediag( func1 ) * takediag( func2 ) def _take( self, index, axis ): func1, func2 = self.funcs return take( func1, index, axis ) * take( func2, index, axis ) def _opposite( self ): func1, func2 = self.funcs return opposite(func1) * opposite(func2) def _negative( self ): func1, func2 = self.funcs negfunc1 = -func1 if not isinstance( negfunc1, Negative ): return multiply( negfunc1, func2 ) negfunc2 = -func2 if not isinstance( negfunc2, Negative ): return multiply( func1, negfunc2 )
[docs]class Negative( ArrayFunc ): 'negate' __slots__ = 'func', def __init__( self, func ): 'constructor' self.func = func ArrayFunc.__init__( self, args=[func], evalf=numpy.negative, shape=func.shape ) @property def blocks( self ): for f, ind in blocks( self.func ): yield negative(f), ind def _add( self, other ): if isinstance( other, Negative ): return negative( self.func + other.func ) def _multiply( self, other ): if isinstance( other, Negative ): return self.func * other.func return negative( self.func * other ) def _negative( self ): return self.func def _elemint( self, weights ): return -elemint( self.func, weights ) def _align( self, axes, ndim ): return -align( self.func, axes, ndim ) def _get( self, i, item ): return -get( self.func, i, item ) def _sum( self, axis ): return -sum( self.func, axis ) def _localgradient( self, ndims ): return -localgradient( self.func, ndims ) def _takediag( self ): return -takediag( self.func ) def _take( self, index, axis ): return -take( self.func, index, axis ) def _opposite( self ): return -opposite(self.func) def _power( self, n ): if n%2 == 0: return power( self.func, n )
[docs]class Add( ArrayFunc ): 'add' __slots__ = 'funcs', def __init__( self, func1, func2 ): 'constructor' self.funcs = func1, func2 shape = _jointshape( func1.shape, func2.shape ) dtype = _jointdtype(func1,func2) ArrayFunc.__init__( self, args=self.funcs, evalf=numpy.add, shape=shape, dtype=dtype )
[docs] def cdef( self ): 'generate C code' arg1 = 'arg1[%s]' % _cindex( self.funcs[0].shape, self.shape ) arg2 = 'arg2[%s]' % _cindex( self.funcs[1].shape, self.shape ) body = 'retval[0] = %s + %s;\nretval++;' % ( arg1, arg2 ) shape = _cshape( self.shape ) for i in reversed( range(self.ndim) ): body = _cfor( varname='i%d'%i, count=shape[i], do=body ) if self.ndim: body = 'int %s;\n' % ', '.join( 'i%d' % idim for idim in range(self.ndim) ) + body args = [ '%s* retval' % _dtypestr(self), '%s* arg1' % _dtypestr(self.funcs[0]), '%s* arg2' % _dtypestr(self.funcs[1]) ] \ + [ 'int %s' % s for s in shape if isinstance(s,str) ] return _cfunc( args, body ), init_args_for( self, *self.funcs )
def __eq__( self, other ): 'compare' return self is other or ( isinstance( other, Add ) and _matchpairs( self.funcs, other.funcs ) ) def _sum( self, axis ): return sum( self.funcs[0], axis ) + sum( self.funcs[1], axis ) def _localgradient( self, ndims ): func1, func2 = self.funcs return localgradient( func1, ndims ) + localgradient( func2, ndims ) def _get( self, i, item ): func1, func2 = self.funcs return get( func1, i, item ) + get( func2, i, item ) def _takediag( self ): func1, func2 = self.funcs return takediag( func1 ) + takediag( func2 ) def _take( self, index, axis ): func1, func2 = self.funcs return take( func1, index, axis ) + take( func2, index, axis ) def _opposite( self ): func1, func2 = self.funcs return opposite(func1) + opposite(func2) def _add( self, other ): func1, func2 = self.funcs func1_other = add( func1, other ) if Add( func1, other ) != func1_other: return add( func1_other, func2 ) func2_other = add( func2, other ) if Add( func2, other ) != func2_other: return add( func1, func2_other )
[docs]class BlockAdd( Add ): 'block addition (used for DG)' __slots__ = () def _multiply( self, other ): func1, func2 = self.funcs return BlockAdd( func1 * other, func2 * other ) def _inflate( self, dofmap, length, axis ): func1, func2 = self.funcs return BlockAdd( inflate( func1, dofmap, length, axis ), inflate( func2, dofmap, length, axis ) ) def _align( self, axes, ndim ): func1, func2 = self.funcs return BlockAdd( align(func1,axes,ndim), align(func2,axes,ndim) ) def _negative( self ): func1, func2 = self.funcs return BlockAdd( negative(func1), negative(func2) ) def _add( self, other ): func1, func2 = self.funcs try1 = func1 + other if try1 != BlockAdd( func1, other ): return try1 + func2 try2 = func2 + other if try2 != BlockAdd( func2, other ): return try2 + func1 return BlockAdd( self, other ) @property def blocks( self ): func1, func2 = self.funcs for f, ind in blocks( func1 ): yield f, ind for f, ind in blocks( func2 ): yield f, ind
[docs]class Dot( ArrayFunc ): 'dot' __slots__ = 'func1', 'func2', 'naxes' def __init__( self, func1, func2, naxes ): 'constructor' assert naxes > 0 self.func1 = func1 self.func2 = func2 self.naxes = naxes shape = _jointshape( func1.shape, func2.shape )[:-naxes] ArrayFunc.__init__( self, args=(func1,func2,naxes), evalf=numeric.contract_fast, shape=shape )
[docs] def cdef( self ): 'generate C code' shape = _jointshape( self.func1.shape, self.func2.shape ) arg1 = 'arg1[%s]' % _cindex( self.func1.shape, shape ) arg2 = 'arg2[%s]' % _cindex( self.func2.shape, shape ) body = 'sum += %s * %s;\n' % ( arg1, arg2 ) cshape = _cshape( shape ) for i in reversed( range(self.ndim+self.naxes) ): body = _cfor( varname='i%d'%i, count=cshape[i], do=body ) if i == self.ndim: body = 'sum = 0;\n%s\nretval[0] = sum;\nretval++;\n' % body body = '%s sum;\n' % _dtypestr(self) + body body = 'int %s;\n' % ', '.join( 'i%d' % idim for idim in range(self.ndim+self.naxes) ) + body args = [ '%s* retval' % _dtypestr(self), '%s* arg1' % _dtypestr(self.func1), '%s* arg2' % _dtypestr(self.func2) ] \ + [ 'int %s' % s for s in cshape if isinstance(s,str) ] return _cfunc( args, body ), init_args_for( self, self.func1, self.func2, None )
@property def axes( self ): return range( self.ndim, self.ndim + self.naxes ) def _get( self, i, item ): return dot( get( self.func1, i, item ), get( self.func2, i, item ), [ ax-1 for ax in self.axes ] ) def _localgradient( self, ndims ): return dot( localgradient( self.func1, ndims ), self.func2[...,_], self.axes ) \ + dot( self.func1[...,_], localgradient( self.func2, ndims ), self.axes ) def _multiply( self, other ): for ax in self.axes: other = insert( other, ax ) assert other.ndim == self.func1.ndim == self.func2.ndim func1_other = multiply( self.func1, other ) if Multiply( self.func1, other ) != func1_other: return dot( func1_other, self.func2, self.axes ) func2_other = multiply( self.func2, other ) if Multiply( self.func2, other ) != func2_other: return dot( self.func1, func2_other, self.axes ) def _add( self, other ): if isinstance( other, Dot ) and self.axes == other.axes: common = _findcommon( (self.func1,self.func2), (other.func1,other.func2) ) if common: f, (g1,g2) = common return dot( f, g1 + g2, self.axes ) def _takediag( self ): n1, n2 = self.ndim-2, self.ndim-1 return dot( takediag( self.func1, n1, n2 ), takediag( self.func2, n1, n2 ), [ ax-2 for ax in self.axes ] ) def _sum( self, axis ): return dot( self.func1, self.func2, self.axes + [axis] ) def _take( self, index, axis ): return dot( take(self.func1,index,axis), take(self.func2,index,axis), self.axes ) def _concatenate( self, other, axis ): if isinstance( other, Dot ) and other.axes == self.axes: common = _findcommon( (self.func1,self.func2), (other.func1,other.func2) ) if common: f, g12 = common tryconcat = concatenate( g12, axis ) if not isinstance( tryconcat, Concatenate ): # avoid inf recursion return dot( f, tryconcat, self.axes ) def _opposite( self ): return dot( opposite(self.func1), opposite(self.func2), self.axes ) def _negative( self ): negfunc1 = -self.func1 if not isinstance( negfunc1, Negative ): return dot( negfunc1, self.func2, self.axes ) negfunc2 = -self.func2 if not isinstance( negfunc2, Negative ): return dot( self.func1, negfunc2, self.axes )
[docs]class Sum( ArrayFunc ): 'sum' __slots__ = 'axis', 'func' def __init__( self, func, axis ): 'constructor' self.axis = axis self.func = func assert 0 <= axis < func.ndim, 'axis out of bounds' shape = func.shape[:axis] + func.shape[axis+1:] ArrayFunc.__init__( self, args=[func,axis-func.ndim], evalf=numpy.sum, shape=shape ) def _sum( self, axis ): trysum = sum( self.func, axis+(axis>=self.axis) ) if not isinstance( trysum, Sum ): # avoid inf recursion return sum( trysum, self.axis ) def _localgradient( self, ndims ): return sum( localgradient( self.func, ndims ), self.axis ) def _opposite( self ): return sum( opposite(self.func), axes=self.axis )
[docs]class Debug( ArrayFunc ): 'debug' __slots__ = 'func', def __init__( self, func ): 'constructor' self.func = func ArrayFunc.__init__( self, args=[func], evalf=self.debug, shape=func.shape ) @staticmethod
[docs] def debug( arr ): 'debug' print arr return arr
def __str__( self ): 'string representation' return '{DEBUG}' def _localgradient( self, ndims ): return Debug( localgradient( self.func, ndims ) )
[docs]class TakeDiag( ArrayFunc ): 'extract diagonal' __slots__ = 'func', def __init__( self, func ): 'constructor' assert func.shape[-1] == func.shape[-2] self.func = func ArrayFunc.__init__( self, args=[func], evalf=numeric.takediag, shape=func.shape[:-1] ) def _localgradient( self, ndims ): return swapaxes( takediag( localgradient( self.func, ndims ), -3, -2 ) ) def _sum( self, axis ): if axis != self.ndim-1: return takediag( sum( self.func, axis ) ) def _opposite( self ): return takediag( opposite(self.func) )
[docs]class Take( ArrayFunc ): 'generalization of numpy.take(), to accept lists, slices, arrays' __slots__ = 'func', 'axis', 'indices' def __init__( self, func, indices, axis ): 'constructor' assert func.shape[axis] != 1 self.func = func self.axis = axis self.indices = indices # try for regular slice start = indices[0] step = indices[1] - start stop = start + step * len(indices) s = [ slice(None) ] * func.ndim s[axis] = slice( start, stop, step ) if numpy.all( numpy.diff(indices) == step ) \ else indices newlen, = numpy.empty( func.shape[axis] )[ indices ].shape assert newlen > 0 shape = func.shape[:axis] + (newlen,) + func.shape[axis+1:] ArrayFunc.__init__( self, args=(func,(Ellipsis,)+tuple(s)), evalf=numpy.ndarray.__getitem__, shape=shape ) def _localgradient( self, ndims ): return take( localgradient( self.func, ndims ), self.indices, self.axis ) def _take( self, index, axis ): if axis == self.axis: if numpy.all( numpy.diff( self.indices ) == 1 ): indices = index + self.indices[0] else: indices = self.indices[index] return take( self.func, indices, axis ) trytake = take( self.func, index, axis ) if not isinstance( trytake, Take ): # avoid inf recursion return take( trytake, self.indices, self.axis )
[docs]class Power( ArrayFunc ): 'power' __slots__ = 'func', 'power' def __init__( self, func, power ): 'constructor' assert _isfunc( func ) assert _isscalar( power ) self.func = func self.power = power ArrayFunc.__init__( self, args=[func,power], evalf=numpy.power, shape=func.shape ) def _localgradient( self, ndims ): return self.power * ( self.func**(self.power-1) )[...,_] * localgradient( self.func, ndims ) def _power( self, n ): func = self.func newpower = n * self.power if self.power % 2 == 0 and newpower % 2 != 0: func = abs( func ) return power( func, newpower ) def _get( self, i, item ): return get( self.func, i, item )**self.power def _sum( self, axis ): if self.power == 2: return dot( self.func, self.func, axis ) def _takediag( self ): return takediag( self.func )**self.power def _take( self, index, axis ): return power( take( self.func, index, axis ), self.power ) def _opposite( self ): return power( opposite(self.func), self.power ) def _multiply( self, other ): if isinstance( other, Power ) and self.func == other.func: return power( self.func, self.power + other.power ) if other == self.func: return power( self.func, self.power + 1 ) def _sign( self ): if self.power % 2 == 0: return expand( 1., self.shape )
[docs]class ElemFunc( ArrayFunc ): 'trivial func' __slots__ = 'domainelem', 'side', 'cascade' def __init__( self, domainelem, side=0 ): 'constructor' self.domainelem = domainelem self.side = side self.cascade = Cascade( domainelem.ndims, side ) ArrayFunc.__init__( self, args=[self.cascade,domainelem], evalf=self.elemfunc, shape=[domainelem.ndims] ) @staticmethod
[docs] def elemfunc( cascade, domainelem ): 'evaluate' for elem, points in cascade: if elem is domainelem: return points raise Exception, '%r not found' % domainelem
def _localgradient( self, ndims ): return self.cascade.transform( ndims ) def _opposite( self ): return ElemFunc( self.domainelem, 1-self.side )
[docs] def find( self, elem, C ): 'find coordinates' assert C.ndim == 2 and C.shape[1] == self.domainelem.ndims assert elem.ndims == self.domainelem.ndims # for now pelem, transform = elem.parent offset = transform.offset Tinv = transform.invtrans while pelem is not self.domainelem: pelem, newtransform = pelem.parent transform = transform.nest( newtransform ) return elem.select_contained( transform.invapply( C ), eps=1e-10 )
[docs]class Pointwise( ArrayFunc ): 'pointwise transformation' __slots__ = 'args', 'evalf', 'deriv' def __init__( self, args, evalf, deriv ): 'constructor' assert _isfunc( args ) shape = args.shape[1:] self.args = args self.evalf = evalf self.deriv = deriv ArrayFunc.__init__( self, args=tuple(args), evalf=evalf, shape=shape ) def _localgradient( self, ndims ): return ( self.deriv( self.args )[...,_] * localgradient( self.args, ndims ) ).sum( 0 ) def _takediag( self ): return pointwise( takediag(self.args), self.evalf, self.deriv ) def _get( self, axis, item ): return pointwise( get( self.args, axis+1, item ), self.evalf, self.deriv ) def _take( self, index, axis ): return pointwise( take( self.args, index, axis+1 ), self.evalf, self.deriv ) def _opposite( self ): # args = [ opposite(arg,side) for arg in self.args ] # @TO args = [ opposite(arg) for arg in self.args ] return pointwise( args, self.evalf, self.deriv )
[docs]class Sign( ArrayFunc ): 'sign' __slots__ = 'func', def __init__( self, func ): 'constructor' assert _isfunc( func ) self.func = func ArrayFunc.__init__( self, args=[func], evalf=numpy.sign, shape=func.shape ) def _localgradient( self, ndims ): return _zeros( self.shape + (ndims,) ) def _takediag( self ): return sign( takediag(self.func) ) def _get( self, axis, item ): return sign( get( self.func, axis, item ) ) def _take( self, index, axis ): return sign( take( self.func, index, axis ) ) def _opposite( self ): return sign( opposite( self.func ) ) def _sign( self ): return self def _power( self, n ): if n % 2 == 0: return expand( 1., self.shape )
class Pointdata( ArrayFunc ): __slots__ = 'data', def __init__ ( self, data, shape ): 'constructor' assert isinstance(data,dict) self.data = data ArrayFunc.__init__( self, args=[ELEM,POINTS,self.data], evalf=self.pointdata, shape=shape ) @staticmethod def pointdata( elem, points, data ): myvals,mypoint = data[elem] assert mypoint is points, 'Illegal point set' return myvals def update_max( self, func ): func = asarray(func) assert func.shape == self.shape data = dict( (elem,(numpy.maximum(func(elem,points),values),points)) for elem,(values,points) in self.data.iteritems() ) return Pointdata( data, self.shape )
[docs]class Eig( Evaluable ): 'Eig' __slots__ = 'func', 'shape', 'symmetric', 'sort' def __init__( self, func, symmetric=False, sort=False ): 'contructor' Evaluable.__init__( self, args=[func, sort], evalf=numeric.eigh if symmetric else numeric.eig ) self.symmetric = symmetric self.sort = sort self.func = func self.shape = func.shape def _opposite( self ): return Eig( opposite(self.func), self.symmetric, self.sort )
class ArrayFromTuple( ArrayFunc ): def __init__( self, arrays, index, shape ): self.arrays = arrays self.index = index ArrayFunc.__init__( self, args=[arrays,index], evalf=self.arrayfromtuple, shape=shape ) @staticmethod def arrayfromtuple( arrays, index ): return arrays[ index ] def _opposite( self ): return ArrayFromTuple( opposite(self.arrays), self.index, self.shape ) # PRIORITY OBJECTS # # Prority objects get priority in situations like A + B, which can be evaluated # as A.__add__(B) and B.__radd__(A), such that they get to decide on how the # operation is performed. The reason is that these objects are wasteful, # generally introducing a lot of zeros, and we would like to see them disappear # by the act of subsequent operations. For this annihilation to work well # priority objects keep themselves at the surface where magic happens. # # Update: "priority objects" as such do not exist anymore, might be # reintroduced later on.
[docs]class Zeros( ArrayFunc ): 'zero' __slots__ = () def __init__( self, shape ): 'constructor' shape = tuple( shape ) ArrayFunc.__init__( self, args=[POINTS,shape], evalf=self.zeros, shape=shape ) @property def blocks( self ): return () @staticmethod
[docs] def zeros( points, shape ): 'prepend point axes' assert not any( sh is None for sh in shape ), 'cannot evaluate zeros for shape %s' % (shape,) shape = points.shape[:-1] + shape strides = [0] * len(shape) return numpy.lib.stride_tricks.as_strided( numpy.array(0.), shape, strides )
def _repeat( self, length, axis ): assert self.shape[axis] == 1 return _zeros( self.shape[:axis] + (length,) + self.shape[axis+1:] ) def _localgradient( self, ndims ): return _zeros( self.shape+(ndims,) ) def _add( self, other ): shape = _jointshape( self.shape, other.shape ) return expand( other, shape ) def _multiply( self, other ): shape = _jointshape( self.shape, other.shape ) return _zeros( shape ) def _dot( self, other, naxes ): shape = _jointshape( self.shape, other.shape ) return _zeros( shape[:-naxes] ) def _cross( self, other, axis ): shape = _jointshape( self.shape, other.shape ) return _zeros( shape ) def _negative( self ): return self def _diagonalize( self ): return _zeros( self.shape + (self.shape[-1],) ) def _sum( self, axis ): return _zeros( self.shape[:axis] + self.shape[axis+1:] ) def _align( self, axes, ndim ): shape = [1] * ndim for ax, sh in zip( axes, self.shape ): shape[ax] = sh return _zeros( shape ) def _get( self, i, item ): return _zeros( self.shape[:i] + self.shape[i+1:] ) def _takediag( self ): sh = max( self.shape[-2], self.shape[-1] ) return _zeros( self.shape[:-2] + (sh,) ) def _take( self, index, axis ): return _zeros( self.shape[:axis] + index.shape + self.shape[axis+1:] ) def _inflate( self, dofmap, length, axis ): assert not isinstance( self.shape[axis], int ) return _zeros( self.shape[:axis] + (length,) + self.shape[axis+1:] ) def _elemint( self, weights ): return numpy.zeros( [1]*self.ndim ) def _power( self, n ): return self def _opposite( self ): return self
[docs]class Inflate( ArrayFunc ): 'inflate' __slots__ = 'func', 'dofmap', 'length', 'axis' def __init__( self, func, dofmap, length, axis ): 'constructor' self.func = func self.dofmap = dofmap self.length = length self.axis = axis shape = func.shape[:axis] + (length,) + func.shape[axis+1:] ArrayFunc.__init__( self, args=[func,dofmap,length,axis-func.ndim], evalf=self.inflate, shape=shape ) @property def blocks( self ): for f, ind in blocks( self.func ): assert ind[self.axis] == None yield f, ind[:self.axis] + (self.dofmap,) + ind[self.axis+1:] @staticmethod
[docs] def inflate( array, indices, length, axis ): 'inflate' warnings.warn( 'using explicit inflation; this is usually a bug.' ) shape = list( array.shape ) shape[axis] = length inflated = numpy.zeros( shape ) inflated[(Ellipsis,indices)+(slice(None),)*(-axis-1)] = array return inflated
def _inflate( self, dofmap, length, axis ): assert axis != self.axis if axis > self.axis: return return inflate( inflate( self.func, dofmap, length, axis ), self.dofmap, self.length, self.axis ) def _localgradient( self, ndims ): return inflate( localgradient(self.func,ndims), self.dofmap, self.length, self.axis ) def _align( self, shuffle, ndims ): return inflate( align(self.func,shuffle,ndims), self.dofmap, self.length, shuffle[self.axis] ) def _get( self, axis, item ): assert axis != self.axis return inflate( get(self.func,axis,item), self.dofmap, self.length, self.axis-(axis<self.axis) ) def _dot( self, other, naxes ): axes = range( self.ndim-naxes, self.ndim ) if isinstance( other, Inflate ) and other.axis == self.axis: assert self.dofmap == other.dofmap other = other.func elif other.shape[self.axis] != 1: other = take( other, self.dofmap, self.axis ) arr = dot( self.func, other, axes ) if self.axis >= self.ndim - naxes: return arr return inflate( arr, self.dofmap, self.length, self.axis ) def _multiply( self, other ): if isinstance( other, Inflate ) and self.axis == other.axis: assert self.dofmap == other.dofmap other = other.func elif other.shape[self.axis] != 1: other = take( other, self.dofmap, self.axis ) return inflate( multiply(self.func,other), self.dofmap, self.length, self.axis ) def _add( self, other ): if isinstance( other, Inflate ) and self.axis == other.axis and self.dofmap == other.dofmap: return inflate( add(self.func,other.func), self.dofmap, self.length, self.axis ) return BlockAdd( self, other ) def _cross( self, other, axis ): if isinstance( other, Inflate ) and self.axis == other.axis: assert self.dofmap == other.dofmap other = other.func elif other.shape[self.axis] != 1: other = take( other, self.dofmap, self.axis ) return inflate( cross(self.func,other,axis), self.dofmap, self.length, self.axis ) def _negative( self ): return inflate( negative(self.func), self.dofmap, self.length, self.axis ) def _power( self, n ): return inflate( power(self.func,n), self.dofmap, self.length, self.axis ) def _takediag( self ): assert self.axis < self.ndim-2 return inflate( takediag(self.func), self.dofmap, self.length, self.axis ) def _take( self, index, axis ): if axis == self.axis: assert index == self.dofmap return self.func return inflate( take( self.func, index, axis ), self.dofmap, self.length, self.axis ) def _diagonalize( self ): assert self.axis < self.ndim-1 return inflate( diagonalize(self.func), self.dofmap, self.length, self.axis ) def _sum( self, axis ): arr = sum( self.func, axis ) if axis == self.axis: return arr return inflate( arr, self.dofmap, self.length, self.axis-(axis<self.axis) ) def _opposite( self ): return inflate( opposite(self.func), opposite(self.dofmap), self.length, self.axis )
[docs]class Diagonalize( ArrayFunc ): 'diagonal matrix' __slots__ = 'func', def __init__( self, func ): 'constructor' n = func.shape[-1] assert n != 1 shape = func.shape + (n,) self.func = func ArrayFunc.__init__( self, args=[func], evalf=numeric.diagonalize, shape=shape ) def _localgradient( self, ndims ): return swapaxes( diagonalize( swapaxes( localgradient( self.func, ndims ), (-2,-1) ) ), (-3,-1) ) def _get( self, i, item ): if i >= self.ndim-2: return kronecker( get( self.func, -1, item ), axis=-1, pos=item, length=self.func.shape[-1] ) return diagonalize( get( self.func, i, item ) ) def _inverse( self ): return diagonalize( reciprocal( self.func ) ) def _determinant( self ): return product( self.func, -1 ) def _multiply( self, other ): return diagonalize( self.func * takediag( other ) ) def _add( self, other ): if isinstance( other, Diagonalize ): return diagonalize( self.func + other.func ) def _negative( self ): return diagonalize( -self.func ) def _sum( self, axis ): if axis >= self.ndim-2: return self.func return diagonalize( sum( self.func, axis ) ) def _align( self, axes, ndim ): if axes[-2:] in [ (ndim-2,ndim-1), (ndim-1,ndim-2) ]: return diagonalize( align( self.func, axes[:-2] + (ndim-2,), ndim-1 ) ) def _opposite( self ): return diagonalize( opposite(self.func) )
[docs]class Repeat( ArrayFunc ): 'repeat singleton axis' __slots__ = 'func', 'axis', 'length' def __init__( self, func, length, axis ): 'constructor' assert func.shape[axis] == 1 self.func = func self.axis = axis self.length = length shape = func.shape[:axis] + (length,) + func.shape[axis+1:] ArrayFunc.__init__( self, args=[func,length,axis-func.ndim], evalf=numeric.fastrepeat, shape=shape ) def _negative( self ): return repeat( -self.func, self.length, self.axis ) def _localgradient( self, ndims ): return repeat( localgradient( self.func, ndims ), self.length, self.axis ) def _get( self, axis, item ): if axis == self.axis: assert 0 <= item < self.length return get( self.func, axis, 0 ) return repeat( get( self.func, axis, item ), self.length, self.axis-(axis<self.axis) ) def _sum( self, axis ): if axis == self.axis: return get( self.func, axis, 0 ) * self.length return repeat( sum( self.func, axis ), self.length, self.axis-(axis<self.axis) ) def _product( self, axis ): if axis == self.axis: return get( self.func, axis, 0 )**self.length return repeat( product( self.func, axis ), self.length, self.axis-(axis<self.axis) ) def _power( self, n ): return repeat( power( self.func, n ), self.length, self.axis ) def _add( self, other ): return repeat( self.func + other, self.length, self.axis ) def _multiply( self, other ): return repeat( self.func * other, self.length, self.axis ) def _align( self, shuffle, ndim ): return repeat( align(self.func,shuffle,ndim), self.length, shuffle[self.axis] ) def _take( self, index, axis ): if axis == self.axis: return repeat( self.func, index.shape[0], self.axis ) return repeat( take(self.func,index,axis), self.length, self.axis ) def _takediag( self ): assert self.axis < self.ndim-2 return repeat( takediag( self.func ), self.length, self.axis ) def _cross( self, other, axis ): if axis != self.axis: return repeat( cross( self.func, other, axis ), self.length, self.axis ) def _dot( self, other, naxes ): axes = range( self.ndim-naxes, self.ndim ) func = dot( self.func, other, axes ) if other.shape[self.axis] != 1: assert other.shape[self.axis] == self.length return func if self.axis >= self.ndim - naxes: return func * self.length return repeat( func, self.length, self.axis ) def _opposite( self ): return repeat( opposite(self.func), self.length, self.axis )
[docs]class Const( ArrayFunc ): 'pointwise transformation' __slots__ = () def __init__( self, func ): 'constructor' func = numpy.asarray( func ) ArrayFunc.__init__( self, args=(POINTS,func), evalf=self.const, shape=func.shape ) @staticmethod
[docs] def const( points, arr ): 'prepend point axes' shape = points.shape[:-1] + arr.shape strides = (0,) * (points.ndim-1) + arr.strides return numpy.lib.stride_tricks.as_strided( arr, shape, strides )
def _localgradient( self, ndims ): return _zeros( self.shape+(ndims,) ) # AUXILIARY FUNCTIONS
def _jointshape( *shapes ): 'determine shape after singleton expansion' ndim = len(shapes[0]) combshape = [1] * ndim for shape in shapes: assert len(shape) == ndim for i, sh in enumerate(shape): if combshape[i] == 1: combshape[i] = sh else: assert sh in ( combshape[i], 1 ), 'incompatible shapes: %s' % ', '.join( str(sh) for sh in shapes ) return tuple(combshape) def _matchndim( *arrays ): 'introduce singleton dimensions to match ndims' arrays = [ asarray(array) for array in arrays ] ndim = _max( array.ndim for array in arrays ) return [ array[(_,)*(ndim-array.ndim)] for array in arrays ] def _isiterable( obj ): 'check for iterability' try: iter(obj) except TypeError: return False return True def _obj2str( obj ): 'convert object to string' if isinstance( obj, numpy.ndarray ): if obj.size < 6: return _obj2str(obj.tolist()) return 'array<%s>' % 'x'.join( map( str, obj.shape ) ) if isinstance( obj, list ): if len(obj) < 6: return '[%s]' % ','.join( _obj2str(o) for o in obj ) return '[#%d]' % len(obj) if isinstance( obj, (tuple,set) ): if len(obj) < 6: return '(%s)' % ','.join( _obj2str(o) for o in obj ) return '(#%d)' % len(obj) if isinstance( obj, dict ): return '{#%d}' % len(obj) if isinstance( obj, slice ): I = '' if obj.start is not None: I += str(obj.start) if obj.step is not None: I += ':' + str(obj.step) I += ':' if obj.stop is not None: I += str(obj.stop) return I if obj is Ellipsis: return '...' if obj is POINTS: return '<points>' if obj is WEIGHTS: return '<weights>' if obj is ELEM: return '<elem>' return str(obj) def _findcommon( (a1,a2), (b1,b2) ): 'find common item in 2x2 data' if _equal( a1, b1 ): return a1, (a2,b2) if _equal( a1, b2 ): return a1, (a2,b1) if _equal( a2, b1 ): return a2, (a1,b2) if _equal( a2, b2 ): return a2, (a1,b1) _matchpairs = lambda (a1,a2), (b1,b2): _equal(a1,b1) and _equal(a2,b2) or _equal(a1,b2) and _equal(a2,b1) _max = max _min = min _sum = sum _isfunc = lambda arg: isinstance( arg, ArrayFunc ) _isscalar = lambda arg: asarray(arg).ndim == 0 _isint = lambda arg: numpy.asarray( arg ).dtype == int _ascending = lambda arg: ( numpy.diff(arg) > 0 ).all() _iszero = lambda arg: isinstance( arg, Zeros ) or isinstance( arg, numpy.ndarray ) and numpy.all( arg == 0 ) _isunit = lambda arg: not _isfunc(arg) and ( numpy.asarray(arg) == 1 ).all() _subsnonesh = lambda shape: tuple( 1 if sh is None else sh for sh in shape ) _normdims = lambda ndim, shapes: tuple( numeric.normdim(ndim,sh) for sh in shapes ) _zeros = lambda shape: Zeros( shape ) _zeros_like = lambda arr: _zeros( arr.shape ) def _call( obj, attr, *args ): 'call method if it exists, return None otherwise' f = getattr( obj, attr, None ) return f and f( *args ) def _norm_and_sort( ndim, args ): 'norm axes, sort, and assert unique' normargs = tuple( sorted( numeric.normdim( ndim, arg ) for arg in args ) ) assert _ascending( normargs ) # strict return normargs def _jointdtype( *args ): 'determine joint dtype' if any( asarray(arg).dtype == float for arg in args ): return float return int def _dtypestr( arg ): if arg.dtype == int: return 'int' if arg.dtype == float: return 'double' raise Exception, 'unknown dtype %s' % arg.dtype def _equal( arg1, arg2 ): 'compare two objects' if arg1 is arg2: return True if isinstance( arg1, dict ) or isinstance( arg2, dict ): return False if isinstance( arg1, (list,tuple) ): if not isinstance( arg2, (list,tuple) ) or len(arg1) != len(arg2): return False return all( _equal(v1,v2) for v1, v2 in zip( arg1, arg2 ) ) if not isinstance( arg1, numpy.ndarray ) and not isinstance( arg2, numpy.ndarray ): return arg1 == arg2 elif isinstance( arg1, numpy.ndarray ) and isinstance( arg2, numpy.ndarray ): return arg1.shape == arg2.shape and numpy.all( arg1 == arg2 ) else: return False
[docs]def asarray( arg ): 'convert to ArrayFunc or numpy.ndarray' if isinstance( arg, numpy.ndarray ) and arg.ndim == 0: arg = arg[...] if _isfunc(arg): return arg arg = numpy.asarray( arg ) if arg.dtype == object: return stack( arg, axis=0 ) elif numpy.all( arg == 0 ): return _zeros( arg.shape ) else: return arg
def _asarray( arg ): warnings.warn( '_asarray is deprecated, use asarray instead', DeprecationWarning ) return asarray( arg ) # FUNCTIONS
[docs]def insert( arg, n ): 'insert axis' arg = asarray( arg ) n = numeric.normdim( arg.ndim+1, n ) I = numpy.arange( arg.ndim ) return align( arg, I + (I>=n), arg.ndim+1 )
[docs]def stack( args, axis=0 ): 'stack functions along new axis' args = [ insert( arg, axis ) for arg in args ] ndim = args[0].ndim assert all( arg.ndim == ndim for arg in args[1:] ), 'arguments have non-matching shapes' return concatenate( args, axis )
[docs]def chain( funcs ): 'chain' funcs = map( asarray, funcs ) shapes = [ func.shape[0] for func in funcs ] return [ concatenate( [ func if i==j else _zeros( (sh,) + func.shape[1:] ) for j, sh in enumerate(shapes) ], axis=0 ) for i, func in enumerate(funcs) ]
[docs]def merge( funcs ): 'Combines unchained funcs into one function object.' cascade = fmap = nmap = None offset = 0 # ndofs = _sum( f.shape[0] for f in funcs ) for inflated_func in funcs: (func, (dofmap,)), = inflated_func.blocks # Returns one scalar function. if fmap is None: fmap = func.stdmap.copy() else: targetlen = len( fmap ) + len( func.stdmap ) fmap.update( func.stdmap ) assert len( fmap ) == targetlen, 'Don`t allow overlap.' if nmap is None: nmap = dofmap.dofmap.copy() else: targetlen = len( nmap ) + len( dofmap.dofmap ) nmap.update( dict( (key, val+offset) for key, val in dofmap.dofmap.iteritems() ) ) assert len( nmap ) == targetlen, 'Don`t allow overlap.' if cascade is None: cascade = func.cascade else: assert func.cascade == cascade, 'Functions have to be defined on domains of same dimension.' offset += inflated_func.shape[0] return function( fmap, nmap, offset, cascade.ndims )
[docs]def vectorize( args ): 'vectorize' return util.sum( kronecker( func, axis=1, length=len(args), pos=ifun ) for ifun, func in enumerate( chain( args ) ) )
[docs]def expand( arg, shape ): 'expand' arg = asarray( arg ) shape = tuple(shape) assert len(shape) == arg.ndim for i, sh in enumerate( shape ): arg = repeat( arg, sh, i ) assert arg.shape == shape return arg
[docs]def repeat( arg, length, axis ): 'repeat' arg = asarray( arg ) axis = numeric.normdim( arg.ndim, axis ) if arg.shape[axis] == length: return arg assert arg.shape[axis] == 1 retval = _call( arg, '_repeat', length, axis ) if retval is not None: shape = arg.shape[:axis] + (length,) + arg.shape[axis+1:] assert retval.shape == shape, 'bug in %s._repeat' % arg return retval return Repeat( arg, length, axis )
[docs]def get( arg, iax, item ): 'get item' assert _isint( item ) and _isscalar( item ) arg = asarray( arg ) iax = numeric.normdim( arg.ndim, iax ) sh = arg.shape[iax] assert isinstance(sh,int), 'cannot get item %r from axis %r' % ( item, sh ) item = 0 if sh == 1 \ else numeric.normdim( sh, item ) if not _isfunc( arg ): return numeric.get( arg, iax, item ) retval = _call( arg, '_get', iax, item ) if retval is not None: assert retval.shape == arg.shape[:iax] + arg.shape[iax+1:], 'bug in %s._get' % arg return retval return Get( arg, iax, item )
[docs]def align( arg, axes, ndim ): 'align' arg = asarray( arg ) assert ndim >= len(axes) assert len(axes) == arg.ndim axes = _normdims( ndim, axes ) assert len(set(axes)) == len(axes), 'duplicate axes in align' if list(axes) == range(ndim): return arg if not _isfunc( arg ): return numeric.align( arg, axes, ndim ) retval = _call( arg, '_align', axes, ndim ) if retval is not None: shape = [1] * ndim for i, axis in enumerate( axes ): shape[axis] = arg.shape[i] assert retval.shape == tuple(shape), 'bug in %s._align' % arg return retval return Align( arg, axes, ndim )
[docs]def bringforward( arg, axis ): 'bring axis forward' arg = asarray(arg) axis = numeric.normdim(arg.ndim,axis) if axis == 0: return arg return transpose( args, [axis] + range(axis) + range(axis+1,args.ndim) )
[docs]def elemint( arg, weights ): 'elementwise integration' arg = asarray( arg ) if not _isfunc( arg ): return arg * ElemArea( weights ) retval = _call( arg, '_elemint', weights ) if retval is not None: return retval return ElemInt( arg, weights )
[docs]def grad( arg, coords, ndims=0 ): 'local derivative' arg = asarray( arg ) if _isfunc( arg ): return arg.grad( coords, ndims ) return _zeros( arg.shape + coords.shape )
[docs]def symgrad( arg, coords, ndims=0 ): 'symmetric gradient' if _isfunc( arg ): return arg.symgrad( coords, ndims ) return _zeros( arg.shape + coords.shape )
[docs]def div( arg, coords, ndims=0 ): 'gradient' if _isfunc( arg ): return arg.div( coords, ndims ) assert arg.shape[-1:] == coords.shape return _zeros( arg.shape[:-1] )
[docs]def sum( arg, axes=-1 ): 'sum over multiply axes' arg = asarray( arg ) if _isiterable(axes): if len(axes) == 0: return arg axes = _norm_and_sort( arg.ndim, axes ) assert numpy.all( numpy.diff(axes) > 0 ), 'duplicate axes in sum' arg = sum( arg, axes[1:] ) axis = axes[0] else: axis = numeric.normdim( arg.ndim, axes ) if arg.shape[axis] == 1: return get( arg, axis, 0 ) if not _isfunc( arg ): return arg.sum( axis ) retval = _call( arg, '_sum', axis ) if retval is not None: assert retval.shape == arg.shape[:axis] + arg.shape[axis+1:], 'bug in %s._sum' % arg return retval return Sum( arg, axis )
[docs]def dot( arg1, arg2, axes ): 'dot product' arg1, arg2 = _matchndim( arg1, arg2 ) shape = _jointshape( arg1.shape, arg2.shape ) if _isiterable(axes): if len(axes) == 0: return arg1 * arg2 axes = _norm_and_sort( len(shape), axes ) assert numpy.all( numpy.diff(axes) > 0 ), 'duplicate axes in sum' else: axes = numeric.normdim( len(shape), axes ), for i, axis in enumerate( axes ): if arg1.shape[axis] == 1 or arg2.shape[axis] == 1: arg1 = sum( arg1, axis ) arg2 = sum( arg2, axis ) axes = axes[:i] + tuple( axis-1 for axis in axes[i+1:] ) return dot( arg1, arg2, axes ) if _isunit( arg1 ): return sum( expand( arg2, shape ), axes ) if _isunit( arg2 ): return sum( expand( arg1, shape ), axes ) if not _isfunc(arg1) and not _isfunc(arg2): return numeric.contract( arg1, arg2, axes ) shuffle = range( len(shape) ) for ax in reversed( axes ): shuffle.append( shuffle.pop(ax) ) arg1 = transpose( arg1, shuffle ) arg2 = transpose( arg2, shuffle ) naxes = len( axes ) shape = tuple( shape[i] for i in shuffle[:-naxes] ) retval = _call( arg1, '_dot', arg2, naxes ) if retval is not None: assert retval.shape == shape, 'bug in %s._dot' % arg1 return retval retval = _call( arg2, '_dot', arg1, naxes ) if retval is not None: assert retval.shape == shape, 'bug in %s._dot' % arg2 return retval return Dot( arg1, arg2, naxes )
[docs]def determinant( arg, axes=(-2,-1) ): 'determinant' arg = asarray( arg ) ax1, ax2 = _norm_and_sort( arg.ndim, axes ) assert ax2 > ax1 # strict n = arg.shape[ax1] assert n == arg.shape[ax2] if n == 1: return get( get( arg, ax2, 0 ), ax1, 0 ) trans = range(ax1) + [-2] + range(ax1,ax2-1) + [-1] + range(ax2-1,arg.ndim-2) arg = align( arg, trans, arg.ndim ) shape = arg.shape[:-2] if not _isfunc( arg ): return numeric.determinant( arg ) retval = _call( arg, '_determinant' ) if retval is not None: assert retval.shape == shape, 'bug in %s._determinant' % arg return retval return Determinant( arg )
[docs]def inverse( arg, axes=(-2,-1) ): 'inverse' arg = asarray( arg ) ax1, ax2 = _norm_and_sort( arg.ndim, axes ) assert ax2 > ax1 # strict n = arg.shape[ax1] assert arg.shape[ax2] == n if n == 1: return reciprocal( arg ) trans = range(ax1) + [-2] + range(ax1,ax2-1) + [-1] + range(ax2-1,arg.ndim-2) arg = align( arg, trans, arg.ndim ) if not _isfunc( arg ): return numeric.inverse( arg ).transpose( trans ) retval = _call( arg, '_inverse' ) if retval is not None: assert retval.shape == arg.shape, 'bug in %s._inverse' % arg return transpose( retval, trans ) return transpose( Inverse(arg), trans )
[docs]def takediag( arg, ax1=-2, ax2=-1 ): 'takediag' arg = asarray( arg ) ax1, ax2 = _norm_and_sort( arg.ndim, (ax1,ax2) ) assert ax2 > ax1 # strict axes = range(ax1) + [-2] + range(ax1,ax2-1) + [-1] + range(ax2-1,arg.ndim-2) arg = align( arg, axes, arg.ndim ) if arg.shape[-1] == 1: return get( arg, -1, 0 ) if arg.shape[-2] == 1: return get( arg, -2, 0 ) assert arg.shape[-1] == arg.shape[-2] shape = arg.shape[:-1] if not _isfunc( arg ): return numeric.takediag( arg ) retval = _call( arg, '_takediag' ) if retval is not None: assert retval.shape == shape, 'bug in %s._takediag' % arg return retval return TakeDiag( arg )
[docs]def localgradient( arg, ndims ): 'local derivative' arg = asarray( arg ) shape = arg.shape + (ndims,) if not _isfunc( arg ): return _zeros( shape ) lgrad = arg._localgradient( ndims ) assert lgrad.shape == shape, 'bug in %s._localgradient' % arg return lgrad
[docs]def dotnorm( arg, coords, ndims=0 ): 'normal component' return sum( arg * coords.normal( ndims-1 ) )
[docs]def kronecker( arg, axis, length, pos ): 'kronecker' axis = numeric.normdim( arg.ndim+1, axis ) arg = insert( arg, axis ) args = [ _zeros_like(arg) ] * length args[pos] = arg return concatenate( args, axis=axis )
[docs]def diagonalize( arg ): 'diagonalize' arg = asarray( arg ) shape = arg.shape + (arg.shape[-1],) if arg.shape[-1] == 1: return arg[...,_] retval = _call( arg, '_diagonalize' ) if retval is not None: assert retval.shape == shape, 'bug in %s._diagonalize' % arg return retval return Diagonalize( arg )
[docs]def concatenate( args, axis=0 ): 'concatenate' args = _matchndim( *args ) axis = numeric.normdim( args[0].ndim, axis ) i = 0 if all( _iszero(arg) for arg in args ): shape = list( args[0].shape ) axis = numeric.normdim( len(shape), axis ) for arg in args[1:]: for i in range( len(shape) ): if i == axis: shape[i] += arg.shape[i] elif shape[i] == 1: shape[i] = arg.shape[i] else: assert arg.shape[i] in (shape[i],1) return _zeros( shape ) while i+1 < len(args): arg1, arg2 = args[i:i+2] if not _isfunc(arg1) and not _isfunc(arg2): arg12 = numpy.concatenate( [ arg1, arg2 ], axis ) else: arg12 = _call( arg1, '_concatenate', arg2, axis ) if arg12 is None: i += 1 continue args = args[:i] + [arg12] + args[i+2:] if len(args) == 1: return args[0] return Concatenate( args, axis )
[docs]def transpose( arg, trans=None ): 'transpose' arg = asarray( arg ) if trans is None: invtrans = range( arg.ndim-1, -1, -1 ) else: trans = _normdims( arg.ndim, trans ) assert sorted(trans) == range(arg.ndim) invtrans = numpy.empty( arg.ndim, dtype=int ) invtrans[ numpy.asarray(trans) ] = numpy.arange( arg.ndim ) return align( arg, invtrans, arg.ndim )
[docs]def product( arg, axis ): 'product' arg = asarray( arg ) axis = numeric.normdim( arg.ndim, axis ) shape = arg.shape[:axis] + arg.shape[axis+1:] if arg.shape[axis] == 1: return get( arg, axis, 0 ) if not _isfunc( arg ): return numpy.product( arg, axis ) retval = _call( arg, '_product', axis ) if retval is not None: assert retval.shape == shape, 'bug in %s._product' % arg return retval return Product( arg, axis )
[docs]def choose( level, choices ): 'choose' choices = _matchndim( *choices ) if _isfunc(level) or any( _isfunc(choice) for choice in choices ): return Choose( level, choices ) return numpy.choose( level, choices )
[docs]def cross( arg1, arg2, axis ): 'cross product' arg1, arg2 = _matchndim( arg1, arg2 ) shape = _jointshape( arg1.shape, arg2.shape ) axis = numeric.normdim( len(shape), axis ) if not _isfunc(arg1) and not _isfunc(arg2): return numeric.cross(arg1,arg2,axis) retval = _call( arg1, '_cross', arg2, axis ) if retval is not None: assert retval.shape == shape, 'bug in %s._cross' % arg1 return retval retval = _call( arg2, '_cross', arg1, axis ) if retval is not None: assert retval.shape == shape, 'bug in %s._cross' % arg2 return -retval return Cross( arg1, arg2, axis )
[docs]def outer( arg1, arg2=None, axis=0 ): 'outer product' arg1, arg2 = _matchndim( arg1, arg2 if arg2 is not None else arg1 ) axis = numeric.normdim( arg1.ndim, axis ) return insert(arg1,axis+1) * insert(arg2,axis)
[docs]def pointwise( args, evalf, deriv ): 'general pointwise operation' args = asarray( _matchndim(*args) ) if _isfunc(args): return Pointwise( args, evalf, deriv ) return evalf( *args )
[docs]def multiply( arg1, arg2 ): 'multiply' arg1, arg2 = _matchndim( arg1, arg2 ) shape = _jointshape( arg1.shape, arg2.shape ) if _isunit( arg1 ): return expand( arg2, shape ) if _isunit( arg2 ): return expand( arg1, shape ) if not _isfunc(arg1) and not _isfunc(arg2): return numpy.multiply( arg1, arg2 ) if _equal( arg1, arg2 ): return power( arg1, 2 ) for idim, sh in enumerate( shape ): if sh == 1: return insert( multiply( get(arg1,idim,0), get(arg2,idim,0) ), idim ) retval = _call( arg1, '_multiply', arg2 ) if retval is not None: assert retval.shape == shape, 'bug in %s._multiply' % arg1 return retval retval = _call( arg2, '_multiply', arg1 ) if retval is not None: assert retval.shape == shape, 'bug in %s._multiply' % arg2 return retval return Multiply( arg1, arg2 )
[docs]def add( arg1, arg2 ): 'add' arg1, arg2 = _matchndim( arg1, arg2 ) shape = _jointshape( arg1.shape, arg2.shape ) if _iszero( arg1 ): return expand( arg2, shape ) if _iszero( arg2 ): return expand( arg1, shape ) if not _isfunc(arg1) and not _isfunc(arg2): return numpy.add( arg1, arg2 ) if _equal( arg1, arg2 ): return arg1 * 2 for idim, sh in enumerate( shape ): if sh == 1: return insert( add( get(arg1,idim,0), get(arg2,idim,0) ), idim ) retval = _call( arg1, '_add', arg2 ) if retval is not None: assert retval.shape == shape, 'bug in %s._add' % arg1 return retval retval = _call( arg2, '_add', arg1 ) if retval is not None: assert retval.shape == shape, 'bug in %s._add' % arg2 return retval return Add( arg1, arg2 )
[docs]def negative( arg ): 'make negative' arg = asarray(arg) if not _isfunc( arg ): return numpy.negative( arg ) retval = _call( arg, '_negative' ) if retval is not None: assert retval.shape == arg.shape, 'bug in %s._negative' % arg return retval return Negative( arg )
[docs]def power( arg, n ): 'power' arg = asarray( arg ) assert _isscalar( n ) if n == 1: return arg if n == 0: return numpy.ones( arg.shape ) if isinstance( arg, numpy.ndarray ): return numpy.power( arg, n ) retval = _call( arg, '_power', n ) if retval is not None: assert retval.shape == arg.shape, 'bug in %s._power' % arg return retval return Power( arg, n )
[docs]def sign( arg ): 'sign' arg = asarray( arg ) if isinstance( arg, numpy.ndarray ): return numpy.sign( arg ) retval = _call( arg, '_sign' ) if retval is not None: assert retval.shape == arg.shape, 'bug in %s._sign' % arg return retval return Sign( arg )
[docs]def eig( arg, axes=(-2,-1), symmetric=False, sort=False ): '''eig( arg, axes [ symmetric ] ) Compute the eigenvalues and vectors of a matrix. The eigenvalues and vectors are positioned on the last axes. * tuple axes The axis on which the eigenvalues and vectors are calculated * bool symmetric Is the matrix symmetric * int sort Sort the eigenvalues and vectors (-1=descending, 0=unsorted, 1=ascending)''' # Sort axis arg = asarray( arg ) ax1, ax2 = _norm_and_sort( arg.ndim, axes ) assert ax2 > ax1 # strict # Check if the matrix is square assert arg.shape[ax1] == arg.shape[ax2] # Move the axis with matrices trans = range(ax1) + [-2] + range(ax1,ax2-1) + [-1] + range(ax2-1,arg.ndim-2) arg = align( arg, trans, arg.ndim ) shapeval = arg.shape[:-1] shapevec = arg.shape # When it's an array calculate directly if not _isfunc(arg): eigval, eigvec = numeric.eigh( arg, sort ) if symmetric else numeric.eig( arg, sort ) return eigval, eigvec # Use _call to see if the object has its own _eig function ret = _call( arg, '_eig' ) if ret is not None: # Check the shapes eigval, eigvec = ret assert eigval.shape == shapeval, 'bug in %s._eig' % arg assert eigvec.shape == shapevec, 'bug in %s._eig' % arg else: eig = Eig( arg, symmetric=symmetric, sort=sort ) eigval = ArrayFromTuple( eig, index=0, shape=arg.shape[:-1] ) eigvec = ArrayFromTuple( eig, index=1, shape=arg.shape ) # Return the evaluable function objects in a tuple like numpy return eigval, eigvec
nsymgrad = lambda arg, coords: ( symgrad(arg,coords) * coords.normal() ).sum() ngrad = lambda arg, coords: ( grad(arg,coords) * coords.normal() ).sum() sin = lambda arg: pointwise( [arg], numpy.sin, cos ) cos = lambda arg: pointwise( [arg], numpy.cos, lambda x: -sin(x) ) tan = lambda arg: pointwise( [arg], numpy.tan, lambda x: cos(x)**-2 ) arcsin = lambda arg: pointwise( [arg], numpy.arcsin, lambda x: reciprocal(sqrt(1-x**2)) ) arccos = lambda arg: pointwise( [arg], numpy.arccos, lambda x: -reciprocal(sqrt(1-x**2)) ) exp = lambda arg: pointwise( [arg], numpy.exp, exp ) ln = lambda arg: pointwise( [arg], numpy.log, reciprocal ) log2 = lambda arg: ln(arg) / ln(2) log10 = lambda arg: ln(arg) / ln(10) sqrt = lambda arg: power( arg, .5 ) reciprocal = lambda arg: power( arg, -1 ) argmin = lambda arg, axis: pointwise( bringforward(arg,axis), lambda *x: numpy.argmin(numeric.stack(x),axis=0), _zeros_like ) argmax = lambda arg, axis: pointwise( bringforward(arg,axis), lambda *x: numpy.argmax(numeric.stack(x),axis=0), _zeros_like ) arctan2 = lambda arg1, arg2=None: pointwise( arg1 if arg2 is None else [arg1,arg2], numpy.arctan2, lambda x: stack([x[1],-x[0]]) / sum(power(x,2),0) ) greater = lambda arg1, arg2=None: pointwise( arg1 if arg2 is None else [arg1,arg2], numpy.greater, _zeros_like ) less = lambda arg1, arg2=None: pointwise( arg1 if arg2 is None else [arg1,arg2], numpy.less, _zeros_like ) min = lambda arg1, *args: choose( argmin( arg1 if not args else (arg1,)+args, axis=0 ), arg1 if not args else (arg1,)+args ) max = lambda arg1, *args: choose( argmax( arg1 if not args else (arg1,)+args, axis=0 ), arg1 if not args else (arg1,)+args ) abs = lambda arg: arg * sign(arg) sinh = lambda arg: .5 * ( exp(arg) - exp(-arg) ) cosh = lambda arg: .5 * ( exp(arg) + exp(-arg) ) tanh = lambda arg: 1 - 2. / ( exp(2*arg) + 1 ) arctanh = lambda arg: .5 * ( ln(1+arg) - ln(1-arg) ) piecewise = lambda level, intervals, *funcs: choose( sum( greater( insert(level,-1), intervals ) ), funcs ) trace = lambda arg, n1=-2, n2=-1: sum( takediag( arg, n1, n2 ) ) eye = lambda n: diagonalize( expand( [1.], (n,) ) ) norm2 = lambda arg, axis=-1: sqrt( sum( arg * arg, axis ) ) heaviside = lambda arg: choose( greater( arg, 0 ), [0.,1.] ) divide = lambda arg1, arg2: multiply( arg1, reciprocal(arg2) ) subtract = lambda arg1, arg2: add( arg1, negative(arg2) ) mean = lambda arg: .5 * ( arg + opposite(arg) ) jump = lambda arg: arg - opposite(arg) add_T = lambda arg, axes=(-2,-1): swapaxes( arg, axes ) + arg
[docs]def swapaxes( arg, axes=(-2,-1) ): 'swap axes' arg = asarray( arg ) n1, n2 = axes trans = numpy.arange( arg.ndim ) trans[n1] = numeric.normdim( arg.ndim, n2 ) trans[n2] = numeric.normdim( arg.ndim, n1 ) return align( arg, trans, arg.ndim )
[docs]def opposite( arg ): 'evaluate jump over interface' if not isinstance( arg, Evaluable ): return arg return arg._opposite()
[docs]def function( fmap, nmap, ndofs, ndims ): 'create function on ndims-element' axis = '~%d' % ndofs cascade = Cascade(ndims) func = Function( cascade, fmap, igrad=0, axis=axis ) dofmap = DofMap( cascade, nmap, axis=axis ) return Inflate( func, dofmap, length=ndofs, axis=0 )
[docs]def take( arg, index, axis ): 'take index' arg = asarray( arg ) axis = numeric.normdim( arg.ndim, axis ) if isinstance( index, slice ): assert index.start == None or index.start >= 0 assert index.stop != None and index.stop > 0 index = numpy.arange( index.start or 0, index.stop, index.step ) assert index.size > 0 elif not _isfunc( index ): index = numpy.asarray( index, dtype=int ) assert numpy.all( index >= 0 ) assert index.size > 0 assert index.ndim == 1 if arg.shape[axis] == 1: return repeat( arg, index.shape[0], axis ) if not _isfunc( index ): allindices = numpy.arange( arg.shape[axis] ) index = allindices[index] if numpy.all( index == allindices ): return arg if index.shape[0] == 1: return insert( get( arg, axis, index[0] ), axis ) retval = _call( arg, '_take', index, axis ) if retval is not None: return retval if _isfunc( index ): return DofIndex( arg, axis, index ) if not _isfunc( arg ): return numpy.take( arg, index, axis ) return Take( arg, index, axis )
[docs]def inflate( arg, dofmap, length, axis ): 'inflate' arg = asarray( arg ) axis = numeric.normdim( arg.ndim, axis ) assert not isinstance( arg.shape[axis], int ) retval = _call( arg, '_inflate', dofmap, length, axis ) if retval is not None: return retval return Inflate( arg, dofmap, length, axis )
def blocks( arg ): arg = asarray( arg ) try: blocks = arg.blocks except AttributeError: blocks = [( arg, tuple( numpy.arange(n) if isinstance(n,int) else None for n in arg.shape ) )] return blocks
[docs]def pointdata ( topo, ischeme, func=None, shape=None, value=None ): 'point data' from . import topology assert isinstance(topo,topology.Topology) if func is not None: assert value is None assert shape is None shape = func.shape else: # func is None if value is not None: assert shape is None value = numpy.asarray( value ) else: # value is None assert shape is not None value = numpy.zeros( shape ) shape = value.shape data = {} for elem in topo: ipoints, iweights = elem.eval( ischeme ) values = numpy.empty( ipoints.shape[:-1]+shape, dtype=float ) values[:] = func(elem,ischeme) if func is not None else value data[ elem ] = values, ipoints return Pointdata( data, shape )
[docs]def fdapprox( func, w, dofs, delta=1.e-5 ): '''Finite difference approximation of the variation of func in directions w around dofs. Input arguments: * func, the functional to differentiate * dofs, DOF vector of linearization point * w, the function space or a tuple of chained spaces * delta, finite difference step scaling of ||dofs||_inf''' log.context( 'FD approx' ) if not isinstance( w, tuple ): w = w, x0 = tuple( wi.dot( dofs ) for wi in w ) step = numpy.linalg.norm( dofs, numpy.inf )*delta ndofs = len( dofs ) dfunc_fd = [] for i in log.iterate( 'dof', range(ndofs) ): pert = dofs.copy() pert[i] += step x1 = tuple( wi.dot( pert ) for wi in w ) dfunc_fd.append( (func( *x1 ) - func( *x0 ))/step ) return dfunc_fd
[docs]def iwscale( coords, ndims ): 'integration weights scale' assert coords.ndim == 1 J = localgradient( coords, ndims ) cndims, = coords.shape assert J.shape == (cndims,ndims), 'wrong jacobian shape: got %s, expected %s' % ( J.shape, (cndims, ndims) ) if cndims == ndims: detJ = determinant( J ) elif ndims == 1: detJ = norm2( J[:,0], axis=0 ) elif cndims == 3 and ndims == 2: detJ = norm2( cross( J[:,0], J[:,1], axis=0 ), axis=0 ) elif ndims == 0: detJ = 1. else: raise NotImplementedError, 'cannot compute determinant for %dx%d jacobian' % J.shape[:2] return detJ
[docs]def supp( funcsp, indices ): 'find support of selection of basis functions' # FRAGILE! makes lots of assumptions on the nature of funcsp supp = [] for func, axes in funcsp.blocks: dofmap = axes[0].dofmap stdmap = func.stdmap for elem, dofs in dofmap.items(): while True: std = stdmap.get( elem ) nshapes = 0 if not std \ else std[1].sum() if isinstance( std, tuple ) \ else std.nshapes if numpy.intersect1d( dofs[:nshapes], indices, assume_unique=True ).size: supp.append( elem ) dofs = dofs[nshapes:] if not elem.parent: break elem, trans = elem.parent assert not dofs.size return supp # CBUILDER # # Set of tools for just-in-time compilation of evaluables. Prototyping stage.
def _cshape( shape ): return [ sh if isinstance(sh,int) else 'n%d' % i for i, sh in enumerate(shape) ] def _cstrides( shape ): if not shape: return [] strides = [1] for sh in _cshape(shape)[:0:-1]: strides.append( '%s*%s' % ( strides[-1], sh ) ) return strides[::-1] def _cindex( myshape, shape ): assert len(shape) == len(myshape) index = [] for i, stride in enumerate( _cstrides(myshape) ): if myshape[i] == shape[i]: index.append( '%s*i%d' % (stride,i) ) else: assert myshape[i] == 1 return '+'.join( index ) or '0' def _cfor( varname, count, do ): return 'for ( %s=0; %s<%s; %s++ ) { %s }' % ( varname, varname, count, varname, do ) def _cfunc( args, body ): return '( %s )\n{ %s }\n' % ( ', '.join( args ), body )
[docs]class CBuilder( object ): 'cbuilder' def __init__( self, cachedir='/tmp/nutils' ): from cffi import FFI self.ffi = FFI() self.codebase = {} self.allcode = [] self.cachedir = cachedir def compile( self ): log.context( 'c-builder' ) if not self.codebase: return import os, hashlib code = '\n'.join( self.allcode ) fname = os.path.join( self.cachedir, hashlib.md5( code ).hexdigest() ) if os.path.isfile( fname+'.so' ): log.info( 'in cache' ) else: log.info( 'compiling %d c-functions' % len(self.codebase) ) if not os.path.isdir( self.cachedir ): os.mkdir( self.cachedir ) open( fname+'.c', 'w' ).write( code ) assert os.system( 'gcc -fpic -g -c -o %s.o -O3 -march=native -mtune=native %s.c' % ((fname,)*2) ) == 0 assert os.system( 'gcc -shared -Wl,-soname,%s.so -o %s.so %s.o -lc' % ((fname,)*3) ) == 0 os.unlink( fname+'.c' ) os.unlink( fname+'.o' ) self.so = self.ffi.dlopen( fname+'.so' ) def cast( self, args, ctypes ): assert len(args) == len(ctypes) for arg, ctype in zip( args, ctypes ): if ctype[-1] == '*': assert isinstance( arg, numpy.ndarray ) # assert arg.flags.c_contiguous # TODO check dtype arg = arg.ctypes.data yield self.ffi.cast( ctype, arg ) def __getitem__( self, (code,init) ): evalf = self.codebase.get(code) if evalf: return evalf arglist, body = code.split( '\n', 1 ) assert arglist[0] == '(' and arglist[-1] == ')' ctypes = [ item.strip().rsplit(' ',1)[0] for item in arglist[1:-1].split( ',' ) ] handle = 'func%d' % len(self.codebase) head = 'void ' + handle self.ffi.cdef( head + arglist + ';' ) def evalf( *args ): cfunc = getattr( self.so, handle ) retval, iterargs = init( args ) for args in iterargs: cfunc( *self.cast( args, ctypes ) ) return retval self.codebase[code] = evalf self.allcode.append( head + code ) return evalf
def init_args_for( RETVAL, *ARGS ): shape = RETVAL.shape ndim = _max( ARG.ndim for ARG in ARGS if ARG is not None ) jointshape = _jointshape( *[ ARG.shape for ARG in ARGS if ARG is not None ] ) def init( args ): assert len(args) == len(ARGS) # 'None' ARGS are filtered out; values are built into the c code args = [ numpy.ascontiguousarray( arg, dtype=ARG.dtype ) if arg.ndim else arg for arg, ARG in zip( args, ARGS ) if ARG is not None ] haspoints = _max( [ arg.ndim for arg in args ] ) - ndim assert haspoints in (0,1) if haspoints: neval = _max( [ arg.shape[0] for arg in args if arg.ndim > ndim ] ) retshape = [ neval ] else: neval = 1 retshape = [] assert all( arg.shape[0] == neval for arg in args if arg.ndim > ndim ) varsizes = [] for idim, sh in enumerate( jointshape ): if not isinstance(sh,int): for arg in args: sh = arg.shape[idim-len(jointshape)] if sh != 1: break varsizes.append( sh ) if idim < len(shape): retshape.append( sh ) retval = numpy.empty( retshape, dtype=RETVAL.dtype ) args = [ retval[...,_] if haspoints else [retval[...,_]] * neval ] \ + [ [arg[...,_]] * neval if arg.ndim == ndim else arg[...,_] for arg in args ] args = [ arg + tuple(varsizes) for arg in zip( *args ) ] return retval, args return init # vim:shiftwidth=2:foldmethod=indent:foldnestmax=2