|
|
import functools, itertools |
|
|
from sympy.core.sympify import _sympify, sympify |
|
|
from sympy.core.expr import Expr |
|
|
from sympy.core import Basic, Tuple |
|
|
from sympy.tensor.array import ImmutableDenseNDimArray |
|
|
from sympy.core.symbol import Symbol |
|
|
from sympy.core.numbers import Integer |
|
|
|
|
|
|
|
|
class ArrayComprehension(Basic): |
|
|
""" |
|
|
Generate a list comprehension. |
|
|
|
|
|
Explanation |
|
|
=========== |
|
|
|
|
|
If there is a symbolic dimension, for example, say [i for i in range(1, N)] where |
|
|
N is a Symbol, then the expression will not be expanded to an array. Otherwise, |
|
|
calling the doit() function will launch the expansion. |
|
|
|
|
|
Examples |
|
|
======== |
|
|
|
|
|
>>> from sympy.tensor.array import ArrayComprehension |
|
|
>>> from sympy import symbols |
|
|
>>> i, j, k = symbols('i j k') |
|
|
>>> a = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3)) |
|
|
>>> a |
|
|
ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3)) |
|
|
>>> a.doit() |
|
|
[[11, 12, 13], [21, 22, 23], [31, 32, 33], [41, 42, 43]] |
|
|
>>> b = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, k)) |
|
|
>>> b.doit() |
|
|
ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, k)) |
|
|
""" |
|
|
def __new__(cls, function, *symbols, **assumptions): |
|
|
if any(len(l) != 3 or None for l in symbols): |
|
|
raise ValueError('ArrayComprehension requires values lower and upper bound' |
|
|
' for the expression') |
|
|
arglist = [sympify(function)] |
|
|
arglist.extend(cls._check_limits_validity(function, symbols)) |
|
|
obj = Basic.__new__(cls, *arglist, **assumptions) |
|
|
obj._limits = obj._args[1:] |
|
|
obj._shape = cls._calculate_shape_from_limits(obj._limits) |
|
|
obj._rank = len(obj._shape) |
|
|
obj._loop_size = cls._calculate_loop_size(obj._shape) |
|
|
return obj |
|
|
|
|
|
@property |
|
|
def function(self): |
|
|
"""The function applied across limits. |
|
|
|
|
|
Examples |
|
|
======== |
|
|
|
|
|
>>> from sympy.tensor.array import ArrayComprehension |
|
|
>>> from sympy import symbols |
|
|
>>> i, j = symbols('i j') |
|
|
>>> a = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3)) |
|
|
>>> a.function |
|
|
10*i + j |
|
|
""" |
|
|
return self._args[0] |
|
|
|
|
|
@property |
|
|
def limits(self): |
|
|
""" |
|
|
The list of limits that will be applied while expanding the array. |
|
|
|
|
|
Examples |
|
|
======== |
|
|
|
|
|
>>> from sympy.tensor.array import ArrayComprehension |
|
|
>>> from sympy import symbols |
|
|
>>> i, j = symbols('i j') |
|
|
>>> a = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3)) |
|
|
>>> a.limits |
|
|
((i, 1, 4), (j, 1, 3)) |
|
|
""" |
|
|
return self._limits |
|
|
|
|
|
@property |
|
|
def free_symbols(self): |
|
|
""" |
|
|
The set of the free_symbols in the array. |
|
|
Variables appeared in the bounds are supposed to be excluded |
|
|
from the free symbol set. |
|
|
|
|
|
Examples |
|
|
======== |
|
|
|
|
|
>>> from sympy.tensor.array import ArrayComprehension |
|
|
>>> from sympy import symbols |
|
|
>>> i, j, k = symbols('i j k') |
|
|
>>> a = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3)) |
|
|
>>> a.free_symbols |
|
|
set() |
|
|
>>> b = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, k+3)) |
|
|
>>> b.free_symbols |
|
|
{k} |
|
|
""" |
|
|
expr_free_sym = self.function.free_symbols |
|
|
for var, inf, sup in self._limits: |
|
|
expr_free_sym.discard(var) |
|
|
curr_free_syms = inf.free_symbols.union(sup.free_symbols) |
|
|
expr_free_sym = expr_free_sym.union(curr_free_syms) |
|
|
return expr_free_sym |
|
|
|
|
|
@property |
|
|
def variables(self): |
|
|
"""The tuples of the variables in the limits. |
|
|
|
|
|
Examples |
|
|
======== |
|
|
|
|
|
>>> from sympy.tensor.array import ArrayComprehension |
|
|
>>> from sympy import symbols |
|
|
>>> i, j, k = symbols('i j k') |
|
|
>>> a = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3)) |
|
|
>>> a.variables |
|
|
[i, j] |
|
|
""" |
|
|
return [l[0] for l in self._limits] |
|
|
|
|
|
@property |
|
|
def bound_symbols(self): |
|
|
"""The list of dummy variables. |
|
|
|
|
|
Note |
|
|
==== |
|
|
|
|
|
Note that all variables are dummy variables since a limit without |
|
|
lower bound or upper bound is not accepted. |
|
|
""" |
|
|
return [l[0] for l in self._limits if len(l) != 1] |
|
|
|
|
|
@property |
|
|
def shape(self): |
|
|
""" |
|
|
The shape of the expanded array, which may have symbols. |
|
|
|
|
|
Note |
|
|
==== |
|
|
|
|
|
Both the lower and the upper bounds are included while |
|
|
calculating the shape. |
|
|
|
|
|
Examples |
|
|
======== |
|
|
|
|
|
>>> from sympy.tensor.array import ArrayComprehension |
|
|
>>> from sympy import symbols |
|
|
>>> i, j, k = symbols('i j k') |
|
|
>>> a = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3)) |
|
|
>>> a.shape |
|
|
(4, 3) |
|
|
>>> b = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, k+3)) |
|
|
>>> b.shape |
|
|
(4, k + 3) |
|
|
""" |
|
|
return self._shape |
|
|
|
|
|
@property |
|
|
def is_shape_numeric(self): |
|
|
""" |
|
|
Test if the array is shape-numeric which means there is no symbolic |
|
|
dimension. |
|
|
|
|
|
Examples |
|
|
======== |
|
|
|
|
|
>>> from sympy.tensor.array import ArrayComprehension |
|
|
>>> from sympy import symbols |
|
|
>>> i, j, k = symbols('i j k') |
|
|
>>> a = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3)) |
|
|
>>> a.is_shape_numeric |
|
|
True |
|
|
>>> b = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, k+3)) |
|
|
>>> b.is_shape_numeric |
|
|
False |
|
|
""" |
|
|
for _, inf, sup in self._limits: |
|
|
if Basic(inf, sup).atoms(Symbol): |
|
|
return False |
|
|
return True |
|
|
|
|
|
def rank(self): |
|
|
"""The rank of the expanded array. |
|
|
|
|
|
Examples |
|
|
======== |
|
|
|
|
|
>>> from sympy.tensor.array import ArrayComprehension |
|
|
>>> from sympy import symbols |
|
|
>>> i, j, k = symbols('i j k') |
|
|
>>> a = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3)) |
|
|
>>> a.rank() |
|
|
2 |
|
|
""" |
|
|
return self._rank |
|
|
|
|
|
def __len__(self): |
|
|
""" |
|
|
The length of the expanded array which means the number |
|
|
of elements in the array. |
|
|
|
|
|
Raises |
|
|
====== |
|
|
|
|
|
ValueError : When the length of the array is symbolic |
|
|
|
|
|
Examples |
|
|
======== |
|
|
|
|
|
>>> from sympy.tensor.array import ArrayComprehension |
|
|
>>> from sympy import symbols |
|
|
>>> i, j = symbols('i j') |
|
|
>>> a = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3)) |
|
|
>>> len(a) |
|
|
12 |
|
|
""" |
|
|
if self._loop_size.free_symbols: |
|
|
raise ValueError('Symbolic length is not supported') |
|
|
return self._loop_size |
|
|
|
|
|
@classmethod |
|
|
def _check_limits_validity(cls, function, limits): |
|
|
|
|
|
new_limits = [] |
|
|
for var, inf, sup in limits: |
|
|
var = _sympify(var) |
|
|
inf = _sympify(inf) |
|
|
|
|
|
|
|
|
if isinstance(sup, list): |
|
|
sup = Tuple(*sup) |
|
|
else: |
|
|
sup = _sympify(sup) |
|
|
new_limits.append(Tuple(var, inf, sup)) |
|
|
if any((not isinstance(i, Expr)) or i.atoms(Symbol, Integer) != i.atoms() |
|
|
for i in [inf, sup]): |
|
|
raise TypeError('Bounds should be an Expression(combination of Integer and Symbol)') |
|
|
if (inf > sup) == True: |
|
|
raise ValueError('Lower bound should be inferior to upper bound') |
|
|
if var in inf.free_symbols or var in sup.free_symbols: |
|
|
raise ValueError('Variable should not be part of its bounds') |
|
|
return new_limits |
|
|
|
|
|
@classmethod |
|
|
def _calculate_shape_from_limits(cls, limits): |
|
|
return tuple([sup - inf + 1 for _, inf, sup in limits]) |
|
|
|
|
|
@classmethod |
|
|
def _calculate_loop_size(cls, shape): |
|
|
if not shape: |
|
|
return 0 |
|
|
loop_size = 1 |
|
|
for l in shape: |
|
|
loop_size = loop_size * l |
|
|
|
|
|
return loop_size |
|
|
|
|
|
def doit(self, **hints): |
|
|
if not self.is_shape_numeric: |
|
|
return self |
|
|
|
|
|
return self._expand_array() |
|
|
|
|
|
def _expand_array(self): |
|
|
res = [] |
|
|
for values in itertools.product(*[range(inf, sup+1) |
|
|
for var, inf, sup |
|
|
in self._limits]): |
|
|
res.append(self._get_element(values)) |
|
|
|
|
|
return ImmutableDenseNDimArray(res, self.shape) |
|
|
|
|
|
def _get_element(self, values): |
|
|
temp = self.function |
|
|
for var, val in zip(self.variables, values): |
|
|
temp = temp.subs(var, val) |
|
|
return temp |
|
|
|
|
|
def tolist(self): |
|
|
"""Transform the expanded array to a list. |
|
|
|
|
|
Raises |
|
|
====== |
|
|
|
|
|
ValueError : When there is a symbolic dimension |
|
|
|
|
|
Examples |
|
|
======== |
|
|
|
|
|
>>> from sympy.tensor.array import ArrayComprehension |
|
|
>>> from sympy import symbols |
|
|
>>> i, j = symbols('i j') |
|
|
>>> a = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3)) |
|
|
>>> a.tolist() |
|
|
[[11, 12, 13], [21, 22, 23], [31, 32, 33], [41, 42, 43]] |
|
|
""" |
|
|
if self.is_shape_numeric: |
|
|
return self._expand_array().tolist() |
|
|
|
|
|
raise ValueError("A symbolic array cannot be expanded to a list") |
|
|
|
|
|
def tomatrix(self): |
|
|
"""Transform the expanded array to a matrix. |
|
|
|
|
|
Raises |
|
|
====== |
|
|
|
|
|
ValueError : When there is a symbolic dimension |
|
|
ValueError : When the rank of the expanded array is not equal to 2 |
|
|
|
|
|
Examples |
|
|
======== |
|
|
|
|
|
>>> from sympy.tensor.array import ArrayComprehension |
|
|
>>> from sympy import symbols |
|
|
>>> i, j = symbols('i j') |
|
|
>>> a = ArrayComprehension(10*i + j, (i, 1, 4), (j, 1, 3)) |
|
|
>>> a.tomatrix() |
|
|
Matrix([ |
|
|
[11, 12, 13], |
|
|
[21, 22, 23], |
|
|
[31, 32, 33], |
|
|
[41, 42, 43]]) |
|
|
""" |
|
|
from sympy.matrices import Matrix |
|
|
|
|
|
if not self.is_shape_numeric: |
|
|
raise ValueError("A symbolic array cannot be expanded to a matrix") |
|
|
if self._rank != 2: |
|
|
raise ValueError('Dimensions must be of size of 2') |
|
|
|
|
|
return Matrix(self._expand_array().tomatrix()) |
|
|
|
|
|
|
|
|
def isLambda(v): |
|
|
LAMBDA = lambda: 0 |
|
|
return isinstance(v, type(LAMBDA)) and v.__name__ == LAMBDA.__name__ |
|
|
|
|
|
class ArrayComprehensionMap(ArrayComprehension): |
|
|
''' |
|
|
A subclass of ArrayComprehension dedicated to map external function lambda. |
|
|
|
|
|
Notes |
|
|
===== |
|
|
|
|
|
Only the lambda function is considered. |
|
|
At most one argument in lambda function is accepted in order to avoid ambiguity |
|
|
in value assignment. |
|
|
|
|
|
Examples |
|
|
======== |
|
|
|
|
|
>>> from sympy.tensor.array import ArrayComprehensionMap |
|
|
>>> from sympy import symbols |
|
|
>>> i, j, k = symbols('i j k') |
|
|
>>> a = ArrayComprehensionMap(lambda: 1, (i, 1, 4)) |
|
|
>>> a.doit() |
|
|
[1, 1, 1, 1] |
|
|
>>> b = ArrayComprehensionMap(lambda a: a+1, (j, 1, 4)) |
|
|
>>> b.doit() |
|
|
[2, 3, 4, 5] |
|
|
|
|
|
''' |
|
|
def __new__(cls, function, *symbols, **assumptions): |
|
|
if any(len(l) != 3 or None for l in symbols): |
|
|
raise ValueError('ArrayComprehension requires values lower and upper bound' |
|
|
' for the expression') |
|
|
|
|
|
if not isLambda(function): |
|
|
raise ValueError('Data type not supported') |
|
|
|
|
|
arglist = cls._check_limits_validity(function, symbols) |
|
|
obj = Basic.__new__(cls, *arglist, **assumptions) |
|
|
obj._limits = obj._args |
|
|
obj._shape = cls._calculate_shape_from_limits(obj._limits) |
|
|
obj._rank = len(obj._shape) |
|
|
obj._loop_size = cls._calculate_loop_size(obj._shape) |
|
|
obj._lambda = function |
|
|
return obj |
|
|
|
|
|
@property |
|
|
def func(self): |
|
|
class _(ArrayComprehensionMap): |
|
|
def __new__(cls, *args, **kwargs): |
|
|
return ArrayComprehensionMap(self._lambda, *args, **kwargs) |
|
|
return _ |
|
|
|
|
|
def _get_element(self, values): |
|
|
temp = self._lambda |
|
|
if self._lambda.__code__.co_argcount == 0: |
|
|
temp = temp() |
|
|
elif self._lambda.__code__.co_argcount == 1: |
|
|
temp = temp(functools.reduce(lambda a, b: a*b, values)) |
|
|
return temp |
|
|
|
