Python Matrix.hstack方法代码示例

# 需要导入模块: from sympy import Matrix [as 别名]
# 或者: from sympy.Matrix import hstack [as 别名]
def potential_energy(Fr, q, u, kde_map, vc_map=None):
    """Returns a potential energy function using the method from Section 5.1
    from Kane 1985.

    'Fr' is a list of the generalized active forces for the system.
    'q' is a list of generalized coordinates.
    'u' is a list of the independent generalized speeds.
    'kde_map' is a dictionary with q dots as keys and the equivalent
        expressions in terms of q's and u's as values.
    'vc_map' is a dictionay with the dependent u's as keys and the expression
        in terms of independent u's as values.
    """
    n = len(q)
    p = len(u)
    m = n - p
    if vc_map is None:
        A_kr = Matrix.zeros(m, p)
    else:
        A_kr, _ = vc_matrix(u, vc_map)
        u_ = u + sorted(vc_map.keys(), cmp=lambda x, y: x.compare(y))
    W_sr, _ = kde_matrix(u_, kde_map)

    dV_dq = symbols('∂V/∂q1:{0}'.format(n + 1))
    dV_eq = Matrix(Fr).T

    for s in range(n):
        dV_eq += dV_dq[s] * (W_sr[s, :p] + W_sr[s, p:]*A_kr[:, :p])

    if vc_map is not None:
        f_arg, non_arg = _f_variables(Fr, q, dV_eq, dV_dq)
        f = map(lambda x: x(*f_arg),
                symbols('f1:{0}'.format(m + 1)))
        dV_eq = subs(dV_eq, dict(zip(dV_dq[-m:], f)))
        dV_dq = dV_dq[:-m]

    dV_dq_map = solve(dV_eq, dV_dq)
    dV_dq_list = map(lambda x: dV_dq_map[x], dV_dq)

    if vc_map is None:
        #print('Checking ∂/∂qr(∂V/∂qs) = ∂/∂qs(∂V/∂qr) for all r, s '
        #      '= 1, ..., n.')
        dV_eq = _equivalent_derivatives(dV_dq_list, q)
        if dV_eq != [0] * (n*(n - 1)//2):
            rs = [(r, s) for r in range(n) for s in range(r + 1, n)]
            for (r, s), x in zip(rs, dV_eq):
                if trigsimp(expand(x)) != 0:
                    print(('∂/∂q{0}(∂V/∂q{1}) != ∂/∂q{1}(∂V/∂q{0}). ' +
                           'V does NOT exist.').format(r + 1, s + 1))
                    print('∂/∂q{0}(∂V/∂q{1}) = {2}'.format(
                            r + 1, s + 1, dV_dq_list[r].diff(q[s])))
                    print('∂/∂q{1}(∂V/∂q{0}) = {2}'.format(
                            r + 1, s + 1, dV_dq_list[s].diff(q[r])))
                    return None
    else:
        dV_dq_list += f
        # Unable to take diff of 'fm.diff(qs)', replace with dummy symbols.
        dfdq = [Dummy('∂f{0}/∂q{1}'.format(i + 1, j + 1))
                for i in range(len(f)) for j in range(n)]
        dfdq_replace = lambda x: reduce(
                lambda y, z: y.replace(z[0], z[1]) if z[0] != 0 else y,
                zip([fm.diff(qs) for fm in f for qs in q], dfdq),
                x)
        dV_eq = map(dfdq_replace,
                    _equivalent_derivatives(dV_dq_list, q))

        X = Matrix(dfdq)
        Z = Matrix([map(lambda x: diff(dV_eqi, x), dfdq)
                    for dV_eqi in dV_eq])
        if Z.rank() == n * (n - 1) / 2:
            print('ρ == n(n - 1)/2')
            print('V may exist but cannot be found by this procedure.')
            return None

        Y = expand(Z*X - Matrix(dV_eq))
        ZI_rref, _ = Matrix.hstack(Z, Matrix.eye(Z.shape[0])).rref()
        # E is the matrix of elementary row operations that gives rref(Z).
        E = ZI_rref[:, Z.shape[1]:]
        f_eq = (E * Y)[Z.rank():]
        f_map = solve(f_eq, f)
        if sorted(f_map.keys(), cmp=lambda x, y: x.compare(y)) != f:
            print('Unable to solve for all f uniquely.')
            return None
        for k, v in f_map.iteritems():
            for qi in non_arg:
                if v.diff(qi) != 0:
                    print('{0} should not be a function of {1}'.format(k, qi))
                    return None
        dV_dq_list = map(trigsimp, (subs(dV_dq_list, f_map)))

    return function_from_partials(dV_dq_list, q)