NvidiaWarp-GarmentCode/warp/sim/integrator_euler.py

# Copyright (c) 2022 NVIDIA CORPORATION.  All rights reserved.
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.

"""This module contains time-integration objects for simulating
models + state forward in time.

"""

import warp as wp

from .collide import triangle_closest_point_barycentric
from .model import PARTICLE_FLAG_ACTIVE, ModelShapeGeometry, ModelShapeMaterials
from .optimizer import Optimizer
from .particles import eval_particle_forces
from .collide import triangle_closest_point_barycentric
from .utils import quat_decompose, quat_twist, compute_vertex_normal, mesh_query_inside_


@wp.kernel
def integrate_particles(
        x: wp.array(dtype=wp.vec3),
        v: wp.array(dtype=wp.vec3),
        f: wp.array(dtype=wp.vec3),
        w: wp.array(dtype=float),
        particle_flags: wp.array(dtype=wp.uint32),
        gravity: wp.vec3,
        viscos_damping: wp.vec3,
        dt: float,
        x_new: wp.array(dtype=wp.vec3),
        v_new: wp.array(dtype=wp.vec3),
):
    tid = wp.tid()
    if (particle_flags[tid] & PARTICLE_FLAG_ACTIVE) == 0:
        return

    x0 = x[tid]
    v0 = v[tid]
    f0 = f[tid]

    inv_mass = w[tid]

    # simple semi-implicit Euler. v1 = v0 + a dt, x1 = x0 + v1 dt
    v1 = v0 + (f0 * inv_mass + gravity * wp.step(0.0 - inv_mass)) * dt
    vel = wp.length(v1)
    # apply viscous damping directly on velocity
    if viscos_damping[0] > 0.0 and vel > viscos_damping[1]:
        v1 *= wp.exp(-viscos_damping[0] * vel * dt)
    # constaint velocity below a max value
    if viscos_damping[2] > 0.0 and vel > viscos_damping[2]:
        v1 = wp.normalize(v1) * viscos_damping[2]

    x1 = x0 + v1 * dt

    x_new[tid] = x1
    v_new[tid] = v1


# semi-implicit Euler integration
@wp.kernel
def integrate_bodies(
        body_q: wp.array(dtype=wp.transform),
        body_qd: wp.array(dtype=wp.spatial_vector),
        body_f: wp.array(dtype=wp.spatial_vector),
        body_com: wp.array(dtype=wp.vec3),
        m: wp.array(dtype=float),
        I: wp.array(dtype=wp.mat33),
        inv_m: wp.array(dtype=float),
        inv_I: wp.array(dtype=wp.mat33),
        gravity: wp.vec3,
        angular_damping: float,
        dt: float,
        # outputs
        body_q_new: wp.array(dtype=wp.transform),
        body_qd_new: wp.array(dtype=wp.spatial_vector),
):
    tid = wp.tid()

    # positions
    q = body_q[tid]
    qd = body_qd[tid]
    f = body_f[tid]

    # masses
    mass = m[tid]
    inv_mass = inv_m[tid]  # 1 / mass

    inertia = I[tid]
    inv_inertia = inv_I[tid]  # inverse of 3x3 inertia matrix

    # unpack transform
    x0 = wp.transform_get_translation(q)
    r0 = wp.transform_get_rotation(q)

    # unpack spatial twist
    w0 = wp.spatial_top(qd)
    v0 = wp.spatial_bottom(qd)

    # unpack spatial wrench
    t0 = wp.spatial_top(f)
    f0 = wp.spatial_bottom(f)

    x_com = x0 + wp.quat_rotate(r0, body_com[tid])

    # linear part
    v1 = v0 + (f0 * inv_mass + gravity * wp.nonzero(inv_mass)) * dt
    x1 = x_com + v1 * dt

    # angular part (compute in body frame)
    wb = wp.quat_rotate_inv(r0, w0)
    tb = wp.quat_rotate_inv(r0, t0) - wp.cross(wb, inertia * wb)  # coriolis forces

    w1 = wp.quat_rotate(r0, wb + inv_inertia * tb * dt)
    r1 = wp.normalize(r0 + wp.quat(w1, 0.0) * r0 * 0.5 * dt)

    # angular damping
    w1 *= 1.0 - angular_damping * dt

    body_q_new[tid] = wp.transform(x1 - wp.quat_rotate(r1, body_com[tid]), r1)
    body_qd_new[tid] = wp.spatial_vector(w1, v1)


@wp.kernel
def eval_springs(
        x: wp.array(dtype=wp.vec3),
        v: wp.array(dtype=wp.vec3),
        spring_indices: wp.array(dtype=int),
        spring_rest_lengths: wp.array(dtype=float),
        spring_stiffness: wp.array(dtype=float),
        spring_damping: wp.array(dtype=float),
        f: wp.array(dtype=wp.vec3),
):
    tid = wp.tid()

    i = spring_indices[tid * 2 + 0]
    j = spring_indices[tid * 2 + 1]

    ke = spring_stiffness[tid]
    kd = spring_damping[tid]
    rest = spring_rest_lengths[tid]

    xi = x[i]
    xj = x[j]

    vi = v[i]
    vj = v[j]

    xij = xi - xj
    vij = vi - vj

    l = wp.length(xij)
    l_inv = 1.0 / l

    # normalized spring direction
    dir = xij * l_inv

    c = l - rest
    dcdt = wp.dot(dir, vij)

    # damping based on relative velocity.
    fs = dir * (ke * c + kd * dcdt)

    wp.atomic_sub(f, i, fs)
    wp.atomic_add(f, j, fs)


@wp.kernel
def eval_triangles(
        x: wp.array(dtype=wp.vec3),
        v: wp.array(dtype=wp.vec3),
        indices: wp.array2d(dtype=int),
        pose: wp.array(dtype=wp.mat22),
        activation: wp.array(dtype=float),
        materials: wp.array2d(dtype=float),
        f: wp.array(dtype=wp.vec3),
):
    tid = wp.tid()

    k_mu = materials[tid, 0]
    k_lambda = materials[tid, 1]
    k_damp = materials[tid, 2]
    k_drag = materials[tid, 3]
    k_lift = materials[tid, 4]

    i = indices[tid, 0]
    j = indices[tid, 1]
    k = indices[tid, 2]

    x0 = x[i]  # point zero
    x1 = x[j]  # point one
    x2 = x[k]  # point two

    v0 = v[i]  # vel zero
    v1 = v[j]  # vel one
    v2 = v[k]  # vel two

    x10 = x1 - x0  # barycentric coordinates (centered at p)
    x20 = x2 - x0

    v10 = v1 - v0
    v20 = v2 - v0

    Dm = pose[tid]

    inv_rest_area = wp.determinant(Dm) * 2.0  # 1 / det(A) = det(A^-1)
    rest_area = 1.0 / inv_rest_area

    # scale stiffness coefficients to account for area
    k_mu = k_mu * rest_area
    k_lambda = k_lambda * rest_area
    k_damp = k_damp * rest_area

    # F = Xs*Xm^-1
    F1 = x10 * Dm[0, 0] + x20 * Dm[1, 0]
    F2 = x10 * Dm[0, 1] + x20 * Dm[1, 1]

    # dFdt = Vs*Xm^-1
    dFdt1 = v10 * Dm[0, 0] + v20 * Dm[1, 0]
    dFdt2 = v10 * Dm[0, 1] + v20 * Dm[1, 1]

    # deviatoric PK1 + damping term
    P1 = F1 * k_mu + dFdt1 * k_damp
    P2 = F2 * k_mu + dFdt2 * k_damp

    # -----------------------------
    # St. Venant-Kirchoff

    # # Green strain, F'*F-I
    # e00 = dot(f1, f1) - 1.0
    # e10 = dot(f2, f1)
    # e01 = dot(f1, f2)
    # e11 = dot(f2, f2) - 1.0

    # E = wp.mat22(e00, e01,
    #              e10, e11)

    # # local forces (deviatoric part)
    # T = wp.mul(E, wp.transpose(Dm))

    # # spatial forces, F*T
    # fq = (f1*T[0,0] + f2*T[1,0])*k_mu*2.0
    # fr = (f1*T[0,1] + f2*T[1,1])*k_mu*2.0
    # alpha = 1.0

    # -----------------------------
    # Baraff & Witkin, note this model is not isotropic

    # c1 = length(f1) - 1.0
    # c2 = length(f2) - 1.0
    # f1 = normalize(f1)*c1*k1
    # f2 = normalize(f2)*c2*k1

    # fq = f1*Dm[0,0] + f2*Dm[0,1]
    # fr = f1*Dm[1,0] + f2*Dm[1,1]

    # -----------------------------
    # Neo-Hookean (with rest stability)

    # force = P*Dm'
    f1 = P1 * Dm[0, 0] + P2 * Dm[0, 1]
    f2 = P1 * Dm[1, 0] + P2 * Dm[1, 1]
    alpha = 1.0 + k_mu / k_lambda

    # -----------------------------
    # Area Preservation

    n = wp.cross(x10, x20)
    area = wp.length(n) * 0.5

    # actuation
    act = activation[tid]

    # J-alpha
    c = area * inv_rest_area - alpha + act

    # dJdx
    n = wp.normalize(n)
    dcdq = wp.cross(x20, n) * inv_rest_area * 0.5
    dcdr = wp.cross(n, x10) * inv_rest_area * 0.5

    f_area = k_lambda * c

    # -----------------------------
    # Area Damping

    dcdt = dot(dcdq, v1) + dot(dcdr, v2) - dot(dcdq + dcdr, v0)
    f_damp = k_damp * dcdt

    f1 = f1 + dcdq * (f_area + f_damp)
    f2 = f2 + dcdr * (f_area + f_damp)
    f0 = f1 + f2

    # -----------------------------
    # Lift + Drag

    vmid = (v0 + v1 + v2) * 0.3333
    vdir = wp.normalize(vmid)

    f_drag = vmid * (k_drag * area * wp.abs(wp.dot(n, vmid)))
    f_lift = n * (k_lift * area * (1.57079 - wp.acos(wp.dot(n, vdir)))) * dot(vmid, vmid)

    # note reversed sign due to atomic_add below.. need to write the unary op -
    f0 = f0 - f_drag - f_lift
    f1 = f1 + f_drag + f_lift
    f2 = f2 + f_drag + f_lift

    # apply forces
    wp.atomic_add(f, i, f0)
    wp.atomic_sub(f, j, f1)
    wp.atomic_sub(f, k, f2)


