def test_two_dof():
    # This is for a 2 d.o.f., 2 particle spring-mass-damper.
    # The first coordinate is the displacement of the first particle, and the
    # second is the relative displacement between the first and second
    # particles. Speeds are defined as the time derivatives of the particles.
    q1, q2, u1, u2 = dynamicsymbols('q1 q2 u1 u2')
    q1d, q2d, u1d, u2d = dynamicsymbols('q1 q2 u1 u2', 1)
    m, c1, c2, k1, k2 = symbols('m c1 c2 k1 k2')
    N = ReferenceFrame('N')
    P1 = Point('P1')
    P2 = Point('P2')
    P1.set_vel(N, u1 * N.x)
    P2.set_vel(N, (u1 + u2) * N.x)
    kd = [q1d - u1, q2d - u2]

    # Now we create the list of forces, then assign properties to each
    # particle, then create a list of all particles.
    FL = [(P1, (-k1 * q1 - c1 * u1 + k2 * q2 + c2 * u2) * N.x), (P2, (-k2 *
        q2 - c2 * u2) * N.x)]
    pa1 = Particle('pa1', P1, m)
    pa2 = Particle('pa2', P2, m)
    BL = [pa1, pa2]

    # Finally we create the KanesMethod object, specify the inertial frame,
    # pass relevant information, and form Fr & Fr*. Then we calculate the mass
    # matrix and forcing terms, and finally solve for the udots.
    KM = KanesMethod(N, q_ind=[q1, q2], u_ind=[u1, u2], kd_eqs=kd)
    KM.kanes_equations(FL, BL)
    MM = KM.mass_matrix
    forcing = KM.forcing
    rhs = MM.inv() * forcing
    assert expand(rhs[0]) == expand((-k1 * q1 - c1 * u1 + k2 * q2 + c2 * u2)/m)
    assert expand(rhs[1]) == expand((k1 * q1 + c1 * u1 - 2 * k2 * q2 - 2 *
                                    c2 * u2) / m)

def test_one_dof():
    # This is for a 1 dof spring-mass-damper case.
    # It is described in more detail in the KanesMethod docstring.
    q, u = dynamicsymbols('q u')
    qd, ud = dynamicsymbols('q u', 1)
    m, c, k = symbols('m c k')
    N = ReferenceFrame('N')
    P = Point('P')
    P.set_vel(N, u * N.x)

    kd = [qd - u]
    FL = [(P, (-k * q - c * u) * N.x)]
    pa = Particle('pa', P, m)
    BL = [pa]

    KM = KanesMethod(N, [q], [u], kd)
    # The old input format raises a deprecation warning, so catch it here so
    # it doesn't cause py.test to fail.
    with warnings.catch_warnings():
        warnings.filterwarnings("ignore", category=SymPyDeprecationWarning)
        KM.kanes_equations(FL, BL)

    MM = KM.mass_matrix
    forcing = KM.forcing
    rhs = MM.inv() * forcing
    assert expand(rhs[0]) == expand(-(q * k + u * c) / m)

    assert simplify(KM.rhs() -
                    KM.mass_matrix_full.LUsolve(KM.forcing_full)) == zeros(2, 1)

    assert (KM.linearize(A_and_B=True, )[0] == Matrix([[0, 1], [-k/m, -c/m]]))

def test_parallel_axis():
    # This is for a 2 dof inverted pendulum on a cart.
    # This tests the parallel axis code in KanesMethod. The inertia of the
    # pendulum is defined about the hinge, not about the center of mass.

    # Defining the constants and knowns of the system
    gravity = symbols('g')
    k, ls = symbols('k ls')
    a, mA, mC = symbols('a mA mC')
    F = dynamicsymbols('F')
    Ix, Iy, Iz = symbols('Ix Iy Iz')

    # Declaring the Generalized coordinates and speeds
    q1, q2 = dynamicsymbols('q1 q2')
    q1d, q2d = dynamicsymbols('q1 q2', 1)
    u1, u2 = dynamicsymbols('u1 u2')
    u1d, u2d = dynamicsymbols('u1 u2', 1)

    # Creating reference frames
    N = ReferenceFrame('N')
    A = ReferenceFrame('A')

    A.orient(N, 'Axis', [-q2, N.z])
    A.set_ang_vel(N, -u2 * N.z)

    # Origin of Newtonian reference frame
    O = Point('O')

    # Creating and Locating the positions of the cart, C, and the
    # center of mass of the pendulum, A
    C = O.locatenew('C', q1 * N.x)
    Ao = C.locatenew('Ao', a * A.y)

    # Defining velocities of the points
    O.set_vel(N, 0)
    C.set_vel(N, u1 * N.x)
    Ao.v2pt_theory(C, N, A)
    Cart = Particle('Cart', C, mC)
    Pendulum = RigidBody('Pendulum', Ao, A, mA, (inertia(A, Ix, Iy, Iz), C))

    # kinematical differential equations

    kindiffs = [q1d - u1, q2d - u2]

    bodyList = [Cart, Pendulum]

    forceList = [(Ao, -N.y * gravity * mA),
                 (C, -N.y * gravity * mC),
                 (C, -N.x * k * (q1 - ls)),
                 (C, N.x * F)]

    km = KanesMethod(N, [q1, q2], [u1, u2], kindiffs)
    with warnings.catch_warnings():
        warnings.filterwarnings("ignore", category=SymPyDeprecationWarning)
        (fr, frstar) = km.kanes_equations(forceList, bodyList)
    mm = km.mass_matrix_full
    assert mm[3, 3] == Iz

