jaxsim 0.2.dev191__py3-none-any.whl → 0.6.1.dev2__py3-none-any.whl

jaxsim/physics/algos/aba_motors.py

@@ -1,284 +0,0 @@
-from typing import Tuple
-
-import jax
-import jax.numpy as jnp
-import numpy as np
-
-import jaxsim.typing as jtp
-from jaxsim.math.adjoint import Adjoint
-from jaxsim.math.cross import Cross
-from jaxsim.physics.model.physics_model import PhysicsModel
-
-from . import utils
-
-
-def aba(
-    model: PhysicsModel,
-    xfb: jtp.Vector,
-    q: jtp.Vector,
-    qd: jtp.Vector,
-    tau: jtp.Vector,
-    f_ext: jtp.Matrix | None = None,
-) -> Tuple[jtp.Vector, jtp.Vector]:
-    """
-    Articulated Body Algorithm (ABA) algorithm with motor dynamics for forward dynamics.
-    """
-
-    x_fb, q, qd, _, tau, f_ext = utils.process_inputs(
-        physics_model=model, xfb=xfb, q=q, qd=qd, tau=tau, f_ext=f_ext
-    )
-
-    # Extract data from the physics model
-    pre_X_λi = model.tree_transforms
-    M = model.spatial_inertias
-    i_X_pre = model.joint_transforms(q=q)
-    S = model.motion_subspaces(q=q)
-    λ = model.parent_array()
-
-    # Extract motor parameters from the physics model
-    Γ = jnp.array([*model._joint_motor_gear_ratio.values()])
-    IM = jnp.array(
-        [jnp.eye(6) * m for m in [*model._joint_motor_inertia.values()]] * model.NB
-    )
-    K̅ᵥ = Γ.T * jnp.array([*model._joint_motor_viscous_friction.values()]) * Γ
-    m_S = jnp.concatenate([S[:1], S[1:] * Γ[:, None, None]], axis=0)
-
-    # Initialize buffers
-    v = jnp.array([jnp.zeros([6, 1])] * model.NB)
-    MA = jnp.array([jnp.zeros([6, 6])] * model.NB)
-    pA = jnp.array([jnp.zeros([6, 1])] * model.NB)
-    c = jnp.array([jnp.zeros([6, 1])] * model.NB)
-    i_X_λi = jnp.zeros_like(i_X_pre)
-
-    m_v = jnp.array([jnp.zeros([6, 1])] * model.NB)
-    m_c = jnp.array([jnp.zeros([6, 1])] * model.NB)
-    pR = jnp.array([jnp.zeros([6, 1])] * model.NB)
-
-    # Base pose B_X_W and velocity
-    base_quat = jnp.vstack(x_fb[0:4])
-    base_pos = jnp.vstack(x_fb[4:7])
-    base_vel = jnp.vstack(jnp.hstack([x_fb[10:13], x_fb[7:10]]))
-
-    # 6D transform of base velocity
-    B_X_W = Adjoint.from_quaternion_and_translation(
-        quaternion=base_quat,
-        translation=base_pos,
-        inverse=True,
-        normalize_quaternion=True,
-    )
-    i_X_λi = i_X_λi.at[0].set(B_X_W)
-
-    # Transforms link -> base
-    i_X_0 = jnp.zeros_like(pre_X_λi)
-    i_X_0 = i_X_0.at[0].set(jnp.eye(6))
-
-    # Initialize base quantities
-    if model.is_floating_base:
-        # Base velocity v₀
-        v_0 = B_X_W @ base_vel
-        v = v.at[0].set(v_0)
-
-        # AB inertia (Mᴬ) and AB bias forces (pᴬ)
-        MA_0 = M[0]
-        MA = MA.at[0].set(MA_0)
-        pA_0 = Cross.vx_star(v[0]) @ MA_0 @ v[0] - Adjoint.inverse(
-            B_X_W
-        ).T @ jnp.vstack(f_ext[0])
-        pA = pA.at[0].set(pA_0)
-
-    Pass1Carry = Tuple[
-        jtp.MatrixJax,
-        jtp.MatrixJax,
-        jtp.MatrixJax,
-        jtp.MatrixJax,
-        jtp.MatrixJax,
-        jtp.MatrixJax,
-        jtp.MatrixJax,
-        jtp.MatrixJax,
-        jtp.MatrixJax,
-    ]
-
-    pass_1_carry = (i_X_λi, v, c, m_v, m_c, MA, pA, pR, i_X_0)
-
-    # Pass 1
-    def loop_body_pass1(carry: Pass1Carry, i: jtp.Int) -> Tuple[Pass1Carry, None]:
-        ii = i - 1
-        i_X_λi, v, c, m_v, m_c, MA, pA, pR, i_X_0 = carry
-
-        # Compute parent-to-child transform
-        i_X_λi_i = i_X_pre[i] @ pre_X_λi[i]
-        i_X_λi = i_X_λi.at[i].set(i_X_λi_i)
-
-        # Propagate link velocity
-        vJ = S[i] * qd[ii] * (qd.size != 0)
-        m_vJ = m_S[i] * qd[ii] * (qd.size != 0)
-
-        v_i = i_X_λi[i] @ v[λ[i]] + vJ
-        v = v.at[i].set(v_i)
-
-        m_v_i = i_X_λi[i] @ v[λ[i]] + m_vJ
-        m_v = m_v.at[i].set(m_v_i)
-
-        c_i = Cross.vx(v[i]) @ vJ
-        c = c.at[i].set(c_i)
-        m_c_i = Cross.vx(m_v[i]) @ m_vJ
-        m_c = m_c.at[i].set(m_c_i)
-
-        # Initialize articulated-body inertia
-        MA_i = jnp.array(M[i])
-        MA = MA.at[i].set(MA_i)
-
-        # Initialize articulated-body bias forces
-        i_X_0_i = i_X_λi[i] @ i_X_0[model.parent[i]]
-        i_X_0 = i_X_0.at[i].set(i_X_0_i)
-        i_Xf_W = Adjoint.inverse(i_X_0[i] @ B_X_W).T
-
-        pA_i = Cross.vx_star(v[i]) @ M[i] @ v[i] - i_Xf_W @ jnp.vstack(f_ext[i])
-        pA = pA.at[i].set(pA_i)
-
-        pR_i = Cross.vx_star(m_v[i]) @ IM[i] @ m_v[i] - K̅ᵥ[i] * m_v[i]
-        pR = pR.at[i].set(pR_i)
-
-        return (i_X_λi, v, c, m_v, m_c, MA, pA, pR, i_X_0), None
-
-    (i_X_λi, v, c, m_v, m_c, MA, pA, pR, i_X_0), _ = jax.lax.scan(
-        f=loop_body_pass1,
-        init=pass_1_carry,
-        xs=np.arange(start=1, stop=model.NB),
-    )
-
-    U = jnp.zeros_like(S)
-    m_U = jnp.zeros_like(S)
-    d = jnp.zeros(shape=(model.NB, 1))
-    u = jnp.zeros(shape=(model.NB, 1))
-    m_u = jnp.zeros(shape=(model.NB, 1))
-
-    Pass2Carry = Tuple[
-        jtp.MatrixJax,
-        jtp.MatrixJax,
-        jtp.MatrixJax,
-        jtp.MatrixJax,
-        jtp.MatrixJax,
-        jtp.MatrixJax,
-        jtp.MatrixJax,
-    ]
-
-    pass_2_carry = (U, m_U, d, u, m_u, MA, pA)
-
-    # Pass 2
-    def loop_body_pass2(carry: Pass2Carry, i: jtp.Int) -> Tuple[Pass2Carry, None]:
-        ii = i - 1
-        U, m_U, d, u, m_u, MA, pA = carry
-
-        # Compute intermediate results
-        u_i = tau[ii] - S[i].T @ pA[i] if tau.size != 0 else -S[i].T @ pA[i]
-        u = u.at[i].set(u_i.squeeze())
-
-        has_motors = jnp.allclose(Γ[i], 1.0)
-
-        m_u_i = (
-            tau[ii] / Γ[i] * has_motors - m_S[i].T @ pR[i]
-            if tau.size != 0
-            else -m_S[i].T @ pR[i]
-        )
-        m_u = m_u.at[i].set(m_u_i.squeeze())
-
-        U_i = MA[i] @ S[i]
-        U = U.at[i].set(U_i)
-
-        m_U_i = IM[i] @ m_S[i]
-        m_U = m_U.at[i].set(m_U_i)
-
-        d_i = S[i].T @ MA[i] @ S[i] + m_S[i].T @ IM[i] @ m_S[i]
-        d = d.at[i].set(d_i.squeeze())
-
-        # Compute the articulated-body inertia and bias forces of this link
-        Ma = MA[i] + IM[i] - U[i] / d[i] @ U[i].T - m_U[i] / d[i] @ m_U[i].T
-        pa = (
-            pA[i]
-            + pR[i]
-            + Ma[i] @ c[i]
-            + IM[i] @ m_c[i]
-            + U[i] / d[i] * u[i]
-            + m_U[i] / d[i] * m_u[i]
-        )
-
-        # Propagate them to the parent, handling the base link
-        def propagate(
-            MA_pA: Tuple[jtp.MatrixJax, jtp.MatrixJax]
-        ) -> Tuple[jtp.MatrixJax, jtp.MatrixJax]:
-            MA, pA = MA_pA
-
-            MA_λi = MA[λ[i]] + i_X_λi[i].T @ Ma @ i_X_λi[i]
-            MA = MA.at[λ[i]].set(MA_λi)
-
-            pA_λi = pA[λ[i]] + i_X_λi[i].T @ pa
-            pA = pA.at[λ[i]].set(pA_λi)
-
-            return MA, pA
-
-        MA, pA = jax.lax.cond(
-            pred=jnp.array([λ[i] != 0, model.is_floating_base]).any(),
-            true_fun=propagate,
-            false_fun=lambda MA_pA: MA_pA,
-            operand=(MA, pA),
-        )
-
-        return (U, m_U, d, u, m_u, MA, pA), None
-
-    (U, m_U, d, u, m_u, MA, pA), _ = jax.lax.scan(
-        f=loop_body_pass2,
-        init=pass_2_carry,
-        xs=np.flip(np.arange(start=1, stop=model.NB)),
-    )
-
-    if model.is_floating_base:
-        a0 = jnp.linalg.solve(-MA[0], pA[0])
-    else:
-        a0 = -B_X_W @ jnp.vstack(model.gravity)
-
-    a = jnp.zeros_like(S)
-    a = a.at[0].set(a0)
-    qdd = jnp.zeros_like(q)
-
-    Pass3Carry = Tuple[jtp.MatrixJax, jtp.VectorJax]
-    pass_3_carry = (a, qdd)
-
-    # Pass 3
-    def loop_body_pass3(carry: Pass3Carry, i: jtp.Int) -> Tuple[Pass3Carry, None]:
-        ii = i - 1
-        a, qdd = carry
-
-        # Propagate link accelerations
-        a_i = i_X_λi[i] @ a[λ[i]] + c[i]
-
-        # Compute joint accelerations
-        qdd_ii = (u[i] + m_u[i] - (U[i].T + m_U[i].T) @ a_i) / d[i]
-        qdd = qdd.at[ii].set(qdd_ii.squeeze()) if qdd.size != 0 else qdd
-
-        a_i = a_i + S[i] * qdd[ii] if qdd.size != 0 else a_i
-        a = a.at[i].set(a_i)
-
-        return (a, qdd), None
-
-    (a, qdd), _ = jax.lax.scan(
-        f=loop_body_pass3,
-        init=pass_3_carry,
-        xs=np.arange(1, model.NB),
-    )
-
-    # Handle 1 DoF models
-    qdd = jnp.atleast_1d(qdd.squeeze())
-    qdd = jnp.vstack(qdd) if qdd.size > 0 else jnp.empty(shape=(0, 1))
-
-    # Get the resulting base acceleration (w/o gravity) in body-fixed representation
-    B_a_WB = a[0]
-
-    # Convert the base acceleration to inertial-fixed representation, and add gravity
-    W_a_WB = jnp.vstack(
-        jnp.linalg.solve(B_X_W, B_a_WB) + jnp.vstack(model.gravity)
-        if model.is_floating_base
-        else jnp.zeros(6)
-    )
-
-    return W_a_WB, qdd