# @wp.func
# def triangle_closest_point(a: wp.vec3, b: wp.vec3, c: wp.vec3, p: wp.vec3):
#     ab = b - a
#     ac = c - a
#     ap = p - a
#
#     d1 = wp.dot(ab, ap)
#     d2 = wp.dot(ac, ap)
#
#     if (d1 <= 0.0 and d2 <= 0.0):
#         return a
#
#     bp = p - b
#     d3 = wp.dot(ab, bp)
#     d4 = wp.dot(ac, bp)
#
#     if (d3 >= 0.0 and d4 <= d3):
#         return b
#
#     vc = d1 * d4 - d3 * d2
#     v = d1 / (d1 - d3)
#     if (vc <= 0.0 and d1 >= 0.0 and d3 <= 0.0):
#         return a + ab * v
#
#     cp = p - c
#     d5 = dot(ab, cp)
#     d6 = dot(ac, cp)
#
#     if (d6 >= 0.0 and d5 <= d6):
#         return c
#
#     vb = d5 * d2 - d1 * d6
#     w = d2 / (d2 - d6)
#     if (vb <= 0.0 and d2 >= 0.0 and d6 <= 0.0):
#         return a + ac * w
#
#     va = d3 * d6 - d5 * d4
#     w = (d4 - d3) / ((d4 - d3) + (d5 - d6))
#     if (va <= 0.0 and (d4 - d3) >= 0.0 and (d5 - d6) >= 0.0):
#         return b + (c - b) * w
#
#     denom = 1.0 / (va + vb + vc)
#     v = vb * denom
#     w = vc * denom
#
#     return a + ab * v + ac * w


@wp.kernel
def eval_triangles_contact(
        # idx : wp.array(dtype=int), # list of indices for colliding particles
        num_particles: int,  # size of particles
        x: wp.array(dtype=wp.vec3),
        v: wp.array(dtype=wp.vec3),
        indices: wp.array2d(dtype=int),
        pose: wp.array(dtype=wp.mat22),
        activation: wp.array(dtype=float),
        materials: wp.array2d(dtype=float),
        f: wp.array(dtype=wp.vec3),
):
    tid = wp.tid()
    face_no = tid // num_particles  # which face
    particle_no = tid % num_particles  # which particle

    k_mu = materials[face_no, 0]
    k_lambda = materials[face_no, 1]
    k_damp = materials[face_no, 2]
    k_drag = materials[face_no, 3]
    k_lift = materials[face_no, 4]

    # at the moment, just one particle
    pos = x[particle_no]

    i = indices[face_no, 0]
    j = indices[face_no, 1]
    k = indices[face_no, 2]

    if i == particle_no or j == particle_no or k == particle_no:
        return

    p = x[i]  # point zero
    q = x[j]  # point one
    r = x[k]  # point two

    # vp = v[i] # vel zero
    # vq = v[j] # vel one
    # vr = v[k] # vel two

    # qp = q-p # barycentric coordinates (centered at p)
    # rp = r-p

    bary = triangle_closest_point_barycentric(p, q, r, pos)
    closest = p * bary[0] + q * bary[1] + r * bary[2]

    diff = pos - closest
    dist = wp.dot(diff, diff)
    n = wp.normalize(diff)
    c = wp.min(dist - 0.01, 0.0)  # 0 unless within 0.01 of surface
    # c = wp.leaky_min(dot(n, x0)-0.01, 0.0, 0.0)
    fn = n * c * 1e5

    wp.atomic_sub(f, particle_no, fn)

    # # apply forces (could do - f / 3 here)
    wp.atomic_add(f, i, fn * bary[0])
    wp.atomic_add(f, j, fn * bary[1])
    wp.atomic_add(f, k, fn * bary[2])


@wp.kernel
def eval_triangles_body_contacts(
        num_particles: int,  # number of particles (size of contact_point)
        x: wp.array(dtype=wp.vec3),  # position of particles
        v: wp.array(dtype=wp.vec3),
        indices: wp.array(dtype=int),  # triangle indices
        body_x: wp.array(dtype=wp.vec3),  # body body positions
        body_r: wp.array(dtype=wp.quat),
        body_v: wp.array(dtype=wp.vec3),
        body_w: wp.array(dtype=wp.vec3),
        contact_body: wp.array(dtype=int),
        contact_point: wp.array(dtype=wp.vec3),  # position of contact points relative to body
        contact_dist: wp.array(dtype=float),
        contact_mat: wp.array(dtype=int),
        materials: wp.array(dtype=float),
        #   body_f : wp.array(dtype=wp.vec3),
        #   body_t : wp.array(dtype=wp.vec3),
        tri_f: wp.array(dtype=wp.vec3),
):
    tid = wp.tid()

    face_no = tid // num_particles  # which face
    particle_no = tid % num_particles  # which particle

    # -----------------------
    # load body body point
    c_body = contact_body[particle_no]
    c_point = contact_point[particle_no]
    c_dist = contact_dist[particle_no]
    c_mat = contact_mat[particle_no]

    # hard coded surface parameter tensor layout (ke, kd, kf, mu)
    ke = materials[c_mat * 4 + 0]  # restitution coefficient
    kd = materials[c_mat * 4 + 1]  # damping coefficient
    kf = materials[c_mat * 4 + 2]  # friction coefficient
    mu = materials[c_mat * 4 + 3]  # coulomb friction

    x0 = body_x[c_body]  # position of colliding body
    r0 = body_r[c_body]  # orientation of colliding body

    v0 = body_v[c_body]
    w0 = body_w[c_body]

    # transform point to world space
    pos = x0 + wp.quat_rotate(r0, c_point)
    # use x0 as center, everything is offset from center of mass

    # moment arm
    r = pos - x0  # basically just c_point in the new coordinates
    rhat = wp.normalize(r)
    pos = pos + rhat * c_dist  # add on 'thickness' of shape, e.g.: radius of sphere/capsule

    # contact point velocity
    dpdt = v0 + wp.cross(w0, r)  # this is body velocity cross offset, so it's the velocity of the contact point.

    # -----------------------
    # load triangle
    i = indices[face_no * 3 + 0]
    j = indices[face_no * 3 + 1]
    k = indices[face_no * 3 + 2]

    p = x[i]  # point zero
    q = x[j]  # point one
    r = x[k]  # point two

    vp = v[i]  # vel zero
    vq = v[j]  # vel one
    vr = v[k]  # vel two

    bary = triangle_closest_point_barycentric(p, q, r, pos)
    closest = p * bary[0] + q * bary[1] + r * bary[2]

    diff = pos - closest  # vector from tri to point
    dist = wp.dot(diff, diff)  # squared distance
    n = wp.normalize(diff)  # points into the object
    c = wp.min(dist - 0.05, 0.0)  # 0 unless within 0.05 of surface
    # c = wp.leaky_min(dot(n, x0)-0.01, 0.0, 0.0)
    # fn = n * c * 1e6    # points towards cloth (both n and c are negative)

    # wp.atomic_sub(tri_f, particle_no, fn)

    fn = c * ke  # normal force (restitution coefficient * how far inside for ground) (negative)

    vtri = vp * bary[0] + vq * bary[1] + vr * bary[2]  # bad approximation for centroid velocity
    vrel = vtri - dpdt

    vn = dot(n, vrel)  # velocity component of body in negative normal direction
    vt = vrel - n * vn  # velocity component not in normal direction

    # contact damping
    fd = 0.0 - wp.max(vn, 0.0) * kd * wp.step(c)  # again, negative, into the ground

    # # viscous friction
    # ft = vt*kf

    # Coulomb friction (box)
    lower = mu * (fn + fd)
    upper = 0.0 - lower  # workaround because no unary ops yet

    nx = cross(n, vec3(0.0, 0.0, 1.0))  # basis vectors for tangent
    nz = cross(n, vec3(1.0, 0.0, 0.0))

    vx = wp.clamp(dot(nx * kf, vt), lower, upper)
    vz = wp.clamp(dot(nz * kf, vt), lower, upper)

    ft = (nx * vx + nz * vz) * (0.0 - wp.step(c))  # wp.vec3(vx, 0.0, vz)*wp.step(c)

    # # Coulomb friction (smooth, but gradients are numerically unstable around |vt| = 0)
    # #ft = wp.normalize(vt)*wp.min(kf*wp.length(vt), 0.0 - mu*c*ke)

    f_total = n * (fn + fd) + ft

    wp.atomic_add(tri_f, i, f_total * bary[0])
    wp.atomic_add(tri_f, j, f_total * bary[1])
    wp.atomic_add(tri_f, k, f_total * bary[2])


