Merge pull request #18 from adenzler-nvidia/dev/adenzler/fix-tile-fn

Run formatter and fix leftover occurrence of tile_fn in io.py

Author: erikfrey

.github/workflows/ruff.yml

@@ -0,0 +1,44 @@
+# Copyright 2025 The Physics-Next Project Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+name: Ruff Format & Lint
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+
+jobs:
+  check:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.x"
+
+      - name: Install Ruff
+        run: pip install ruff
+
+      # Check formatting (fail if files aren't formatted)
+      - name: Check Formatting
+        run: ruff format --check .
+
+      # Check linting (fail on rule violations)
+      #- name: Check Linting
+      #  run: ruff check --output-format=github .

mujoco/mjx/_src/forward_test.py

@@ -1,3 +1,18 @@
+# Copyright 2025 The Physics-Next Project Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
 """Tests for forward dynamics functions."""
 
 from absl.testing import absltest
@@ -14,14 +29,13 @@
 
 def _assert_eq(a, b, name):
   tol = _TOLERANCE * 10  # avoid test noise
-  err_msg = f'mismatch: {name}'
+  err_msg = f"mismatch: {name}"
   np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol)
 
 
 class ForwardTest(absltest.TestCase):
-
   def _load(self, fname: str, is_sparse: bool = True):
-    path = epath.resource_path('mujoco.mjx') / 'test_data' / fname
+    path = epath.resource_path("mujoco.mjx") / "test_data" / fname
     mjm = mujoco.MjModel.from_xml_path(path.as_posix())
     mjm.opt.jacobian = is_sparse
     mjd = mujoco.MjData(mjm)
@@ -34,28 +48,30 @@ def _load(self, fname: str, is_sparse: bool = True):
 
   def test_fwd_velocity(self):
     """Tests MJX fwd_velocity."""
-    _, mjd, m, d = self._load('humanoid/humanoid.xml')
-    
+    _, mjd, m, d = self._load("humanoid/humanoid.xml")
+
     d.actuator_velocity.zero_()
     mjx.fwd_velocity(m, d)
 
-    _assert_eq(d.actuator_velocity.numpy()[0], mjd.actuator_velocity, 'actuator_velocity')
-    _assert_eq(d.qfrc_bias.numpy()[0], mjd.qfrc_bias, 'qfrc_bias')
+    _assert_eq(
+      d.actuator_velocity.numpy()[0], mjd.actuator_velocity, "actuator_velocity"
+    )
+    _assert_eq(d.qfrc_bias.numpy()[0], mjd.qfrc_bias, "qfrc_bias")
 
   def test_fwd_acceleration(self):
     """Tests MJX fwd_acceleration."""
-    _, mjd, m, d = self._load('humanoid/humanoid.xml', is_sparse=False)
+    _, mjd, m, d = self._load("humanoid/humanoid.xml", is_sparse=False)
 
     for arr in (d.qfrc_smooth, d.qacc_smooth):
       arr.zero_()
 
-    mjx.factor_m(m, d) # for dense, get tile cholesky factorization
+    mjx.factor_m(m, d)  # for dense, get tile cholesky factorization
     mjx.fwd_acceleration(m, d)
 
-    _assert_eq(d.qfrc_smooth.numpy()[0], mjd.qfrc_smooth, 'qfrc_smooth')
-    _assert_eq(d.qacc_smooth.numpy()[0], mjd.qacc_smooth, 'qacc_smooth')
+    _assert_eq(d.qfrc_smooth.numpy()[0], mjd.qfrc_smooth, "qfrc_smooth")
+    _assert_eq(d.qacc_smooth.numpy()[0], mjd.qacc_smooth, "qacc_smooth")
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
   wp.init()
   absltest.main()

mujoco/mjx/_src/io.py

@@ -1,3 +1,18 @@
+# Copyright 2025 The Physics-Next Project Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
 import warp as wp
 import mujoco
 import numpy as np
@@ -19,7 +34,7 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
   m.nM = mjm.nM
   m.opt.gravity = wp.vec3(mjm.opt.gravity)
   m.opt.is_sparse = support.is_sparse(mjm)
- 
+
   m.qpos0 = wp.array(mjm.qpos0, dtype=wp.float32, ndim=1)
   m.qpos_spring = wp.array(mjm.qpos_spring, dtype=wp.float32, ndim=1)
 
@@ -75,10 +90,12 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
     qLD_tilesize = np.array(sorted(tiles.keys()))
 
   m.qLD_update_tree = wp.array(qLD_update_tree, dtype=wp.vec3i, ndim=1)
