Python sympy.flatten函数代码示例

def Tensor(*tensor_products):
    """Create a new tensor"""
    ## First, go through and make the data nice
    if(len(tensor_products)==0):
        # Since Tensor objects are additive, the empty object should be zero
        return sympify(0)
    if(len(tensor_products)==1 and isinstance(tensor_products[0], TensorFunction)) :
        tensor_products = list(t_p for t_p in tensor_products[0].tensor_products if t_p!=0)
    elif(len(tensor_products)==1 and isinstance(tensor_products[0], TensorProductFunction)) :
        tensor_products = [tensor_products[0],]
    else:
        if(len(tensor_products)==1 and isinstance(tensor_products[0], list)):
            tensor_products = tensor_products[0]
        tensor_products = flatten(list(t_p for t_p in tensor_products if t_p!=0))
        if(len(tensor_products)>0):
            rank=tensor_products[0].rank
            for t_p in tensor_products:
                if(t_p.rank != rank):
                    raise ValueError("Cannot add rank-{0} tensor to rank-{1} tensors.".format(t_p.rank, rank))

    ## Now, create the object and set its data.  Because of sympy's
    ## caching, tensors with different data need to be created as
    ## classes with different names.  So we just create a lighweight
    ## subclass with a unique name (the number at the end gets
    ## incremented every time we construct a tensor).
    global _Tensor_count
    ThisTensorFunction = type('TensorFunction_'+str(_Tensor_count),
                             (TensorFunction,), {})
    _Tensor_count += 1
    T = ThisTensorFunction( *tuple( set( flatten( [t_p.args for t_p in tensor_products] ) ) ) )
    T.tensor_products = tensor_products
    return T

def test_solve_biquadratic():
    x0, y0, x1, y1, r = symbols('x0 y0 x1 y1 r')

    f_1 = (x - 1)**2 + (y - 1)**2 - r**2
    f_2 = (x - 2)**2 + (y - 2)**2 - r**2

    assert solve_poly_system([f_1, f_2], x, y) == \
        [(S(3)/2 + sqrt(-1 + 2*r**2)/2, S(3)/2 - sqrt(-1 + 2*r**2)/2),
         (S(3)/2 - sqrt(-1 + 2*r**2)/2, S(3)/2 + sqrt(-1 + 2*r**2)/2)]

    f_1 = (x - 1)**2 + (y - 2)**2 - r**2
    f_2 = (x - 1)**2 + (y - 1)**2 - r**2

    assert solve_poly_system([f_1, f_2], x, y) == \
        [(1 + sqrt(((2*r - 1)*(2*r + 1)))/2, S(3)/2),
         (1 - sqrt(((2*r - 1)*(2*r + 1)))/2, S(3)/2)]

    query = lambda expr: expr.is_Pow and expr.exp is S.Half

    f_1 = (x - 1 )**2 + (y - 2)**2 - r**2
    f_2 = (x - x1)**2 + (y - 1)**2 - r**2

    result = solve_poly_system([f_1, f_2], x, y)

    assert len(result) == 2 and all(len(r) == 2 for r in result)
    assert all(r.count(query) == 1 for r in flatten(result))

    f_1 = (x - x0)**2 + (y - y0)**2 - r**2
    f_2 = (x - x1)**2 + (y - y1)**2 - r**2

    result = solve_poly_system([f_1, f_2], x, y)

    assert len(result) == 2 and all(len(r) == 2 for r in result)
    assert all(len(r.find(query)) == 1 for r in flatten(result))