def test_aux():
    # Same as above, except we have 2 auxiliary speeds for the ground contact
    # point, which is known to be zero. In one case, we go through then
    # substitute the aux. speeds in at the end (they are zero, as well as their
    # derivative), in the other case, we use the built-in auxiliary speed part
    # of KanesMethod. The equations from each should be the same.
    q1, q2, q3, u1, u2, u3 = dynamicsymbols('q1 q2 q3 u1 u2 u3')
    q1d, q2d, q3d, u1d, u2d, u3d = dynamicsymbols('q1 q2 q3 u1 u2 u3', 1)
    u4, u5, f1, f2 = dynamicsymbols('u4, u5, f1, f2')
    u4d, u5d = dynamicsymbols('u4, u5', 1)
    r, m, g = symbols('r m g')

    N = ReferenceFrame('N')
    Y = N.orientnew('Y', 'Axis', [q1, N.z])
    L = Y.orientnew('L', 'Axis', [q2, Y.x])
    R = L.orientnew('R', 'Axis', [q3, L.y])
    w_R_N_qd = R.ang_vel_in(N)
    R.set_ang_vel(N, u1 * L.x + u2 * L.y + u3 * L.z)

    C = Point('C')
    C.set_vel(N, u4 * L.x + u5 * (Y.z ^ L.x))
    Dmc = C.locatenew('Dmc', r * L.z)
    Dmc.v2pt_theory(C, N, R)
    Dmc.a2pt_theory(C, N, R)

    I = inertia(L, m / 4 * r**2, m / 2 * r**2, m / 4 * r**2)

    kd = [dot(R.ang_vel_in(N) - w_R_N_qd, uv) for uv in L]

    ForceList = [(Dmc, - m * g * Y.z), (C, f1 * L.x + f2 * (Y.z ^ L.x))]
    BodyD = RigidBody('BodyD', Dmc, R, m, (I, Dmc))
    BodyList = [BodyD]

    KM = KanesMethod(N, q_ind=[q1, q2, q3], u_ind=[u1, u2, u3, u4, u5],
                     kd_eqs=kd)
    with warnings.catch_warnings():
        warnings.filterwarnings("ignore", category=SymPyDeprecationWarning)
        (fr, frstar) = KM.kanes_equations(ForceList, BodyList)
    fr = fr.subs({u4d: 0, u5d: 0}).subs({u4: 0, u5: 0})
    frstar = frstar.subs({u4d: 0, u5d: 0}).subs({u4: 0, u5: 0})

    KM2 = KanesMethod(N, q_ind=[q1, q2, q3], u_ind=[u1, u2, u3], kd_eqs=kd,
                      u_auxiliary=[u4, u5])
    with warnings.catch_warnings():
        warnings.filterwarnings("ignore", category=SymPyDeprecationWarning)
        (fr2, frstar2) = KM2.kanes_equations(ForceList, BodyList)
    fr2 = fr2.subs({u4d: 0, u5d: 0}).subs({u4: 0, u5: 0})
    frstar2 = frstar2.subs({u4d: 0, u5d: 0}).subs({u4: 0, u5: 0})

    frstar.simplify()
    frstar2.simplify()

    assert (fr - fr2).expand() == Matrix([0, 0, 0, 0, 0])
    assert (frstar - frstar2).expand() == Matrix([0, 0, 0, 0, 0])

def get_equations(m_val, g_val, l_val):
    # This function body is copyied from:
    # http://www.pydy.org/examples/double_pendulum.html
    # Retrieved 2015-09-29
    from sympy import symbols
    from sympy.physics.mechanics import (
        dynamicsymbols, ReferenceFrame, Point, Particle, KanesMethod
    )

    q1, q2 = dynamicsymbols('q1 q2')
    q1d, q2d = dynamicsymbols('q1 q2', 1)
    u1, u2 = dynamicsymbols('u1 u2')
    u1d, u2d = dynamicsymbols('u1 u2', 1)
    l, m, g = symbols('l m g')

    N = ReferenceFrame('N')
    A = N.orientnew('A', 'Axis', [q1, N.z])
    B = N.orientnew('B', 'Axis', [q2, N.z])

    A.set_ang_vel(N, u1 * N.z)
    B.set_ang_vel(N, u2 * N.z)

    O = Point('O')
    P = O.locatenew('P', l * A.x)
    R = P.locatenew('R', l * B.x)

    O.set_vel(N, 0)
    P.v2pt_theory(O, N, A)
    R.v2pt_theory(P, N, B)

    ParP = Particle('ParP', P, m)
    ParR = Particle('ParR', R, m)

    kd = [q1d - u1, q2d - u2]
    FL = [(P, m * g * N.x), (R, m * g * N.x)]
    BL = [ParP, ParR]

    KM = KanesMethod(N, q_ind=[q1, q2], u_ind=[u1, u2], kd_eqs=kd)

    (fr, frstar) = KM.kanes_equations(FL, BL)
    kdd = KM.kindiffdict()
    mm = KM.mass_matrix_full
    fo = KM.forcing_full
    qudots = mm.inv() * fo
    qudots = qudots.subs(kdd)
    qudots.simplify()
    # Edit:
    depv = [q1, q2, u1, u2]
    subs = list(zip([m, g, l], [m_val, g_val, l_val]))
    return zip(depv, [expr.subs(subs) for expr in qudots])

