jaxsim 0.5.1.dev37__tar.gz → 0.5.1.dev46__tar.gz

{jaxsim-0.5.1.dev37 → jaxsim-0.5.1.dev46}/.github/workflows/pixi.yml

@@ -35,7 +35,7 @@ jobs:
           pixi update --json | pixi exec pixi-diff-to-markdown --explicit-column > diff.md
 
       - name: Test project against updated pixi
-        run: pixi run --environment tasks-cpu tests
+        run: pixi run --environment tasks-cpu test
         env:
           PY_COLORS: "1"
           JAX_PLATFORM_NAME: cpu

{jaxsim-0.5.1.dev37 → jaxsim-0.5.1.dev46}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.5.1.dev37
+Version: 0.5.1.dev46
 Summary: A differentiable physics engine and multibody dynamics library for control and robot learning.
 Author-email: Diego Ferigo <dgferigo@gmail.com>, Filippo Luca Ferretti <filippoluca.ferretti@outlook.com>
 Maintainer-email: Filippo Luca Ferretti <filippo.ferretti@iit.it>, Alessandro Croci <alessandro.croci@iit.it>
@@ -78,6 +78,7 @@ Requires-Dist: pre-commit; extra == "style"
 Provides-Extra: testing
 Requires-Dist: idyntree>=12.2.1; extra == "testing"
 Requires-Dist: pytest>=6.0; extra == "testing"
+Requires-Dist: pytest-benchmark; extra == "testing"
 Requires-Dist: pytest-icdiff; extra == "testing"
 Requires-Dist: robot-descriptions; extra == "testing"
 Provides-Extra: viz

{jaxsim-0.5.1.dev37 → jaxsim-0.5.1.dev46}/pyproject.toml