@wp.kernel
def eval_bending(
        x: wp.array(dtype=wp.vec3),
        v: wp.array(dtype=wp.vec3),
        indices: wp.array2d(dtype=int),
        rest: wp.array(dtype=float),
        bending_properties: wp.array2d(dtype=float),
        f: wp.array(dtype=wp.vec3),
):
    tid = wp.tid()
    ke = bending_properties[tid, 0]
    kd = bending_properties[tid, 1]

    i = indices[tid, 0]
    j = indices[tid, 1]
    k = indices[tid, 2]
    l = indices[tid, 3]

    rest_angle = rest[tid]

    x1 = x[i]
    x2 = x[j]
    x3 = x[k]
    x4 = x[l]

    v1 = v[i]
    v2 = v[j]
    v3 = v[k]
    v4 = v[l]

    n1 = wp.cross(x3 - x1, x4 - x1)  # normal to face 1
    n2 = wp.cross(x4 - x2, x3 - x2)  # normal to face 2

    n1_length = wp.length(n1)
    n2_length = wp.length(n2)

    if n1_length < 1.0e-6 or n2_length < 1.0e-6:
        return

    rcp_n1 = 1.0 / n1_length
    rcp_n2 = 1.0 / n2_length

    cos_theta = wp.dot(n1, n2) * rcp_n1 * rcp_n2

    n1 = n1 * rcp_n1 * rcp_n1
    n2 = n2 * rcp_n2 * rcp_n2

    e = x4 - x3
    e_hat = wp.normalize(e)
    e_length = wp.length(e)

    s = wp.sign(wp.dot(wp.cross(n2, n1), e_hat))
    angle = wp.acos(cos_theta) * s

    d1 = n1 * e_length
    d2 = n2 * e_length
    d3 = n1 * wp.dot(x1 - x4, e_hat) + n2 * wp.dot(x2 - x4, e_hat)
    d4 = n1 * wp.dot(x3 - x1, e_hat) + n2 * wp.dot(x3 - x2, e_hat)

    # elastic
    f_elastic = ke * (angle - rest_angle)

    # damping
    f_damp = kd * (wp.dot(d1, v1) + wp.dot(d2, v2) + wp.dot(d3, v3) + wp.dot(d4, v4))

    # total force, proportional to edge length
    f_total = 0.0 - e_length * (f_elastic + f_damp)

    wp.atomic_add(f, i, d1 * f_total)
    wp.atomic_add(f, j, d2 * f_total)
    wp.atomic_add(f, k, d3 * f_total)
    wp.atomic_add(f, l, d4 * f_total)


@wp.kernel
def eval_tetrahedra(
        x: wp.array(dtype=wp.vec3),
        v: wp.array(dtype=wp.vec3),
        indices: wp.array2d(dtype=int),
        pose: wp.array(dtype=wp.mat33),
        activation: wp.array(dtype=float),
        materials: wp.array2d(dtype=float),
        f: wp.array(dtype=wp.vec3),
):
    tid = wp.tid()

    i = indices[tid, 0]
    j = indices[tid, 1]
    k = indices[tid, 2]
    l = indices[tid, 3]

    act = activation[tid]

    k_mu = materials[tid, 0]
    k_lambda = materials[tid, 1]
    k_damp = materials[tid, 2]

    x0 = x[i]
    x1 = x[j]
    x2 = x[k]
    x3 = x[l]

    v0 = v[i]
    v1 = v[j]
    v2 = v[k]
    v3 = v[l]

    x10 = x1 - x0
    x20 = x2 - x0
    x30 = x3 - x0

    v10 = v1 - v0
    v20 = v2 - v0
    v30 = v3 - v0

    Ds = wp.mat33(x10, x20, x30)
    Dm = pose[tid]

    inv_rest_volume = wp.determinant(Dm) * 6.0
    rest_volume = 1.0 / inv_rest_volume

    alpha = 1.0 + k_mu / k_lambda - k_mu / (4.0 * k_lambda)

    # scale stiffness coefficients to account for area
    k_mu = k_mu * rest_volume
    k_lambda = k_lambda * rest_volume
    k_damp = k_damp * rest_volume

    # F = Xs*Xm^-1
    F = Ds * Dm
    dFdt = wp.mat33(v10, v20, v30) * Dm

    col1 = wp.vec3(F[0, 0], F[1, 0], F[2, 0])
    col2 = wp.vec3(F[0, 1], F[1, 1], F[2, 1])
    col3 = wp.vec3(F[0, 2], F[1, 2], F[2, 2])

    # -----------------------------
    # Neo-Hookean (with rest stability [Smith et al 2018])

    Ic = dot(col1, col1) + dot(col2, col2) + dot(col3, col3)

    # deviatoric part
    P = F * k_mu * (1.0 - 1.0 / (Ic + 1.0)) + dFdt * k_damp
    H = P * wp.transpose(Dm)

    f1 = wp.vec3(H[0, 0], H[1, 0], H[2, 0])
    f2 = wp.vec3(H[0, 1], H[1, 1], H[2, 1])
    f3 = wp.vec3(H[0, 2], H[1, 2], H[2, 2])

    # -----------------------------
    # C_sqrt

    # alpha = 1.0

    # r_s = wp.sqrt(wp.abs(dot(col1, col1) + dot(col2, col2) + dot(col3, col3) - 3.0))

    # f1 = wp.vec3()
    # f2 = wp.vec3()
    # f3 = wp.vec3()

    # if (r_s > 0.0):
    #     r_s_inv = 1.0/r_s

    #     C = r_s
    #     dCdx = F*wp.transpose(Dm)*r_s_inv*wp.sign(r_s)

    #     grad1 = vec3(dCdx[0,0], dCdx[1,0], dCdx[2,0])
    #     grad2 = vec3(dCdx[0,1], dCdx[1,1], dCdx[2,1])
    #     grad3 = vec3(dCdx[0,2], dCdx[1,2], dCdx[2,2])

    #     f1 = grad1*C*k_mu
    #     f2 = grad2*C*k_mu
    #     f3 = grad3*C*k_mu

    # -----------------------------
    # C_spherical

    # alpha = 1.0

    # r_s = wp.sqrt(dot(col1, col1) + dot(col2, col2) + dot(col3, col3))
    # r_s_inv = 1.0/r_s

    # C = r_s - wp.sqrt(3.0)
    # dCdx = F*wp.transpose(Dm)*r_s_inv

    # grad1 = vec3(dCdx[0,0], dCdx[1,0], dCdx[2,0])
    # grad2 = vec3(dCdx[0,1], dCdx[1,1], dCdx[2,1])
    # grad3 = vec3(dCdx[0,2], dCdx[1,2], dCdx[2,2])

    # f1 = grad1*C*k_mu
    # f2 = grad2*C*k_mu
    # f3 = grad3*C*k_mu

    # ----------------------------
    # C_D

    # alpha = 1.0

    # r_s = wp.sqrt(dot(col1, col1) + dot(col2, col2) + dot(col3, col3))

    # C = r_s*r_s - 3.0
    # dCdx = F*wp.transpose(Dm)*2.0

    # grad1 = vec3(dCdx[0,0], dCdx[1,0], dCdx[2,0])
    # grad2 = vec3(dCdx[0,1], dCdx[1,1], dCdx[2,1])
    # grad3 = vec3(dCdx[0,2], dCdx[1,2], dCdx[2,2])

    # f1 = grad1*C*k_mu
    # f2 = grad2*C*k_mu
    # f3 = grad3*C*k_mu

    # ----------------------------
    # Hookean

    # alpha = 1.0

    # I = wp.mat33(wp.vec3(1.0, 0.0, 0.0),
    #              wp.vec3(0.0, 1.0, 0.0),
    #              wp.vec3(0.0, 0.0, 1.0))

    # P = (F + wp.transpose(F) + I*(0.0-2.0))*k_mu
    # H = P * wp.transpose(Dm)

    # f1 = wp.vec3(H[0, 0], H[1, 0], H[2, 0])
    # f2 = wp.vec3(H[0, 1], H[1, 1], H[2, 1])
    # f3 = wp.vec3(H[0, 2], H[1, 2], H[2, 2])

    # hydrostatic part
    J = wp.determinant(F)

    # print(J)
    s = inv_rest_volume / 6.0
    dJdx1 = wp.cross(x20, x30) * s
    dJdx2 = wp.cross(x30, x10) * s
    dJdx3 = wp.cross(x10, x20) * s

    f_volume = (J - alpha + act) * k_lambda
    f_damp = (wp.dot(dJdx1, v1) + wp.dot(dJdx2, v2) + wp.dot(dJdx3, v3)) * k_damp

    f_total = f_volume + f_damp

    f1 = f1 + dJdx1 * f_total
    f2 = f2 + dJdx2 * f_total
    f3 = f3 + dJdx3 * f_total
    f0 = (f1 + f2 + f3) * (0.0 - 1.0)

    # apply forces
    wp.atomic_sub(f, i, f0)
    wp.atomic_sub(f, j, f1)
    wp.atomic_sub(f, k, f2)
    wp.atomic_sub(f, l, f3)