jaxsim/physics/algos/crba.py

@@ -1,154 +0,0 @@
-from typing import Tuple
-
-import jax
-import jax.numpy as jnp
-import numpy as np
-
-import jaxsim.typing as jtp
-from jaxsim.physics.model.physics_model import PhysicsModel
-
-from . import utils
-
-
-def crba(model: PhysicsModel, q: jtp.Vector) -> jtp.Matrix:
-    """
-    Compute the Composite Rigid-Body Inertia Matrix (CRBA) for an articulated body or robot given joint positions.
-
-    Args:
-        model (PhysicsModel): The physics model of the articulated body or robot.
-        q (jtp.Vector): Joint positions (Generalized coordinates).
-
-    Returns:
-        jtp.Matrix: The Composite Rigid-Body Inertia Matrix (CRBA) of the articulated body or robot.
-    """
-
-    _, q, _, _, _, _ = utils.process_inputs(
-        physics_model=model, xfb=None, q=q, qd=None, tau=None, f_ext=None
-    )
-
-    Xtree = model.tree_transforms
-    Mc = model.spatial_inertias
-    S = model.motion_subspaces(q=q)
-    Xj = model.joint_transforms(q=q)
-
-    Xup = jnp.zeros_like(Xtree)
-    i_X_0 = jnp.zeros_like(Xtree)
-    i_X_0 = i_X_0.at[0].set(jnp.eye(6))
-
-    # Parent array mapping: i -> λ(i).
-    # Exception: λ(0) must not be used, it's initialized to -1.
-    λ = model.parent
-
-    # ====================
-    # Propagate kinematics
-    # ====================
-
-    ForwardPassCarry = Tuple[jtp.MatrixJax, jtp.MatrixJax]
-    forward_pass_carry = (Xup, i_X_0)
-
-    def propagate_kinematics(
-        carry: ForwardPassCarry, i: jtp.Int
-    ) -> Tuple[ForwardPassCarry, None]:
-        Xup, i_X_0 = carry
-
-        Xup_i = Xj[i] @ Xtree[i]
-        Xup = Xup.at[i].set(Xup_i)
-
-        i_X_0_i = Xup[i] @ i_X_0[λ[i]]
-        i_X_0 = i_X_0.at[i].set(i_X_0_i)
-
-        return (Xup, i_X_0), None
-
-    (Xup, i_X_0), _ = jax.lax.scan(
-        f=propagate_kinematics,
-        init=forward_pass_carry,
-        xs=np.arange(start=1, stop=model.NB),
-    )
-
-    # ===================
-    # Compute mass matrix
-    # ===================
-
-    M = jnp.zeros(shape=(6 + model.dofs(), 6 + model.dofs()))
-
-    BackwardPassCarry = Tuple[jtp.MatrixJax, jtp.MatrixJax]
-    backward_pass_carry = (Mc, M)
-
-    def backward_pass(
-        carry: BackwardPassCarry, i: jtp.Int
-    ) -> Tuple[BackwardPassCarry, None]:
-        ii = i - 1
-        Mc, M = carry
-
-        Mc_λi = Mc[λ[i]] + Xup[i].T @ Mc[i] @ Xup[i]
-        Mc = Mc.at[λ[i]].set(Mc_λi)
-
-        Fi = Mc[i] @ S[i]
-        M_ii = S[i].T @ Fi
-        M = M.at[ii + 6, ii + 6].set(M_ii.squeeze())
-
-        j = i
-
-        CarryInnerFn = Tuple[jtp.Int, jtp.MatrixJax, jtp.MatrixJax]
-        carry_inner_fn = (j, Fi, M)
-
-        def while_loop_body(carry: CarryInnerFn) -> CarryInnerFn:
-            j, Fi, M = carry
-
-            Fi = Xup[j].T @ Fi
-            j = λ[j]
-            jj = j - 1
-
-            M_ij = Fi.T @ S[j]
-
-            M = M.at[ii + 6, jj + 6].set(M_ij.squeeze())
-            M = M.at[jj + 6, ii + 6].set(M_ij.squeeze())
-
-            return j, Fi, M
-
-        # The following functions are part of a (rather messy) workaround for computing
-        # a while loop using a for loop with fixed number of iterations.
-        def inner_fn(carry: CarryInnerFn, k: jtp.Int) -> Tuple[CarryInnerFn, None]:
-            def compute_inner(carry: CarryInnerFn) -> Tuple[CarryInnerFn, None]:
-                j, Fi, M = carry
-                out = jax.lax.cond(
-                    pred=(λ[j] > 0),
-                    true_fun=while_loop_body,
-                    false_fun=lambda carry: carry,
-                    operand=carry,
-                )
-                return out, None
-
-            j, Fi, M = carry
-            return jax.lax.cond(
-                pred=(k == j),
-                true_fun=compute_inner,
-                false_fun=lambda carry: (carry, None),
-                operand=carry,
-            )
-
-        (j, Fi, M), _ = jax.lax.scan(
-            f=inner_fn,
-            init=carry_inner_fn,
-            xs=np.flip(np.arange(start=1, stop=model.NB)),
-        )
-
-        Fi = i_X_0[j].T @ Fi
-
-        M = M.at[0:6, ii + 6].set(Fi.squeeze())
-        M = M.at[ii + 6, 0:6].set(Fi.squeeze())
-
-        return (Mc, M), None
-
-    # This scan performs the backward pass to compute Mbj, Mjb and Mjj, that
-    # also includes a fake while loop implemented with a scan and two cond.
-    (Mc, M), _ = jax.lax.scan(
-        f=backward_pass,
-        init=backward_pass_carry,
-        xs=np.flip(np.arange(start=1, stop=model.NB)),
-    )
-
-    # Store the locked 6D rigid-body inertia matrix Mbb ∈ ℝ⁶ˣ⁶
-    M = M.at[0:6, 0:6].set(Mc[0])
-
-    return M

