Skip to content

preconditioner

Jacobi preconditioner for the system matrix A = J M⁻¹ Jᵀ + C.

JacobiPreconditioner

Bases: LinearOperator

A Jacobi (diagonal) preconditioner for the system matrix A = J M⁻¹ Jᵀ + C.

This class provides a .matvec() method that applies the inverse of the diagonal of A, for use with Warp's iterative solvers.

Source code in src/axion/optim/preconditioner.py 
 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
148
149
class JacobiPreconditioner(LinearOperator):
    """
    A Jacobi (diagonal) preconditioner for the system matrix A = J M⁻¹ Jᵀ + C.

    This class provides a .matvec() method that applies the inverse of the
    diagonal of A, for use with Warp's iterative solvers.
    """

    def __init__(self, engine, regularization):
        super().__init__(
            shape=(engine.dims.N_w, engine.dims.N_c, engine.dims.N_c),
            dtype=wp.float32,
            device=engine.device,
            matvec=None,  # Will be set later
        )
        self.engine = engine
        self.regularization = regularization

        # Storage for the inverse diagonal elements
        self._P_inv_diag = wp.zeros(
            (engine.dims.N_w, engine.dims.N_c), dtype=wp.float32, device=self.device
        )

    def update(self):
        """
        Re-computes the preconditioner's data. This must be called each time
        the Jacobian (J) or compliance (C) values change.
        """
        wp.launch(
            kernel=compute_inv_diag_kernel,
            dim=(self.engine.dims.N_w, self.engine.dims.N_c),
            inputs=[
                self.engine.data.world_M_inv,
                self.engine.data.J_values.full,
                self.engine.data.C_values.full,
                self.engine.data.constraint_body_idx.full,
                self.engine.data.constraint_active_mask.full,
                self.regularization,
            ],
            outputs=[
                self._P_inv_diag,
                self.engine.data.system_diag,
            ],
            device=self.device,
        )

    def matvec(self, x, y, z, alpha, beta):
        """
        Performs the preconditioning operation z = beta*y + alpha*(M⁻¹@x),
        where M⁻¹ is the inverse diagonal matrix stored in `_P_inv_diag`.
        """
        wp.launch(
            kernel=apply_preconditioner_kernel,
            dim=(self.engine.dims.N_w, self.engine.dims.N_c),
            inputs=[
                self._P_inv_diag,
                self.engine.data.constraint_active_mask.full,
                x,
                y,
                alpha,
                beta,
                z,
            ],
            device=self.device,
        )

matvec

Performs the preconditioning operation z = betay + alpha(M⁻¹@x), where M⁻¹ is the inverse diagonal matrix stored in _P_inv_diag.

Source code in src/axion/optim/preconditioner.py 
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
def matvec(self, x, y, z, alpha, beta):
    """
    Performs the preconditioning operation z = beta*y + alpha*(M⁻¹@x),
    where M⁻¹ is the inverse diagonal matrix stored in `_P_inv_diag`.
    """
    wp.launch(
        kernel=apply_preconditioner_kernel,
        dim=(self.engine.dims.N_w, self.engine.dims.N_c),
        inputs=[
            self._P_inv_diag,
            self.engine.data.constraint_active_mask.full,
            x,
            y,
            alpha,
            beta,
            z,
        ],
        device=self.device,
    )

update

Re-computes the preconditioner's data. This must be called each time the Jacobian (J) or compliance (C) values change.

Source code in src/axion/optim/preconditioner.py 
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
def update(self):
    """
    Re-computes the preconditioner's data. This must be called each time
    the Jacobian (J) or compliance (C) values change.
    """
    wp.launch(
        kernel=compute_inv_diag_kernel,
        dim=(self.engine.dims.N_w, self.engine.dims.N_c),
        inputs=[
            self.engine.data.world_M_inv,
            self.engine.data.J_values.full,
            self.engine.data.C_values.full,
            self.engine.data.constraint_body_idx.full,
            self.engine.data.constraint_active_mask.full,
            self.regularization,
        ],
        outputs=[
            self._P_inv_diag,
            self.engine.data.system_diag,
        ],
        device=self.device,
    )

