Skip to content

Commit

Permalink
Save some kindyn computation in JaxSimModelData
Browse files Browse the repository at this point in the history
  • Loading branch information
flferretti committed Dec 20, 2024
1 parent be7f69a commit f257982
Show file tree
Hide file tree
Showing 24 changed files with 497 additions and 211 deletions.
4 changes: 2 additions & 2 deletions src/jaxsim/api/com.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ def com_position(

m = js.model.total_mass(model=model)

W_H_L = js.model.forward_kinematics(model=model, data=data)
W_H_L = data.kyn_dyn.forward_kinematics
W_H_B = data.base_transform()
B_H_W = jaxsim.math.Transform.inverse(transform=W_H_B)

Expand Down Expand Up @@ -269,7 +269,7 @@ def bias_acceleration(
"""

# Compute the pose of all links with forward kinematics.
W_H_L = js.model.forward_kinematics(model=model, data=data)
W_H_L = data.kyn_dyn.forward_kinematics

# Compute the bias acceleration of all links by zeroing the generalized velocity
# in the active representation.
Expand Down
49 changes: 49 additions & 0 deletions src/jaxsim/api/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -232,3 +232,52 @@ def other_representation_to_inertial(

case _:
raise ValueError(other_representation)


def convert_mass_matrix(
M: jtp.Matrix,
base_transform: jtp.Matrix,
dofs: jtp.Int,
velocity_representation: VelRepr,
):

# The mass matrix is always save in body-fixed representation.

match velocity_representation:
case VelRepr.Body:
return M

case VelRepr.Inertial:

B_X_W = Adjoint.from_transform(transform=base_transform, inverse=True)
invT = jax.scipy.linalg.block_diag(B_X_W, jnp.eye(dofs))

return invT.T @ M @ invT

case VelRepr.Mixed:

BW_H_B = base_transform.at[0:3, 3].set(jnp.zeros(3))
B_X_BW = Adjoint.from_transform(transform=BW_H_B, inverse=True)
invT = jax.scipy.linalg.block_diag(B_X_BW, jnp.eye(dofs))

return invT.T @ M @ invT


def convert_jacobian(
J: jtp.Matrix,
base_transform: jtp.Matrix,
dofs: jtp.Int,
velocity_representation: VelRepr,
):
# TODO (flferretti): save actual Jacobian instead of full doubly left and perform conversion.
return J


def convert_jacobian_derivative(
Jd: jtp.Matrix,
base_transform: jtp.Matrix,
dofs: jtp.Int,
velocity_representation: VelRepr,
):
# TODO (flferretti): save actual Jacobian derivative instead of full doubly left and perform conversion.
return Jd
8 changes: 5 additions & 3 deletions src/jaxsim/api/contact.py
Original file line number Diff line number Diff line change
Expand Up @@ -42,12 +42,14 @@ def collidable_point_kinematics(

W_p_Ci, W_ṗ_Ci = jaxsim.rbda.collidable_points.collidable_points_pos_vel(
model=model,
base_position=data.base_position(),
base_position=data.base_position,
base_quaternion=data.base_orientation(dcm=False),
joint_positions=data.joint_positions(model=model),
base_linear_velocity=data.base_velocity()[0:3],
base_angular_velocity=data.base_velocity()[3:6],
joint_velocities=data.joint_velocities(model=model),
joint_transforms=data.kyn_dyn.joint_transforms,
motion_subspaces=data.kyn_dyn.motion_subspaces,
)

return W_p_Ci, W_ṗ_Ci
Expand Down Expand Up @@ -460,7 +462,7 @@ def transforms(model: js.model.JaxSimModel, data: js.data.JaxSimModelData) -> jt
)[indices_of_enabled_collidable_points]

# Get the transforms of the parent link of all collidable points.
W_H_L = js.model.forward_kinematics(model=model, data=data)[
W_H_L = data.kyn_dyn.forward_kinematics[
parent_link_idx_of_enabled_collidable_points
]

Expand Down Expand Up @@ -612,7 +614,7 @@ def jacobian_derivative(
]

# Get the transforms of all the parent links.
W_H_Li = js.model.forward_kinematics(model=model, data=data)
W_H_Li = data.kyn_dyn.forward_kinematics

# =====================================================
# Compute quantities to adjust the input representation
Expand Down
Loading

0 comments on commit f257982

Please sign in to comment.