def test_solve_biquadratic():
    x0, y0, x1, y1, r = symbols('x0 y0 x1 y1 r')

    f_1 = (x - 1)**2 + (y - 1)**2 - r**2
    f_2 = (x - 2)**2 + (y - 2)**2 - r**2
    s = sqrt(2*r**2 - 1)
    a = (3 - s)/2
    b = (3 + s)/2
    assert solve_poly_system([f_1, f_2], x, y) == [(a, b), (b, a)]

    f_1 = (x - 1)**2 + (y - 2)**2 - r**2
    f_2 = (x - 1)**2 + (y - 1)**2 - r**2

    assert solve_poly_system([f_1, f_2], x, y) == \
        [(1 - sqrt(((2*r - 1)*(2*r + 1)))/2, S(3)/2),
         (1 + sqrt(((2*r - 1)*(2*r + 1)))/2, S(3)/2)]

    query = lambda expr: expr.is_Pow and expr.exp is S.Half

    f_1 = (x - 1 )**2 + (y - 2)**2 - r**2
    f_2 = (x - x1)**2 + (y - 1)**2 - r**2

    result = solve_poly_system([f_1, f_2], x, y)

    assert len(result) == 2 and all(len(r) == 2 for r in result)
    assert all(r.count(query) == 1 for r in flatten(result))

    f_1 = (x - x0)**2 + (y - y0)**2 - r**2
    f_2 = (x - x1)**2 + (y - y1)**2 - r**2

    result = solve_poly_system([f_1, f_2], x, y)

    assert len(result) == 2 and all(len(r) == 2 for r in result)
    assert all(len(r.find(query)) == 1 for r in flatten(result))

    s1 = (x*y - y, x**2 - x)
    assert solve(s1) == [{x: 1}, {x: 0, y: 0}]
    s2 = (x*y - x, y**2 - y)
    assert solve(s2) == [{y: 1}, {x: 0, y: 0}]
    gens = (x, y)
    for seq in (s1, s2):
        (f, g), opt = parallel_poly_from_expr(seq, *gens)
        raises(SolveFailed, lambda: solve_biquadratic(f, g, opt))
    seq = (x**2 + y**2 - 2, y**2 - 1)
    (f, g), opt = parallel_poly_from_expr(seq, *gens)
    assert solve_biquadratic(f, g, opt) == [
        (-1, -1), (-1, 1), (1, -1), (1, 1)]
    ans = [(0, -1), (0, 1)]
    seq = (x**2 + y**2 - 1, y**2 - 1)
    (f, g), opt = parallel_poly_from_expr(seq, *gens)
    assert solve_biquadratic(f, g, opt) == ans
    seq = (x**2 + y**2 - 1, x**2 - x + y**2 - 1)
    (f, g), opt = parallel_poly_from_expr(seq, *gens)
    assert solve_biquadratic(f, g, opt) == ans

 def _eval_expand_basic(self, deep=True, **hints):
     from sympy import flatten
     if not deep:
         return self
     else:
         return Integral(self.function.expand(deep=deep, **hints),\
         flatten(*self.limits))

    def _preprocess(self, args, expr):
        """Preprocess args, expr to replace arguments that do not map
        to valid Python identifiers.

        Returns string form of args, and updated expr.
        """
        from sympy import Dummy, Function, flatten, Derivative, ordered, Basic
        from sympy.matrices import DeferredVector

        # Args of type Dummy can cause name collisions with args
        # of type Symbol.  Force dummify of everything in this
        # situation.
        dummify = self._dummify or any(
            isinstance(arg, Dummy) for arg in flatten(args))

        argstrs = [None]*len(args)
        for arg, i in reversed(list(ordered(zip(args, range(len(args)))))):
            if iterable(arg):
                s, expr = self._preprocess(arg, expr)
            elif isinstance(arg, DeferredVector):
                s = str(arg)
            elif isinstance(arg, Basic) and arg.is_symbol:
                s = self._argrepr(arg)
                if dummify or not self._is_safe_ident(s):
                    dummy = Dummy()
                    s = self._argrepr(dummy)
                    expr = self._subexpr(expr, {arg: dummy})
            elif dummify or isinstance(arg, (Function, Derivative)):
                dummy = Dummy()
                s = self._argrepr(dummy)
                expr = self._subexpr(expr, {arg: dummy})
            else:
                s = str(arg)
            argstrs[i] = s
        return argstrs, expr

def explode_term_respect_to(term, cls, deep=False, container=list):

    exploded = [term] # we start with the given term since we've to build a list, eventually

    if isinstance(term, cls): 
        exploded = flatten(term.expand().args, cls=cls) if deep else term.args

    return container(exploded)

 def _new(cls, *args, **kwargs):
     shape, flat_list = cls._handle_ndarray_creation_inputs(*args, **kwargs)
     flat_list = flatten(flat_list)
     self = object.__new__(cls)
     self._shape = shape
     self._array = list(flat_list)
     self._rank = len(shape)
     self._loop_size = functools.reduce(lambda x,y: x*y, shape) if shape else 0
     return self

    def variables(self):
        """Model variables definition."""
        v = super().variables
        ncol = self.collocation.n

        x = [xi.name for xi in v['x']]
        u = [ui.name for ui in v['u']]
        
        # Piece states and controls
        xp = [[f'{n}_piece_{k}' for n in x] for k in range(ncol)]
        up = [[f'{n}_piece_{k}' for n in u] for k in range(ncol)]
        
        additional_vars = sym2num.var.make_dict(
            [sym2num.var.SymbolArray('piece_len'),
             sym2num.var.SymbolArray('xp', xp),
             sym2num.var.SymbolArray('up', up),
             sym2num.var.SymbolArray('xp_flat', sympy.flatten(xp)),
             sym2num.var.SymbolArray('up_flat', sympy.flatten(up))]
        )
        return collections.OrderedDict([*v.items(), *additional_vars.items()])