@wp.kernel
def eval_particle_ground_contacts(
        particle_x: wp.array(dtype=wp.vec3),
        particle_v: wp.array(dtype=wp.vec3),
        particle_radius: wp.array(dtype=float),
        particle_flags: wp.array(dtype=wp.uint32),
        ke: float,
        kd: float,
        kf: float,
        mu: float,
        ground: wp.array(dtype=float),
        # outputs
        f: wp.array(dtype=wp.vec3),
):
    tid = wp.tid()
    if (particle_flags[tid] & PARTICLE_FLAG_ACTIVE) == 0:
        return

    x = particle_x[tid]
    v = particle_v[tid]
    radius = particle_radius[tid]

    n = wp.vec3(ground[0], ground[1], ground[2])
    c = wp.min(wp.dot(n, x) + ground[3] - radius, 0.0)

    vn = wp.dot(n, v)
    jn = c * ke

    if c >= 0.0:
        return

    jd = min(vn, 0.0) * kd

    # contact force
    fn = jn + jd

    # friction force
    vt = v - n * vn
    vs = wp.length(vt)

    if vs > 0.0:
        vt = vt / vs

    # Coulomb condition
    ft = wp.min(vs * kf, mu * wp.abs(fn))

    # total force
    f[tid] = f[tid] - n * fn - vt * ft


@wp.kernel
def eval_particle_contacts(
    particle_x: wp.array(dtype=wp.vec3),
    particle_v: wp.array(dtype=wp.vec3),
    body_q: wp.array(dtype=wp.transform),
    body_qd: wp.array(dtype=wp.spatial_vector),
    particle_radius: wp.array(dtype=float),
    particle_flags: wp.array(dtype=wp.uint32),
    body_com: wp.array(dtype=wp.vec3),
    shape_body: wp.array(dtype=int),
    shape_materials: ModelShapeMaterials,
    particle_ke: float,
    particle_kd: float,
    particle_kf: float,
    particle_mu: float,
    particle_ka: float,
    contact_count: wp.array(dtype=int),
    contact_particle: wp.array(dtype=int),
    contact_shape: wp.array(dtype=int),
    contact_body_pos: wp.array(dtype=wp.vec3),
    contact_body_vel: wp.array(dtype=wp.vec3),
    contact_normal: wp.array(dtype=wp.vec3),
    contact_max: int,
    # outputs
    particle_f: wp.array(dtype=wp.vec3),
    body_f: wp.array(dtype=wp.spatial_vector),
):
    tid = wp.tid()

    count = min(contact_max, contact_count[0])
    if tid >= count:
        return

    shape_index = contact_shape[tid]
    body_index = shape_body[shape_index]
    particle_index = contact_particle[tid]
    if (particle_flags[particle_index] & PARTICLE_FLAG_ACTIVE) == 0:
        return

    px = particle_x[particle_index]
    pv = particle_v[particle_index]

    X_wb = wp.transform_identity()
    X_com = wp.vec3()
    body_v_s = wp.spatial_vector()

    if body_index >= 0:
        X_wb = body_q[body_index]
        X_com = body_com[body_index]
        body_v_s = body_qd[body_index]

    # body position in world space
    bx = wp.transform_point(X_wb, contact_body_pos[tid])
    r = bx - wp.transform_point(X_wb, X_com)

    n = contact_normal[tid]
    c = wp.dot(n, px - bx) - particle_radius[tid]

    if c > particle_ka:
        return

    # take average material properties of shape and particle parameters
    ke = 0.5 * (particle_ke + shape_materials.ke[shape_index])
    kd = 0.5 * (particle_kd + shape_materials.kd[shape_index])
    kf = 0.5 * (particle_kf + shape_materials.kf[shape_index])
    mu = 0.5 * (particle_mu + shape_materials.mu[shape_index])

    body_w = wp.spatial_top(body_v_s)
    body_v = wp.spatial_bottom(body_v_s)

    # compute the body velocity at the particle position
    bv = body_v + wp.cross(body_w, r) + wp.transform_vector(X_wb, contact_body_vel[tid])

    # relative velocity
    v = pv - bv

    # decompose relative velocity
    vn = wp.dot(n, v)
    vt = v - n * vn

    # contact elastic
    fn = n * c * ke

    # contact damping
    fd = n * wp.min(vn, 0.0) * kd

    # viscous friction
    # ft = vt*kf

    # Coulomb friction (box)
    # lower = mu * c * ke
    # upper = 0.0 - lower

    # vx = wp.clamp(wp.dot(wp.vec3(kf, 0.0, 0.0), vt), lower, upper)
    # vz = wp.clamp(wp.dot(wp.vec3(0.0, 0.0, kf), vt), lower, upper)

    # ft = wp.vec3(vx, 0.0, vz)

    # Coulomb friction (smooth, but gradients are numerically unstable around |vt| = 0)
    ft = wp.normalize(vt) * wp.min(kf * wp.length(vt), abs(mu * c * ke))

    f_total = fn + (fd + ft)
    t_total = wp.cross(r, f_total)

    wp.atomic_sub(particle_f, particle_index, f_total)

    if body_index >= 0:
        wp.atomic_add(body_f, body_index, wp.spatial_vector(t_total, f_total))


@wp.kernel
def eval_rigid_contacts(
        body_q: wp.array(dtype=wp.transform),
        body_qd: wp.array(dtype=wp.spatial_vector),
        body_com: wp.array(dtype=wp.vec3),
        shape_materials: ModelShapeMaterials,
        geo: ModelShapeGeometry,
        contact_count: wp.array(dtype=int),
        contact_body0: wp.array(dtype=int),
        contact_body1: wp.array(dtype=int),
        contact_point0: wp.array(dtype=wp.vec3),
        contact_point1: wp.array(dtype=wp.vec3),
        contact_normal: wp.array(dtype=wp.vec3),
        contact_shape0: wp.array(dtype=int),
        contact_shape1: wp.array(dtype=int),
        # outputs
        body_f: wp.array(dtype=wp.spatial_vector),
):
    tid = wp.tid()
    if contact_shape0[tid] == contact_shape1[tid]:
        return

    count = contact_count[0]
    if tid >= count:
        return

    # retrieve contact thickness, compute average contact material properties
    ke = 0.0  # restitution coefficient
    kd = 0.0  # damping coefficient
    kf = 0.0  # friction coefficient
    mu = 0.0  # coulomb friction
    mat_nonzero = 0
    thickness_a = 0.0
    thickness_b = 0.0
    shape_a = contact_shape0[tid]
    shape_b = contact_shape1[tid]
    if shape_a >= 0:
        mat_nonzero += 1
        ke += shape_materials.ke[shape_a]
        kd += shape_materials.kd[shape_a]
        kf += shape_materials.kf[shape_a]
        mu += shape_materials.mu[shape_a]
        thickness_a = geo.thickness[shape_a]
    if shape_b >= 0:
        mat_nonzero += 1
        ke += shape_materials.ke[shape_b]
        kd += shape_materials.kd[shape_b]
        kf += shape_materials.kf[shape_b]
        mu += shape_materials.mu[shape_b]
        thickness_b = geo.thickness[shape_b]
    if mat_nonzero > 0:
        ke = ke / float(mat_nonzero)
        kd = kd / float(mat_nonzero)
        kf = kf / float(mat_nonzero)
        mu = mu / float(mat_nonzero)

    body_a = contact_body0[tid]
    body_b = contact_body1[tid]

    # body position in world space
    n = contact_normal[tid]
    bx_a = contact_point0[tid]
    bx_b = contact_point1[tid]
    if body_a >= 0:
        X_wb_a = body_q[body_a]
        X_com_a = body_com[body_a]
        bx_a = wp.transform_point(X_wb_a, bx_a) - thickness_a * n
        r_a = bx_a - wp.transform_point(X_wb_a, X_com_a)

    if body_b >= 0:
        X_wb_b = body_q[body_b]
        X_com_b = body_com[body_b]
        bx_b = wp.transform_point(X_wb_b, bx_b) + thickness_b * n
        r_b = bx_b - wp.transform_point(X_wb_b, X_com_b)

    d = wp.dot(n, bx_a - bx_b)

    if d >= 0.0:
        return

    # compute contact point velocity
    bv_a = wp.vec3(0.0)
    bv_b = wp.vec3(0.0)
    if body_a >= 0:
        body_v_s_a = body_qd[body_a]
        body_w_a = wp.spatial_top(body_v_s_a)
        body_v_a = wp.spatial_bottom(body_v_s_a)
        bv_a = body_v_a + wp.cross(body_w_a, r_a)

    if body_b >= 0:
        body_v_s_b = body_qd[body_b]
        body_w_b = wp.spatial_top(body_v_s_b)
        body_v_b = wp.spatial_bottom(body_v_s_b)
        bv_b = body_v_b + wp.cross(body_w_b, r_b)

    # relative velocity
    v = bv_a - bv_b

    # print(v)

    # decompose relative velocity
    vn = wp.dot(n, v)
    vt = v - n * vn

    # contact elastic
    fn = d * ke

    # contact damping
    fd = wp.min(vn, 0.0) * kd * wp.step(d)

    # viscous friction
    # ft = vt*kf

    # Coulomb friction (box)
    # lower = mu * d * ke
    # upper = 0.0 - lower

    # vx = wp.clamp(wp.dot(wp.vec3(kf, 0.0, 0.0), vt), lower, upper)
    # vz = wp.clamp(wp.dot(wp.vec3(0.0, 0.0, kf), vt), lower, upper)

    # ft = wp.vec3(vx, 0.0, vz)

    # Coulomb friction (smooth, but gradients are numerically unstable around |vt| = 0)
    # ft = wp.normalize(vt)*wp.min(kf*wp.length(vt), abs(mu*d*ke))
    ft = wp.normalize(vt) * wp.min(kf * wp.length(vt), 0.0 - mu * (fn + fd))

    # f_total = fn + (fd + ft)
    f_total = n * (fn + fd) + ft
    # t_total = wp.cross(r, f_total)

    # print("apply contact force")
    # print(f_total)

    if body_a >= 0:
        wp.atomic_sub(body_f, body_a, wp.spatial_vector(wp.cross(r_a, f_total), f_total))
    if body_b >= 0:
        wp.atomic_add(body_f, body_b, wp.spatial_vector(wp.cross(r_b, f_total), f_total))


