rename MixedDim to MixedDimKinematic

Author: victorpoughon

src/torchlensmaker/__init__.py

@@ -61,7 +61,7 @@
 
 from torchlensmaker.kinematics.kinematics_elements import (
     KinematicElement,
-    MixedDim,
+    MixedDimKinematic,
     AbsolutePosition,
     Gap,
     Rotate3D,

src/torchlensmaker/kinematics/kinematics_elements.py

@@ -176,8 +176,7 @@ def forward(
         return self.func.forward(dfk, ifk, self.y, self.z)
 
 
-class MixedDim(KinematicElement):
-    # TODO true mixed dim that works with any element?
+class MixedDimKinematic(KinematicElement):
     def __init__(self, module_2d: nn.Module, module_3d: nn.Module):
         super().__init__()
         self.module_2d = module_2d
@@ -195,7 +194,7 @@ def __init__(self, offset: Float[torch.Tensor, ""] | float | int):
         super().__init__()
         translate_2d = Translate2D(x=offset)
         translate_3d = Translate3D(x=offset)
-        self.mixed_dim = MixedDim(translate_2d, translate_3d)
+        self.mixed_dim = MixedDimKinematic(translate_2d, translate_3d)
 
     def forward(self, dfk: HomMatrix, ifk: HomMatrix) -> tuple[HomMatrix, HomMatrix]:
         return self.mixed_dim(dfk, ifk)
@@ -206,7 +205,7 @@ def __init__(
         self, angles: tuple[float | int, float | int] | Float[torch.Tensor, "2"]
     ):
         super().__init__()
-        self.mixed_dim = MixedDim(Rotate2D(angles[0]), Rotate3D(angles[1], angles[0]))
+        self.mixed_dim = MixedDimKinematic(Rotate2D(angles[0]), Rotate3D(angles[1], angles[0]))
 
     def forward(self, dfk: HomMatrix, ifk: HomMatrix) -> tuple[HomMatrix, HomMatrix]:
         return self.mixed_dim(dfk, ifk)
@@ -220,7 +219,7 @@ def __init__(
         z: Float[torch.Tensor, ""] | float | int = 0.0,
     ):
         super().__init__()
-        self.mixed_dim = MixedDim(Translate2D(x, y), Translate3D(x, y, z))
+        self.mixed_dim = MixedDimKinematic(Translate2D(x, y), Translate3D(x, y, z))
 
     def forward(self, dfk: HomMatrix, ifk: HomMatrix) -> tuple[HomMatrix, HomMatrix]:
         return self.mixed_dim(dfk, ifk)

src/torchlensmaker/kinematics/tests/test_kinematics_elements.py

@@ -38,7 +38,7 @@
     Rotate3D,
     AbsolutePosition,
     AbsolutePositionVec3D,
-    MixedDim,
+    MixedDimKinematic,
     Gap,
     Rotate,
     Translate,
@@ -243,7 +243,7 @@ def test_elements_mixed() -> None:
 
     chain_model_mixed = nn.ModuleList(
         [
-            MixedDim(
+            MixedDimKinematic(
                 Rotate2D(0.1),
                 Rotate3D(y=0.1),
             ),