jaxsim/physics/algos/forward_kinematics.py

@@ -1,79 +0,0 @@
-from typing import Tuple
-
-import jax
-import jax.numpy as jnp
-import numpy as np
-
-import jaxsim.typing as jtp
-from jaxsim.math.adjoint import Adjoint
-from jaxsim.physics.model.physics_model import PhysicsModel
-
-from . import utils
-
-
-def forward_kinematics_model(
-    model: PhysicsModel, q: jtp.Vector, xfb: jtp.Vector
-) -> jtp.Array:
-    """
-    Compute the forward kinematics transformations for all links in an articulated body or robot.
-
-    Args:
-        model (PhysicsModel): The physics model of the articulated body or robot.
-        q (jtp.Vector): Joint positions (Generalized coordinates).
-        xfb (jtp.Vector): The base pose vector, including the quaternion (first 4 elements) and translation (last 3 elements).
-
-    Returns:
-        jtp.Array: A 3D array containing the forward kinematics transformations for all links.
-    """
-
-    x_fb, q, _, _, _, _ = utils.process_inputs(
-        physics_model=model, xfb=xfb, q=q, qd=None, tau=None, f_ext=None
-    )
-
-    W_X_0 = Adjoint.from_quaternion_and_translation(
-        quaternion=x_fb[0:4], translation=x_fb[4:7]
-    )
-
-    # This is the 6D velocity transform from i-th link frame to the world frame
-    W_X_i = jnp.zeros(shape=[model.NB, 6, 6])
-    W_X_i = W_X_i.at[0].set(W_X_0)
-
-    i_X_pre = model.joint_transforms(q=q)
-    pre_X_λi = model.tree_transforms
-
-    # This is the parent-to-child 6D velocity transforms of all links
-    i_X_λi = jnp.zeros_like(i_X_pre)
-
-    # Parent array mapping: i -> λ(i).
-    # Exception: λ(0) must not be used, it's initialized to -1.
-    λ = model.parent
-
-    PropagateKinematicsCarry = Tuple[jtp.MatrixJax, jtp.MatrixJax]
-    propagate_kinematics_carry = (i_X_λi, W_X_i)
-
-    def propagate_kinematics(
-        carry: PropagateKinematicsCarry, i: jtp.Int
-    ) -> Tuple[PropagateKinematicsCarry, None]:
-        i_X_λi, W_X_i = carry
-
-        i_X_λi_i = i_X_pre[i] @ pre_X_λi[i]
-        i_X_λi = i_X_λi.at[i].set(i_X_λi_i)
-
-        W_X_i_i = W_X_i[λ[i]] @ Adjoint.inverse(i_X_λi[i])
-        W_X_i = W_X_i.at[i].set(W_X_i_i)
-
-        return (i_X_λi, W_X_i), None
-
-    (_, W_X_i), _ = jax.lax.scan(
-        f=propagate_kinematics,
-        init=propagate_kinematics_carry,
-        xs=np.arange(start=1, stop=model.NB),
-    )
-
-    return jnp.stack([Adjoint.to_transform(adjoint=X) for X in list(W_X_i)])
-
-
-def forward_kinematics(
-    model: PhysicsModel, body_index: jtp.Int, q: jtp.Vector, xfb: jtp.Vector
-) -> jtp.Matrix:
-    return forward_kinematics_model(model=model, q=q, xfb=xfb)[body_index]