def test_input_format():
    # 1 dof problem from test_one_dof
    q, u = dynamicsymbols('q u')
    qd, ud = dynamicsymbols('q u', 1)
    m, c, k = symbols('m c k')
    N = ReferenceFrame('N')
    P = Point('P')
    P.set_vel(N, u * N.x)

    kd = [qd - u]
    FL = [(P, (-k * q - c * u) * N.x)]
    pa = Particle('pa', P, m)
    BL = [pa]

    KM = KanesMethod(N, [q], [u], kd)
    # test for input format kane.kanes_equations((body1, body2, particle1))
    assert KM.kanes_equations(BL)[0] == Matrix([0])
    # test for input format kane.kanes_equations(bodies=(body1, body 2), loads=(load1,load2))
    assert KM.kanes_equations(bodies=BL, loads=None)[0] == Matrix([0])
    # test for input format kane.kanes_equations(bodies=(body1, body 2), loads=None)
    assert KM.kanes_equations(BL, loads=None)[0] == Matrix([0])
    # test for input format kane.kanes_equations(bodies=(body1, body 2))
    assert KM.kanes_equations(BL)[0] == Matrix([0])
    # test for error raised when a wrong force list (in this case a string) is provided
    from sympy.utilities.pytest import raises
    raises(ValueError, lambda: KM._form_fr('bad input'))

    # 2 dof problem from test_two_dof
    q1, q2, u1, u2 = dynamicsymbols('q1 q2 u1 u2')
    q1d, q2d, u1d, u2d = dynamicsymbols('q1 q2 u1 u2', 1)
    m, c1, c2, k1, k2 = symbols('m c1 c2 k1 k2')
    N = ReferenceFrame('N')
    P1 = Point('P1')
    P2 = Point('P2')
    P1.set_vel(N, u1 * N.x)
    P2.set_vel(N, (u1 + u2) * N.x)
    kd = [q1d - u1, q2d - u2]

    FL = ((P1, (-k1 * q1 - c1 * u1 + k2 * q2 + c2 * u2) * N.x), (P2, (-k2 *
        q2 - c2 * u2) * N.x))
    pa1 = Particle('pa1', P1, m)
    pa2 = Particle('pa2', P2, m)
    BL = (pa1, pa2)

    KM = KanesMethod(N, q_ind=[q1, q2], u_ind=[u1, u2], kd_eqs=kd)
    # test for input format
    # kane.kanes_equations((body1, body2), (load1, load2))
    KM.kanes_equations(BL, FL)
    MM = KM.mass_matrix
    forcing = KM.forcing
    rhs = MM.inv() * forcing
    assert expand(rhs[0]) == expand((-k1 * q1 - c1 * u1 + k2 * q2 + c2 * u2)/m)
    assert expand(rhs[1]) == expand((k1 * q1 + c1 * u1 - 2 * k2 * q2 - 2 *
                                    c2 * u2) / m)

def test_pend():
    q, u = dynamicsymbols('q u')
    qd, ud = dynamicsymbols('q u', 1)
    m, l, g = symbols('m l g')
    N = ReferenceFrame('N')
    P = Point('P')
    P.set_vel(N, -l * u * sin(q) * N.x + l * u * cos(q) * N.y)
    kd = [qd - u]

    FL = [(P, m * g * N.x)]
    pa = Particle('pa', P, m)
    BL = [pa]

    KM = KanesMethod(N, [q], [u], kd)
    KM.kanes_equations(FL, BL)
    MM = KM.mass_matrix
    forcing = KM.forcing
    rhs = MM.inv() * forcing
    rhs.simplify()
    assert expand(rhs[0]) == expand(-g / l * sin(q))

 def mass_matrix(self):
     #if not (self._kanes_method and self.up_to_date):
     #    self._init_kanes_method()
     # TODO move the creation of Kane's Method somewhere else.
     self._kanes_method = KanesMethod(self.root.frame,
             q_ind=self.independent_coordinates(),
             u_ind=self.independent_speeds(),
             kd_eqs=self.kinematic_differential_equations()
             )
     # TODO must make this call to get the mass matrix, etc.?
     self._kanes_method.kanes_equations(self.force_list(), self.body_list());
     return self._kanes_method.mass_matrix