@wp.func
def eval_joint_force(
        q: float,
        qd: float,
        target: float,
        target_ke: float,
        target_kd: float,
        act: float,
        limit_lower: float,
        limit_upper: float,
        limit_ke: float,
        limit_kd: float,
        axis: wp.vec3,
):
    limit_f = 0.0

    # compute limit forces, damping only active when limit is violated
    if q < limit_lower:
        limit_f = limit_ke * (limit_lower - q) - limit_kd * min(qd, 0.0)

    if q > limit_upper:
        limit_f = limit_ke * (limit_upper - q) - limit_kd * max(qd, 0.0)

    # joint dynamics
    total_f = (target_ke * (q - target) + target_kd * qd + act - limit_f) * axis

    return total_f


@wp.kernel
def eval_body_joints(
        body_q: wp.array(dtype=wp.transform),
        body_qd: wp.array(dtype=wp.spatial_vector),
        body_com: wp.array(dtype=wp.vec3),
        joint_q_start: wp.array(dtype=int),
        joint_qd_start: wp.array(dtype=int),
        joint_type: wp.array(dtype=int),
        joint_enabled: wp.array(dtype=int),
        joint_child: wp.array(dtype=int),
        joint_parent: wp.array(dtype=int),
        joint_X_p: wp.array(dtype=wp.transform),
        joint_X_c: wp.array(dtype=wp.transform),
        joint_axis: wp.array(dtype=wp.vec3),
        joint_axis_start: wp.array(dtype=int),
        joint_axis_dim: wp.array(dtype=int, ndim=2),
        joint_target: wp.array(dtype=float),
        joint_act: wp.array(dtype=float),
        joint_target_ke: wp.array(dtype=float),
        joint_target_kd: wp.array(dtype=float),
        joint_limit_lower: wp.array(dtype=float),
        joint_limit_upper: wp.array(dtype=float),
        joint_limit_ke: wp.array(dtype=float),
        joint_limit_kd: wp.array(dtype=float),
        joint_attach_ke: float,
        joint_attach_kd: float,
        body_f: wp.array(dtype=wp.spatial_vector),
):
    tid = wp.tid()
    type = joint_type[tid]

    # early out for free joints
    if joint_enabled[tid] == 0 or type == wp.sim.JOINT_FREE:
        return

    c_child = joint_child[tid]
    c_parent = joint_parent[tid]

    X_pj = joint_X_p[tid]
    X_cj = joint_X_c[tid]

    X_wp = X_pj
    r_p = wp.vec3()
    w_p = wp.vec3()
    v_p = wp.vec3()

    # parent transform and moment arm
    if c_parent >= 0:
        X_wp = body_q[c_parent] * X_wp
        r_p = wp.transform_get_translation(X_wp) - wp.transform_point(body_q[c_parent], body_com[c_parent])

        twist_p = body_qd[c_parent]

        w_p = wp.spatial_top(twist_p)
        v_p = wp.spatial_bottom(twist_p) + wp.cross(w_p, r_p)

    # child transform and moment arm
    X_wc = body_q[c_child] * X_cj
    r_c = wp.transform_get_translation(X_wc) - wp.transform_point(body_q[c_child], body_com[c_child])

    twist_c = body_qd[c_child]

    w_c = wp.spatial_top(twist_c)
    v_c = wp.spatial_bottom(twist_c) + wp.cross(w_c, r_c)

    # joint properties (for 1D joints)
    q_start = joint_q_start[tid]
    qd_start = joint_qd_start[tid]
    axis_start = joint_axis_start[tid]

    target = joint_target[axis_start]
    target_ke = joint_target_ke[axis_start]
    target_kd = joint_target_kd[axis_start]
    limit_ke = joint_limit_ke[axis_start]
    limit_kd = joint_limit_kd[axis_start]
    limit_lower = joint_limit_lower[axis_start]
    limit_upper = joint_limit_upper[axis_start]

    act = joint_act[qd_start]

    x_p = wp.transform_get_translation(X_wp)
    x_c = wp.transform_get_translation(X_wc)

    q_p = wp.transform_get_rotation(X_wp)
    q_c = wp.transform_get_rotation(X_wc)

    # translational error
    x_err = x_c - x_p
    r_err = wp.quat_inverse(q_p) * q_c
    v_err = v_c - v_p
    w_err = w_c - w_p

    # total force/torque on the parent
    t_total = wp.vec3()
    f_total = wp.vec3()

    # reduce angular damping stiffness for stability
    angular_damping_scale = 0.01

    if type == wp.sim.JOINT_FIXED:
        ang_err = wp.normalize(wp.vec3(r_err[0], r_err[1], r_err[2])) * wp.acos(r_err[3]) * 2.0

        f_total += x_err * joint_attach_ke + v_err * joint_attach_kd
        t_total += (
                wp.transform_vector(X_wp, ang_err) * joint_attach_ke + w_err * joint_attach_kd * angular_damping_scale
        )

    if type == wp.sim.JOINT_PRISMATIC:
        axis = joint_axis[axis_start]

        # world space joint axis
        axis_p = wp.transform_vector(X_wp, axis)

        # evaluate joint coordinates
        q = wp.dot(x_err, axis_p)
        qd = wp.dot(v_err, axis_p)

        f_total = eval_joint_force(
            q, qd, target, target_ke, target_kd, act, limit_lower, limit_upper, limit_ke, limit_kd, axis_p
        )

        # attachment dynamics
        ang_err = wp.normalize(wp.vec3(r_err[0], r_err[1], r_err[2])) * wp.acos(r_err[3]) * 2.0

        # project off any displacement along the joint axis
        f_total += (x_err - q * axis_p) * joint_attach_ke + (v_err - qd * axis_p) * joint_attach_kd
        t_total += (
                wp.transform_vector(X_wp, ang_err) * joint_attach_ke + w_err * joint_attach_kd * angular_damping_scale
        )

    if type == wp.sim.JOINT_REVOLUTE:
        axis = joint_axis[axis_start]

        axis_p = wp.transform_vector(X_wp, axis)
        axis_c = wp.transform_vector(X_wc, axis)

        # swing twist decomposition
        twist = quat_twist(axis, r_err)

        q = wp.acos(twist[3]) * 2.0 * wp.sign(wp.dot(axis, wp.vec3(twist[0], twist[1], twist[2])))
        qd = wp.dot(w_err, axis_p)

        t_total = eval_joint_force(
            q, qd, target, target_ke, target_kd, act, limit_lower, limit_upper, limit_ke, limit_kd, axis_p
        )

        # attachment dynamics
        swing_err = wp.cross(axis_p, axis_c)

        f_total += x_err * joint_attach_ke + v_err * joint_attach_kd
        t_total += swing_err * joint_attach_ke + (w_err - qd * axis_p) * joint_attach_kd * angular_damping_scale

    if type == wp.sim.JOINT_BALL:
        ang_err = wp.normalize(wp.vec3(r_err[0], r_err[1], r_err[2])) * wp.acos(r_err[3]) * 2.0

        # todo: joint limits
        t_total += target_kd * w_err + target_ke * wp.transform_vector(X_wp, ang_err)
        f_total += x_err * joint_attach_ke + v_err * joint_attach_kd

    if type == wp.sim.JOINT_COMPOUND:
        q_pc = wp.quat_inverse(q_p) * q_c

        # decompose to a compound rotation each axis
        angles = quat_decompose(q_pc)

        # reconstruct rotation axes
        axis_0 = wp.vec3(1.0, 0.0, 0.0)
        q_0 = wp.quat_from_axis_angle(axis_0, angles[0])

        axis_1 = wp.quat_rotate(q_0, wp.vec3(0.0, 1.0, 0.0))
        q_1 = wp.quat_from_axis_angle(axis_1, angles[1])

        axis_2 = wp.quat_rotate(q_1 * q_0, wp.vec3(0.0, 0.0, 1.0))
        q_2 = wp.quat_from_axis_angle(axis_2, angles[2])

        q_w = q_p

        axis_0 = wp.transform_vector(X_wp, axis_0)
        axis_1 = wp.transform_vector(X_wp, axis_1)
        axis_2 = wp.transform_vector(X_wp, axis_2)

        # joint dynamics
        t_total = wp.vec3()
        # # TODO remove wp.quat_rotate(q_w, ...)?
        # t_total += eval_joint_force(angles[0], wp.dot(wp.quat_rotate(q_w, axis_0), w_err), joint_target[axis_start+0], joint_target_ke[axis_start+0],joint_target_kd[axis_start+0], joint_act[axis_start+0], joint_limit_lower[axis_start+0], joint_limit_upper[axis_start+0], joint_limit_ke[axis_start+0], joint_limit_kd[axis_start+0], wp.quat_rotate(q_w, axis_0))
        # t_total += eval_joint_force(angles[1], wp.dot(wp.quat_rotate(q_w, axis_1), w_err), joint_target[axis_start+1], joint_target_ke[axis_start+1],joint_target_kd[axis_start+1], joint_act[axis_start+1], joint_limit_lower[axis_start+1], joint_limit_upper[axis_start+1], joint_limit_ke[axis_start+1], joint_limit_kd[axis_start+1], wp.quat_rotate(q_w, axis_1))
        # t_total += eval_joint_force(angles[2], wp.dot(wp.quat_rotate(q_w, axis_2), w_err), joint_target[axis_start+2], joint_target_ke[axis_start+2],joint_target_kd[axis_start+2], joint_act[axis_start+2], joint_limit_lower[axis_start+2], joint_limit_upper[axis_start+2], joint_limit_ke[axis_start+2], joint_limit_kd[axis_start+2], wp.quat_rotate(q_w, axis_2))

        t_total += eval_joint_force(
            angles[0],
            wp.dot(axis_0, w_err),
            joint_target[axis_start + 0],
            joint_target_ke[axis_start + 0],
            joint_target_kd[axis_start + 0],
            joint_act[axis_start + 0],
            joint_limit_lower[axis_start + 0],
            joint_limit_upper[axis_start + 0],
            joint_limit_ke[axis_start + 0],
            joint_limit_kd[axis_start + 0],
            axis_0,
        )
        t_total += eval_joint_force(
            angles[1],
            wp.dot(axis_1, w_err),
            joint_target[axis_start + 1],
            joint_target_ke[axis_start + 1],
            joint_target_kd[axis_start + 1],
            joint_act[axis_start + 1],
            joint_limit_lower[axis_start + 1],
            joint_limit_upper[axis_start + 1],
            joint_limit_ke[axis_start + 1],
            joint_limit_kd[axis_start + 1],
            axis_1,
        )
        t_total += eval_joint_force(
            angles[2],
            wp.dot(axis_2, w_err),
            joint_target[axis_start + 2],
            joint_target_ke[axis_start + 2],
            joint_target_kd[axis_start + 2],
            joint_act[axis_start + 2],
            joint_limit_lower[axis_start + 2],
            joint_limit_upper[axis_start + 2],
            joint_limit_ke[axis_start + 2],
            joint_limit_kd[axis_start + 2],
            axis_2,
        )

        f_total += x_err * joint_attach_ke + v_err * joint_attach_kd

    if type == wp.sim.JOINT_UNIVERSAL:
        q_pc = wp.quat_inverse(q_p) * q_c

        # decompose to a compound rotation each axis
        angles = quat_decompose(q_pc)

        # reconstruct rotation axes
        axis_0 = wp.vec3(1.0, 0.0, 0.0)
        q_0 = wp.quat_from_axis_angle(axis_0, angles[0])

        axis_1 = wp.quat_rotate(q_0, wp.vec3(0.0, 1.0, 0.0))
        q_1 = wp.quat_from_axis_angle(axis_1, angles[1])

        axis_2 = wp.quat_rotate(q_1 * q_0, wp.vec3(0.0, 0.0, 1.0))
        q_2 = wp.quat_from_axis_angle(axis_2, angles[2])

        q_w = q_p

        axis_0 = wp.transform_vector(X_wp, axis_0)
        axis_1 = wp.transform_vector(X_wp, axis_1)
        axis_2 = wp.transform_vector(X_wp, axis_2)

        # joint dynamics
        t_total = wp.vec3()

        # free axes
        # # TODO remove wp.quat_rotate(q_w, ...)?
        # t_total += eval_joint_force(angles[0], wp.dot(wp.quat_rotate(q_w, axis_0), w_err), joint_target[axis_start+0], joint_target_ke[axis_start+0],joint_target_kd[axis_start+0], joint_act[axis_start+0], joint_limit_lower[axis_start+0], joint_limit_upper[axis_start+0], joint_limit_ke[axis_start+0], joint_limit_kd[axis_start+0], wp.quat_rotate(q_w, axis_0))
        # t_total += eval_joint_force(angles[1], wp.dot(wp.quat_rotate(q_w, axis_1), w_err), joint_target[axis_start+1], joint_target_ke[axis_start+1],joint_target_kd[axis_start+1], joint_act[axis_start+1], joint_limit_lower[axis_start+1], joint_limit_upper[axis_start+1], joint_limit_ke[axis_start+1], joint_limit_kd[axis_start+1], wp.quat_rotate(q_w, axis_1))

        # # last axis (fixed)
        # t_total += eval_joint_force(angles[2], wp.dot(wp.quat_rotate(q_w, axis_2), w_err), 0.0, joint_attach_ke, joint_attach_kd*angular_damping_scale, 0.0, 0.0, 0.0, 0.0, 0.0, wp.quat_rotate(q_w, axis_2))

        # TODO remove wp.quat_rotate(q_w, ...)?
        t_total += eval_joint_force(
            angles[0],
            wp.dot(axis_0, w_err),
            joint_target[axis_start + 0],
            joint_target_ke[axis_start + 0],
            joint_target_kd[axis_start + 0],
            joint_act[axis_start + 0],
            joint_limit_lower[axis_start + 0],
            joint_limit_upper[axis_start + 0],
            joint_limit_ke[axis_start + 0],
            joint_limit_kd[axis_start + 0],
            axis_0,
        )
        t_total += eval_joint_force(
            angles[1],
            wp.dot(axis_1, w_err),
            joint_target[axis_start + 1],
            joint_target_ke[axis_start + 1],
            joint_target_kd[axis_start + 1],
            joint_act[axis_start + 1],
            joint_limit_lower[axis_start + 1],
            joint_limit_upper[axis_start + 1],
            joint_limit_ke[axis_start + 1],
            joint_limit_kd[axis_start + 1],
            axis_1,
        )

        # last axis (fixed)
        t_total += eval_joint_force(
            angles[2],
            wp.dot(axis_2, w_err),
            0.0,
            joint_attach_ke,
            joint_attach_kd * angular_damping_scale,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            axis_2,
        )

        f_total += x_err * joint_attach_ke + v_err * joint_attach_kd

    # write forces
    if c_parent >= 0:
        wp.atomic_add(body_f, c_parent, wp.spatial_vector(t_total + wp.cross(r_p, f_total), f_total))

    wp.atomic_sub(body_f, c_child, wp.spatial_vector(t_total + wp.cross(r_c, f_total), f_total))