apply_preconditioner_kernel

Applies the Jacobi preconditioner: z = beta*y + alpha * P⁻¹ * x

Source code in src/axion/optim/preconditioner.py 
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
@wp.kernel
def apply_preconditioner_kernel(
    P_inv_diag: wp.array(dtype=wp.float32, ndim=2),
    constraint_active_mask: wp.array(dtype=wp.float32, ndim=2),
    vec_x: wp.array(dtype=wp.float32, ndim=2),
    vec_y: wp.array(dtype=wp.float32, ndim=2),
    alpha: float,
    beta: float,
    out_vec_z: wp.array(dtype=wp.float32, ndim=2),
):
    """Applies the Jacobi preconditioner: z = beta*y + alpha * P⁻¹ * x"""
    world_idx, constraint_idx = wp.tid()

    is_active = constraint_active_mask[world_idx, constraint_idx]

    # Calculate the preconditioned value (M⁻¹ x)
    # If inactive, the result of the matrix operation is 0.0
    preconditioned_x = 0.0
    if is_active > 0.0:
        preconditioned_x = P_inv_diag[world_idx, constraint_idx] * vec_x[world_idx, constraint_idx]

    # Combine with beta * y and write to output.
    if beta == 0.0:
        out_vec_z[world_idx, constraint_idx] = alpha * preconditioned_x
    else:
        out_vec_z[world_idx, constraint_idx] = (
            beta * vec_y[world_idx, constraint_idx] + alpha * preconditioned_x
        )

compute_inv_diag_kernel

Computes the inverse of the diagonal of the system matrix A = J M⁻¹ Jᵀ + C. The result P_inv_diag[i] = 1.0 / A[i,i] is stored.

Source code in src/axion/optim/preconditioner.py 
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
@wp.kernel
def compute_inv_diag_kernel(
    body_M_inv: wp.array(dtype=SpatialInertia, ndim=2),
    J_values: wp.array(dtype=wp.spatial_vector, ndim=3),
    C_values: wp.array(dtype=wp.float32, ndim=2),
    constraint_body_idx: wp.array(dtype=wp.int32, ndim=3),
    constraint_active_mask: wp.array(dtype=wp.float32, ndim=2),
    regularization: wp.float32,
    # Output array
    P_inv_diag: wp.array(dtype=wp.float32, ndim=2),
    system_diag: wp.array(dtype=wp.float32, ndim=2),
):
    """
    Computes the inverse of the diagonal of the system matrix A = J M⁻¹ Jᵀ + C.
    The result P_inv_diag[i] = 1.0 / A[i,i] is stored.
    """
    world_idx, constraint_idx = wp.tid()

    is_active = constraint_active_mask[world_idx, constraint_idx]
    if is_active == 0.0:
        return

    body_1 = constraint_body_idx[world_idx, constraint_idx, 0]
    body_2 = constraint_body_idx[world_idx, constraint_idx, 1]

    result = 0.0
    if body_1 >= 0:
        M_inv_1 = body_M_inv[world_idx, body_1]
        J_1 = J_values[world_idx, constraint_idx, 0]
        result += wp.dot(J_1, to_spatial_momentum(M_inv_1, J_1))
    if body_2 >= 0:
        M_inv_2 = body_M_inv[world_idx, body_2]
        J_2 = J_values[world_idx, constraint_idx, 1]
        result += wp.dot(J_2, to_spatial_momentum(M_inv_2, J_2))

    # Add diagonal compliance term C[i,i]
    diag_A = result + C_values[world_idx, constraint_idx] + regularization

    # Store raw diagonal
    system_diag[world_idx, constraint_idx] = diag_A

    # Compute and store inverse, with stabilization
    P_inv_diag[world_idx, constraint_idx] = 1.0 / (diag_A + 1e-6)