def test_one_dof():
    # This is for a 1 dof spring-mass-damper case.
    # It is described in more detail in the KanesMethod docstring.
    q, u = dynamicsymbols('q u')
    qd, ud = dynamicsymbols('q u', 1)
    m, c, k = symbols('m c k')
    N = ReferenceFrame('N')
    P = Point('P')
    P.set_vel(N, u * N.x)

    kd = [qd - u]
    FL = [(P, (-k * q - c * u) * N.x)]
    pa = Particle('pa', P, m)
    BL = [pa]

    KM = KanesMethod(N, [q], [u], kd)
    KM.kanes_equations(FL, BL)
    MM = KM.mass_matrix
    forcing = KM.forcing
    rhs = MM.inv() * forcing
    assert expand(rhs[0]) == expand(-(q * k + u * c) / m)
    assert (KM.linearize(A_and_B=True, new_method=True)[0] ==
            Matrix([[0, 1], [-k/m, -c/m]]))

    # Ensure that the old linearizer still works and that the new linearizer
    # gives the same results. The old linearizer is deprecated and should be
    # removed in >= 0.7.7.
    M_old = KM.mass_matrix_full
    # The old linearizer raises a deprecation warning, so catch it here so
    # it doesn't cause py.test to fail.
    with warnings.catch_warnings():
        warnings.filterwarnings("ignore", category=SymPyDeprecationWarning)
        F_A_old, F_B_old, r_old = KM.linearize()
    M_new, F_A_new, F_B_new, r_new = KM.linearize(new_method=True)
    assert simplify(M_new.inv() * F_A_new - M_old.inv() * F_A_old) == zeros(2)

def test_one_dof():
    # This is for a 1 dof spring-mass-damper case.
    # It is described in more detail in the KanesMethod docstring.
    q, u = dynamicsymbols('q u')
    qd, ud = dynamicsymbols('q u', 1)
    m, c, k = symbols('m c k')
    N = ReferenceFrame('N')
    P = Point('P')
    P.set_vel(N, u * N.x)

    kd = [qd - u]
    FL = [(P, (-k * q - c * u) * N.x)]
    pa = Particle('pa', P, m)
    BL = [pa]

    KM = KanesMethod(N, [q], [u], kd)
    KM.kanes_equations(FL, BL)
    MM = KM.mass_matrix
    forcing = KM.forcing
    rhs = MM.inv() * forcing
    assert expand(rhs[0]) == expand(-(q * k + u * c) / m)
    assert KM.linearize() == (Matrix([[0, 1], [-k, -c]]), Matrix([]), Matrix([]))

def test_linearize_pendulum_kane_minimal():
    q1 = dynamicsymbols('q1')                     # angle of pendulum
    u1 = dynamicsymbols('u1')                     # Angular velocity
    q1d = dynamicsymbols('q1', 1)                 # Angular velocity
    L, m, t = symbols('L, m, t')
    g = 9.8

    # Compose world frame
    N = ReferenceFrame('N')
    pN = Point('N*')
    pN.set_vel(N, 0)

    # A.x is along the pendulum
    A = N.orientnew('A', 'axis', [q1, N.z])
    A.set_ang_vel(N, u1*N.z)

    # Locate point P relative to the origin N*
    P = pN.locatenew('P', L*A.x)
    P.v2pt_theory(pN, N, A)
    pP = Particle('pP', P, m)

    # Create Kinematic Differential Equations
    kde = Matrix([q1d - u1])

    # Input the force resultant at P
    R = m*g*N.x

    # Solve for eom with kanes method
    KM = KanesMethod(N, q_ind=[q1], u_ind=[u1], kd_eqs=kde)
    with warnings.catch_warnings():
        warnings.filterwarnings("ignore", category=SymPyDeprecationWarning)
        (fr, frstar) = KM.kanes_equations([(P, R)], [pP])

    # Linearize
    A, B, inp_vec = KM.linearize(A_and_B=True, new_method=True, simplify=True)

    assert A == Matrix([[0, 1], [-9.8*cos(q1)/L, 0]])
    assert B == Matrix([])