@wp.func
def compute_muscle_force(
        i: int,
        body_X_s: wp.array(dtype=wp.transform),
        body_v_s: wp.array(dtype=wp.spatial_vector),
        body_com: wp.array(dtype=wp.vec3),
        muscle_links: wp.array(dtype=int),
        muscle_points: wp.array(dtype=wp.vec3),
        muscle_activation: float,
        body_f_s: wp.array(dtype=wp.spatial_vector),
):
    link_0 = muscle_links[i]
    link_1 = muscle_links[i + 1]

    if link_0 == link_1:
        return 0

    r_0 = muscle_points[i]
    r_1 = muscle_points[i + 1]

    xform_0 = body_X_s[link_0]
    xform_1 = body_X_s[link_1]

    pos_0 = wp.transform_point(xform_0, r_0 - body_com[link_0])
    pos_1 = wp.transform_point(xform_1, r_1 - body_com[link_1])

    n = wp.normalize(pos_1 - pos_0)

    # todo: add passive elastic and viscosity terms
    f = n * muscle_activation

    wp.atomic_sub(body_f_s, link_0, wp.spatial_vector(f, wp.cross(pos_0, f)))
    wp.atomic_add(body_f_s, link_1, wp.spatial_vector(f, wp.cross(pos_1, f)))

    return 0