def indexed_terms_appearing_in(term, indexed, only_subscripts=False, do_traversal=False):

    indexed_terms_set = set()

    for subterm in preorder_traversal(term) if do_traversal else flatten(term.args, cls=Add):
        try:
            with bind_Mul_indexed(subterm, indexed) as (_, subscripts):
                indexed_terms_set.add(subscripts if only_subscripts else indexed[subscripts])
        except DestructuringError:
            continue

    return list(indexed_terms_set)

    def dynparms(self, parm_order=None):
        """Return list of RobotDef symbolic dynamic parameters."""

        if not parm_order:
            parm_order = self._dyn_parms_order
        parm_order = parm_order.lower()

        parms = []
        for i in range(0, self.dof):

            if parm_order == 'khalil' or parm_order == 'tensor first':
                # Lxx Lxy Lxz Lyy Lyz Lzz lx ly lz m
                parms += self.Le[i]
                parms += sympy.flatten(self.l[i])
                parms += [self.m[i]]

            elif parm_order == 'siciliano' or parm_order == 'mass first':
                # m lx ly lz Lxx Lxy Lxz Lyy Lyz Lzz
                parms += [self.m[i]]
                parms += sympy.flatten(self.l[i])
                parms += self.Le[i]

            else:
                raise Exception(
                    'RobotDef.Parms(): dynamic parameters order \''
                    + parm_order + '\' not know.')

            if self.driveinertiamodel == 'simplified':
                parms += [self.Ia[i]]

            if self.frictionmodel is not None:
                if 'viscous' in self.frictionmodel:
                    parms += [self.fv[i]]
                if 'Coulomb' in self.frictionmodel:
                    parms += [self.fc[i]]
                if 'offset' in self.frictionmodel:
                    parms += [self.fo[i]]

        return parms

def test_scara_dh_sym_geo_kin():

    pi = sympy.pi
    q = sympybotics.robotdef.q

    a1, a2, d3, d4 = sympy.symbols('a1, a2, d3, d4')

    scara = sympybotics.robotdef.RobotDef(
        'SCARA - Spong',
        [( 0, a1,  0, q),
         (pi, a2,  0, q),
         ( 0,  0,  q, 0),
         ( 0,  0, d4, q)],
        dh_convention='standard')

    scara_geo = sympybotics.geometry.Geometry(scara)
    scara_kin = sympybotics.kinematics.Kinematics(scara, scara_geo)

    cos, sin = sympy.cos, sympy.sin
    q1, q2, q3, q4 = sympy.flatten(scara.q)

    T_spong = sympy.Matrix([
        [(-sin(q1)*sin(q2) + cos(q1)*cos(q2))*cos(q4) + (sin(q1)*cos(q2) +
         sin(q2)*cos(q1))*sin(q4), -(-sin(q1)*sin(q2) +
         cos(q1)*cos(q2))*sin(q4) + (sin(q1)*cos(q2) +
         sin(q2)*cos(q1))*cos(q4), 0, a1*cos(q1) - a2*sin(q1)*sin(q2) +
         a2*cos(q1)*cos(q2)],
        [(sin(q1)*sin(q2) - cos(q1)*cos(q2))*sin(q4) + (sin(q1)*cos(q2) +
         sin(q2)*cos(q1))*cos(q4), (sin(q1)*sin(q2) -
         cos(q1)*cos(q2))*cos(q4) - (sin(q1)*cos(q2) +
         sin(q2)*cos(q1))*sin(q4), 0, a1*sin(q1) + a2*sin(q1)*cos(q2) +
         a2*sin(q2)*cos(q1)],
        [0, 0, -1, -d4 - q3],
        [0, 0, 0, 1]])

    J_spong = sympy.Matrix([[-a1*sin(q1) - a2*sin(q1)*cos(q2) -
                             a2*sin(q2)*cos(q1), -a2*sin(q1)*cos(q2) -
                             a2*sin(q2)*cos(q1), 0, 0],
                            [a1*cos(q1) - a2*sin(q1)*sin(q2) +
                             a2*cos(q1)*cos(q2), -a2*sin(q1)*sin(q2) +
                             a2*cos(q1)*cos(q2), 0, 0],
                            [0, 0, -1, 0],
                            [0, 0, 0, 0],
                            [0, 0, 0, 0],
                            [1, 1, 0, -1]])

    assert (scara_geo.T[-1] - T_spong).expand() == sympy.zeros(4)
    assert (scara_kin.J[-1] - J_spong).expand() == sympy.zeros(6, 4)

 def add_derivative(self, name, fname, wrt, flatten_wrt=False):
     # Test if we have only one wrt item
     if isinstance(wrt, (str, sympy.NDimArray)):
         wrt = (wrt,)
     
     out = self.default_function_output(fname)
     for wrt_array in wrt:
         if utils.isstr(wrt_array):
             wrt_array = self.variables[wrt_array]
         if flatten_wrt:
             wrt_array = sympy.flatten(wrt_array)
         out = sympy.derive_by_array(out, wrt_array)
     
     args = self.function_codegen_arguments(fname)
     deriv = function.SymbolicSubsFunction(args, out)
     setattr(self, name, deriv)

 def sub_args(args, dummies_dict):
     if isinstance(args, string_types):
         return args
     elif isinstance(args, DeferredVector):
         return str(args)
     elif iterable(args):
         dummies = flatten([sub_args(a, dummies_dict) for a in args])
         return ",".join(str(a) for a in dummies)
     else:
         # replace these with Dummy symbols
         if isinstance(args, (Function, Symbol, Derivative)):
             dummies = Dummy()
             dummies_dict.update({args : dummies})
             return str(dummies)
         else:
             return str(args)