jaxsim/physics/algos/jacobian.py

@@ -1,98 +0,0 @@
-from typing import Tuple
-
-import jax
-import jax.numpy as jnp
-import numpy as np
-
-import jaxsim.typing as jtp
-from jaxsim.math.adjoint import Adjoint
-from jaxsim.physics.model.physics_model import PhysicsModel
-
-from . import utils
-
-
-def jacobian(model: PhysicsModel, body_index: jtp.Int, q: jtp.Vector) -> jtp.Matrix:
-    """
-    Compute the Jacobian matrix for a specific link in an articulated body or robot.
-
-    Args:
-        model (PhysicsModel): The physics model of the articulated body or robot.
-        body_index (jtp.Int): The index of the link for which to compute the Jacobian matrix.
-        q (jtp.Vector): Joint positions (Generalized coordinates).
-
-    Returns:
-        jtp.Matrix: The Jacobian matrix for the specified link.
-    """
-    _, q, _, _, _, _ = utils.process_inputs(physics_model=model, q=q)
-
-    S = model.motion_subspaces(q=q)
-    i_X_pre = model.joint_transforms(q=q)
-    pre_X_λi = model.tree_transforms
-    i_X_λi = jnp.zeros_like(i_X_pre)
-
-    i_X_0 = jnp.zeros_like(i_X_pre)
-    i_X_0 = i_X_0.at[0].set(jnp.eye(6))
-
-    # Parent array mapping: i -> λ(i).
-    # Exception: λ(0) must not be used, it's initialized to -1.
-    λ = model.parent
-
-    # ====================
-    # Propagate kinematics
-    # ====================
-
-    PropagateKinematicsCarry = Tuple[jtp.MatrixJax, jtp.MatrixJax]
-    propagate_kinematics_carry = (i_X_λi, i_X_0)
-
-    def propagate_kinematics(
-        carry: PropagateKinematicsCarry, i: jtp.Int
-    ) -> Tuple[PropagateKinematicsCarry, None]:
-        i_X_λi, i_X_0 = carry
-
-        # For each body (i), compute the parent (λi) to body (i) adjoint matrix
-        i_X_λi_i = i_X_pre[i] @ pre_X_λi[i]
-        i_X_λi = i_X_λi.at[i].set(i_X_λi_i)
-
-        # Compute the base (0) to body (i) adjoint matrix.
-        # This works fine since we traverse the kinematic tree following the link
-        # indices assigned with BFS.
-        i_X_0_i = i_X_λi[i] @ i_X_0[λ[i]]
-        i_X_0 = i_X_0.at[i].set(i_X_0_i)
-
-        return (i_X_λi, i_X_0), None
-
-    (i_X_λi, i_X_0), _ = jax.lax.scan(
-        f=propagate_kinematics,
-        init=propagate_kinematics_carry,
-        xs=np.arange(start=1, stop=model.NB),
-    )
-
-    # ============================
-    # Compute doubly-left Jacobian
-    # ============================
-
-    J = jnp.zeros(shape=(6, 6 + model.dofs()))
-
-    Jb = i_X_0[body_index]
-    J = J.at[0:6, 0:6].set(Jb)
-
-    # To make JIT happy, we operate on a boolean version of κ(i).
-    # Checking if j ∈ κ(i) is equivalent to: κ_bool(j) is True.
-    κ_bool = model.support_body_array_bool(body_index=body_index)
-
-    def compute_jacobian(J: jtp.MatrixJax, i: jtp.Int) -> Tuple[jtp.MatrixJax, None]:
-        def update_jacobian(J: jtp.MatrixJax, i: jtp.Int) -> jtp.MatrixJax:
-            ii = i - 1
-            Js_i = i_X_0[body_index] @ Adjoint.inverse(i_X_0[i]) @ S[i]
-            J = J.at[0:6, 6 + ii].set(Js_i.squeeze())
-
-            return J
-
-        J = jax.lax.select(pred=κ_bool[i], on_true=update_jacobian(J, i), on_false=J)
-        return J, None
-
-    J, _ = jax.lax.scan(
-        f=compute_jacobian, init=J, xs=np.arange(start=1, stop=model.NB)
-    )
-
-    return J