@wp.kernel
def eval_muscles(
        body_X_s: wp.array(dtype=wp.transform),
        body_v_s: wp.array(dtype=wp.spatial_vector),
        body_com: wp.array(dtype=wp.vec3),
        muscle_start: wp.array(dtype=int),
        muscle_params: wp.array(dtype=float),
        muscle_links: wp.array(dtype=int),
        muscle_points: wp.array(dtype=wp.vec3),
        muscle_activation: wp.array(dtype=float),
        # output
        body_f_s: wp.array(dtype=wp.spatial_vector),
):
    tid = wp.tid()

    m_start = muscle_start[tid]
    m_end = muscle_start[tid + 1] - 1

    activation = muscle_activation[tid]

    for i in range(m_start, m_end):
        compute_muscle_force(i, body_X_s, body_v_s, body_com, muscle_links, muscle_points, activation, body_f_s)


@wp.kernel
def eval_reference_shape_drag(
        particle_x: wp.array(dtype=wp.vec3),
        particle_v: wp.array(dtype=wp.vec3),
        particle_radius: wp.array(dtype=float),
        particle_flags: wp.array(dtype=wp.uint32),
        particle_reference_label: wp.array(dtype=int),
        particle_normal: wp.array(dtype=wp.vec3),
        reference_shapes: wp.array(dtype=wp.uint64),
        reference_transform: wp.array(dtype=wp.transform),
        reference_scale: wp.array(dtype=wp.vec3),
        reference_marigin: float,
        reference_k: float,
        model_shapes: wp.array(dtype=wp.uint64),
        watertight_shape: int,
        cloth_reference_drag_particles: wp.array(dtype=int),
        # outputs
        particle_f: wp.array(dtype=wp.vec3),
):
    """
    Dim: num_particles
    inputs=[
            state.particle_q,
            state.particle_qd,
            model.particle_radius,
            model.particle_flags,
            model.particle_reference_label,
            model.cloth_reference_shape_ids,
            model.cloth_reference_transforms,
            model.cloth_reference_scale,
            model.cloth_reference_margin
        ],
    outputs=[particle_f],
    """
    particle_index = wp.tid()
    shape_index = particle_reference_label[particle_index]
    if (cloth_reference_drag_particles[particle_index] == 0 or 
        cloth_reference_drag_particles[particle_index] == 1):   # NEW: Don't do anything for self-intersections
        return
    if (particle_flags[particle_index] & PARTICLE_FLAG_ACTIVE) == 0 or shape_index < 0:
        return

    px = particle_x[particle_index]
    v = particle_v[particle_index]
    radius = particle_radius[particle_index]

    X_ws = reference_transform[shape_index]
    X_sw = wp.transform_inverse(X_ws)

    # transform particle position to shape local space
    x_local = wp.transform_point(X_sw, px)

    # geo description
    geo_scale = reference_scale[shape_index]

    # evaluate shape sdf
    d = 1.0e6
    n = wp.vec3()

    mesh = reference_shapes[shape_index]

    face_index = int(0)
    face_u = float(0.0)
    face_v = float(0.0)
    sign = float(0.0)

    point = wp.cw_div(x_local, geo_scale)

    if wp.mesh_query_point(mesh, point, 1000.0, sign, face_index, face_u, face_v):
        pass
        sign = mesh_query_inside_(model_shapes[watertight_shape], wp.cw_div(x_local, geo_scale))
        shape_p = wp.mesh_eval_position(mesh, face_index, face_u, face_v)
        # shape_v = wp.mesh_eval_velocity(mesh, face_index, face_u, face_v)

        shape_p = wp.cw_mul(shape_p, geo_scale)
        # shape_v = wp.cw_mul(shape_v, geo_scale)

        delta = wp.transform_vector(X_ws, shape_p - x_local)
        d = wp.length(delta) * sign
        n = wp.normalize(delta) * sign

        # if (is_global or cloth_reference_drag_particles[particle_index] != 0) and d > reference_marigin + radius:
        if d > reference_marigin + radius:
            # Used in cloth self-collision resolution
            f = n * (d - reference_marigin - radius) * reference_k
            wp.atomic_add(particle_f, particle_index, f)

        if d < 0:
            f = n * d * reference_k
            wp.atomic_add(particle_f, particle_index, f)


def compute_forces(model, state, particle_f, body_f, requires_grad):
    # damped springs
    if model.spring_count:
        wp.launch(
            kernel=eval_springs,
            dim=model.spring_count,
            inputs=[
                state.particle_q,
                state.particle_qd,
                model.spring_indices,
                model.spring_rest_length,
                model.spring_stiffness,
                model.spring_damping,
            ],
            outputs=[particle_f],
            device=model.device,
        )

    # particle-particle interactions
    if model.particle_count:
        eval_particle_forces(model, state, particle_f)

    if model.cloth_reference_drag and model.particle_count:
        wp.launch(
            kernel=eval_reference_shape_drag,
            dim=model.particle_count,
            inputs=[
                state.particle_q,
                state.particle_qd,
                model.particle_radius,
                model.particle_flags,
                model.particle_reference_label,
                model.particle_normal,
                model.cloth_reference_shape_ids,
                model.cloth_reference_transforms,
                model.cloth_reference_scale,
                model.cloth_reference_margin,
                model.cloth_reference_k,
                model.shape_geo.source,
                model.cloth_reference_watertight_whole_shape_index,
                model.cloth_reference_drag_particles
            ],
            outputs=[particle_f],
            device=model.device,
        )

    # triangle elastic and lift/drag forces
    if model.tri_count:
        wp.launch(
            kernel=eval_triangles,
            dim=model.tri_count,
            inputs=[
                state.particle_q,
                state.particle_qd,
                model.tri_indices,
                model.tri_poses,
                model.tri_activations,
                model.tri_materials,
            ],
            outputs=[particle_f],
            device=model.device,
        )

    # triangle/triangle contacts
    if model.enable_tri_collisions and model.tri_count:
        wp.launch(
            kernel=eval_triangles_contact,
            dim=model.tri_count * model.particle_count,
            inputs=[
                model.particle_count,
                state.particle_q,
                state.particle_qd,
                model.tri_indices,
                model.tri_poses,
                model.tri_activations,
                model.tri_materials,
            ],
            outputs=[particle_f],
            device=model.device,
        )

    # triangle bending
    if model.edge_count:
        wp.launch(
            kernel=eval_bending,
            dim=model.edge_count,
            inputs=[
                state.particle_q,
                state.particle_qd,
                model.edge_indices,
                model.edge_rest_angle,
                model.edge_bending_properties,
            ],
            outputs=[particle_f],
            device=model.device,
        )

    # particle ground contact
    if model.ground and model.particle_count:
        wp.launch(
            kernel=eval_particle_ground_contacts,
            dim=model.particle_count,
            inputs=[
                state.particle_q,
                state.particle_qd,
                model.particle_radius,
                model.particle_flags,
                model.soft_contact_ke,
                model.soft_contact_kd,
                model.soft_contact_kf,
                model.soft_contact_mu,
                model.ground_plane,
            ],
            outputs=[particle_f],
            device=model.device,
        )

    # tetrahedral FEM
    if model.tet_count:
        wp.launch(
            kernel=eval_tetrahedra,
            dim=model.tet_count,
            inputs=[
                state.particle_q,
                state.particle_qd,
                model.tet_indices,
                model.tet_poses,
                model.tet_activations,
                model.tet_materials,
            ],
            outputs=[particle_f],
            device=model.device,
        )

    if model.rigid_contact_max and (
            model.ground and model.shape_ground_contact_pair_count or model.shape_contact_pair_count
    ):
        wp.launch(
            kernel=eval_rigid_contacts,
            dim=model.rigid_contact_max,
            inputs=[
                state.body_q,
                state.body_qd,
                model.body_com,
                model.shape_materials,
                model.shape_geo,
                model.rigid_contact_count,
                model.rigid_contact_body0,
                model.rigid_contact_body1,
                model.rigid_contact_point0,
                model.rigid_contact_point1,
                model.rigid_contact_normal,
                model.rigid_contact_shape0,
                model.rigid_contact_shape1,
            ],
            outputs=[body_f],
            device=model.device,
        )

    if model.joint_count:
        wp.launch(
            kernel=eval_body_joints,
            dim=model.joint_count,
            inputs=[
                state.body_q,
                state.body_qd,
                model.body_com,
                model.joint_q_start,
                model.joint_qd_start,
                model.joint_type,
                model.joint_enabled,
                model.joint_child,
                model.joint_parent,
                model.joint_X_p,
                model.joint_X_c,
                model.joint_axis,
                model.joint_axis_start,
                model.joint_axis_dim,
                model.joint_target,
                model.joint_act,
                model.joint_target_ke,
                model.joint_target_kd,
                model.joint_limit_lower,
                model.joint_limit_upper,
                model.joint_limit_ke,
                model.joint_limit_kd,
                model.joint_attach_ke,
                model.joint_attach_kd,
            ],
            outputs=[body_f],
            device=model.device,
        )

    # particle shape contact
    if model.particle_count and model.shape_count > 1:
        wp.launch(
            kernel=eval_particle_contacts,
            dim=model.soft_contact_max,
            inputs=[
                state.particle_q,
                state.particle_qd,
                state.body_q,
                state.body_qd,
                model.particle_radius,
                model.particle_flags,
                model.body_com,
                model.shape_body,
                model.shape_materials,
                model.soft_contact_ke,
                model.soft_contact_kd,
                model.soft_contact_kf,
                model.soft_contact_mu,
                model.particle_adhesion,
                model.soft_contact_count,
                model.soft_contact_particle,
                model.soft_contact_shape,
                model.soft_contact_body_pos,
                model.soft_contact_body_vel,
                model.soft_contact_normal,
                model.soft_contact_max,
            ],
            # outputs
            outputs=[particle_f, body_f],
            device=model.device,
        )

    # evaluate muscle actuation
    if False and model.muscle_count:
        wp.launch(
            kernel=eval_muscles,
            dim=model.muscle_count,
            inputs=[
                state.body_q,
                state.body_qd,
                model.body_com,
                model.muscle_start,
                model.muscle_params,
                model.muscle_bodies,
                model.muscle_points,
                model.muscle_activation,
            ],
            outputs=[body_f],
            device=model.device,
        )

    # if (model.articulation_count):

    #     # evaluate joint torques
    #     wp.launch(
    #         kernel=eval_body_tau,
    #         dim=model.articulation_count,
    #         inputs=[
    #             model.articulation_joint_start,
    #             model.joint_type,
    #             model.joint_parent,
    #             model.joint_q_start,
    #             model.joint_qd_start,
    #             state.joint_q,
    #             state.joint_qd,
    #             state.joint_act,
    #             model.joint_target,
    #             model.joint_target_ke,
    #             model.joint_target_kd,
    #             model.joint_limit_lower,
    #             model.joint_limit_upper,
    #             model.joint_limit_ke,
    #             model.joint_limit_kd,
    #             model.joint_axis,
    #             state.joint_S_s,
    #             state.body_f_s
    #         ],
    #         outputs=[
    #             state.body_ft_s,
    #             state.joint_tau
    #         ],
    #         device=model.device,
    #         preserve_output=True)

    state.particle_f = particle_f
    state.body_f = body_f


