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