
    sg/                         d dl Z d dlZd dlmZmZ d dlmZ d dlmZm	Z	 d dl
mZ d dlmZ d dlmZ  G d d	e      Zd
 Z G d de      Zy)    N)_sympifysympify)Expr)BasicTuple)ImmutableDenseNDimArray)Symbol)Integerc                       e Zd ZdZd Zed        Zed        Zed        Zed        Z	ed        Z
ed        Zed	        Zd
 Zd Zed        Zed        Zed        Zd Zd Zd Zd Zd Zy)ArrayComprehensiona  
    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))
    c                    t        d |D              rt        d      t        |      g}|j                  | j	                  ||             t        j                  | g|i |}|j                  dd  |_        | j                  |j                        |_
        t        |j                        |_        | j                  |j                        |_        |S )Nc              3   @   K   | ]  }t        |      d k7  xs d  yw   Nlen.0ls     Y/var/www/html/venv/lib/python3.12/site-packages/sympy/tensor/array/array_comprehension.py	<genexpr>z-ArrayComprehension.__new__.<locals>.<genexpr>%         4qs1v{"d"4   KArrayComprehension requires values lower and upper bound for the expression   )any
ValueErrorr   extend_check_limits_validityr   __new___args_limits_calculate_shape_from_limits_shaper   _rank_calculate_loop_size
_loop_sizeclsfunctionsymbolsassumptionsarglistobjs         r   r    zArrayComprehension.__new__$   s    4G44 4 5 58$%s11(GDEmmC9'9[9iim55ckkB


O	11#**=
    c                      | j                   d   S )aA  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
        r   )r!   selfs    r   r*   zArrayComprehension.function1   s     zz!}r/   c                     | j                   S )au  
        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))
        r"   r1   s    r   limitszArrayComprehension.limitsA   s     ||r/   c                     | j                   j                  }| j                  D ]M  \  }}}|j                  |       |j                  j	                  |j                        }|j	                  |      }O |S )a)  
        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}
        )r*   free_symbolsr"   discardunion)r2   expr_free_symvarinfsupcurr_free_symss         r   r7   zArrayComprehension.free_symbolsR   sp    ( 22!\\ 	@MCc!!#& --33C4D4DEN)//?M	@ r/   c                 F    | j                   D cg c]  }|d   	 c}S c c}w )aL  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]
        r   r4   r2   r   s     r   	variableszArrayComprehension.variablesm   s      #ll+!+++s   c                 d    | j                   D cg c]  }t        |      dk7  s|d    c}S c c}w )zThe list of dummy variables.

        Note
        ====

        Note that all variables are dummy variables since a limit without
        lower bound or upper bound is not accepted.
        r   r   )r"   r   r@   s     r   bound_symbolsz ArrayComprehension.bound_symbols}   s*     #ll:c!fk!:::s   --c                     | j                   S )aE  
        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)
        )r$   r1   s    r   shapezArrayComprehension.shape   s    0 {{r/   c                 p    | j                   D ]'  \  }}}t        ||      j                  t              s' y y)a  
        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
        FT)r"   r   atomsr	   )r2   _r<   r=   s       r   is_shape_numericz#ArrayComprehension.is_shape_numeric   s9    &  << 	KAsCS#$$V,	 r/   c                     | j                   S )a9  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
        )r%   r1   s    r   rankzArrayComprehension.rank   s     zzr/   c                 \    | j                   j                  rt        d      | j                   S )a  
        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
        z Symbolic length is not supported)r'   r7   r   r1   s    r   __len__zArrayComprehension.__len__   s'    ( ??''?@@r/   c                 ~   g }|D ]  \  }}}t        |      }t        |      }t        |t              r	t        | }nt        |      }|j	                  t        |||             t        d ||fD              rt        d      ||kD  dk(  rt        d      ||j                  v s||j                  v st        d       |S )Nc              3      K   | ]B  }t        |t               xs+ |j                  t        t              |j                         k7   D y wN)
isinstancer   rG   r	   r
   )r   is     r   r   z<ArrayComprehension._check_limits_validity.<locals>.<genexpr>   sB      UDE #1d++U0HAGGI0UU Us   AA