def test_two_dof():
    # This is for a 2 d.o.f., 2 particle spring-mass-damper.
    # The first coordinate is the displacement of the first particle, and the
    # second is the relative displacement between the first and second
    # particles. Speeds are defined as the time derivatives of the particles.
    q1, q2, u1, u2 = dynamicsymbols("q1 q2 u1 u2")
    q1d, q2d, u1d, u2d = dynamicsymbols("q1 q2 u1 u2", 1)
    m, c1, c2, k1, k2 = symbols("m c1 c2 k1 k2")
    N = ReferenceFrame("N")
    P1 = Point("P1")
    P2 = Point("P2")
    P1.set_vel(N, u1 * N.x)
    P2.set_vel(N, (u1 + u2) * N.x)
    kd = [q1d - u1, q2d - u2]

    # Now we create the list of forces, then assign properties to each
    # particle, then create a list of all particles.
    FL = [(P1, (-k1 * q1 - c1 * u1 + k2 * q2 + c2 * u2) * N.x), (P2, (-k2 * q2 - c2 * u2) * N.x)]
    pa1 = Particle("pa1", P1, m)
    pa2 = Particle("pa2", P2, m)
    BL = [pa1, pa2]

    # Finally we create the KanesMethod object, specify the inertial frame,
    # pass relevant information, and form Fr & Fr*. Then we calculate the mass
    # matrix and forcing terms, and finally solve for the udots.
    KM = KanesMethod(N, q_ind=[q1, q2], u_ind=[u1, u2], kd_eqs=kd)
    # The old input format raises a deprecation warning, so catch it here so
    # it doesn't cause py.test to fail.
    with warnings.catch_warnings():
        warnings.filterwarnings("ignore", category=SymPyDeprecationWarning)
        KM.kanes_equations(FL, BL)
    MM = KM.mass_matrix
    forcing = KM.forcing
    rhs = MM.inv() * forcing
    assert expand(rhs[0]) == expand((-k1 * q1 - c1 * u1 + k2 * q2 + c2 * u2) / m)
    assert expand(rhs[1]) == expand((k1 * q1 + c1 * u1 - 2 * k2 * q2 - 2 * c2 * u2) / m)

    assert simplify(KM.rhs() - KM.mass_matrix_full.LUsolve(KM.forcing_full)) == zeros(4, 1)

def test_pend():
    q, u = dynamicsymbols("q u")
    qd, ud = dynamicsymbols("q u", 1)
    m, l, g = symbols("m l g")
    N = ReferenceFrame("N")
    P = Point("P")
    P.set_vel(N, -l * u * sin(q) * N.x + l * u * cos(q) * N.y)
    kd = [qd - u]

    FL = [(P, m * g * N.x)]
    pa = Particle("pa", P, m)
    BL = [pa]

    KM = KanesMethod(N, [q], [u], kd)
    with warnings.catch_warnings():
        warnings.filterwarnings("ignore", category=SymPyDeprecationWarning)
        KM.kanes_equations(FL, BL)
    MM = KM.mass_matrix
    forcing = KM.forcing
    rhs = MM.inv() * forcing
    rhs.simplify()
    assert expand(rhs[0]) == expand(-g / l * sin(q))
    assert simplify(KM.rhs() - KM.mass_matrix_full.LUsolve(KM.forcing_full)) == zeros(2, 1)

def test_linearize_pendulum_kane_minimal():
    q1 = dynamicsymbols("q1")  # angle of pendulum
    u1 = dynamicsymbols("u1")  # Angular velocity
    q1d = dynamicsymbols("q1", 1)  # Angular velocity
    L, m, t = symbols("L, m, t")
    g = 9.8

    # Compose world frame
    N = ReferenceFrame("N")
    pN = Point("N*")
    pN.set_vel(N, 0)

    # A.x is along the pendulum
    A = N.orientnew("A", "axis", [q1, N.z])
    A.set_ang_vel(N, u1 * N.z)

    # Locate point P relative to the origin N*
    P = pN.locatenew("P", L * A.x)
    P.v2pt_theory(pN, N, A)
    pP = Particle("pP", P, m)

    # Create Kinematic Differential Equations
    kde = Matrix([q1d - u1])

    # Input the force resultant at P
    R = m * g * N.x

    # Solve for eom with kanes method
    KM = KanesMethod(N, q_ind=[q1], u_ind=[u1], kd_eqs=kde)
    (fr, frstar) = KM.kanes_equations([(P, R)], [pP])

    # Linearize
    A, B, inp_vec = KM.linearize(A_and_B=True, new_method=True, simplify=True)

    assert A == Matrix([[0, 1], [-9.8 * cos(q1) / L, 0]])
    assert B == Matrix([])

