Add static typing with Python type annotations

.github/workflows/python-tests.yml

@@ -1,35 +1,41 @@
 name: Python Tests
-
 on:
   pull_request:
     branches: [ master ]
 
 jobs:
-  test:
+  test-and-typecheck:
     runs-on: ubuntu-latest
     strategy:
       matrix:
         python-version: ['3.10', '3.11', '3.12']
 
     steps:
-    - uses: actions/checkout@v4
+      - uses: actions/checkout@v4
+
+      - name: Set up Miniconda
+        uses: conda-incubator/setup-miniconda@v3
+        with:
+          auto-update-conda: true
+          python-version: ${{ matrix.python-version }}
 
-    - name: Set up Miniconda
-      uses: conda-incubator/setup-miniconda@v3
-      with:
-        auto-update-conda: true
-        python-version: ${{ matrix.python-version }}
+      - name: Install dependencies
+        shell: bash -l {0}
+        run: |
+          conda create -n tracepyci python=${{ matrix.python-version }} --yes
+          conda activate tracepyci
+          conda install --yes numpy scipy matplotlib scikit-learn pandas pytest pytest-cov
+          pip install mypy types-PyYAML pandas-stubs
+          pip install .
 
-    - name: Install dependencies
-      shell: bash -l {0}
-      run: |
-        conda create -n tracepyci python=${{ matrix.python-version }} --yes
-        conda activate tracepyci
-        conda install --yes numpy scipy matplotlib scikit-learn pandas pytest pytest-cov
-        pip install .
+      - name: Run tests
+        shell: bash -l {0}
+        run: |
+          conda activate tracepyci
+          pytest tests/
 
-    - name: Run tests
-      shell: bash -l {0}
-      run: |
-        conda activate tracepyci
-        pytest tests/
+      - name: Run type checking
+        shell: bash -l {0}
+        run: |
+          conda activate tracepyci
+          mypy tracepy/ --ignore-missing-imports

tests/test_hyperbolic.py

@@ -24,6 +24,6 @@
 def test_rms_hyperbolic():
     geo = [back_lens, lens, stop]
     ray_group = tp.ray_plane(geo, [0., 0., 0.], 1.1, d=[0.,0.,1.], nrays=100)
-    rms = tp.spotdiagram(geo, ray_group, optimizer=True)
+    rms = tp.spot_rms(geo, ray_group)
     assert rms == 0.
 

tests/test_parabolic.py

@@ -17,5 +17,5 @@
 def test_rms_parabolic():
     geo = [mirror, stop]
     ray_group = tp.ray_plane(geo, [0., 0., -1.5], 1.1, d=[0.,0.,1.], nrays=100)
-    rms = tp.spotdiagram(geo, ray_group, optimizer=True)
+    rms = tp.spot_rms(geo, ray_group)
     assert rms == 0.

tracepy/__init__.py

@@ -14,8 +14,8 @@
 from .geometry import geometry
 from .optimize import optimize
 from .geoplot import plotxz, plotyz, plot2d
-from .optplot import spotdiagram, plotobject, rayaberration
-from .iotables import *
-from .transforms import *
-from .raygroup import *
-from.index import *
+from .optplot import spotdiagram, plotobject, rayaberration, spot_rms
+from .iotables import save_optics
+from .raygroup import ray_plane
+from .index import cauchy_two_term, glass_index
+from .utils import gen_rot

tracepy/constants.py

@@ -0,0 +1,12 @@
+# Constants used by optimizer
+SURV_CONST = 100 # Weight of failed propagation.
+MAX_RMS = 999 # Maximum RMS penalty for trace error.
+
+# Constants used for ray objects
+MAX_INTERSECTION_ITERATIONS = 1e4 # Max iter before failed intersection search.
+MAX_REFRACTION_ITERATIONS = 1e5 # Max iter before failed refraction.
+INTERSECTION_CONVERGENCE_TOLERANCE = 1e-6 # Tolerance for intersection search.
+REFRACTION_CONVERGENCE_TOLERANCE = 1e-15 # Tolerance for refraction.
+
+# Constants used for plotting
+PLOT_ROUNDING_ACC = 14 # Rounding accuracy for spot diagrams and plots.

tracepy/exceptions.py

@@ -6,3 +6,6 @@ class NotOnSurfaceError(Exception):
 
 class TraceError(Exception):
     """ Custom error for lens systems where no rays survive being traced. """
+
+class InvalidGeometry(Exception):
+    """ Invalid parameters were given to define a geometry object. """

tracepy/geometry.py

@@ -1,9 +1,11 @@
 import numpy as np
 
-from .transforms import gen_rot
-from .exceptions import NotOnSurfaceError
+from .utils import gen_rot
+from .exceptions import NotOnSurfaceError, InvalidGeometry
 from .index import glass_index
 
+from typing import Dict, List, Tuple, Union, Optional
+
 class geometry:
     """Class for the different surfaces in an optical system.
 
@@ -39,55 +41,48 @@ class geometry:
         If c is 0 then the surface is planar.
     name (optional): str
         Name of the surface, used for optimization
-    R (generated): np.matrix((3,3))
+    R (generated): np.array((3,3))
         Rotation matrix for the surface from rotation angles D.
 
     """
 
