import and typing improvements

Author: victorpoughon

src/torchlensmaker/__init__.py

@@ -14,51 +14,130 @@
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
+# ruff: noqa: F401
+
+######
 # Core
-from torchlensmaker.core.physics import *
-from torchlensmaker.core.outline import *
-from torchlensmaker.surfaces.sphere_r import *
-from torchlensmaker.core.sag_functions import *
-from torchlensmaker.core.transforms import *
+######
+
+from torchlensmaker.core.physics import reflection, refraction
+from torchlensmaker.core.transforms import (
+    TransformBase,
+    LinearTransform,
+    ComposeTransform,
+    TranslateTransform,
+    IdentityTransform,
+)
 from torchlensmaker.core.intersect import *
 from torchlensmaker.core.full_forward import *
 from torchlensmaker.core.collision_detection import *
 from torchlensmaker.core.parameter import *
 from torchlensmaker.core.geometry import *
+from torchlensmaker.core.sag_functions import (
+    Spherical,
+    Parabolic,
+    Aspheric,
+    XYPolynomial,
+    Conical,
+    SagSum,
+    SagFunction,
+)
+# from torchlensmaker.core.outline import *
 
+##########
 # Surfaces
-from torchlensmaker.surfaces.conics import *
-from torchlensmaker.surfaces.sphere_r import *
-from torchlensmaker.surfaces.implicit_surface import *
-from torchlensmaker.surfaces.local_surface import *
-from torchlensmaker.surfaces.plane import *
-from torchlensmaker.surfaces.implicit_cylinder import *
-
-# Optics
+##########
+
+from torchlensmaker.surfaces.sphere_r import SphereR
+from torchlensmaker.surfaces.conics import Sphere, Parabola, Conic, Asphere
+from torchlensmaker.surfaces.implicit_surface import ImplicitSurface
+from torchlensmaker.surfaces.local_surface import LocalSurface
+from torchlensmaker.surfaces.plane import Plane, CircularPlane, SquarePlane
+from torchlensmaker.surfaces.implicit_cylinder import ImplicitCylinder
+from torchlensmaker.surfaces.sag_surface import SagSurface
+
+##################
+# Optical elements
+##################
+
 from torchlensmaker.elements.sequential import Sequential
-from torchlensmaker.elements.utils import *
-from torchlensmaker.elements.kinematics import *
-from torchlensmaker.elements.optical_surfaces import *
-from torchlensmaker.elements.light_sources import *
-
-from torchlensmaker.optical_data import *
-from torchlensmaker.lenses import *
-from torchlensmaker.materials import *
+from torchlensmaker.elements.utils import MixedDim
+from torchlensmaker.elements.kinematics import (
+    SubChain,
+    AbsoluteTransform,
+    RelativeTransform,
+    Gap,
+    Rotate3D,
+    Rotate2D,
+    Translate2D,
+    Translate3D,
+    Rotate,
+    Translate,
+)
+from torchlensmaker.elements.optical_surfaces import (
+    CollisionSurface,
+    ReflectiveSurface,
+    RefractiveSurface,
+    Aperture,
+    FocalPoint,
+    ImagePlane,
+    linear_magnification,
+)
+from torchlensmaker.elements.light_sources import (
+    LightSourceBase,
+    RaySource,
+    PointSourceAtInfinity,
+    PointSource,
+    ObjectAtInfinity,
+    Object,
+    Wavelength,
+)
+
+# Top level stuff - to be reorganized
+from torchlensmaker.optical_data import OpticalData, default_input
+from torchlensmaker.lenses import LensBase, BiLens, Lens, PlanoLens
+from torchlensmaker.materials import (
+    MaterialModel,
+    NonDispersiveMaterial,
+    CauchyMaterial,
+    SellmeierMaterial,
+)
+
+##########
+# Sampling
+##########
+
 from torchlensmaker.sampling import *
 
+##############
 # Optimization
-from torchlensmaker.optimize import *
+##############
+
+import torch.optim as optim
+from torchlensmaker.optimize import (
+    optimize,
+    OptimizationRecord,
+    plot_optimization_record,
+)
 
