1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147

#![allow(missing_docs)] // For downcast.

use crate::dynamics::joint::MultibodyLink;
use crate::dynamics::solver::{
    AnyJointVelocityConstraint, JointGenericVelocityGroundConstraint, WritebackId,
};
use crate::dynamics::{IntegrationParameters, JointMotor, Multibody};
use crate::math::Real;
use na::DVector;

/// Initializes and generate the velocity constraints applicable to the multibody links attached
/// to this multibody_joint.
pub fn unit_joint_limit_constraint(
    params: &IntegrationParameters,
    multibody: &Multibody,
    link: &MultibodyLink,
    limits: [Real; 2],
    curr_pos: Real,
    dof_id: usize,
    j_id: &mut usize,
    jacobians: &mut DVector<Real>,
    constraints: &mut Vec<AnyJointVelocityConstraint>,
    insert_at: &mut Option<usize>,
) {
    let ndofs = multibody.ndofs();
    let joint_velocity = multibody.joint_velocity(link);

    let min_enabled = curr_pos < limits[0];
    let max_enabled = limits[1] < curr_pos;
    let erp_inv_dt = params.joint_erp_inv_dt();
    let cfm_coeff = params.joint_cfm_coeff();
    let rhs_bias = ((curr_pos - limits[1]).max(0.0) - (limits[0] - curr_pos).max(0.0)) * erp_inv_dt;
    let rhs_wo_bias = joint_velocity[dof_id];

    let dof_j_id = *j_id + dof_id + link.assembly_id;
    jacobians.rows_mut(*j_id, ndofs * 2).fill(0.0);
    jacobians[dof_j_id] = 1.0;
    jacobians[dof_j_id + ndofs] = 1.0;
    multibody
        .inv_augmented_mass()
        .solve_mut(&mut jacobians.rows_mut(*j_id + ndofs, ndofs));

    let lhs = jacobians[dof_j_id + ndofs]; // = J^t * M^-1 J
    let impulse_bounds = [
        min_enabled as u32 as Real * -Real::MAX,
        max_enabled as u32 as Real * Real::MAX,
    ];

    let constraint = JointGenericVelocityGroundConstraint {
        mj_lambda2: multibody.solver_id,
        ndofs2: ndofs,
        j_id2: *j_id,
        joint_id: usize::MAX,
        impulse: 0.0,
        impulse_bounds,
        inv_lhs: crate::utils::inv(lhs),
        rhs: rhs_wo_bias + rhs_bias,
        rhs_wo_bias,
        cfm_coeff,
        cfm_gain: 0.0,
        writeback_id: WritebackId::Limit(dof_id),
    };

    if let Some(at) = insert_at {
        constraints[*at] = AnyJointVelocityConstraint::JointGenericGroundConstraint(constraint);
        *at += 1;
    } else {
        constraints.push(AnyJointVelocityConstraint::JointGenericGroundConstraint(
            constraint,
        ));
    }
    *j_id += 2 * ndofs;
}

/// Initializes and generate the velocity constraints applicable to the multibody links attached
/// to this multibody_joint.
pub fn unit_joint_motor_constraint(
    params: &IntegrationParameters,
    multibody: &Multibody,
    link: &MultibodyLink,
    motor: &JointMotor,
    curr_pos: Real,
    limits: Option<[Real; 2]>,
    dof_id: usize,
    j_id: &mut usize,
    jacobians: &mut DVector<Real>,
    constraints: &mut Vec<AnyJointVelocityConstraint>,
    insert_at: &mut Option<usize>,
) {
    let inv_dt = params.inv_dt();
    let ndofs = multibody.ndofs();
    let joint_velocity = multibody.joint_velocity(link);

    let motor_params = motor.motor_params(params.dt);

    let dof_j_id = *j_id + dof_id + link.assembly_id;
    jacobians.rows_mut(*j_id, ndofs * 2).fill(0.0);
    jacobians[dof_j_id] = 1.0;
    jacobians[dof_j_id + ndofs] = 1.0;
    multibody
        .inv_augmented_mass()
        .solve_mut(&mut jacobians.rows_mut(*j_id + ndofs, ndofs));

    let lhs = jacobians[dof_j_id + ndofs]; // = J^t * M^-1 J
    let impulse_bounds = [-motor_params.max_impulse, motor_params.max_impulse];

    let mut rhs_wo_bias = 0.0;
    if motor_params.erp_inv_dt != 0.0 {
        rhs_wo_bias += (curr_pos - motor_params.target_pos) * motor_params.erp_inv_dt;
    }

    let mut target_vel = motor_params.target_vel;
    if let Some(limits) = limits {
        target_vel = target_vel.clamp(
            (limits[0] - curr_pos) * inv_dt,
            (limits[1] - curr_pos) * inv_dt,
        );
    };

    let dvel = joint_velocity[dof_id];
    rhs_wo_bias += dvel - target_vel;

    let constraint = JointGenericVelocityGroundConstraint {
        mj_lambda2: multibody.solver_id,
        ndofs2: ndofs,
        j_id2: *j_id,
        joint_id: usize::MAX,
        impulse: 0.0,
        impulse_bounds,
        cfm_coeff: motor_params.cfm_coeff,
        cfm_gain: motor_params.cfm_gain,
        inv_lhs: crate::utils::inv(lhs),
        rhs: rhs_wo_bias,
        rhs_wo_bias,
        writeback_id: WritebackId::Limit(dof_id),
    };

    if let Some(at) = insert_at {
        constraints[*at] = AnyJointVelocityConstraint::JointGenericGroundConstraint(constraint);
        *at += 1;
    } else {
        constraints.push(AnyJointVelocityConstraint::JointGenericGroundConstraint(
            constraint,
        ));
    }
    *j_id += 2 * ndofs;
}