def test_sub_qdot2():
    # This test solves exercises 8.3 from Kane 1985 and defines
    # all velocities in terms of q, qdot. We check that the generalized active
    # forces are correctly computed if u terms are only defined in the
    # kinematic differential equations.
    #
    # This functionality was added in PR 8948. Without qdot/u substitution, the
    # KanesMethod constructor will fail during the constraint initialization as
    # the B matrix will be poorly formed and inversion of the dependent part
    # will fail.

    g, m, Px, Py, Pz, R, t = symbols('g m Px Py Pz R t')
    q = dynamicsymbols('q:5')
    qd = dynamicsymbols('q:5', level=1)
    u = dynamicsymbols('u:5')

    ## Define inertial, intermediate, and rigid body reference frames
    A = ReferenceFrame('A')
    B_prime = A.orientnew('B_prime', 'Axis', [q[0], A.z])
    B = B_prime.orientnew('B', 'Axis', [pi/2 - q[1], B_prime.x])
    C = B.orientnew('C', 'Axis', [q[2], B.z])

    ## Define points of interest and their velocities
    pO = Point('O')
    pO.set_vel(A, 0)

    # R is the point in plane H that comes into contact with disk C.
    pR = pO.locatenew('R', q[3]*A.x + q[4]*A.y)
    pR.set_vel(A, pR.pos_from(pO).diff(t, A))
    pR.set_vel(B, 0)

    # C^ is the point in disk C that comes into contact with plane H.
    pC_hat = pR.locatenew('C^', 0)
    pC_hat.set_vel(C, 0)

    # C* is the point at the center of disk C.
    pCs = pC_hat.locatenew('C*', R*B.y)
    pCs.set_vel(C, 0)
    pCs.set_vel(B, 0)

    # calculate velocites of points C* and C^ in frame A
    pCs.v2pt_theory(pR, A, B) # points C* and R are fixed in frame B
    pC_hat.v2pt_theory(pCs, A, C) # points C* and C^ are fixed in frame C

    ## Define forces on each point of the system
    R_C_hat = Px*A.x + Py*A.y + Pz*A.z
    R_Cs = -m*g*A.z
    forces = [(pC_hat, R_C_hat), (pCs, R_Cs)]

    ## Define kinematic differential equations
    # let ui = omega_C_A & bi (i = 1, 2, 3)
    # u4 = qd4, u5 = qd5
    u_expr = [C.ang_vel_in(A) & uv for uv in B]
    u_expr += qd[3:]
    kde = [ui - e for ui, e in zip(u, u_expr)]
    km1 = KanesMethod(A, q, u, kde)
    with warnings.catch_warnings():
        warnings.filterwarnings("ignore", category=SymPyDeprecationWarning)
        fr1, _ = km1.kanes_equations(forces, [])

    ## Calculate generalized active forces if we impose the condition that the
    # disk C is rolling without slipping
    u_indep = u[:3]
    u_dep = list(set(u) - set(u_indep))
    vc = [pC_hat.vel(A) & uv for uv in [A.x, A.y]]
    km2 = KanesMethod(A, q, u_indep, kde,
                      u_dependent=u_dep, velocity_constraints=vc)
    with warnings.catch_warnings():
        warnings.filterwarnings("ignore", category=SymPyDeprecationWarning)
        fr2, _ = km2.kanes_equations(forces, [])

    fr1_expected = Matrix([
        -R*g*m*sin(q[1]),
        -R*(Px*cos(q[0]) + Py*sin(q[0]))*tan(q[1]),
        R*(Px*cos(q[0]) + Py*sin(q[0])),
        Px,
        Py])
    fr2_expected = Matrix([
        -R*g*m*sin(q[1]),
        0,
        0])
    assert (trigsimp(fr1.expand()) == trigsimp(fr1_expected.expand()))
    assert (trigsimp(fr2.expand()) == trigsimp(fr2_expected.expand()))

#!/usr/bin/env python

from sympy.physics.mechanics import KanesMethod

from .kinetics import *

# Equations of Motion
# ===================

coordinates = [theta1, theta2, theta3]

speeds = [omega1, omega2, omega3]

kane = KanesMethod(inertial_frame,
                   coordinates,
                   speeds,
                   kinematical_differential_equations)

loads = [lower_leg_grav_force,
         upper_leg_grav_force,
         torso_grav_force,
         lower_leg_torque,
         upper_leg_torque,
         torso_torque]

bodies = [lower_leg, upper_leg, torso]

fr, frstar = kane.kanes_equations(bodies, loads)

mass_matrix = kane.mass_matrix_full
forcing_vector = kane.forcing_full

def test_rolling_disc():
    # Rolling Disc Example
    # Here the rolling disc is formed from the contact point up, removing the
    # need to introduce generalized speeds. Only 3 configuration and three
    # speed variables are need to describe this system, along with the disc's
    # mass and radius, and the local gravity (note that mass will drop out).
    q1, q2, q3, u1, u2, u3 = dynamicsymbols('q1 q2 q3 u1 u2 u3')
    q1d, q2d, q3d, u1d, u2d, u3d = dynamicsymbols('q1 q2 q3 u1 u2 u3', 1)
    r, m, g = symbols('r m g')

    # The kinematics are formed by a series of simple rotations. Each simple
    # rotation creates a new frame, and the next rotation is defined by the new
    # frame's basis vectors. This example uses a 3-1-2 series of rotations, or
    # Z, X, Y series of rotations. Angular velocity for this is defined using
    # the second frame's basis (the lean frame).
    N = ReferenceFrame('N')
    Y = N.orientnew('Y', 'Axis', [q1, N.z])
    L = Y.orientnew('L', 'Axis', [q2, Y.x])
    R = L.orientnew('R', 'Axis', [q3, L.y])
    w_R_N_qd = R.ang_vel_in(N)
    R.set_ang_vel(N, u1 * L.x + u2 * L.y + u3 * L.z)

    # This is the translational kinematics. We create a point with no velocity
    # in N; this is the contact point between the disc and ground. Next we form
    # the position vector from the contact point to the disc's center of mass.
    # Finally we form the velocity and acceleration of the disc.
    C = Point('C')
    C.set_vel(N, 0)
    Dmc = C.locatenew('Dmc', r * L.z)
    Dmc.v2pt_theory(C, N, R)

    # This is a simple way to form the inertia dyadic.
    I = inertia(L, m / 4 * r**2, m / 2 * r**2, m / 4 * r**2)

    # Kinematic differential equations; how the generalized coordinate time
    # derivatives relate to generalized speeds.
    kd = [dot(R.ang_vel_in(N) - w_R_N_qd, uv) for uv in L]

    # Creation of the force list; it is the gravitational force at the mass
    # center of the disc. Then we create the disc by assigning a Point to the
    # center of mass attribute, a ReferenceFrame to the frame attribute, and mass
    # and inertia. Then we form the body list.
    ForceList = [(Dmc, - m * g * Y.z)]
    BodyD = RigidBody('BodyD', Dmc, R, m, (I, Dmc))
    BodyList = [BodyD]

    # Finally we form the equations of motion, using the same steps we did
    # before. Specify inertial frame, supply generalized speeds, supply
    # kinematic differential equation dictionary, compute Fr from the force
    # list and Fr* from the body list, compute the mass matrix and forcing
    # terms, then solve for the u dots (time derivatives of the generalized
    # speeds).
    KM = KanesMethod(N, q_ind=[q1, q2, q3], u_ind=[u1, u2, u3], kd_eqs=kd)
    KM.kanes_equations(ForceList, BodyList)
    MM = KM.mass_matrix
    forcing = KM.forcing
    rhs = MM.inv() * forcing
    kdd = KM.kindiffdict()
    rhs = rhs.subs(kdd)
    rhs.simplify()
    assert rhs.expand() == Matrix([(6*u2*u3*r - u3**2*r*tan(q2) +
        4*g*sin(q2))/(5*r), -2*u1*u3/3, u1*(-2*u2 + u3*tan(q2))]).expand()

    # This code tests our output vs. benchmark values. When r=g=m=1, the
    # critical speed (where all eigenvalues of the linearized equations are 0)
    # is 1 / sqrt(3) for the upright case.
    A = KM.linearize(A_and_B=True, new_method=True)[0]
    A_upright = A.subs({r: 1, g: 1, m: 1}).subs({q1: 0, q2: 0, q3: 0, u1: 0, u3: 0})
    assert A_upright.subs(u2, 1 / sqrt(3)).eigenvals() == {S(0): 6}

