sympy/core/multidimensional.py

   1 """
   2 Provides functionality for multidimensional usage of scalar-functions.
   3
   4 Read the vectorize docstring for more details.
   5 """
   6
   7 def apply_on_element(f, args, kwargs, n):
   8     """
   9     Returns a structure with the same dimension as the specified argument,
  10     where each basic element is replaced by the function f applied on it. All
  11     other arguments stay the same.
  12     """
  13     # Get the specified argument.
  14     if isinstance(n, int):
  15         structure = args[n]
  16         is_arg = True
  17     elif isinstance(n, str):
  18         structure = kwargs[n]
  19         is_arg = False
  20
  21     # Define reduced function that is only dependend of the specified argument.
  22     def f_reduced(x):
  23         if hasattr(x, "__iter__"):
  24             return map(f_reduced, x)
  25         else:
  26             if is_arg:
  27                 args[n] = x
  28             else:
  29                 kwargs[n] = x
  30             return f(*args, **kwargs)
  31
  32     # f_reduced will call itself recursively so that in the end f is applied to
  33     # all basic elements.
  34     return map(f_reduced, structure)
  35
  36 def iter_copy(structure):
  37     """
  38     Returns a copy of an iterable object (also copying all embedded iterables).
  39     """
  40     l = []
  41     for i in structure:
  42         if hasattr(i, "__iter__"):
  43             l.append(iter_copy(i))
  44         else:
  45             l.append(i)
  46     return l
  47
  48 def structure_copy(structure):
  49     """
  50     Returns a copy of the given structure (numpy-array, list, iterable, ..).
  51     """
  52     if hasattr(structure, "copy"):
  53         return structure.copy()
  54     return iter_copy(structure)
  55
  56
  57 class vectorize:
  58     """
  59     Generalizes a function taking scalars to accept multidimensional arguments.
  60
  61     For example:
  62         (1)
  63             @vectorize(0)
  64             def sin(x):
  65                 ....
  66
  67             >>sin([1, x, y])
  68             [sin(1), sin(x), sin(y)]
  69
  70         (2)
  71             @vectorize(0,1)
  72             def diff(f(y), y)
  73                 ....
  74
  75             >>diff([f(x,y,z),g(x,y,z),h(x,y,z)], [x,y,z])
  76             [[d/dx f, d/dy f, d/dz f],
  77              [d/dx g, d/dy g, d/dz g],
  78              [d/dx h, d/dy h, d/dz h]]
  79     """
  80     def __init__(self, *mdargs):
  81         """
  82         The given numbers and strings characterize the arguments that will be
  83         treated as data structures, where the decorated function will be applied
  84         to every single element.
  85         If no argument is given, everything is treated multidimensional.
  86         """
  87         for a in mdargs:
  88             assert isinstance(a, (int,str))
  89         self.mdargs = mdargs
  90
  91     def __call__(self, f):
  92         """
  93         Returns a wrapper for the one-dimensional function that can handle
  94         multidimensional arguments.
  95         """
  96         def wrapper(*args, **kwargs):
  97             # Get arguments that should be treated multidimensional
  98             if self.mdargs:
  99                 mdargs = self.mdargs
 100             else:
 101                 mdargs = range(len(args)) + kwargs.keys()
 102
 103             arglength = len(args)
 104
 105             for n in mdargs:
 106                 if isinstance(n, int):
 107                     if n>=arglength:
 108                         continue
 109                     entry = args[n]
 110                     is_arg = True
 111                 elif isinstance(n, str):
 112                     try:
 113                         entry = kwargs[n]
 114                     except KeyError:
 115                         continue
 116                     is_arg = False
 117                 if hasattr(entry, "__iter__"):
 118                     # Create now a copy of the given array and manipulate then
 119                     # the entries directly.
 120                     if is_arg:
 121                         args = list(args)
 122                         args[n] = structure_copy(entry)
 123                     else:
 124                         kwargs[n] = structure_copy(entry)
 125                     result = apply_on_element(wrapper, args, kwargs, n)
 126                     return result
 127             return f(*args, **kwargs)
 128         wrapper.__doc__ = f.__doc__
 129         wrapper.__name__ = f.__name__
 130         return wrapper