add types standalone module

Author: victorpoughon

src/torchlensmaker/__init__.py

@@ -18,6 +18,21 @@
 
 import torch
 
+#######
+# Types
+#######
+
+from torchlensmaker.types import (
+    ScalarTensor,
+    BatchTensor,
+    Batch2DTensor,
+    Batch3DTensor,
+    BatchNDTensor,
+    HomMatrix2D,
+    HomMatrix3D,
+    HomMatrix,
+)
+
 ######
 # Core
 ######
@@ -51,9 +66,6 @@
 ############
 
 from torchlensmaker.kinematics.homogeneous_geometry import (
-    HomMatrix2D,
-    HomMatrix3D,
-    HomMatrix,
     transform_points,
     transform_vectors,
     hom_identity,

src/torchlensmaker/kinematics/homogeneous_geometry.py

@@ -22,9 +22,7 @@
 
 from torchlensmaker.core.rot3d import euler_angles_to_matrix
 
-HomMatrix2D: TypeAlias = Float[torch.Tensor, "3 3"]
-HomMatrix3D: TypeAlias = Float[torch.Tensor, "4 4"]
-HomMatrix: TypeAlias = HomMatrix2D | HomMatrix3D
+from torchlensmaker.types import HomMatrix2D, HomMatrix3D, HomMatrix
 
 
 def hom_matrix_2d(M: Float[torch.Tensor, "2 2"]) -> HomMatrix2D:

src/torchlensmaker/kinematics/kinematics_elements.py

@@ -23,7 +23,7 @@
 from torchlensmaker.core.tensor_manip import to_tensor, init_param, expand_bool_tuple
 from torchlensmaker.optical_data import OpticalData
 
-from .homogeneous_geometry import (
+from torchlensmaker.types import (
     HomMatrix2D,
     HomMatrix3D,
     HomMatrix,

src/torchlensmaker/kinematics/kinematics_kernels.py

@@ -17,9 +17,12 @@
 import torch
 from jaxtyping import Float
 
-from .homogeneous_geometry import (
+from torchlensmaker.types import (
     HomMatrix2D,
     HomMatrix3D,
+)
+
+from .homogeneous_geometry import (
     hom_rotate_2d,
     hom_rotate_3d,
     hom_translate_2d,

src/torchlensmaker/kinematics/tests/test_kinematics_elements.py

@@ -21,12 +21,11 @@
 import torch
 import torch.nn as nn
 
+from torchlensmaker.types import HomMatrix2D, HomMatrix3D
 
 from torchlensmaker.kinematics.homogeneous_geometry import (
     hom_identity_2d,
     hom_identity_3d,
-    HomMatrix2D,
-    HomMatrix3D,
 )
 
 from torchlensmaker.kinematics.kinematics_elements import (

src/torchlensmaker/types.py

@@ -0,0 +1,30 @@
+# This file is part of Torch Lens Maker
+# Copyright (C) 2024-present Victor Poughon
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+
+from typing import TypeAlias
+from jaxtyping import Float
+import torch
+
+ScalarTensor: TypeAlias = Float[torch.Tensor, ""]
+BatchTensor: TypeAlias = Float[torch.Tensor, "..."]
+Batch2DTensor: TypeAlias = Float[torch.Tensor, "... 2"]
+Batch3DTensor: TypeAlias = Float[torch.Tensor, "... 3"]
+BatchNDTensor: TypeAlias = Float[torch.Tensor, "... D"]
+
+HomMatrix2D: TypeAlias = Float[torch.Tensor, "3 3"]
+HomMatrix3D: TypeAlias = Float[torch.Tensor, "4 4"]
+HomMatrix: TypeAlias = HomMatrix2D | HomMatrix3D