+########
 # Viewer
+########
+
 import torchlensmaker.viewer.tlmviewer as viewer
 from torchlensmaker.viewer.render_sequence import *
 
-
+##########
 # Analysis
+##########
+
 from torchlensmaker.analysis.plot_magnification import plot_magnification
 from torchlensmaker.analysis.plot_material_model import plot_material_models
 from torchlensmaker.analysis.spot_diagram import spot_diagram
 
+##################
 # Export build123d
+##################
+
 import torchlensmaker.export_build123d as export
 from torchlensmaker.export_build123d import show_part

src/torchlensmaker/core/collision_detection.py

@@ -17,12 +17,10 @@
 from __future__ import annotations
 
 import torch
-import math
-from functools import partial
 
 from dataclasses import dataclass
 
-from typing import TYPE_CHECKING, Optional, Callable, Any
+from typing import TYPE_CHECKING, Optional, Any
 
 if TYPE_CHECKING:
     from torchlensmaker.surfaces.implicit_cylinder import ImplicitSurface

src/torchlensmaker/core/tensorframe.py

@@ -50,7 +50,7 @@ def get(self, names: str | list[str]) -> torch.Tensor:
             else:
                 idx = [self.columns.index(n) for n in names]
             return self.data[:, idx]
-        except:
+        except ValueError:
             raise KeyError(f"TensorFrame doesn't have column(s): {names}")
 
     def masked(self, mask: torch.Tensor) -> TensorFrame:

src/torchlensmaker/core/transforms.py

@@ -18,8 +18,6 @@
 import functools
 from collections.abc import Sequence
 
-from torchlensmaker.core.rot2d import rotation_matrix_2D
-from torchlensmaker.core.rot3d import euler_angles_to_matrix
 
 # for shorter type annotations
 Tensor = torch.Tensor

src/torchlensmaker/elements/utils.py

@@ -17,8 +17,6 @@
 import torch.nn as nn
 from torchlensmaker.optical_data import OpticalData
 
-from typing import Any
-
 
 class Marker(nn.Module):
     "WIP"

src/torchlensmaker/export_build123d.py

@@ -21,7 +21,6 @@
 from torchlensmaker.elements.sequential import Sequential
 
 import os
-import torch.nn as nn
 import build123d as bd
 from os.path import join
 

src/torchlensmaker/sampling/samplers.py

@@ -18,11 +18,10 @@
 import torch.nn as nn
 import numpy as np
 
-from typing import Any, Sequence
-
-Tensor = torch.Tensor
-
 from torchlensmaker.core.tensor_manip import to_tensor
+from typing import Any, Sequence, TypeAlias
+
+Tensor: TypeAlias = torch.Tensor
 
 
 class Sampler:

src/torchlensmaker/testing/check_local_collide.py

@@ -24,8 +24,6 @@ def check_local_collide(
 ) -> None:
     "Call surface.local_collide() and performs tests on the output"
 
-    dim = P.shape[-1]
-
     # Check that rays are the correct dtype
     assert P.dtype == surface.dtype
     assert V.dtype == surface.dtype
@@ -63,7 +61,6 @@ def check_local_collide(
     )
 
     if isinstance(surface, tlm.ImplicitSurface):
-        N = sum(P.shape[:-1])
         rmse = surface.rmse(local_points)
     else:
         rmse = None

tests/test_tensorframe.py

@@ -1,3 +1,4 @@
+import pytest
 import torch
 from torchlensmaker.core.tensorframe import TensorFrame
 
@@ -89,3 +90,19 @@ def test_stack() -> None:
 
     assert torch.all(tf3.data == torch.tensor([[1, 2], [1, 2], [3, 4], [3, 4]]))
     assert tf3.columns == tf1.columns == tf2.columns
+
+
+def test_missing_column() -> None:
+    N = 2
+    tf1 = TensorFrame(
+        torch.column_stack(
+            (
+                torch.full((N,), 1),
+                torch.full((N,), 2),
+            )
+        ),
+        ["a", "b"],
+    )
+
+    with pytest.raises(KeyError):
+        tf1.get("c")