@@ -67,6 +67,7 @@ style = [
 testing = [
     "idyntree >= 12.2.1",
     "pytest >=6.0",
+    "pytest-benchmark",
     "pytest-icdiff",
     "robot-descriptions"
 ]
@@ -121,7 +122,7 @@ multi_line_output = 3
 profile = "black"
 
 [tool.pytest.ini_options]
-addopts = "-rsxX -v --strict-markers"
+addopts = "-rsxX -v --strict-markers --benchmark-skip"
 minversion = "6.0"
 testpaths = [
     "tests",
@@ -233,7 +234,8 @@ jaxsim = { path = "./", editable = true }
 
 [tool.pixi.feature.test.tasks]
 pipcheck = "pip check"
-tests = { cmd = "pytest", depends_on = ["pipcheck"] }
+test = { cmd = "pytest", depends_on = ["pipcheck"] }
+benchmark = { cmd = "pytest --benchmark-only", depends_on = ["pipcheck"] }
 
 [tool.pixi.feature.test.dependencies]
 black = "24.*"
@@ -241,6 +243,7 @@ idyntree = "*"
 isort = "*"
 pre-commit = "*"
 pytest = "*"
+pytest-benchmark = "*"
 pytest-icdiff = "*"
 robot_descriptions = "*"
 

{jaxsim-0.5.1.dev37 → jaxsim-0.5.1.dev46}/src/jaxsim/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.5.1.dev37'
-__version_tuple__ = version_tuple = (0, 5, 1, 'dev37')
+__version__ = version = '0.5.1.dev46'
+__version_tuple__ = version_tuple = (0, 5, 1, 'dev46')

{jaxsim-0.5.1.dev37 → jaxsim-0.5.1.dev46}/src/jaxsim.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.5.1.dev37
+Version: 0.5.1.dev46
 Summary: A differentiable physics engine and multibody dynamics library for control and robot learning.
 Author-email: Diego Ferigo <dgferigo@gmail.com>, Filippo Luca Ferretti <filippoluca.ferretti@outlook.com>
 Maintainer-email: Filippo Luca Ferretti <filippo.ferretti@iit.it>, Alessandro Croci <alessandro.croci@iit.it>
@@ -78,6 +78,7 @@ Requires-Dist: pre-commit; extra == "style"
 Provides-Extra: testing
 Requires-Dist: idyntree>=12.2.1; extra == "testing"
 Requires-Dist: pytest>=6.0; extra == "testing"
+Requires-Dist: pytest-benchmark; extra == "testing"
 Requires-Dist: pytest-icdiff; extra == "testing"
 Requires-Dist: robot-descriptions; extra == "testing"
 Provides-Extra: viz

{jaxsim-0.5.1.dev37 → jaxsim-0.5.1.dev46}/src/jaxsim.egg-info/SOURCES.txt

@@ -121,6 +121,7 @@ tests/test_api_joint.py
 tests/test_api_link.py
 tests/test_api_model.py
 tests/test_automatic_differentiation.py
+tests/test_benchmark.py
 tests/test_contact.py
 tests/test_exceptions.py
 tests/test_meshes.py

{jaxsim-0.5.1.dev37 → jaxsim-0.5.1.dev46}/src/jaxsim.egg-info/requires.txt

@@ -23,6 +23,7 @@ pre-commit
 [testing]
 idyntree>=12.2.1
 pytest>=6.0
+pytest-benchmark
 pytest-icdiff
 robot-descriptions
 

{jaxsim-0.5.1.dev37 → jaxsim-0.5.1.dev46}/tests/conftest.py

@@ -20,6 +20,21 @@ def pytest_addoption(parser):
         help="Run tests only if GPU is available and utilized",
     )
 
+    parser.addoption(
+        "--batch-size",
+        action="store",
+        default="None",
+        help="Batch size for vectorized benchmarks (only applies to benchmark tests)",
+    )
+
+
+def pytest_generate_tests(metafunc):
+    if (
+        "batch_size" in metafunc.fixturenames
+        and (batch_size := metafunc.config.getoption("--batch-size")) != "None"
+    ):
+        metafunc.parametrize("batch_size", [1, int(batch_size)])
+
 
 def check_gpu_usage():
     # Set environment variable to prioritize GPU.
@@ -109,6 +124,18 @@ def velocity_representation(request) -> jaxsim.VelRepr:
     return request.param
 
 
+@pytest.fixture(scope="session")
+def batch_size(request) -> int:
+    """
+    Fixture providing the batch size for vectorized benchmarks.
+
+    Returns:
+        The batch size for vectorized benchmarks.
+    """
+
+    return 1
+
+
 # ================================
 # Fixtures providing JaxSim models
 # ================================

jaxsim-0.5.1.dev46/tests/test_benchmark.py

@@ -0,0 +1,127 @@
+from collections.abc import Callable
+
+import jax
+import pytest
+
+import jaxsim
+import jaxsim.api as js
+
+
+def vectorize_data(model: js.model.JaxSimModel, batch_size: int):
+    key = jax.random.PRNGKey(seed=0)
+
+    return jax.vmap(
+        lambda key: js.data.random_model_data(
+            model=model,
+            key=key,
+        )
+    )(jax.numpy.repeat(key[None, :], repeats=batch_size, axis=0))
+
+
+def benchmark_test_function(
+    func: Callable, model: js.model.JaxSimModel, benchmark, batch_size
+):
+    """Reusability wrapper for benchmark tests."""
+    data = vectorize_data(model=model, batch_size=batch_size)
+
+    # Warm-up call to avoid including compilation time
+    jax.vmap(func, in_axes=(None, 0))(model, data)
+    benchmark(jax.vmap(func, in_axes=(None, 0)), model, data)
+
+
+@pytest.mark.benchmark
+def test_forward_dynamics_aba(
+    jaxsim_model_ergocub_reduced: js.model.JaxSimModel, benchmark, batch_size
+):
+    model = jaxsim_model_ergocub_reduced
+
+    benchmark_test_function(js.model.forward_dynamics_aba, model, benchmark, batch_size)
+
+
+@pytest.mark.benchmark
+def test_free_floating_bias_forces(
+    jaxsim_model_ergocub_reduced: js.model.JaxSimModel, benchmark, batch_size
+):
+    model = jaxsim_model_ergocub_reduced
+
+    benchmark_test_function(
+        js.model.free_floating_bias_forces, model, benchmark, batch_size
+    )
+
+
+@pytest.mark.benchmark
+def test_forward_kinematics(
+    jaxsim_model_ergocub_reduced: js.model.JaxSimModel, benchmark, batch_size
+):
+    model = jaxsim_model_ergocub_reduced
+
+    benchmark_test_function(js.model.forward_kinematics, model, benchmark, batch_size)
+
+
+@pytest.mark.benchmark
+def test_free_floating_mass_matrix(
+    jaxsim_model_ergocub_reduced: js.model.JaxSimModel, benchmark, batch_size
+):
+    model = jaxsim_model_ergocub_reduced
+
+    benchmark_test_function(
+        js.model.free_floating_mass_matrix, model, benchmark, batch_size
+    )
+
+
+@pytest.mark.benchmark
+def test_free_floating_jacobian(
+    jaxsim_model_ergocub_reduced: js.model.JaxSimModel, benchmark, batch_size
+):
+    model = jaxsim_model_ergocub_reduced
+
+    benchmark_test_function(
+        js.model.generalized_free_floating_jacobian, model, benchmark, batch_size
+    )
+
+
+@pytest.mark.benchmark
+def test_free_floating_jacobian_derivative(
+    jaxsim_model_ergocub_reduced: js.model.JaxSimModel, benchmark, batch_size
+):
+    model = jaxsim_model_ergocub_reduced
+
+    benchmark_test_function(
+        js.model.generalized_free_floating_jacobian_derivative,
+        model,
+        benchmark,
+        batch_size,
+    )
+
+
+@pytest.mark.benchmark
+def test_soft_contact_model(
+    jaxsim_model_ergocub_reduced: js.model.JaxSimModel, benchmark, batch_size
+):
+    model = jaxsim_model_ergocub_reduced
+
+    benchmark_test_function(js.ode.system_dynamics, model, benchmark, batch_size)
+
+
+@pytest.mark.benchmark
+def test_rigid_contact_model(
+    jaxsim_model_ergocub_reduced: js.model.JaxSimModel, benchmark, batch_size
+):
+    model = jaxsim_model_ergocub_reduced
+
+    with model.editable(validate=False) as model:
+        model.contact_model = jaxsim.rbda.contacts.RigidContacts()
+
+    benchmark_test_function(js.ode.system_dynamics, model, benchmark, batch_size)
+
+
+@pytest.mark.benchmark
+def test_relaxed_rigid_contact_model(
+    jaxsim_model_ergocub_reduced: js.model.JaxSimModel, benchmark, batch_size
+):
+    model = jaxsim_model_ergocub_reduced
+
+    with model.editable(validate=False) as model:
+        model.contact_model = jaxsim.rbda.contacts.RelaxedRigidContacts()
+
+    benchmark_test_function(js.ode.system_dynamics, model, benchmark, batch_size)