two_link_torque_dp1 = (two_frame_dp1, two_torque_dp1 * inertial_frame_dp1.z + one_torque_dp2 * inertial_frame_dp1.z)

one_link_torque_dp2 = (one_frame_dp2, one_torque_dp2 * inertial_frame_dp2.z - two_torque_dp2 * inertial_frame_dp2.z)

two_link_torque_dp2 = (two_frame_dp2, two_torque_dp2 * inertial_frame_dp2.z)

# Equations of Motion
# ===================

coordinates_dp1 = [theta1_dp1, theta2_dp1]
coordinates_dp2 = [theta1_dp2, theta2_dp2]

speeds_dp1 = [omega1_dp1, omega2_dp1]
speeds_dp2 = [omega1_dp2, omega2_dp2]

kane_dp1 = KanesMethod(inertial_frame_dp1, coordinates_dp1, speeds_dp1, kinematical_differential_equations_dp1)
kane_dp2 = KanesMethod(inertial_frame_dp2, coordinates_dp2, speeds_dp2, kinematic_differential_equations_dp2)

loads_dp1 = [one_grav_force_dp1, two_grav_force_dp1, one_link_torque_dp1, two_link_torque_dp1]
loads_dp2 = [one_grav_force_dp2, two_grav_force_dp2, two_link_torque_dp2, two_link_torque_dp2]

bodies_dp1 = [oneB_dp1, twoB_dp1]
particles_dp2 = [oneP_dp2, twoP_dp2]

fr_dp1, frstar_dp1 = kane_dp1.kanes_equations(loads, bodies)

mass_matrix_dp1 = kane_dp1.mass_matrix
forcing_vector_dp1 = kane_dp1.forcing

mass_matrix_dp2 = kane_dp2.mass_matrix
forcing_vector_dp2 = kane_dp2.forcing

# =============

l_ankle_torque, l_hip_torque, r_hip_torque = dynamicsymbols("T_a, T_k, T_h")

l_leg_torque = (l_leg_frame, l_ankle_torque * inertial_frame.z - l_hip_torque * inertial_frame.z)

body_torque = (body_frame, l_hip_torque * inertial_frame.z)

# Equations of Motion
# ===================

coordinates = [theta1, theta2]

speeds = [omega1, omega2]

kane = KanesMethod(inertial_frame, coordinates, speeds, kinematical_differential_equations)

loads = [l_leg_grav_force, body_grav_force, r_leg_grav_force, l_leg_torque, body_torque]

bodies = [l_leg, body, r_leg]

fr, frstar = kane.kanes_equations(loads, bodies)

mass_matrix = kane.mass_matrix_full
forcing_vector = kane.forcing_full

# List the symbolic arguments
# ===========================

# Constants
# ---------