# FIXME Inputs/outputs should be swapped
@wp.kernel
def replace_mesh_points(
        points: wp.array(dtype=wp.vec3),
        new_points: wp.array(dtype=wp.vec3),
):
    tid = wp.tid()
    points[tid] = new_points[tid]


class SemiImplicitIntegrator:
    """A semi-implicit integrator using symplectic Euler

    After constructing `Model` and `State` objects this time-integrator
    may be used to advance the simulation state forward in time.

    Semi-implicit time integration is a variational integrator that
    preserves energy, however it not unconditionally stable, and requires a time-step
    small enough to support the required stiffness and damping forces.

    See: https://en.wikipedia.org/wiki/Semi-implicit_Euler_method

    Example
    -------

    .. code-block:: python

        integrator = wp.SemiImplicitIntegrator()

        # simulation loop
        for i in range(100):
            state = integrator.simulate(model, state_in, state_out, dt)

    """

    def __init__(self, angular_damping=0.05):
        self.angular_damping = angular_damping

    def simulate(self, model, state_in, state_out, dt, requires_grad=False):
        with wp.ScopedTimer("simulate", False):
            particle_f = None
            body_f = None

            if state_in.particle_count:
                particle_f = state_in.particle_f

            if state_in.body_count:
                body_f = state_in.body_f

            compute_forces(model, state_in, particle_f, body_f, requires_grad=requires_grad)

            # -------------------------------------
            # integrate bodies

            if model.body_count:
                wp.launch(
                    kernel=integrate_bodies,
                    dim=model.body_count,
                    inputs=[
                        state_in.body_q,
                        state_in.body_qd,
                        state_in.body_f,
                        model.body_com,
                        model.body_mass,
                        model.body_inertia,
                        model.body_inv_mass,
                        model.body_inv_inertia,
                        model.gravity,
                        self.angular_damping,
                        dt,
                    ],
                    outputs=[state_out.body_q, state_out.body_qd],
                    device=model.device,
                )

            # ----------------------------
            # integrate particles

            if model.particle_count:
                wp.launch(
                    kernel=integrate_particles,
                    dim=model.particle_count,
                    inputs=[
                        state_in.particle_q,
                        state_in.particle_qd,
                        state_in.particle_f,
                        model.particle_inv_mass,
                        model.particle_flags,
                        model.gravity,
                        model.global_viscous_damping,
                        dt
                    ],
                    outputs=[state_out.particle_q, state_out.particle_qd],
                    device=model.device,
                )

                wp.launch(
                    kernel=replace_mesh_points,
                    dim=model.particle_count,
                    inputs=[model.particle_shape.points],
                    outputs=[state_out.particle_q],
                    device=model.device,
                )
                model.particle_shape.refit()
                compute_vertex_normal(model.particle_shape, model.particle_normal)

            return state_out


@wp.kernel
def compute_particle_residual(
        particle_qd_0: wp.array(dtype=wp.vec3),
        particle_qd_1: wp.array(dtype=wp.vec3),
        particle_f: wp.array(dtype=wp.vec3),
        particle_m: wp.array(dtype=float),
        particle_flags: wp.array(dtype=wp.uint32),
        gravity: wp.vec3,
        dt: float,
        residual: wp.array(dtype=wp.vec3),
):
    tid = wp.tid()
    if (particle_flags[tid] & PARTICLE_FLAG_ACTIVE) == 0:
        return

    m = particle_m[tid]
    v1 = particle_qd_1[tid]
    v0 = particle_qd_0[tid]
    f = particle_f[tid]

    err = wp.vec3()

    if m > 0.0:
        err = (v1 - v0) * m - f * dt - gravity * dt * m

    residual[tid] = err


@wp.kernel
def update_particle_position(
        particle_q_0: wp.array(dtype=wp.vec3),
        particle_q_1: wp.array(dtype=wp.vec3),
        particle_qd_1: wp.array(dtype=wp.vec3),
        x: wp.array(dtype=wp.vec3),
        particle_flags: wp.array(dtype=wp.uint32),
        dt: float,
):
    tid = wp.tid()
    if (particle_flags[tid] & PARTICLE_FLAG_ACTIVE) == 0:
        return

    qd_1 = x[tid]

    q_0 = particle_q_0[tid]
    q_1 = q_0 + qd_1 * dt

    particle_q_1[tid] = q_1
    particle_qd_1[tid] = qd_1


def compute_residual(model, state_in, state_out, particle_f, residual, dt):
    wp.launch(
        kernel=compute_particle_residual,
        dim=model.particle_count,
        inputs=[
            state_in.particle_qd,
            state_out.particle_qd,
            particle_f,
            model.particle_mass,
            model.particle_flags,
            model.gravity,
            dt,
            residual.astype(dtype=wp.vec3),
        ],
        device=model.device,
    )


def init_state(model, state_in, state_out, dt):
    wp.launch(
        kernel=integrate_particles,
        dim=model.particle_count,
        inputs=[
            state_in.particle_q,
            state_in.particle_qd,
            state_in.particle_f,
            model.particle_inv_mass,
            model.particle_flags,
            model.gravity,
            model.global_viscous_damping,
            dt
        ],
        outputs=[state_out.particle_q, state_out.particle_qd],
        device=model.device,
    )


# compute the final positions given output velocity (x)
def update_state(model, state_in, state_out, x, dt):
    wp.launch(
        kernel=update_particle_position,
        dim=model.particle_count,
        inputs=[state_in.particle_q, state_out.particle_q, state_out.particle_qd, x, model.particle_flags, dt],
        device=model.device,
    )


class VariationalImplicitIntegrator:
    def __init__(self, model, solver="gd", alpha=0.1, max_iters=32, report=False):
        self.solver = solver
        self.alpha = alpha
        self.max_iters = max_iters
        self.report = report

        self.opt = Optimizer(model.particle_count * 3, mode=self.solver, device=model.device)

        # allocate temporary space for evaluating particle forces
        self.particle_f = wp.zeros(model.particle_count, dtype=wp.vec3, device=model.device)

    def simulate(self, model, state_in, state_out, dt, requires_grad=False):
        if state_in is state_out:
            raise RuntimeError("Implicit integrators require state objects to not alias each other")

        with wp.ScopedTimer("simulate", False):
            # alloc particle force buffer
            if model.particle_count:
                def residual_func(x, dfdx):
                    self.particle_f.zero_()

                    update_state(model, state_in, state_out, x.astype(wp.vec3), dt)
                    compute_forces(model, state_out, self.particle_f, None)
                    compute_residual(model, state_in, state_out, self.particle_f, dfdx, dt)

                # initialize output state using the input velocity to create 'predicted state'
                init_state(model, state_in, state_out, dt)

                # our optimization variable
                x = state_out.particle_qd.astype(dtype=float)

                self.opt.solve(
                    x=x, grad_func=residual_func, max_iters=self.max_iters, alpha=self.alpha, report=self.report
                )

                # final update to output state with solved velocity
                update_state(model, state_in, state_out, x.astype(wp.vec3), dt)

            return state_out