zABounds should be an Expression(combination of Integer and Symbol)Tz-Lower bound should be inferior to upper boundz)Variable should not be part of its bounds)	r   rQ   listr   appendr   	TypeErrorr   r7   )r)   r*   r5   
new_limitsr;   r<   r=   s          r   r   z)ArrayComprehension._check_limits_validity   s     
# 	NMCc3-C3-C #t$SksmeCc23 UJMsU U cddc	d" !PQQc&&&#1A1A*A !LMM!	N" r/   c           
      ^    t        |D cg c]  \  }}}||z
  dz    c}}}      S c c}}}w Nr   )tuple)r)   r5   rH   r<   r=   s        r   r#   z/ArrayComprehension._calculate_shape_from_limits   s,    v>>3cCi!m>??>s   (c                 (    |syd}|D ]  }||z  }	 |S )Nr   r    )r)   rE   	loop_sizer   s       r   r&   z'ArrayComprehension._calculate_loop_size   s-    	 	&A!AI	& r/   c                 >    | j                   s| S | j                         S rP   )rI   _expand_array)r2   hintss     r   doitzArrayComprehension.doit  s    $$K!!##r/   c                    g }t        j                  | j                  D cg c]  \  }}}t        ||dz          c}}} D ]"  }|j	                  | j                  |             $ t        || j                        S c c}}}w rX   )	itertoolsproductr"   rangerT   _get_elementr   rE   )r2   resr;   r<   r=   valuess         r   r^   z ArrayComprehension._expand_array  s    ''+/<<*9 *9,9Cc +0SU*; *9 : 	2F JJt((01	2
 'sDJJ77*9s   A=c                 ~    | j                   }t        | j                  |      D ]  \  }}|j                  ||      } |S rP   )r*   ziprA   subs)r2   rg   tempr;   vals        r   re   zArrayComprehension._get_element  s>    }}DNNF3 	'HC99S#&D	'r/   c                 l    | j                   r| j                         j                         S t        d      )a  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]]
        z-A symbolic array cannot be expanded to a list)rI   r^   tolistr   r1   s    r   rn   zArrayComprehension.tolist  s1    $   %%'..00HIIr/   c                     ddl m} | j                  st        d      | j                  dk7  rt        d       || j                         j                               S )aE  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]])
        r   )Matrixz/A symbolic array cannot be expanded to a matrix   zDimensions must be of size of 2)sympy.matricesrp   rI   r   r%   r^   tomatrix)r2   rp   s     r   rs   zArrayComprehension.tomatrix3  sP    . 	*$$NOO::?>??d((*33566r/   N)__name__
__module____qualname____doc__r    propertyr*   r5   r7   rA   rC   rE   rI   rK   rM   classmethodr   r#   r&   r`   r^   re   rn   rs   r[   r/   r   r   r   
   s    2       4 , , 	; 	;  2  .0  , @ @  $8J.7r/   r   c                 h    d }t        | t        |            xr | j                  |j                  k(  S )Nc                       y)Nr   r[   r[   r/   r   <lambda>zisLambda.<locals>.<lambda>U  s    r/   )rQ   typert   )vLAMBDAs     r   isLambdar   T  s*    Faf&H1::+HHr/   c                   ,    e Zd ZdZd Zed        Zd Zy)ArrayComprehensionMapa[  
    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]

    c                    t        d |D              rt        d      t        |      st        d      | j                  ||      }t	        j
                  | g|i |}|j                  |_        | j                  |j                        |_	        t        |j                        |_        | j                  |j                        |_        ||_        |S )Nc              3   @   K   | ]  }t        |      d k7  xs d  ywr   r   r   s     r   r   z0ArrayComprehensionMap.__new__.<locals>.<genexpr>r  r   r   r   zData type not supported)r   r   r   r   r   r    r!   r"   r#   r$   r   r%   r&   r'   _lambdar(   s         r   r    zArrayComprehensionMap.__new__q  s    4G44 4 5 5 !677,,Xw?mmC9'9[9ii55ckkB


O	11#**=
r/   c                 *      G  fddt               }|S )Nc                       e Zd Z fdZy)%ArrayComprehensionMap.func.<locals>._c                 6    t        j                  g|i |S rP   )r   r   )r)   argskwargsr2   s      r   r    z-ArrayComprehensionMap.func.<locals>._.__new__  s    ,T\\KDKFKKr/   N)rt   ru   rv   r    r1   s   r   rH   r     s	    Lr/   rH   )r   )r2   rH   s   ` r   funczArrayComprehensionMap.func  s    	L% 	L r/   c                     | j                   }| j                   j                  j                  dk(  r	 |       }|S | j                   j                  j                  dk(  r |t        j                  d |            }|S )Nr   r   c                     | |z  S rP   r[   )abs     r   r|   z4ArrayComprehensionMap._get_element.<locals>.<lambda>  s
    ac r/   )r   __code__co_argcount	functoolsreduce)r2   rg   rk   s      r   re   z"ArrayComprehensionMap._get_element  sh    ||<<  ,,16D  \\""..!3	(()96BCDr/   N)rt   ru   rv   rw   r    rx   r   re   r[   r/   r   r   r   X  s%    0"  r/   r   )r   rb   sympy.core.sympifyr   r   sympy.core.exprr   
sympy.corer   r   sympy.tensor.arrayr   sympy.core.symbolr	   sympy.core.numbersr
   r   r   r   r[   r/   r   <module>r      s<     0   # 6 $ &G7 G7T
I7. 7r/   