# Copyright 2025 The Newton Developers
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
import warp as wp
from mujoco_warp._src import math
from mujoco_warp._src import support
from mujoco_warp._src.types import MJ_MINVAL
from mujoco_warp._src.types import Data
from mujoco_warp._src.types import DisableBit
from mujoco_warp._src.types import GeomType
from mujoco_warp._src.types import JointType
from mujoco_warp._src.types import Model
from mujoco_warp._src.warp_util import event_scope
wp.set_module_options({"enable_backward": False})
@wp.func
def _pow2(val: float) -> float:
return val * val
@wp.func
def _pow4(val: float) -> float:
sq = val * val
return sq * sq
@wp.func
def _geom_semiaxes(size: wp.vec3, geom_type: int) -> wp.vec3: # kernel_analyzer: ignore
if geom_type == GeomType.SPHERE:
r = size[0]
return wp.vec3(r, r, r)
if geom_type == GeomType.CAPSULE:
radius = size[0]
half_length = size[1]
return wp.vec3(radius, radius, half_length + radius)
if geom_type == GeomType.CYLINDER:
radius = size[0]
half_length = size[1]
return wp.vec3(radius, radius, half_length)
# ellipsoid, box, mesh, sdf -> use size directly
return size
@wp.func
def _ellipsoid_max_moment(size: wp.vec3, dir: int) -> float:
d0 = size[dir]
d1 = size[(dir + 1) % 3]
d2 = size[(dir + 2) % 3]
return wp.static(8.0 / 15.0 * wp.pi) * d0 * _pow4(wp.max(d1, d2))
@wp.kernel
def _spring_damper_dof_passive(
# Model:
opt_disableflags: int,
qpos_spring: wp.array2d(dtype=float),
jnt_type: wp.array(dtype=int),
jnt_qposadr: wp.array(dtype=int),
jnt_dofadr: wp.array(dtype=int),
jnt_stiffness: wp.array2d(dtype=float),
dof_damping: wp.array2d(dtype=float),
# Data in:
qpos_in: wp.array2d(dtype=float),
qvel_in: wp.array2d(dtype=float),
# Data out:
qfrc_spring_out: wp.array2d(dtype=float),
qfrc_damper_out: wp.array2d(dtype=float),
):
worldid, jntid = wp.tid()
dofid = jnt_dofadr[jntid]
stiffness = jnt_stiffness[worldid % jnt_stiffness.shape[0], jntid]
damping = dof_damping[worldid % dof_damping.shape[0], dofid]
has_stiffness = stiffness != 0.0 and not opt_disableflags & DisableBit.SPRING
has_damping = damping != 0.0 and not opt_disableflags & DisableBit.DAMPER
if not has_stiffness:
qfrc_spring_out[worldid, dofid] = 0.0
if not has_damping:
qfrc_damper_out[worldid, dofid] = 0.0
if not (has_stiffness or has_damping):
return
jnttype = jnt_type[jntid]
qposid = jnt_qposadr[jntid]
qpos_spring_id = worldid % qpos_spring.shape[0]
if jnttype == JointType.FREE:
# spring
if has_stiffness:
dif = wp.vec3(
qpos_in[worldid, qposid + 0] - qpos_spring[qpos_spring_id, qposid + 0],
qpos_in[worldid, qposid + 1] - qpos_spring[qpos_spring_id, qposid + 1],
qpos_in[worldid, qposid + 2] - qpos_spring[qpos_spring_id, qposid + 2],
)
qfrc_spring_out[worldid, dofid + 0] = -stiffness * dif[0]
qfrc_spring_out[worldid, dofid + 1] = -stiffness * dif[1]
qfrc_spring_out[worldid, dofid + 2] = -stiffness * dif[2]
rot = wp.quat(
qpos_in[worldid, qposid + 3],
qpos_in[worldid, qposid + 4],
qpos_in[worldid, qposid + 5],
qpos_in[worldid, qposid + 6],
)
rot = wp.normalize(rot)
ref = wp.quat(
qpos_spring[qpos_spring_id, qposid + 3],
qpos_spring[qpos_spring_id, qposid + 4],
qpos_spring[qpos_spring_id, qposid + 5],
qpos_spring[qpos_spring_id, qposid + 6],
)
dif = math.quat_sub(rot, ref)
qfrc_spring_out[worldid, dofid + 3] = -stiffness * dif[0]
qfrc_spring_out[worldid, dofid + 4] = -stiffness * dif[1]
qfrc_spring_out[worldid, dofid + 5] = -stiffness * dif[2]
# damper
if has_damping:
qfrc_damper_out[worldid, dofid + 0] = -damping * qvel_in[worldid, dofid + 0]
qfrc_damper_out[worldid, dofid + 1] = -damping * qvel_in[worldid, dofid + 1]
qfrc_damper_out[worldid, dofid + 2] = -damping * qvel_in[worldid, dofid + 2]
qfrc_damper_out[worldid, dofid + 3] = -damping * qvel_in[worldid, dofid + 3]
qfrc_damper_out[worldid, dofid + 4] = -damping * qvel_in[worldid, dofid + 4]
qfrc_damper_out[worldid, dofid + 5] = -damping * qvel_in[worldid, dofid + 5]
elif jnttype == JointType.BALL:
# spring
if has_stiffness:
rot = wp.quat(
qpos_in[worldid, qposid + 0],
qpos_in[worldid, qposid + 1],
qpos_in[worldid, qposid + 2],
qpos_in[worldid, qposid + 3],
)
rot = wp.normalize(rot)
ref = wp.quat(
qpos_spring[qpos_spring_id, qposid + 0],
qpos_spring[qpos_spring_id, qposid + 1],
qpos_spring[qpos_spring_id, qposid + 2],
qpos_spring[qpos_spring_id, qposid + 3],
)
dif = math.quat_sub(rot, ref)
qfrc_spring_out[worldid, dofid + 0] = -stiffness * dif[0]
qfrc_spring_out[worldid, dofid + 1] = -stiffness * dif[1]
qfrc_spring_out[worldid, dofid + 2] = -stiffness * dif[2]
# damper
if has_damping:
qfrc_damper_out[worldid, dofid + 0] = -damping * qvel_in[worldid, dofid + 0]
qfrc_damper_out[worldid, dofid + 1] = -damping * qvel_in[worldid, dofid + 1]
qfrc_damper_out[worldid, dofid + 2] = -damping * qvel_in[worldid, dofid + 2]
else: # mjJNT_SLIDE, mjJNT_HINGE
# spring
if has_stiffness:
fdif = qpos_in[worldid, qposid] - qpos_spring[qpos_spring_id, qposid]
qfrc_spring_out[worldid, dofid] = -stiffness * fdif
# damper
if has_damping:
qfrc_damper_out[worldid, dofid] = -damping * qvel_in[worldid, dofid]
@wp.kernel
def _spring_damper_tendon_passive(
# Model:
tendon_stiffness: wp.array2d(dtype=float),
tendon_damping: wp.array2d(dtype=float),
tendon_lengthspring: wp.array2d(dtype=wp.vec2),
# Data in:
ten_J_in: wp.array3d(dtype=float),
ten_length_in: wp.array2d(dtype=float),
ten_velocity_in: wp.array2d(dtype=float),
# In:
dsbl_spring: bool,
dsbl_damper: bool,
# Data out:
qfrc_spring_out: wp.array2d(dtype=float),
qfrc_damper_out: wp.array2d(dtype=float),
):
worldid, tenid, dofid = wp.tid()
stiffness = tendon_stiffness[worldid % tendon_stiffness.shape[0], tenid]
damping = tendon_damping[worldid % tendon_damping.shape[0], tenid]
has_stiffness = stiffness != 0.0 and not dsbl_spring
has_damping = damping != 0.0 and not dsbl_damper
if not has_stiffness and not has_damping:
return
J = ten_J_in[worldid, tenid, dofid]
if has_stiffness:
# compute spring force along tendon
length = ten_length_in[worldid, tenid]
lengthspring = tendon_lengthspring[worldid % tendon_lengthspring.shape[0], tenid]
lower = lengthspring[0]
upper = lengthspring[1]
if length > upper:
frc_spring = stiffness * (upper - length)
elif length < lower:
frc_spring = stiffness * (lower - length)
else:
frc_spring = 0.0
# transform to joint torque
wp.atomic_add(qfrc_spring_out[worldid], dofid, J * frc_spring)
if has_damping:
# compute damper linear force along tendon
frc_damper = -damping * ten_velocity_in[worldid, tenid]
# transform to joint torque
wp.atomic_add(qfrc_damper_out[worldid], dofid, J * frc_damper)
@wp.kernel
def _gravity_force(
# Model:
opt_gravity: wp.array(dtype=wp.vec3),
body_parentid: wp.array(dtype=int),
body_rootid: wp.array(dtype=int),
body_mass: wp.array2d(dtype=float),
body_gravcomp: wp.array2d(dtype=float),
dof_bodyid: wp.array(dtype=int),
# Data in:
xipos_in: wp.array2d(dtype=wp.vec3),
subtree_com_in: wp.array2d(dtype=wp.vec3),
cdof_in: wp.array2d(dtype=wp.spatial_vector),
# Data out:
qfrc_gravcomp_out: wp.array2d(dtype=float),
):
worldid, bodyid, dofid = wp.tid()
bodyid += 1 # skip world body
gravcomp = body_gravcomp[worldid % body_gravcomp.shape[0], bodyid]
gravity = opt_gravity[worldid % opt_gravity.shape[0]]
if gravcomp:
force = -gravity * body_mass[worldid % body_mass.shape[0], bodyid] * gravcomp
pos = xipos_in[worldid, bodyid]
jac, _ = support.jac_dof(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos, bodyid, dofid, worldid)
wp.atomic_add(qfrc_gravcomp_out[worldid], dofid, wp.dot(jac, force))
@wp.kernel
def _fluid_force(
# Model:
opt_density: wp.array(dtype=float),
opt_viscosity: wp.array(dtype=float),
opt_wind: wp.array(dtype=wp.vec3),
body_rootid: wp.array(dtype=int),
body_geomnum: wp.array(dtype=int),
body_geomadr: wp.array(dtype=int),
body_mass: wp.array2d(dtype=float),
body_inertia: wp.array2d(dtype=wp.vec3),
geom_type: wp.array(dtype=int),
geom_size: wp.array2d(dtype=wp.vec3),
geom_fluid: wp.array2d(dtype=float),
body_fluid_ellipsoid: wp.array(dtype=bool),
# Data in:
xipos_in: wp.array2d(dtype=wp.vec3),
ximat_in: wp.array2d(dtype=wp.mat33),
geom_xpos_in: wp.array2d(dtype=wp.vec3),
geom_xmat_in: wp.array2d(dtype=wp.mat33),
subtree_com_in: wp.array2d(dtype=wp.vec3),
cvel_in: wp.array2d(dtype=wp.spatial_vector),
# Out:
fluid_applied_out: wp.array2d(dtype=wp.spatial_vector),
):
"""Computes body-space fluid forces for both inertia-box and ellipsoid models."""
worldid, bodyid = wp.tid()
zero_force = wp.spatial_vector(wp.vec3(0.0), wp.vec3(0.0))
if bodyid == 0:
fluid_applied_out[worldid, bodyid] = zero_force
return
# skip bodies with negligible mass
mass = body_mass[worldid % body_mass.shape[0], bodyid]
if mass < MJ_MINVAL:
fluid_applied_out[worldid, bodyid] = zero_force
return
wind = opt_wind[worldid % opt_wind.shape[0]]
density = opt_density[worldid % opt_density.shape[0]]
viscosity = opt_viscosity[worldid % opt_viscosity.shape[0]]
# Body kinematics
xipos = xipos_in[worldid, bodyid]
rot = ximat_in[worldid, bodyid]
rotT = wp.transpose(rot)
cvel = cvel_in[worldid, bodyid]
ang_global = wp.spatial_top(cvel)
lin_global = wp.spatial_bottom(cvel)
subtree_root = subtree_com_in[worldid, body_rootid[bodyid]]
lin_com = lin_global - wp.cross(xipos - subtree_root, ang_global)
if body_fluid_ellipsoid[bodyid]:
force_global = wp.vec3(0.0)
torque_global = wp.vec3(0.0)
start = body_geomadr[bodyid]
count = body_geomnum[bodyid]
for i in range(count):
geomid = start + i
coef = geom_fluid[geomid, 0]
if coef <= 0.0:
continue
size = geom_size[worldid % geom_size.shape[0], geomid]
semiaxes = _geom_semiaxes(size, geom_type[geomid])
geom_rot = geom_xmat_in[worldid, geomid]
geom_rotT = wp.transpose(geom_rot)
geom_pos = geom_xpos_in[worldid, geomid]
lin_point = lin_com + wp.cross(ang_global, geom_pos - xipos)
l_ang = geom_rotT @ ang_global
l_lin = geom_rotT @ lin_point
if wind[0] or wind[1] or wind[2]:
l_lin -= geom_rotT @ wind
lfrc_torque = wp.vec3(0.0)
lfrc_force = wp.vec3(0.0)
if density > 0.0:
# added-mass forces and torques
virtual_mass = wp.vec3(geom_fluid[geomid, 6], geom_fluid[geomid, 7], geom_fluid[geomid, 8])
virtual_inertia = wp.vec3(geom_fluid[geomid, 9], geom_fluid[geomid, 10], geom_fluid[geomid, 11])
virtual_lin_mom = wp.vec3(
density * virtual_mass[0] * l_lin[0],
density * virtual_mass[1] * l_lin[1],
density * virtual_mass[2] * l_lin[2],
)
virtual_ang_mom = wp.vec3(
density * virtual_inertia[0] * l_ang[0],
density * virtual_inertia[1] * l_ang[1],
density * virtual_inertia[2] * l_ang[2],
)
added_mass_force = wp.cross(virtual_lin_mom, l_ang)
added_mass_torque = wp.cross(virtual_lin_mom, l_lin) + wp.cross(virtual_ang_mom, l_ang)
lfrc_force += added_mass_force
lfrc_torque += added_mass_torque
# lift force orthogonal to velocity from Kutta-Joukowski theorem
magnus_coef = geom_fluid[geomid, 5]
kutta_coef = geom_fluid[geomid, 4]
blunt_drag_coef = geom_fluid[geomid, 1]
slender_drag_coef = geom_fluid[geomid, 2]
ang_drag_coef = geom_fluid[geomid, 3]
volume = wp.static(4.0 / 3.0 * wp.pi) * semiaxes[0] * semiaxes[1] * semiaxes[2]
d_max = wp.max(wp.max(semiaxes[0], semiaxes[1]), semiaxes[2])
d_min = wp.min(wp.min(semiaxes[0], semiaxes[1]), semiaxes[2])
d_mid = semiaxes[0] + semiaxes[1] + semiaxes[2] - d_max - d_min
A_max = wp.pi * d_max * d_mid
lin_speed = wp.length(l_lin)
magnus_force = wp.cross(l_ang, l_lin) * (magnus_coef * density * volume)
s12 = semiaxes[1] * semiaxes[2]
s20 = semiaxes[2] * semiaxes[0]
s01 = semiaxes[0] * semiaxes[1]
proj_denom = _pow4(s12) * _pow2(l_lin[0]) + _pow4(s20) * _pow2(l_lin[1]) + _pow4(s01) * _pow2(l_lin[2])
proj_num = _pow2(s12 * l_lin[0]) + _pow2(s20 * l_lin[1]) + _pow2(s01 * l_lin[2])
A_proj = 0.0
cos_alpha = 0.0
if proj_num > MJ_MINVAL and proj_denom > MJ_MINVAL:
A_proj = wp.pi * wp.sqrt(proj_denom / wp.max(MJ_MINVAL, proj_num))
if lin_speed > MJ_MINVAL:
cos_alpha = proj_num / wp.max(MJ_MINVAL, lin_speed * proj_denom)
norm = wp.vec3(
_pow2(s12) * l_lin[0],
_pow2(s20) * l_lin[1],
_pow2(s01) * l_lin[2],
)
kutta_force = wp.vec3(0.0)
if density > 0.0 and kutta_coef != 0.0 and lin_speed > MJ_MINVAL:
kutta_circ = wp.cross(norm, l_lin) * (kutta_coef * density * cos_alpha * A_proj)
kutta_force = wp.cross(kutta_circ, l_lin)
eq_sphere_D = wp.static(2.0 / 3.0) * (semiaxes[0] + semiaxes[1] + semiaxes[2])
lin_visc_force_coef = wp.static(3.0 * wp.pi) * eq_sphere_D
lin_visc_torq_coef = wp.pi * eq_sphere_D * eq_sphere_D * eq_sphere_D
I_max = wp.static(8.0 / 15.0 * wp.pi) * d_mid * _pow4(d_max)
II0 = _ellipsoid_max_moment(semiaxes, 0)
II1 = _ellipsoid_max_moment(semiaxes, 1)
II2 = _ellipsoid_max_moment(semiaxes, 2)
mom_visc = wp.vec3(
l_ang[0] * (ang_drag_coef * II0 + slender_drag_coef * (I_max - II0)),
l_ang[1] * (ang_drag_coef * II1 + slender_drag_coef * (I_max - II1)),
l_ang[2] * (ang_drag_coef * II2 + slender_drag_coef * (I_max - II2)),
)
drag_lin_coef = viscosity * lin_visc_force_coef + density * lin_speed * (
A_proj * blunt_drag_coef + slender_drag_coef * (A_max - A_proj)
)
drag_ang_coef = viscosity * lin_visc_torq_coef + density * wp.length(mom_visc)
lfrc_torque -= drag_ang_coef * l_ang
lfrc_force += magnus_force + kutta_force - drag_lin_coef * l_lin
lfrc_torque *= coef
lfrc_force *= coef
# map force/torque from local to world frame: lfrc -> bfrc
torque_global += geom_rot @ lfrc_torque
force_global += geom_rot @ lfrc_force
fluid_applied_out[worldid, bodyid] = wp.spatial_vector(force_global, torque_global)
return
l_ang = rotT @ ang_global
l_lin = rotT @ lin_com
if wind[0] or wind[1] or wind[2]:
l_lin -= rotT @ wind
lfrc_torque = wp.vec3(0.0)
lfrc_force = wp.vec3(0.0)
has_viscosity = viscosity > 0.0
has_density = density > 0.0
if has_viscosity or has_density:
inertia = body_inertia[worldid % body_inertia.shape[0], bodyid]
mass = body_mass[worldid % body_mass.shape[0], bodyid]
scl = 6.0 / mass
box0 = wp.sqrt(wp.max(MJ_MINVAL, inertia[1] + inertia[2] - inertia[0]) * scl)
box1 = wp.sqrt(wp.max(MJ_MINVAL, inertia[0] + inertia[2] - inertia[1]) * scl)
box2 = wp.sqrt(wp.max(MJ_MINVAL, inertia[0] + inertia[1] - inertia[2]) * scl)
if has_viscosity:
diam = (box0 + box1 + box2) / 3.0
lfrc_torque = -l_ang * wp.pow(diam, 3.0) * wp.pi * viscosity
lfrc_force = -3.0 * l_lin * diam * wp.pi * viscosity
if has_density:
lfrc_force -= wp.vec3(
0.5 * density * box1 * box2 * wp.abs(l_lin[0]) * l_lin[0],
0.5 * density * box0 * box2 * wp.abs(l_lin[1]) * l_lin[1],
0.5 * density * box0 * box1 * wp.abs(l_lin[2]) * l_lin[2],
)
scl = density / 64.0
box0_pow4 = wp.pow(box0, 4.0)
box1_pow4 = wp.pow(box1, 4.0)
box2_pow4 = wp.pow(box2, 4.0)
lfrc_torque -= wp.vec3(
box0 * (box1_pow4 + box2_pow4) * wp.abs(l_ang[0]) * l_ang[0] * scl,
box1 * (box0_pow4 + box2_pow4) * wp.abs(l_ang[1]) * l_ang[1] * scl,
box2 * (box0_pow4 + box1_pow4) * wp.abs(l_ang[2]) * l_ang[2] * scl,
)
torque_global = rot @ lfrc_torque
force_global = rot @ lfrc_force
fluid_applied_out[worldid, bodyid] = wp.spatial_vector(force_global, torque_global)
def _fluid(m: Model, d: Data):
fluid_applied = wp.empty((d.nworld, m.nbody), dtype=wp.spatial_vector)
wp.launch(
_fluid_force,
dim=(d.nworld, m.nbody),
inputs=[
m.opt.density,
m.opt.viscosity,
m.opt.wind,
m.body_rootid,
m.body_geomnum,
m.body_geomadr,
m.body_mass,
m.body_inertia,
m.geom_type,
m.geom_size,
m.geom_fluid,
m.body_fluid_ellipsoid,
d.xipos,
d.ximat,
d.geom_xpos,
d.geom_xmat,
d.subtree_com,
d.cvel,
],
outputs=[fluid_applied],
)
support.apply_ft(m, d, fluid_applied, d.qfrc_fluid, False)
@wp.kernel
def _qfrc_passive(
# Model:
jnt_actgravcomp: wp.array(dtype=int),
dof_jntid: wp.array(dtype=int),
has_fluid: bool,
# Data in:
qfrc_spring_in: wp.array2d(dtype=float),
qfrc_damper_in: wp.array2d(dtype=float),
qfrc_gravcomp_in: wp.array2d(dtype=float),
qfrc_fluid_in: wp.array2d(dtype=float),
# In:
gravcomp: bool,
# Data out:
qfrc_passive_out: wp.array2d(dtype=float),
):
worldid, dofid = wp.tid()
qfrc_passive = qfrc_spring_in[worldid, dofid]
qfrc_passive += qfrc_damper_in[worldid, dofid]
# add gravcomp unless added by actuators
if gravcomp and not jnt_actgravcomp[dof_jntid[dofid]]:
qfrc_passive += qfrc_gravcomp_in[worldid, dofid]
# add fluid force
if has_fluid:
qfrc_passive += qfrc_fluid_in[worldid, dofid]
qfrc_passive_out[worldid, dofid] = qfrc_passive
@wp.kernel
def _flex_elasticity(
# Model:
nflex: int,
opt_timestep: wp.array(dtype=float),
body_dofadr: wp.array(dtype=int),
flex_dim: wp.array(dtype=int),
flex_vertadr: wp.array(dtype=int),
flex_edgeadr: wp.array(dtype=int),
flex_elemadr: wp.array(dtype=int),
flex_elemnum: wp.array(dtype=int),
flex_elemedgeadr: wp.array(dtype=int),
flex_vertbodyid: wp.array(dtype=int),
flex_elem: wp.array(dtype=int),
flex_elemedge: wp.array(dtype=int),
flexedge_length0: wp.array(dtype=float),
flex_stiffness: wp.array2d(dtype=float),
flex_damping: wp.array(dtype=float),
# Data in:
flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
flexedge_length_in: wp.array2d(dtype=float),
flexedge_velocity_in: wp.array2d(dtype=float),
# In:
dsbl_damper: bool,
# Data out:
qfrc_spring_out: wp.array2d(dtype=float),
):
worldid, elemid = wp.tid()
timestep = opt_timestep[worldid % opt_timestep.shape[0]]
for i in range(nflex):
locid = elemid - flex_elemadr[i]
if locid >= 0 and locid < flex_elemnum[i]:
f = i
break
dim = flex_dim[f]
nvert = dim + 1
nedge = nvert * (nvert - 1) / 2
edges = wp.where(
dim == 3,
wp.matrix(0, 1, 1, 2, 2, 0, 2, 3, 0, 3, 1, 3, shape=(6, 2), dtype=int),
wp.matrix(1, 2, 2, 0, 0, 1, 0, 0, 0, 0, 0, 0, shape=(6, 2), dtype=int),
)
if timestep > 0.0 and not dsbl_damper:
kD = flex_damping[f] / timestep
else:
kD = 0.0
gradient = wp.matrix(0.0, shape=(6, 6))
for e in range(nedge):
vert0 = flex_elem[(dim + 1) * elemid + edges[e, 0]]
vert1 = flex_elem[(dim + 1) * elemid + edges[e, 1]]
xpos0 = flexvert_xpos_in[worldid, vert0]
xpos1 = flexvert_xpos_in[worldid, vert1]
for i in range(3):
gradient[e, 0 + i] = xpos0[i] - xpos1[i]
gradient[e, 3 + i] = xpos1[i] - xpos0[i]
elongation = wp.spatial_vectorf(0.0)
for e in range(nedge):
idx = flex_elemedge[elemid * nedge + e]
vel = flexedge_velocity_in[worldid, flex_edgeadr[f] + idx]
deformed = flexedge_length_in[worldid, flex_edgeadr[f] + idx]
reference = flexedge_length0[flex_edgeadr[f] + idx]
previous = deformed - vel * timestep
elongation[e] = deformed * deformed - reference * reference + (deformed * deformed - previous * previous) * kD
metric = wp.matrix(0.0, shape=(6, 6))
id = int(0)
for ed1 in range(nedge):
for ed2 in range(ed1, nedge):
metric[ed1, ed2] = flex_stiffness[elemid, id]
metric[ed2, ed1] = flex_stiffness[elemid, id]
id += 1
force = wp.matrix(0.0, shape=(6, 3))
for ed1 in range(nedge):
for ed2 in range(nedge):
for i in range(2):
for x in range(3):
force[edges[ed2, i], x] -= elongation[ed1] * gradient[ed2, 3 * i + x] * metric[ed1, ed2]
for v in range(nvert):
vert = flex_elem[(dim + 1) * elemid + v]
bodyid = flex_vertbodyid[flex_vertadr[f] + vert]
for x in range(3):
wp.atomic_add(qfrc_spring_out, worldid, body_dofadr[bodyid] + x, force[v, x])
@wp.kernel
def _flex_bending(
# Model:
nflex: int,
body_dofadr: wp.array(dtype=int),
flex_dim: wp.array(dtype=int),
flex_vertadr: wp.array(dtype=int),
flex_edgeadr: wp.array(dtype=int),
flex_edgenum: wp.array(dtype=int),
flex_vertbodyid: wp.array(dtype=int),
flex_edge: wp.array(dtype=wp.vec2i),
flex_edgeflap: wp.array(dtype=wp.vec2i),
flex_bending: wp.array2d(dtype=float),
# Data in:
flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
# Data out:
qfrc_spring_out: wp.array2d(dtype=float),
):
worldid, edgeid = wp.tid()
nvert = 4
for i in range(nflex):
locid = edgeid - flex_edgeadr[i]
if locid >= 0 and locid < flex_edgenum[i]:
f = i
break
if flex_dim[f] != 2:
return
if flex_edgeflap[edgeid][1] == -1:
return
v = wp.vec4i(
flex_vertadr[f] + flex_edge[edgeid][0],
flex_vertadr[f] + flex_edge[edgeid][1],
flex_vertadr[f] + flex_edgeflap[edgeid][0],
flex_vertadr[f] + flex_edgeflap[edgeid][1],
)
frc = wp.matrix(0.0, shape=(4, 3))
if flex_bending[edgeid, 16]:
v0 = flexvert_xpos_in[worldid, v[0]]
v1 = flexvert_xpos_in[worldid, v[1]]
v2 = flexvert_xpos_in[worldid, v[2]]
v3 = flexvert_xpos_in[worldid, v[3]]
frc[1] = wp.cross(v2 - v0, v3 - v0)
frc[2] = wp.cross(v3 - v0, v1 - v0)
frc[3] = wp.cross(v1 - v0, v2 - v0)
frc[0] = -(frc[1] + frc[2] + frc[3])
force = wp.matrix(0.0, shape=(nvert, 3))
for i in range(nvert):
for x in range(3):
for j in range(nvert):
force[i, x] -= flex_bending[edgeid, 4 * i + j] * flexvert_xpos_in[worldid, v[j]][x]
force[i, x] -= flex_bending[edgeid, 16] * frc[i, x]
for i in range(nvert):
bodyid = flex_vertbodyid[v[i]]
for x in range(3):
wp.atomic_add(qfrc_spring_out, worldid, body_dofadr[bodyid] + x, force[i, x])
[docs]
@event_scope
def passive(m: Model, d: Data):
"""Adds all passive forces."""
dsbl_spring = m.opt.disableflags & DisableBit.SPRING
dsbl_damper = m.opt.disableflags & DisableBit.DAMPER
if dsbl_spring and dsbl_damper:
d.qfrc_spring.zero_()
d.qfrc_damper.zero_()
d.qfrc_gravcomp.zero_()
d.qfrc_fluid.zero_()
d.qfrc_passive.zero_()
return
wp.launch(
_spring_damper_dof_passive,
dim=(d.nworld, m.njnt),
inputs=[
m.opt.disableflags,
m.qpos_spring,
m.jnt_type,
m.jnt_qposadr,
m.jnt_dofadr,
m.jnt_stiffness,
m.dof_damping,
d.qpos,
d.qvel,
],
outputs=[d.qfrc_spring, d.qfrc_damper],
)
if m.ntendon:
wp.launch(
_spring_damper_tendon_passive,
dim=(d.nworld, m.ntendon, m.nv),
inputs=[
m.tendon_stiffness,
m.tendon_damping,
m.tendon_lengthspring,
d.ten_J,
d.ten_length,
d.ten_velocity,
dsbl_spring,
dsbl_damper,
],
outputs=[
d.qfrc_spring,
d.qfrc_damper,
],
)
if not dsbl_spring:
wp.launch(
_flex_elasticity,
dim=(d.nworld, m.nflexelem),
inputs=[
m.nflex,
m.opt.timestep,
m.body_dofadr,
m.flex_dim,
m.flex_vertadr,
m.flex_edgeadr,
m.flex_elemadr,
m.flex_elemnum,
m.flex_elemedgeadr,
m.flex_vertbodyid,
m.flex_elem,
m.flex_elemedge,
m.flexedge_length0,
m.flex_stiffness,
m.flex_damping,
d.flexvert_xpos,
d.flexedge_length,
d.flexedge_velocity,
dsbl_damper,
],
outputs=[d.qfrc_spring],
)
wp.launch(
_flex_bending,
dim=(d.nworld, m.nflexedge),
inputs=[
m.nflex,
m.body_dofadr,
m.flex_dim,
m.flex_vertadr,
m.flex_edgeadr,
m.flex_edgenum,
m.flex_vertbodyid,
m.flex_edge,
m.flex_edgeflap,
m.flex_bending,
d.flexvert_xpos,
],
outputs=[d.qfrc_spring],
)
gravcomp = m.ngravcomp and not (m.opt.disableflags & DisableBit.GRAVITY)
if gravcomp:
d.qfrc_gravcomp.zero_()
wp.launch(
_gravity_force,
dim=(d.nworld, m.nbody - 1, m.nv),
inputs=[
m.opt.gravity,
m.body_parentid,
m.body_rootid,
m.body_mass,
m.body_gravcomp,
m.dof_bodyid,
d.xipos,
d.subtree_com,
d.cdof,
],
outputs=[d.qfrc_gravcomp],
)
if m.has_fluid:
_fluid(m, d)
wp.launch(
_qfrc_passive,
dim=(d.nworld, m.nv),
inputs=[
m.jnt_actgravcomp,
m.dof_jntid,
m.has_fluid,
d.qfrc_spring,
d.qfrc_damper,
d.qfrc_gravcomp,
d.qfrc_fluid,
gravcomp,
],
outputs=[
d.qfrc_passive,
],
)