dynamics uav payload

gistfile1.txt

import numpy as np
import jax.numpy as jnp
from jax import jit, jacfwd
import pandas as pd
from casadi import *

def attitude_jacobian_bar(Thetab):
  phi, theta, psi = Thetab[0][0], Thetab[1][0], Thetab[2][0]
  H = jnp.array([
    [jnp.cos(phi) / jnp.cos(theta),   jnp.sin(phi) / jnp.cos(theta),     0 ],
    [               -jnp.sin(psi),                   jnp.cos(psi),     0 ],
    [jnp.cos(psi) * jnp.tan(theta),   jnp.sin(psi) / jnp.tan(theta),     1 ]
  ])
  return H

def attitude_jacobian(Theta):
  phi, theta, psi = Theta[0][0], Theta[1][0], Theta[2][0]
  H = np.array([
    [1,              jnp.sin(phi) * jnp.tan(theta),   jnp.cos(phi) * jnp.tan(theta)],
    [0,              jnp.cos(phi),                  -jnp.sin(phi)],
    [0,              jnp.sin(phi) / jnp.cos(theta),   jnp.cos(phi) / jnp.cos(theta)]
  ])
  return H


def rotation_matrix(Theta):
  phi, theta, psi = Theta[0][0], Theta[1][0], Theta[2][0]
  R_x = jnp.array([[1, 0, 0],
                  [0, jnp.cos(phi), -jnp.sin(phi)],
                  [0, jnp.sin(phi), jnp.cos(phi)]])
  R_y = jnp.array([[jnp.cos(theta), 0, jnp.sin(theta)],
                  [0, 1, 0],
                  [-jnp.sin(theta), 0, jnp.cos(theta)]])
  R_z = jnp.array([[jnp.cos(psi), -jnp.sin(psi), 0],
                  [jnp.sin(psi), jnp.cos(psi), 0],
                  [0, 0, 1]])
  R = jnp.dot(R_z, jnp.dot(R_y, R_x))
  return R


# Quadrotor parameters
mass = 1.75
L = 0.238537
J = jnp.diag(jnp.array([0.01252, 0.01743, 0.0238]))
kf = 8.54858e-06,
km = 1.6e-2

# Bar parameters
Lb = 1.5       
d1 = Lb / 2
d2 = Lb / 2      
massb = 0.5  
Jb = massb*Lb*Lb / 12        

# Cable parameters
Lc = 0.6
l1 = Lc
l2 = Lc

g = 9.81

