working static types for geometry.py, move optplot constant to constant.py

Author: GNiendorf

tracepy/constants.py

@@ -7,3 +7,6 @@
 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,10 +1,10 @@
 import numpy as np
 
 from .utils import gen_rot
-from .exceptions import NotOnSurfaceError
+from .exceptions import NotOnSurfaceError, InvalidGeometry
 from .index import glass_index
 
-from typing import Dict, List, Tuple
+from typing import Dict, List, Tuple, Union, Optional
 
 class geometry:
     """Class for the different surfaces in an optical system.
@@ -47,16 +47,24 @@ class geometry:
     """
 
     def __init__(self, params: Dict):
-        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)
+        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 = np.array(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()
@@ -66,22 +74,15 @@ def check_params(self) -> None:
 
         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.")
@@ -131,6 +132,9 @@ def conics(self, point: np.ndarray) -> Tuple[float, List[float]]:
         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.
@@ -160,5 +164,8 @@ def conics_plot(self, point: np.ndarray) -> np.ndarray:
         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

tracepy/optplot.py

@@ -2,6 +2,7 @@
 import matplotlib.pyplot as plt
 
 from .ray import ray
+from .constants import PLOT_ROUNDING_ACC
 from .geometry import geometry
 from .transforms import transform_points
 from .exceptions import TraceError
@@ -40,7 +41,7 @@ def _gen_object_points(surface: geometry,
     points_obj = transform_points(surface.R, surface, points)
 
     # Round arrays to upper bound on accuracy.
-    return np.around(points_obj, 14)
+    return np.around(points_obj, PLOT_ROUNDING_ACC)
 
 def calculate_rms(points: np.ndarray) -> float:
     """Calculates the RMS of the given points.