Skip to content

Commit a81fdf9

Browse files
A-HayasakaA-Hayasaka
committed
vectorize dynamics jacobian calculation
1 parent ac2c5a8 commit a81fdf9

File tree

1 file changed

+64
-78
lines changed

1 file changed

+64
-78
lines changed

‎constraints.py

Lines changed: 64 additions & 78 deletions
Original file line numberDiff line numberDiff line change
@@ -594,81 +594,78 @@ def dynamics(mass, pos, vel, quat, t):
594594
t_nodes,
595595
)
596596

597+
# mass
598+
mass_i_[1:] += dx
599+
f_p = dynamics(
600+
mass_i_[1:],
601+
pos_i_[1:],
602+
vel_i_[1:],
603+
quat_i_[1:],
604+
t_nodes,
605+
)
606+
mass_i_[1:] -= dx
607+
608+
rh_mass = -(f_p - f_center) / dx * (tf - to) * unit_t / 2.0 # rh acc mass
597609
for j in range(n):
598-
mass_i_[j + 1] += dx
610+
jac["mass"]["coo"][0].extend(list(range((a + j) * 3, (a + j + 1) * 3)))
611+
jac["mass"]["coo"][1].extend([(a + i + j + 1)] * 3)
612+
jac["mass"]["coo"][2].extend(rh_mass[j].tolist())
613+
614+
# position
615+
for k in range(3):
616+
pos_i_[1:, k] += dx
599617
f_p = dynamics(
600618
mass_i_[1:],
601619
pos_i_[1:],
602620
vel_i_[1:],
603621
quat_i_[1:],
604622
t_nodes,
605623
)
606-
mass_i_[j + 1] -= dx
607-
rh_mass = (
608-
-(f_p[j] - f_center[j]) / dx * (tf - to) * unit_t / 2.0
609-
) # rh acc mass
610-
jac["mass"]["coo"][0].extend(list(range((a + j) * 3, (a + j + 1) * 3)))
611-
jac["mass"]["coo"][1].extend([(a + i + j + 1)] * 3)
612-
jac["mass"]["coo"][2].extend(rh_mass.tolist())
624+
pos_i_[1:, k] -= dx
625+
rh_pos = -(f_p - f_center) / dx * (tf - to) * unit_t / 2.0 # rh acc pos
613626

614-
for k in range(3):
615-
pos_i_[j + 1, k] += dx
616-
f_p = dynamics(
617-
mass_i_[1:],
618-
pos_i_[1:],
619-
vel_i_[1:],
620-
quat_i_[1:],
621-
t_nodes,
622-
)
623-
pos_i_[j + 1, k] -= dx
624-
rh_pos = (
625-
-(f_p[j] - f_center[j]) / dx * (tf - to) * unit_t / 2.0
626-
) # rh acc pos
627+
for j in range(n):
627628
jac["position"]["coo"][0].extend(
628629
list(range((a + j) * 3, (a + j + 1) * 3))
629630
)
630631
jac["position"]["coo"][1].extend([(a + i + j + 1) * 3 + k] * 3)
631-
jac["position"]["coo"][2].extend(rh_pos.tolist())
632-
633-
if param[2] > 0.0:
634-
for k in range(3):
635-
vel_i_[j + 1, k] += dx
636-
f_p = dynamics(
637-
mass_i_[1:],
638-
pos_i_[1:],
639-
vel_i_[1:],
640-
quat_i_[1:],
641-
t_nodes,
642-
)
643-
vel_i_[j + 1, k] -= dx
644-
submat_vel[j * 3, (j + 1) * 3 + k] += (
645-
-(f_p[j, 0] - f_center[j, 0]) / dx * (tf - to) * unit_t / 2.0
646-
) # rh acc_x vel
647-
submat_vel[j * 3 + 1, (j + 1) * 3 + k] += (
648-
-(f_p[j, 1] - f_center[j, 1]) / dx * (tf - to) * unit_t / 2.0
649-
) # rh acc_x vel
650-
submat_vel[j * 3 + 2, (j + 1) * 3 + k] += (
651-
-(f_p[j, 2] - f_center[j, 2]) / dx * (tf - to) * unit_t / 2.0
652-
) # rh acc_x vel
632+
jac["position"]["coo"][2].extend(rh_pos[j].tolist())
653633