-  m.qLD_update_treeadr = wp.array(qLD_update_treeadr, dtype=wp.int32, ndim=1, device="cpu")
+  m.qLD_update_treeadr = wp.array(
+    qLD_update_treeadr, dtype=wp.int32, ndim=1, device="cpu"
+  )
   m.qLD_tile = wp.array(qLD_tile, dtype=wp.int32, ndim=1)
-  m.qLD_tileadr = wp.array(qLD_tileadr, dtype=wp.int32, ndim=1, device='cpu')
-  m.qLD_tilesize = wp.array(qLD_tilesize, dtype=wp.int32, ndim=1, device='cpu')
+  m.qLD_tileadr = wp.array(qLD_tileadr, dtype=wp.int32, ndim=1, device="cpu")
+  m.qLD_tilesize = wp.array(qLD_tilesize, dtype=wp.int32, ndim=1, device="cpu")
   m.body_dofadr = wp.array(mjm.body_dofadr, dtype=wp.int32, ndim=1)
   m.body_dofnum = wp.array(mjm.body_dofnum, dtype=wp.int32, ndim=1)
   m.body_jntadr = wp.array(mjm.body_jntadr, dtype=wp.int32, ndim=1)
@@ -179,7 +196,13 @@ def tile(x):
 
   # TODO(taylorhowell): sparse actuator_moment
   actuator_moment = np.zeros((mjm.nu, mjm.nv))
-  mujoco.mju_sparse2dense(actuator_moment, mjd.actuator_moment, mjd.moment_rownnz, mjd.moment_rowadr, mjd.moment_colind)
+  mujoco.mju_sparse2dense(
+    actuator_moment,
+    mjd.actuator_moment,
+    mjd.moment_rownnz,
+    mjd.moment_rowadr,
+    mjd.moment_colind,
+  )
 
   d.qpos = wp.array(tile(mjd.qpos), dtype=wp.float32, ndim=2)
   d.qvel = wp.array(tile(mjd.qvel), dtype=wp.float32, ndim=2)

mujoco/mjx/_src/math.py

@@ -1,7 +1,23 @@
+# Copyright 2025 The Physics-Next Project Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
 import warp as wp
 
 from . import types
 
+
 @wp.func
 def mul_quat(u: wp.quat, v: wp.quat) -> wp.quat:
   return wp.quat(
@@ -91,7 +107,7 @@ def motion_cross_force(v: wp.spatial_vector, f: wp.spatial_vector) -> wp.spatial
 
 @wp.func
 def quat_to_vel(quat: wp.quat) -> wp.vec3:
-  axis  = wp.vec3(quat[1], quat[2], quat[3])
+  axis = wp.vec3(quat[1], quat[2], quat[3])
   sin_a_2 = wp.norm_l2(axis)
 
   if sin_a_2 == 0.0:

mujoco/mjx/_src/passive_test.py

@@ -1,3 +1,18 @@
+# Copyright 2025 The Physics-Next Project Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
 """Tests for passive force functions."""
 
 from absl.testing import absltest
@@ -11,28 +26,28 @@
 # due to float precision
 _TOLERANCE = 5e-5
 
+
 def _assert_eq(a, b, name):
   tol = _TOLERANCE * 10  # avoid test noise
-  err_msg = f'mismatch: {name}'
+  err_msg = f"mismatch: {name}"
   np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol)
 
 
 class PassiveTest(absltest.TestCase):
-
   def test_passive(self):
     """Tests MJX passive."""
-    _, mjd, m, d = test_util.fixture('pendula.xml')
+    _, mjd, m, d = test_util.fixture("pendula.xml")
 
     for arr in (d.qfrc_spring, d.qfrc_damper, d.qfrc_passive):
       arr.zero_()
 
     mjx.passive(m, d)
 
-    _assert_eq(d.qfrc_spring.numpy()[0], mjd.qfrc_spring, 'qfrc_spring')
-    _assert_eq(d.qfrc_damper.numpy()[0], mjd.qfrc_damper, 'qfrc_damper')
-    _assert_eq(d.qfrc_passive.numpy()[0], mjd.qfrc_passive, 'qfrc_passive')
+    _assert_eq(d.qfrc_spring.numpy()[0], mjd.qfrc_spring, "qfrc_spring")
+    _assert_eq(d.qfrc_damper.numpy()[0], mjd.qfrc_damper, "qfrc_damper")
+    _assert_eq(d.qfrc_passive.numpy()[0], mjd.qfrc_passive, "qfrc_passive")
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
   wp.init()
   absltest.main()