def dynamicsC(x, u):
    # State variables
    pb = x[0:3].reshape(3, 1)
    vb = x[3:6].reshape(3, 1)
    thetab = x[6:9].reshape(3, 1)
    omegab = x[9:12].reshape(3, 1)
    Hb = attitude_jacobian_bar(thetab)

    p1 = x[12:15].reshape(3, 1)
    v1 = x[15:18].reshape(3, 1)
    theta1 = x[18:21].reshape(3, 1)
    omega1 = x[21:24].reshape(3, 1)
    R1 = rotation_matrix(theta1)
    H1 = attitude_jacobian(theta1)

    p2 = x[24:27].reshape(3, 1)
    v2 = x[27:30].reshape(3, 1)
    theta2 = x[30:33].reshape(3, 1)
    omega2 = x[33:36].reshape(3, 1)
    R2 = rotation_matrix(theta2)
    H2 = attitude_jacobian(theta2)
    
    nb = (p2 - p1) / jnp.linalg.norm(p2 - p1)
    n1 = (p1 - pb + d1*nb) / l1
    n2 = (p2 - pb - d2*nb) / l2

    #d1/d2 for tension calculation
    d1_ = -d1
    d2_ = d2

    S_w_n = jnp.cross(omegab.reshape(3,), nb.reshape(3,)).reshape(3,1)
    a = jnp.cross(nb.reshape(3,), n1.reshape(3,)).reshape(3,1)
    b = jnp.cross(nb.reshape(3,), n2.reshape(3,)).reshape(3,1)

    M_z = jnp.array([[massb/mass, 0], [0, massb/mass]]) 
    M_z = M_z + jnp.array([[1, (n1.T @ n2)[0][0]], [(n1.T @ n2)[0][0], 1]]) 
    M_z = M_z + (massb*d1_*d2 / Jb) * jnp.array([[d1_/d2 * jnp.linalg.norm(a)**2, (a.T @ b)[0][0]], [(a.T @ b)[0][0], d2/d1_ * jnp.linalg.norm(a)**2]])
    
    T_z = jnp.vstack(
              (
                 jnp.hstack( (massb/mass * n1.T, jnp.zeros((1,3))) ), 
                 jnp.hstack( (jnp.zeros((1,3)), massb/mass * n2.T) )
              )
            ) @ jnp.array([0, 0, u[0], 0, 0, u[4]]).reshape(6,1)  
    T_z = T_z +  massb * jnp.array([
                              [jnp.linalg.norm(v1 - (vb + d1_ * S_w_n))**2 / l1], 
                              [jnp.linalg.norm(v2 - (vb + d2 * S_w_n))**2 / l2]])
    T_z = T_z + massb * jnp.linalg.norm(omegab)**2 * jnp.array([[d1_ * (nb.T @ n1)[0][0]], [d2 * (nb.T @ n2)[0][0]]])

    T12 = jnp.linalg.inv(M_z) @ T_z

    T1 = T12[0]
    T2 = T12[1]

    # Control inputs
    F1 = u[0]
    tau1 = jnp.array([u[1], u[2], u[3]])
    F2 = u[4]
    tau2 = jnp.array([u[5], u[6], u[7]])

    # Dynamics
    pb_d = vb
    vb_d = (T1*n1 + T2*n2 - massb*g*jnp.array([0, 0, 1]).reshape((3,1))) / massb
    thetab_d = Hb @ omegab
    cross_pr = jnp.cross(nb.reshape(3,), (d2*T2*n2 - d1*T1*n1).reshape(3,)).reshape(3,1) 
    omegab_d, _, _, _ = jnp.linalg.lstsq( jnp.diag(jnp.array([0.0, 0.0, Jb])) , cross_pr, rcond=None)

    p1_d = v1
    v1_d = (F1* R1 @ jnp.array([0, 0, 1]).reshape(3,1) - mass*g*jnp.array([0, 0, 1]).reshape(3,1) - T1*n1) / mass
    theta1_d = H1@omega1
    cross_pr = jnp.cross(omega1.reshape(3,), (J@omega1).reshape(3,)).reshape(3,1)
    omega1_d, _, _, _ = jnp.linalg.lstsq(J, tau1.reshape(3,1) - cross_pr, rcond=None)

    p2_d = v2
    v2_d = (F2* R2 @ jnp.array([0, 0, 1]).reshape(3,1) - mass*g*jnp.array([0, 0, 1]).reshape((3,1)) - T2*n2) / mass
    theta2_d = H2@omega2
    cross_pr = jnp.cross(omega2.reshape(3,), (J@omega2).reshape(3,)).reshape(3,1)
    omega2_d, _, _, _ = jnp.linalg.lstsq(J, tau2.reshape((3,1)) - cross_pr, rcond=None)

    ret = jnp.vstack((pb_d, vb_d, thetab_d, omegab_d, p1_d, v1_d, theta1_d, omega1_d, p2_d, v2_d, theta2_d, omega2_d)).reshape(36,)  # x_dot
    return ret


def get_eqb(pb = [0, 0, 2], yawb = 0, yaw1 = 0, yaw2 = 0, thetaf = 0, thetal = 0):
  p1 = np.array(pb) + np.array([(d1 + l1*jnp.sin(thetal))*jnp.sin(yawb), (- d1 - l1*jnp.sin(thetal))*jnp.cos(yawb), l1*jnp.cos(thetal)])
  p2 = np.array(pb) + np.array([(-d2 - l2*jnp.sin(thetal))*jnp.sin(yawb), (d2 + l2*jnp.sin(thetal))*jnp.cos(yawb), l2*jnp.cos(thetal)])

  pitch1 = thetaf
  pitch2 = -thetaf

  x1 = np.hstack((p1, [0, 0, 0, 0, pitch1, yaw1, 0, 0, 0]))
  x2 = np.hstack((p2, [0, 0, 0, 0, pitch2, yaw2, 0, 0, 0]))
  xb = np.hstack((pb, [0, 0, 0, 0, 0, yawb, 0, 0, 0]))

  M = 2*mass + massb
  F = g*M / (2*jnp.cos(thetaf)) # eqb thrust

  x = jnp.hstack((xb, x1, x2))
  u = jnp.array([F, 0, 0, 0, F, 0, 0, 0]).reshape(8,)

  return jnp.array(x), jnp.array(u)

def controller():
  pass

if __name__ == "__main__":
    x, u = get_eqb()
    x_dot = dynamicsC(x, u)
    print(x_dot)