#.........这里部分代码省略.........
    #                                is q[3].
    # Velocity constraints: f_v, for u3, u4 and u5.
    # Acceleration constraints: f_a.
    f_c = Matrix([dot(-r*B.z, A.z) - q[3]])
    f_v = Matrix([dot(O.vel(N) - (P.vel(N) + cross(C.ang_vel_in(N),
        O.pos_from(P))), ai).expand() for ai in A])
    v_o_n = cross(C.ang_vel_in(N), O.pos_from(P))
    a_o_n = v_o_n.diff(t, A) + cross(A.ang_vel_in(N), v_o_n)
    f_a = Matrix([dot(O.acc(N) - a_o_n, ai) for ai in A])

    # Solve for constraint equations in the form of
    #                           u_dependent = A_rs * [u_i; u_aux].
    # First, obtain constraint coefficient matrix:  M_v * [u; ua] = 0;
    # Second, taking u[0], u[1], u[2] as independent,
    #         taking u[3], u[4], u[5] as dependent,
    #         rearranging the matrix of M_v to be A_rs for u_dependent.
    # Third, u_aux ==0 for u_dep, and resulting dictionary of u_dep_dict.
    M_v = zeros(3, 9)
    for i in range(3):
        for j, ui in enumerate(u + ua):
            M_v[i, j] = f_v[i].diff(ui)

    M_v_i = M_v[:, :3]
    M_v_d = M_v[:, 3:6]
    M_v_aux = M_v[:, 6:]
    M_v_i_aux = M_v_i.row_join(M_v_aux)
    A_rs = - M_v_d.inv() * M_v_i_aux

    u_dep = A_rs[:, :3] * Matrix(u[:3])
    u_dep_dict = dict(zip(u[3:], u_dep))
    #u_dep_dict = {udi : u_depi[0] for udi, u_depi in zip(u[3:], u_dep.tolist())}

    # Active forces: F_O acting on point O; F_P acting on point P.
    # Generalized active forces (unconstrained): Fr_u = F_point * pv_point.
    F_O = m*g*A.z
    F_P = Fx * A.x + Fy * A.y + Fz * A.z
    Fr_u = Matrix([dot(F_O, pv_o) + dot(F_P, pv_p) for pv_o, pv_p in
            zip(partial_v_O, partial_v_P)])

    # Inertia force: R_star_O.
    # Inertia of disc: I_C_O, where J is a inertia component about principal axis.
    # Inertia torque: T_star_C.
    # Generalized inertia forces (unconstrained): Fr_star_u.
    R_star_O = -m*O.acc(N)
    I_C_O = inertia(B, I, J, I)
    T_star_C = -(dot(I_C_O, C.ang_acc_in(N)) \
                 + cross(C.ang_vel_in(N), dot(I_C_O, C.ang_vel_in(N))))
    Fr_star_u = Matrix([dot(R_star_O, pv) + dot(T_star_C, pav) for pv, pav in
                        zip(partial_v_O, partial_w_C)])

    # Form nonholonomic Fr: Fr_c, and nonholonomic Fr_star: Fr_star_c.
    # Also, nonholonomic Fr_star in steady turning condition: Fr_star_steady.
    Fr_c = Fr_u[:3, :].col_join(Fr_u[6:, :]) + A_rs.T * Fr_u[3:6, :]
    Fr_star_c = Fr_star_u[:3, :].col_join(Fr_star_u[6:, :])\
                + A_rs.T * Fr_star_u[3:6, :]
    Fr_star_steady = Fr_star_c.subs(ud_zero).subs(u_dep_dict)\
            .subs(steady_conditions).subs({q[3]: -r*cos(q[1])}).expand()


    # Second, using KaneMethod in mechanics for fr, frstar and frstar_steady.

    # Rigid Bodies: disc, with inertia I_C_O.
    iner_tuple = (I_C_O, O)
    disc = RigidBody('disc', O, C, m, iner_tuple)
    bodyList = [disc]

    # Generalized forces: Gravity: F_o; Auxiliary forces: F_p.
    F_o = (O, F_O)
    F_p = (P, F_P)
    forceList = [F_o,  F_p]

    # KanesMethod.
    kane = KanesMethod(
        N, q_ind= q[:3], u_ind= u[:3], kd_eqs=kindiffs,
        q_dependent=q[3:], configuration_constraints = f_c,
        u_dependent=u[3:], velocity_constraints= f_v,
        u_auxiliary=ua
        )

    # fr, frstar, frstar_steady and kdd(kinematic differential equations).
    (fr, frstar)= kane.kanes_equations(forceList, bodyList)
    frstar_steady = frstar.subs(ud_zero).subs(u_dep_dict).subs(steady_conditions)\
                    .subs({q[3]: -r*cos(q[1])}).expand()
    kdd = kane.kindiffdict()


    # Test
    # First try Fr_c == fr;
    # Second try Fr_star_c == frstar;
    # Third try Fr_star_steady == frstar_steady.
    # Both signs are checked in case the equations were found with an inverse
    # sign.
    assert ((Matrix(Fr_c).expand() == fr.expand()) or
             (Matrix(Fr_c).expand() == (-fr).expand()))

    assert ((Matrix(Fr_star_c).expand() == frstar.expand()) or
             (Matrix(Fr_star_c).expand() == (-frstar).expand()))

    assert ((Matrix(Fr_star_steady).expand() == frstar_steady.expand()) or
             (Matrix(Fr_star_steady).expand() == (-frstar_steady).expand()))