def condense_list_of_dicts(dicts,key):
    if not all([key in d for d in dicts]):
        raise ValueError('{key} not in every dict in list'.format(key=key))
    
    keyvals = set([d[key] for d in dicts])
    result = []
    for val in keyvals:
        temp_dict = {}
        temp_dict[key] = val
        matching_dicts = [d for d in dicts if d[key] == val]
        all_keys = set(sympy.flatten([list(d.keys()) for d in matching_dicts]))
        for k in all_keys:
            if k != key:
                temp_dict[k] = set([d[k] for d in matching_dicts])
        result.append(temp_dict)
    return result

def TensorProduct(*input_vectors, **kwargs):
    if('coefficient' in kwargs and kwargs['coefficient']==0):
        return sympify(0)

    ## First, go through and make the data nice
    if(len(input_vectors)==0):
        # Since TensorProducts are multiplicative, the empty object
        # should be 1 (or whatever coefficient was passed, if any)
        return kwargs.get('coefficient', sympify(1))
    if(len(input_vectors)==1 and isinstance(input_vectors[0], TensorProductFunction)) :
        vectors = list(input_vectors[0].vectors)
        coefficient = deepcopy(input_vectors[0].coefficient)
        symmetric = bool(input_vectors[0].symmetric)
    else:
        if(len(input_vectors)==1 and isinstance(input_vectors[0], list)):
            input_vectors = input_vectors[0]
        vectors = list(input_vectors)
        coefficient = deepcopy(kwargs.get('coefficient', 1))
        symmetric = bool(kwargs.get('symmetric', True))

    ## Now, make sure none of the input vectors are zero
    for v in vectors:
        if v==0:
            return sympify(0)

    ## Finally, create the object and set its data.  Because of
    ## sympy's caching, tensor products with different data need to be
    ## created as classes with different names.  So we just create a
    ## lighweight subclass with a unique name (the number at the end
    ## gets incremented every time we construct a tensor product).
    global _TensorProduct_count
    ThisTensorProductFunction = type('TensorProductFunction_'+str(_TensorProduct_count),
                                     (TensorProductFunction,), {})
    _TensorProduct_count += 1
    # print('About to construct a tensor with args ',
    #       tuple( set( flatten( [v.args for v in vectors] ) ) ),
    #       input_vectors,
    #       kwargs,
    #       vectors,
    #       [v.args for v in vectors] )
    TP = ThisTensorProductFunction( *tuple( set( flatten( [v.args for v in vectors] ) ) ) )
    TP.vectors = vectors
    TP.coefficient = coefficient
    TP.symmetric = symmetric
    return TP