Use bitwise operators instead of jnp.logical_*

ami-iit · Oct 23, 2024 · c8d6b53 · c8d6b53
1 parent 4742ae4
commit c8d6b53
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 17 deletions.
diff --git a/src/jaxsim/api/model.py b/src/jaxsim/api/model.py
@@ -2027,15 +2027,15 @@ def step(
             # Raise runtime error for not supported case in which Rigid contacts and
             # Baumgarte stabilization are enabled and used with ForwardEuler integrator.
             jaxsim.exceptions.raise_runtime_error_if(
-                condition=jnp.logical_and(
-                    isinstance(
-                        integrator,
-                        jaxsim.integrators.fixed_step.ForwardEuler
-                        | jaxsim.integrators.fixed_step.ForwardEulerSO3,
-                    ),
-                    jnp.array(
-                        [data_tf.contacts_params.K, data_tf.contacts_params.D]
-                    ).any(),
+                condition=isinstance(
+                    integrator,
+                    jaxsim.integrators.fixed_step.ForwardEuler
+                    | jaxsim.integrators.fixed_step.ForwardEulerSO3,
+                )
+                & (
+                    data_tf.contacts_params.K
+                    is not None | data_tf.contacts_params.D
+                    is not None
                 ),
                 msg="Baumgarte stabilization is not supported with ForwardEuler integrators",
             )

diff --git a/src/jaxsim/integrators/variable_step.py b/src/jaxsim/integrators/variable_step.py
@@ -446,14 +446,10 @@ def reject_step():
                 params_next,
                 discarded_steps,
             ) = jax.lax.cond(
-                pred=jnp.array(
-                    [
-                        discarded_steps >= self.max_step_rejections,
-                        local_error <= 1.0,
-                        Δt_next < self.dt_min,
-                        integrator_init,
-                    ]
-                ).any(),
+                pred=discarded_steps
+                >= self.max_step_rejections | local_error
+                <= 1.0 | Δt_next
+                < self.dt_min | integrator_init,
                 true_fun=accept_step,
                 false_fun=reject_step,
             )