-    def __init__(self, params):
-        self.P = params['P']
-        self.D = np.array(params.get('D', [0., 0., 0.]))
-        self.action = params['action']
-        self.Diam = params['Diam']
-        self.N = params.get('N', 1.)
-        self.kappa = params.get('kappa', None)
-        self.diam = params.get('diam', 0.)
-        self.c = params.get('c', 0.)
-        self.name = params.get('name', None)
-        self.R = gen_rot(self.D)
+    def __init__(self, params: Dict):
+        P = params.get('P')
+        # Allow on axis integer for P.
+        if isinstance(P, float) or isinstance(P, int):
+            P = np.array([0., 0., P])
+        elif isinstance(P, List):
+            P = np.array(P)
+        else:
+            raise InvalidGeometry()
+        self.P: np.ndarray = P
+        self.D: np.ndarray = np.array(params.get('D', [0., 0., 0.]))
+        self.action: str = params['action']
+        self.Diam: Union[float, int] = params['Diam']
+        self.N: Union[float, int] = params.get('N', 1.)
+        self.kappa: Optional[Union[float, int]] = params.get('kappa')
+        self.diam: Union[float, int] = params.get('diam', 0.)
+        self.c: Union[float, int] = params.get('c', 0.)
+        self.name: str = params.get('name', None)
+        self.R: np.ndarray = gen_rot(self.D)
         if params.get('glass'):
             self.glass = glass_index(params.get('glass'))
         self.check_params()
 
-    def __getitem__(self, item):
-        """ Return attribute of geometry. """
-        return getattr(self, item)
-
-    def __setitem__(self, item, value):
-        """ Set attribute of geometry. """
-        return setattr(self, item, value)
-
-    def check_params(self):
+    def check_params(self) -> None:
         """Check that required parameters are given and update needed parameters.
 
         Summary
         -------
-            If P is given as a float/int then it is converted to a np array
-            with that float/int in the Z direction. If c != 0 (in the case
-            of a conic) then kappa must be specified, and if kappa is greater
-            than 0 then the value of c is redundant by boundary conditions of
-            the conic equation. Lastly, if c == 0 in the case of a planar
-            surface the None value of kappa needs to be set to a dummy value
-            to avoid exceptions in calculating the conic equation. Note that
-            this does not affect the calculation since c is 0.
+            If c != 0 (in the case of a conic) then kappa must be specified,
+            and if kappa is greater than 0 then the value of c is redundant
+            by boundary conditions of the conic equation. Lastly, if c == 0 in
+            the case of a planar surface the None value of kappa needs to be set
+            to a dummy value to avoid exceptions in calculating the conic equation.
+            Note that this does not affect the calculation since c is 0.
 
         """
 
-        if isinstance(self.P, float) or isinstance(self.P, int):
-            #Allow on axis integer for P.
-            self.P = np.array([0., 0., self.P])
-        else:
-            self.P = np.array(self.P)
         if self.c != 0:
             if self.kappa is None:
                 raise Exception("Specify a kappa for this conic.")
@@ -98,15 +93,15 @@ def check_params(self):
             #Used for planes, does not affect calculations.
             self.kappa = 1.
 
-    def get_surface(self, point):
+    def get_surface(self, point: np.ndarray) -> Tuple[float, List[float]]:
         """ Returns the function and derivitive of a surface for a point. """
         return self.conics(point)
 
-    def get_surface_plot(self, points):
+    def get_surface_plot(self, points: np.ndarray) -> np.ndarray:
         """ Returns the function value for an array of points. """
         return self.conics_plot(points)
 
-    def conics(self, point):
+    def conics(self, point: np.ndarray) -> Tuple[float, List[float]]:
         """Returns function value and derivitive list for conics and sphere surfaces.
 
         Note
@@ -137,14 +132,17 @@ def conics(self, point):
         rho = np.sqrt(pow(X,2) + pow(Y, 2))
         if rho > self.Diam/2. or rho < self.diam/2.:
             raise NotOnSurfaceError()
+        # Ensure kappa is not None before using it in calculations
+        if self.kappa is None:
+            raise ValueError("kappa must not be None for conic calculations")
         #Conic equation.
         function = Z - self.c*pow(rho, 2)/(1 + pow((1-self.kappa*pow(self.c, 2)*pow(rho,2)), 0.5))
         #See Spencer, Murty section on rotational surfaces for definition of E.
         E = self.c / pow((1-self.kappa*pow(self.c, 2)*pow(rho,2)), 0.5)
         derivitive = [-X*E, -Y*E, 1.]
         return function, derivitive
 
-    def conics_plot(self, point):
+    def conics_plot(self, point: np.ndarray) -> np.ndarray:
         """Returns Z values for an array of points for plotting conics.
 
         Parameters
@@ -166,5 +164,8 @@ def conics_plot(self, point):
         nan_idx = (rho > self.Diam/2.) + (rho < self.diam/2.)
         rho = np.sqrt(pow(X[~nan_idx],2) + pow(Y[~nan_idx], 2))
         function[nan_idx] = np.nan
+        # Ensure kappa is not None before using it in calculations
+        if self.kappa is None:
+            raise ValueError("kappa must not be None for conic plot calculations")
         function[~nan_idx] = self.c*pow(rho, 2)/(1 + pow((1-self.kappa*pow(self.c, 2)*pow(rho,2)), 0.5))
         return function