654-
for k in range(4):
655-
quat_i_[j + 1, k] += dx
634+
# velocity
635+
if param[2] > 0.0:
636+
for k in range(3):
637+
vel_i_[1:, k] += dx
656638
f_p = dynamics(
657639
mass_i_[1:],
658640
pos_i_[1:],
659641
vel_i_[1:],
660642
quat_i_[1:],
661643
t_nodes,
662644
)
663-
quat_i_[j + 1, k] -= dx
664-
rh_quat = (
665-
-(f_p[j] - f_center[j]) / dx * (tf - to) * unit_t / 2.0
666-
) # rh acc quat
645+
vel_i_[1:, k] -= dx
646+
rh_vel = -(f_p - f_center) / dx * (tf - to) * unit_t / 2.0
647+
for j in range(n):
648+
submat_vel[j * 3 : j * 3 + 3, (j + 1) * 3 + k] += rh_vel[j]
649+
650+
# quaternion
651+
for k in range(4):
652+
quat_i_[1:, k] += dx
653+
f_p = dynamics(
654+
mass_i_[1:],
655+
pos_i_[1:],
656+
vel_i_[1:],
657+
quat_i_[1:],
658+
t_nodes,
659+
)
660+
quat_i_[1:, k] -= dx
661+
rh_quat = -(f_p - f_center) / dx * (tf - to) * unit_t / 2.0 # rh acc quat
662+
663+
for j in range(n):
667664
jac["quaternion"]["coo"][0].extend(
668665
list(range((a + j) * 3, (a + j + 1) * 3))
669666
)
670667
jac["quaternion"]["coo"][1].extend([(a + i + j + 1) * 4 + k] * 3)
671-
jac["quaternion"]["coo"][2].extend(rh_quat.tolist())
668+
jac["quaternion"]["coo"][2].extend(rh_quat[j].tolist())
672669

673670
rh_to = f_center.ravel() * unit_t / 2.0 # rh to
674671
rh_tf = -rh_to # rh tf
@@ -783,36 +780,25 @@ def equality_jac_dynamics_quaternion(xdict, pdict, unitdict, condition):
783780

784781
f_center = dynamics_quaternion(quat_i_[1:], u_i_, unit_u)
785782

786-
for j in range(n):
787-
788-
for k in range(4):
789-
quat_i_[j + 1, k] += dx
790-
f_p = dynamics_quaternion(quat_i_[1:], u_i_, unit_u)
791-
submat_quat[j * 4, (j + 1) * 4 + k] += (
792-
-(f_p[j, 0] - f_center[j, 0]) / dx * (tf - to) * unit_t / 2.0
793-
) # rh q0 quat
794-
submat_quat[j * 4 + 1, (j + 1) * 4 + k] += (
795-
-(f_p[j, 1] - f_center[j, 1]) / dx * (tf - to) * unit_t / 2.0
796-
) # rh q1 quat
797-
submat_quat[j * 4 + 2, (j + 1) * 4 + k] += (
798-
-(f_p[j, 2] - f_center[j, 2]) / dx * (tf - to) * unit_t / 2.0
799-
) # rh q2 quat
800-
submat_quat[j * 4 + 3, (j + 1) * 4 + k] += (
801-
-(f_p[j, 3] - f_center[j, 3]) / dx * (tf - to) * unit_t / 2.0
802-
) # rh q3 quat
803-
quat_i_[j + 1, k] -= dx
804-
805-
for k in range(3):
806-
u_i_[j, k] += dx
807-
f_p = dynamics_quaternion(quat_i_[1:], u_i_, unit_u)
808-
u_i_[j, k] -= dx
809-
810-
rh_pos = (
811-
-(f_p[j] - f_center[j]) / dx * (tf - to) * unit_t / 2.0
812-
) # rh q0 quat
783+
# quaternion
784+
for k in range(4):
785+
quat_i_[1:, k] += dx
786+
f_p = dynamics_quaternion(quat_i_[1:], u_i_, unit_u)
787+
quat_i_[1:, k] -= dx
788+
rh_quat = -(f_p - f_center) / dx * (tf - to) * unit_t / 2.0
789+
for j in range(n):
790+
submat_quat[j * 4 : j * 4 + 4, (j + 1) * 4 + k] += rh_quat[j]
791+
792+
# u (angular velocity)
793+
for k in range(3):
794+
u_i_[:, k] += dx
795+
f_p = dynamics_quaternion(quat_i_[1:], u_i_, unit_u)
796+
u_i_[:, k] -= dx
797+
rh_u = -(f_p - f_center) / dx * (tf - to) * unit_t / 2.0
798+
for j in range(n):
813799
jac["u"]["coo"][0].extend(list(range((a + j) * 4, (a + j + 1) * 4)))
814800
jac["u"]["coo"][1].extend([(a + j) * 3 + k] * 4)
815-
jac["u"]["coo"][2].extend(rh_pos.tolist())
801+
jac["u"]["coo"][2].extend(rh_u[j].tolist())
816802

817803
rh_to = f_center.ravel() * unit_t / 2.0 # rh to
818804
rh_tf = -rh_to # rh tf

0 commit comments

Comments
 (0)