From edc101268652ec9b64bb97452190f590af9d47cb Mon Sep 17 00:00:00 2001
From: GNiendorf <gavinniendorf@gmail.com>
Date: Tue, 3 Sep 2024 21:12:21 -0400
Subject: [PATCH 1/6] Add type checking to tests

---
 .github/workflows/python-tests.yml | 46 +++++++++++++++++-------------
 1 file changed, 26 insertions(+), 20 deletions(-)

diff --git a/.github/workflows/python-tests.yml b/.github/workflows/python-tests.yml
index f24dd28..ec0c310 100644
--- a/.github/workflows/python-tests.yml
+++ b/.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
\ No newline at end of file

From d84265007be45baa0338879d3cb66a874f3e612d Mon Sep 17 00:00:00 2001
From: GNiendorf <gavinniendorf@gmail.com>
Date: Fri, 6 Sep 2024 21:49:16 -0400
Subject: [PATCH 2/6] working static type defs for transform.py, raygroup.py,
 and ray.py

---
 tracepy/__init__.py   |  3 ++-
 tracepy/geometry.py   |  2 +-
 tracepy/geoplot.py    |  4 ++--
 tracepy/optimize.py   |  2 +-
 tracepy/optplot.py    |  2 +-
 tracepy/ray.py        | 45 ++++++++++++++++++++++++++-----------------
 tracepy/raygroup.py   |  9 ++++++++-
 tracepy/transforms.py | 44 ++++++++++++------------------------------
 tracepy/utils.py      | 31 +++++++++++++++++++++++++++++
 9 files changed, 85 insertions(+), 57 deletions(-)
 create mode 100644 tracepy/utils.py

diff --git a/tracepy/__init__.py b/tracepy/__init__.py
index 85c96b5..f587c24 100644
--- a/tracepy/__init__.py
+++ b/tracepy/__init__.py
@@ -18,4 +18,5 @@
 from .iotables import *
 from .transforms import *
 from .raygroup import *
-from.index import *
\ No newline at end of file
+from .index import *
+from .utils import gen_rot
\ No newline at end of file
diff --git a/tracepy/geometry.py b/tracepy/geometry.py
index c8d0027..cea4a65 100644
--- a/tracepy/geometry.py
+++ b/tracepy/geometry.py
@@ -1,6 +1,6 @@
 import numpy as np
 
-from .transforms import gen_rot
+from .utils import gen_rot
 from .exceptions import NotOnSurfaceError
 from .index import glass_index
 
diff --git a/tracepy/geoplot.py b/tracepy/geoplot.py
index d34d91d..0cf44d7 100644
--- a/tracepy/geoplot.py
+++ b/tracepy/geoplot.py
@@ -161,7 +161,7 @@ def plotxz(geo_params, rays, pltparams={'c': 'red', 'alpha': 0.3 }, both=None):
 
     """
 
-    rays = np.array([rayiter for rayiter in rays if rayiter.P is not None])
+    rays = np.array([rayiter for rayiter in rays if rayiter.active != 0])
     #Override 1,1,1 subplot if displaying side-by-side.
     if both is None:
         #Keep aspect ratio equal.
@@ -187,7 +187,7 @@ def plotyz(geo_params, rays, pltparams={'c': 'red', 'alpha': 0.3 }, both=None):
 
     """
 
-    rays = np.array([rayiter for rayiter in rays if rayiter.P is not None])
+    rays = np.array([rayiter for rayiter in rays if rayiter.active != 0])
     #Override 1,1,1 subplot if displaying side-by-side.
     if both is None:
         #Keep aspect ratio equal.
diff --git a/tracepy/optimize.py b/tracepy/optimize.py
index 7b629ec..8b8d92b 100644
--- a/tracepy/optimize.py
+++ b/tracepy/optimize.py
@@ -62,7 +62,7 @@ def get_rms(inputs, geoparams, vary_dicts):
 
     params_iter = update_geometry(inputs, geoparams, vary_dicts)
     raygroup_iter = ray_plane(params_iter, [0., 0., 0.], 1.1, d=[0.,0.,1.], nrays=50)
-    ratio_surv = np.sum([1 for ray in raygroup_iter if ray.P is not None])/len(raygroup_iter)
+    ratio_surv = np.sum([1 for ray in raygroup_iter if ray.active != 0])/len(raygroup_iter)
     try:
         rms = spotdiagram(params_iter, raygroup_iter, optimizer=True)
     except TraceError:
diff --git a/tracepy/optplot.py b/tracepy/optplot.py
index 61ccaf1..f76dd67 100644
--- a/tracepy/optplot.py
+++ b/tracepy/optplot.py
@@ -29,7 +29,7 @@ def _gen_object_points(surface, surface_idx, rays):
 
     """
 
-    points = np.array([rayiter.P_hist[surface_idx] for rayiter in rays if rayiter.P is not None])
+    points = np.array([rayiter.P_hist[surface_idx] for rayiter in rays if rayiter.active != 0])
     if points.size == 0:
         #No rays survived
         raise TraceError()
diff --git a/tracepy/ray.py b/tracepy/ray.py
index 22e8348..f69a586 100644
--- a/tracepy/ray.py
+++ b/tracepy/ray.py
@@ -3,6 +3,10 @@
 from .transforms import *
 from .exceptions import NormalizationError, NotOnSurfaceError
 
+from .geometry import geometry
+
+from typing import Union, List, Optional
+
 class ray:
     """Class for rays and their propagation through surfaces.
 
@@ -24,24 +28,27 @@ class ray:
         Index of refraction of current material.
     wvl: float/int
         Wavelength of the ray in microns 550nm --> 0.55.
+    active: int
+        Is the ray still being propagated? 1 if so, else 0 for failed.
     """
 
-    def __init__(self, params, N_0=1):
+    def __init__(self, params: dict, N_0: Union[float, int] = 1):
         self.P = np.array(params['P'])
         self.D = np.array(params['D'])
         self.P_hist = [self.P]
         self.D_hist = [self.D]
         self.N = N_0
         self.wvl = params.get('wvl',0.55) #Added default wavelength 550nm
+        self.active = 1 #Is the ray still active?
         if abs(np.linalg.norm(self.D)-1.) > .01:
             #Ray direction cosines are not normalized.
             raise NormalizationError()
 
-    def transform(self, surface):
+    def transform(self, surface: geometry):
         """ Updates position and direction of a ray to obj coordinate system. """
         self.P, self.D = transform(surface.R, surface, np.array([self.P]), np.array([self.D]))
 
-    def find_intersection(self, surface):
+    def find_intersection(self, surface: geometry):
         """Finds the intersection point of a ray with a surface.
 
         Note
@@ -76,20 +83,20 @@ def find_intersection(self, surface):
                 func, normal= surface.get_surface([X, Y, Z])
                 deriv = np.dot(normal, self.D)
                 #Newton-raphson method
-                s_j = s_j[1], s_j[1]-func/deriv
+                s_j = [s_j[1], s_j[1]-func/deriv]
             except NotOnSurfaceError:
-                self.P = None
+                self.active = 0
                 return None
             #Error is how far f(X, Y, Z) is from 0.
             error = abs(func)
             n_iter += 1
         if n_iter == n_max or s_0+s_j[0] < 0 or np.dot(([X, Y, Z]-self.P), self.D) < 0.:
-            self.P = None
+            self.active = 0
         else:
             self.normal = normal
             self.P = np.array([X, Y, Z])
 
-    def interact(self, surface, typeof):
+    def interact(self, surface: geometry, typeof: str):
         """Updates new direction of a ray for a given interaction type.
 
         Note
@@ -122,7 +129,7 @@ def interact(self, surface, typeof):
         elif typeof == 'refraction':
             self.refraction(surface, mu, a, b)
 
-    def reflection(self, surface, a):
+    def reflection(self, surface: geometry, a: Union[float, int]):
         """Reflects the ray off a surface and updates the ray's direction.
 
         Note
@@ -143,7 +150,11 @@ def reflection(self, surface, a):
         K, L, M = self.normal
         self.D = np.array([k-2.*a*K, l-2.*a*L, m-2.*a*M])
 
-    def refraction(self, surface, mu, a, b):
+    def refraction(self,
+                   surface: geometry,
+                   mu: Union[float, int],
+                   a: Union[float, int],
+                   b: Union[float, int]):
         """Simulates refraction of a ray into a surface and updates the ray's direction.
 
         Note
@@ -177,12 +188,12 @@ def refraction(self, surface, mu, a, b):
         nmax = 1e5
         while error > 1e-15 and niter < nmax:
             #Newton-raphson method
-            G = G[1], (pow(G[1],2)-b)/(2*(G[1]+a))
+            G = [G[1], (pow(G[1],2)-b)/(2*(G[1]+a))]
             #See Spencer, Murty for where this is inspired by.
             error = abs(pow(G[1],2)+2*a*G[1]+b)
             niter += 1
         if niter==nmax:
-            self.P = None
+            self.active = 0
             return 0.
         #Update direction and index of refraction of the current material.
         self.D = np.array([mu*k+G[1]*K,mu*l+G[1]*L,mu*m+G[1]*M])
@@ -191,7 +202,7 @@ def refraction(self, surface, mu, a, b):
         else:
             self.N = surface.N
 
-    def ray_lab_frame(self, surface):
+    def ray_lab_frame(self, surface: geometry):
         """ Updates position and direction of a ray in the lab frame. """
         self.P, self.D = lab_frame(surface.R, surface, np.array([self.P]), np.array([self.D]))
 
@@ -200,15 +211,13 @@ def update(self):
         self.P_hist.append(self.P)
         self.D_hist.append(self.D)
 
-    def propagate(self, surfaces):
+    def propagate(self, surfaces: List[geometry]):
         """Propagates a ray through a given surfaces list.
 
         Note
         ----
-        If self.P is None then the ray failed to converge or
+        If self.active is 0 then the ray failed to converge or
         took too many iterations to meet the required accuracy.
-        Note that this is used (self.P is None) as a flag in
-        many other functions in TracePy.
 
         Parameters
         ----------
@@ -221,11 +230,11 @@ def propagate(self, surfaces):
             self.transform(surface)
             self.find_intersection(surface)
             #Results from failure to converge.
-            if self.P is None:
+            if self.active == 0:
                 break
             self.interact(surface, surface.action)
             #Results from too many iterations.
-            if self.P is None:
+            if self.active == 0:
                 break
             self.ray_lab_frame(surface)
             #Update current to history arrays.
diff --git a/tracepy/raygroup.py b/tracepy/raygroup.py
index dc48ec7..e8cb0c5 100644
--- a/tracepy/raygroup.py
+++ b/tracepy/raygroup.py
@@ -3,7 +3,14 @@
 from .ray import ray
 from .geometry import geometry
 
-def ray_plane(geo_params, pos, radius, d, nrays=100, wvl=0.55):
+from typing import Union, List, Dict
+
+def ray_plane(geo_params: List[Dict],
+              pos: Union[List[float], float, int],
+              radius: Union[float, int],
+              d: List[float],
+              nrays: int = 100,
+              wvl: Union[float, int] = 0.55):
     """Creates a plane of rays and propagates them through geometry.
 
     Note
diff --git a/tracepy/transforms.py b/tracepy/transforms.py
index e5ed8a6..dbe7084 100644
--- a/tracepy/transforms.py
+++ b/tracepy/transforms.py
@@ -1,36 +1,13 @@
 import numpy as np
 
-def gen_rot(ang):
-    """Returns a rotation matrix from 3 rotation angles.
+from .geometry import geometry
 
-    Parameters
-    ----------
-    ang : np.array of length 3
-        Euler angles alpha, beta, gamma in the lab frame.
-
-    Returns
-    -------
-    np.array((3,3))
-        Returns the rotation matrix.
-
-    """
-
-    alpha, beta, gamma = ang
-    R_11 = np.cos(alpha)*np.cos(gamma)+np.sin(alpha)*np.sin(beta)*np.sin(gamma)
-    R_12 = -np.cos(beta)*np.sin(gamma)
-    R_13 = -np.sin(alpha)*np.cos(gamma)+np.cos(alpha)*np.sin(beta)*np.sin(gamma)
-    R_21 = np.cos(alpha)*np.sin(gamma)-np.sin(alpha)*np.sin(beta)*np.cos(gamma)
-    R_22 = np.cos(beta)*np.cos(gamma)
-    R_23 = -np.sin(alpha)*np.sin(gamma)-np.cos(alpha)*np.sin(beta)*np.cos(gamma)
-    R_31 = np.sin(alpha)*np.cos(beta)
-    R_32 = np.sin(beta)
-    R_33 = np.cos(alpha)*np.cos(beta)
-    R = np.array([[R_11, R_12, R_13],\
-                   [R_21, R_22, R_23],\
-                   [R_31, R_32, R_33]])
-    return R
+from typing import Optional, Tuple, Union
 
-def transform(R, surface, points, D=None):
+def transform(R: np.ndarray,
+              surface: geometry,
+              points: np.ndarray,
+              D: Optional[np.ndarray] = None) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
     """Transforms points into the reference frame of surface.
 
     Note
@@ -39,7 +16,7 @@ def transform(R, surface, points, D=None):
 
     Parameters
     ----------
-    R : np.matrix((3,3))
+    R : np.array((3,3))
         Rotation matrix for surface.
     surface : geometry object
         Surface whos reference frame to transform into.
@@ -68,7 +45,10 @@ def transform(R, surface, points, D=None):
         tran_points = tran_points.flatten()
     return tran_points
 
-def lab_frame(R, surface, points, D=None):
+def lab_frame(R: np.ndarray,
+              surface: geometry,
+              points: np.ndarray,
+              D: Optional[np.ndarray] = None) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
     """Transforms points into the lab frame.
 
     Note
@@ -77,7 +57,7 @@ def lab_frame(R, surface, points, D=None):
 
     Parameters
     ----------
-    R : np.matrix((3,3))
+    R : np.array((3,3))
         Rotation matrix for surface.
     surface : geometry object
         Surface whos reference frame to transform from.
diff --git a/tracepy/utils.py b/tracepy/utils.py
new file mode 100644
index 0000000..0f373ab
--- /dev/null
+++ b/tracepy/utils.py
@@ -0,0 +1,31 @@
+import numpy as np
+
+def gen_rot(ang: np.ndarray) -> np.ndarray:
+    """Returns a rotation matrix from 3 rotation angles.
+
+    Parameters
+    ----------
+    ang : np.array of length 3
+        Euler angles alpha, beta, gamma in the lab frame.
+
+    Returns
+    -------
+    np.array((3,3))
+        Returns the rotation matrix.
+
+    """
+
+    alpha, beta, gamma = ang
+    R_11 = np.cos(alpha)*np.cos(gamma)+np.sin(alpha)*np.sin(beta)*np.sin(gamma)
+    R_12 = -np.cos(beta)*np.sin(gamma)
+    R_13 = -np.sin(alpha)*np.cos(gamma)+np.cos(alpha)*np.sin(beta)*np.sin(gamma)
+    R_21 = np.cos(alpha)*np.sin(gamma)-np.sin(alpha)*np.sin(beta)*np.cos(gamma)
+    R_22 = np.cos(beta)*np.cos(gamma)
+    R_23 = -np.sin(alpha)*np.sin(gamma)-np.cos(alpha)*np.sin(beta)*np.cos(gamma)
+    R_31 = np.sin(alpha)*np.cos(beta)
+    R_32 = np.sin(beta)
+    R_33 = np.cos(alpha)*np.cos(beta)
+    R = np.array([[R_11, R_12, R_13],\
+                   [R_21, R_22, R_23],\
+                   [R_31, R_32, R_33]])
+    return R

From 0f4612983ae2a0527b6a9b1dcff7ad13e7ea8f40 Mon Sep 17 00:00:00 2001
From: GNiendorf <gavinniendorf@gmail.com>
Date: Sat, 7 Sep 2024 17:25:43 -0400
Subject: [PATCH 3/6] working static types for all files but geometry.py,
 index.py

---
 tests/test_hyperbolic.py |  2 +-
 tests/test_parabolic.py  |  2 +-
 tracepy/__init__.py      |  9 ++--
 tracepy/constants.py     |  9 ++++
 tracepy/geometry.py      | 24 ++++------
 tracepy/geoplot.py       | 50 +++++++++++++--------
 tracepy/index.py         | 26 ++++++++---
 tracepy/iotables.py      |  4 +-
 tracepy/optimize.py      | 28 +++++++-----
 tracepy/optplot.py       | 65 ++++++++++++++++++---------
 tracepy/ray.py           | 66 ++++++++++++++--------------
 tracepy/raygroup.py      |  2 +-
 tracepy/transforms.py    | 94 ++++++++++++++++++++++++++++------------
 13 files changed, 242 insertions(+), 139 deletions(-)
 create mode 100644 tracepy/constants.py

diff --git a/tests/test_hyperbolic.py b/tests/test_hyperbolic.py
index 6c7a9bb..9da47aa 100644
--- a/tests/test_hyperbolic.py
+++ b/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.
 
diff --git a/tests/test_parabolic.py b/tests/test_parabolic.py
index 5d540bd..efae9cf 100644
--- a/tests/test_parabolic.py
+++ b/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.
diff --git a/tracepy/__init__.py b/tracepy/__init__.py
index f587c24..cf89ee9 100644
--- a/tracepy/__init__.py
+++ b/tracepy/__init__.py
@@ -14,9 +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
\ No newline at end of file
diff --git a/tracepy/constants.py b/tracepy/constants.py
new file mode 100644
index 0000000..b430121
--- /dev/null
+++ b/tracepy/constants.py
@@ -0,0 +1,9 @@
+# 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.
diff --git a/tracepy/geometry.py b/tracepy/geometry.py
index cea4a65..40b28b7 100644
--- a/tracepy/geometry.py
+++ b/tracepy/geometry.py
@@ -4,6 +4,8 @@
 from .exceptions import NotOnSurfaceError
 from .index import glass_index
 
+from typing import Dict, List, Tuple
+
 class geometry:
     """Class for the different surfaces in an optical system.
 
@@ -39,12 +41,12 @@ 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):
+    def __init__(self, params: Dict):
         self.P = params['P']
         self.D = np.array(params.get('D', [0., 0., 0.]))
         self.action = params['action']
@@ -59,15 +61,7 @@ def __init__(self, params):
             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
@@ -98,15 +92,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
@@ -144,7 +138,7 @@ def conics(self, point):
         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
diff --git a/tracepy/geoplot.py b/tracepy/geoplot.py
index 0cf44d7..9d8e477 100644
--- a/tracepy/geoplot.py
+++ b/tracepy/geoplot.py
@@ -1,17 +1,19 @@
-import matplotlib.pyplot as plt
 import numpy as np
-from numpy import pi
+import matplotlib.pyplot as plt
 
+from .ray import ray
 from .geometry import geometry
-from .transforms import lab_frame
+from .transforms import lab_frame_points
 
-def _gen_points(surfaces):
+from typing import List, Dict, Optional
+
+def _gen_points(surfaces: List[geometry]) -> np.ndarray:
     """Generates the mesh points for each surface in the obj frame.
 
     Parameters
     ----------
-    surface : geometry object
-        Surface whos reference frame to generate the points in.
+    surfaces : list of geometry objects
+        List of surfaces whos reference frames to generate the points in.
 
     Returns
     -------
@@ -37,12 +39,14 @@ def _gen_points(surfaces):
         surfpoints.append(meshpoints)
     return np.array(surfpoints)
 
-def _plot_rays(rays, axes, pltparams):
+def _plot_rays(rays: np.ndarray,
+               axes: List[int],
+               pltparams: Dict) -> None:
     """Plots 2d ray history points. Takes list axes to specify axes (0, 1, 2) to plot.
 
     Parameters
     ----------
-    rays : list of ray objects
+    rays : np.array of ray objects
         Rays that are going to be plotted.
     axes : list of length 2 with integers from range [0,2]
         Axes (X, Y, Z) to plot from ray points.
@@ -61,7 +65,7 @@ def _plot_rays(rays, axes, pltparams):
         #Plot direction of ray after stop.
         plt.plot([G_p, G_p+I_p],[F_p, F_p+H_p], **pltparams)
 
-def _clip_lens(surfaces, surfpoints, idx):
+def _clip_lens(surfaces: List[geometry], surfpoints: np.ndarray, idx: int) -> np.ndarray:
     """Clips points ouside of a lens intersection point.
 
     Parameters
@@ -89,7 +93,7 @@ def _clip_lens(surfaces, surfpoints, idx):
     surfpoints[idx+1][:,2][clipped_idx] = np.nan
     return surfpoints
 
-def _plot_surfaces(geo_params, axes):
+def _plot_surfaces(geo_params: List[Dict], axes: List[int]) -> None:
     """Plots 2d surface cross sections. Takes list axes to specify axes (0, 1, 2) to plot.
 
     Note
@@ -115,7 +119,7 @@ def _plot_surfaces(geo_params, axes):
         if lens_condition:
             surfpoints = _clip_lens(surfaces, surfpoints, idx)
         with np.errstate(invalid='ignore'):
-            if np.any(np.mod(surf.D/pi, 1) != 0) and surf.c == 0 and surf.diam == 0:
+            if np.any(np.mod(surf.D/np.pi, 1) != 0) and surf.c == 0 and surf.diam == 0:
                     #Find cross section points.
                     cross_idx = abs(surfpoints[idx][:,axes[1]]) == 0
             else:
@@ -123,7 +127,7 @@ def _plot_surfaces(geo_params, axes):
                     cross_idx = abs(surfpoints[idx][:,1-axes[0]]) == 0
         cross_points = surfpoints[idx][cross_idx]
         #Transform to lab frame.
-        points = lab_frame(surf.R, surf, cross_points)
+        points = lab_frame_points(surf.R, surf, cross_points)
         F, G = points[:,axes[0]], points[:,axes[1]]
         #Connect the surfaces in a lens
         if surfaces[idx].action == surfaces[idx-1].action == 'refraction' and start is not None:
@@ -145,14 +149,17 @@ def _plot_surfaces(geo_params, axes):
             end = np.array([F[-1], G[-1]])
         plt.plot(G, F, 'k')
 
-def plotxz(geo_params, rays, pltparams={'c': 'red', 'alpha': 0.3 }, both=None):
+def plotxz(geo_params: List[Dict],
+           ray_list: List[ray],
+           pltparams: Dict = {'c': 'red', 'alpha': 0.3 },
+           both: Optional[bool] = None) -> None:
     """Plots the xz coordinates of all rays and surface cross sections.
 
     Parameters
     ----------
     geo_params : list of dictionaries
         Surfaces in propagation order to plot.
-    rays : list of ray objects
+    ray_list : list of ray objects
         Rays that are going to be plotted.
     pltparams : dictionary
         Plot characteristics of rays such as colors and alpha.
@@ -161,7 +168,7 @@ def plotxz(geo_params, rays, pltparams={'c': 'red', 'alpha': 0.3 }, both=None):
 
     """
 
-    rays = np.array([rayiter for rayiter in rays if rayiter.active != 0])
+    rays = np.array([rayiter for rayiter in ray_list if rayiter.active])
     #Override 1,1,1 subplot if displaying side-by-side.
     if both is None:
         #Keep aspect ratio equal.
@@ -171,14 +178,17 @@ def plotxz(geo_params, rays, pltparams={'c': 'red', 'alpha': 0.3 }, both=None):
     plt.xlabel("Z")
     plt.ylabel("X")
 
-def plotyz(geo_params, rays, pltparams={'c': 'red', 'alpha': 0.3 }, both=None):
+def plotyz(geo_params: List[Dict],
+           ray_list: List[ray],
+           pltparams: Dict = {'c': 'red', 'alpha': 0.3 },
+           both: Optional[bool] = None) -> None:
     """Plots the yz coordinates of all rays and surface cross sections.
 
     Parameters
     ----------
     geo_params : list of dictionaries
         Surfaces in propagation order to plot.
-    rays : list of ray objects
+    ray_list : list of ray objects
         Rays that are going to be plotted.
     pltparams : dictionary
         Plot characteristics of rays such as colors and alpha.
@@ -187,7 +197,7 @@ def plotyz(geo_params, rays, pltparams={'c': 'red', 'alpha': 0.3 }, both=None):
 
     """
 
-    rays = np.array([rayiter for rayiter in rays if rayiter.active != 0])
+    rays = np.array([rayiter for rayiter in ray_list if rayiter.active])
     #Override 1,1,1 subplot if displaying side-by-side.
     if both is None:
         #Keep aspect ratio equal.
@@ -197,7 +207,9 @@ def plotyz(geo_params, rays, pltparams={'c': 'red', 'alpha': 0.3 }, both=None):
     plt.xlabel("Z")
     plt.ylabel("Y")
 
-def plot2d(geo_params, rays, pltparams={'c': 'red', 'alpha': 0.3 }):
+def plot2d(geo_params: List[Dict],
+           rays: List[ray],
+           pltparams: Dict = {'c': 'red', 'alpha': 0.3 }) -> None:
     """Plots both xz and yz side-by-side.
 
     Parameters
diff --git a/tracepy/index.py b/tracepy/index.py
index a793262..e9a1b28 100755
--- a/tracepy/index.py
+++ b/tracepy/index.py
@@ -5,15 +5,29 @@
 import numpy as np
 from scipy.optimize import curve_fit
 
-def cauchy_two_term(x,B,C):
-    '''
-    This is a simple two term cauchy for index of refraction as function of wavelength.
+def cauchy_two_term(x: float, B: float, C: float) -> float:
+    """
+    This is a simple two-term Cauchy equation for the index of refraction as a function of wavelength.
     https://en.wikipedia.org/wiki/Cauchy%27s_equation
     It is used to create a function for refractive index database entries when only tabulated data is available.
-    '''
-    return B + (C/(x**2))
+    
+    Parameters
+    ----------
+    x : float
+        Wavelength (in micrometers).
+    B : float
+        First term of the Cauchy equation.
+    C : float
+        Second term of the Cauchy equation.
+        
+    Returns
+    -------
+    float
+        Refractive index at the given wavelength.
+    """
+    return B + (C / (x ** 2))
 
-def glass_index (glass):
+def glass_index(glass):
     '''
     Given a glass name this function will return a function that takes wavelength in microns and returns index of refraction.
     - All values were taken from refractiveindexinfo's database of different optical glasses.
diff --git a/tracepy/iotables.py b/tracepy/iotables.py
index 7c363d6..fa7a0d9 100644
--- a/tracepy/iotables.py
+++ b/tracepy/iotables.py
@@ -1,6 +1,8 @@
 import pandas as pd
 
-def save_optics(geo_params, filename):
+# TODO: Rewrite save_optics to iterate over geo_params to handle edge
+#       cases, and convert to dataframe at end before exporting to csv.
+def save_optics(geo_params, filename: str):
     """Save geometry to an optics table in csv format.
 
     Note
diff --git a/tracepy/optimize.py b/tracepy/optimize.py
index 8b8d92b..8c91024 100644
--- a/tracepy/optimize.py
+++ b/tracepy/optimize.py
@@ -1,11 +1,16 @@
 import numpy as np
 from scipy.optimize import least_squares
 
-from .optplot import spotdiagram
+from .optplot import spot_rms
 from .raygroup import ray_plane
+from .constants import SURV_CONST, MAX_RMS
 from .exceptions import TraceError
 
-def update_geometry(inputs, geoparams, vary_dicts):
+from typing import List, Union, Dict, Optional
+
+def update_geometry(inputs: List[Union[float, int]],
+                    geoparams: List[Dict],
+                    vary_dicts: List[Dict]) -> List[Dict]:
     """Return the geometry requested by the optimization algorithm.
 
     Parameters
@@ -35,7 +40,9 @@ def update_geometry(inputs, geoparams, vary_dicts):
                     vary_idxs += 1
     return geoparams
 
-def get_rms(inputs, geoparams, vary_dicts):
+def get_rms(inputs: List[Union[float, int]],
+            geoparams: List[Dict],
+            vary_dicts: List[Dict]) -> float:
     """Return the rms of an updated geometry.
 
     Note
@@ -62,16 +69,17 @@ def get_rms(inputs, geoparams, vary_dicts):
 
     params_iter = update_geometry(inputs, geoparams, vary_dicts)
     raygroup_iter = ray_plane(params_iter, [0., 0., 0.], 1.1, d=[0.,0.,1.], nrays=50)
-    ratio_surv = np.sum([1 for ray in raygroup_iter if ray.active != 0])/len(raygroup_iter)
+    ratio_surv = np.sum([1 for ray in raygroup_iter if ray.active])/len(raygroup_iter)
     try:
-        rms = spotdiagram(params_iter, raygroup_iter, optimizer=True)
+        rms = spot_rms(params_iter, raygroup_iter)
     except TraceError:
-        rms = 999.
-    #Weight of failed propagation.
-    surv_const = 100
-    return rms + (1-ratio_surv)*surv_const
+        rms = MAX_RMS
+    return rms + (1 - ratio_surv) * SURV_CONST
 
-def optimize(geoparams, vary_dicts, typeof='least_squares', max_iter=None):
+def optimize(geoparams:List[Dict],
+             vary_dicts: List[Dict],
+             typeof: str = 'least_squares',
+             max_iter: Optional[int] = None) -> List[Dict]:
     """Optimize a given geometry for a given varylist and return the new geometry.
 
     Parameters
diff --git a/tracepy/optplot.py b/tracepy/optplot.py
index f76dd67..8323cf8 100644
--- a/tracepy/optplot.py
+++ b/tracepy/optplot.py
@@ -1,12 +1,18 @@
 import numpy as np
-
 import matplotlib.pyplot as plt
 
+from .ray import ray
 from .geometry import geometry
-from .transforms import transform
+from .transforms import transform_points
 from .exceptions import TraceError
 
-def _gen_object_points(surface, surface_idx, rays):
+from typing import List, Dict, Tuple
+
+# TODO: Decouple plotting the spot diagram from the RMS calculation used by optimizer.
+# This involves breaking up the three functions below to decouple the two tasks.
+def _gen_object_points(surface: geometry,
+                       surface_idx: int,
+                       rays: List[ray]) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
     """Transform intersection points into a surfaces' reference frame.
 
     Parameters
@@ -29,12 +35,12 @@ def _gen_object_points(surface, surface_idx, rays):
 
     """
 
-    points = np.array([rayiter.P_hist[surface_idx] for rayiter in rays if rayiter.active != 0])
+    points = np.array([rayiter.P_hist[surface_idx] for rayiter in rays if rayiter.active])
     if points.size == 0:
         #No rays survived
         raise TraceError()
     #Get X,Y points in obj. reference frame.
-    points_obj = transform(surface.R, surface, points)
+    points_obj = transform_points(surface.R, surface, points)
     #Round arrays to upper bound on accuracy.
     points_obj = np.around(points_obj, 14)
     if points_obj.ndim == 2:
@@ -44,12 +50,31 @@ def _gen_object_points(surface, surface_idx, rays):
         points_obj = np.array([points_obj])
     return X, Y, points_obj
 
-def spotdiagram(geo_params, rays, pltparams = {'color': 'red'}, optimizer=False):
-    """Plots the transformed intersection points of rays on the stop surface.
+def spot_rms(geo_params: List[Dict], rays: List[ray]) -> float:
+    """Calculates the RMS of the spot diagram points.
 
-    Note
-    ----
-    Directly plots the spot diagram rather than returning something.
+    Parameters
+    ----------
+    geo_params : list of dictionaries
+        Surfaces in order of propagation.
+    rays : list of ray objects
+        Rays that were propagated through the geometry list.
+
+    Returns
+    -------
+    float
+        The calculated RMS value.
+
+    """
+
+    stop = geometry(geo_params[-1])
+    _, _, points_obj = _gen_object_points(stop, -1, rays)
+    return np.std(points_obj[:, [0, 1]] - points_obj[:, [0, 1]].mean(axis=0))
+
+def spotdiagram(geo_params: List[Dict],
+                rays: List[ray],
+                pltparams: Dict = {'color': 'red'}) -> None:
+    """Plots the transformed intersection points of rays on the stop surface.
 
     Parameters
     ----------
@@ -59,24 +84,22 @@ def spotdiagram(geo_params, rays, pltparams = {'color': 'red'}, optimizer=False)
         Rays that were propagated through the geometry list.
     pltparams : dictionary
         Plotting attributes of the spot diagram.
-    optimizer : bool
-        Flag for whether the optimizer is calling the function to get
-        the rms of the spot diagram.
 
     """
 
     stop = geometry(geo_params[-1])
     X, Y, points_obj = _gen_object_points(stop, -1, rays)
-    rms = np.std(points_obj[:,[0,1]] - points_obj[:,[0,1]].mean(axis=0))
-    if optimizer:
-        return rms
-    plt.subplot(1,1,1, aspect='equal')
+    rms = np.std(points_obj[:, [0, 1]] - points_obj[:, [0, 1]].mean(axis=0))
+
+    plt.subplot(1, 1, 1, aspect='equal')
     plt.locator_params(axis='x', nbins=8)
-    plt.ticklabel_format(style='sci', axis='both', scilimits=(0,0))
+    plt.ticklabel_format(style='sci', axis='both', scilimits=(0, 0))
     plt.plot(X, Y, '+', **pltparams)
     plt.text(0.65, 1.015, 'RMS=%.2E' % rms, transform=plt.gca().transAxes)
 
-def plotobject(geo_params, rays, pltparams = {'color': 'blue'}):
+def plotobject(geo_params: List[Dict],
+               rays: List[ray],
+               pltparams: Dict = {'color': 'blue'}) -> None:
     """Plots the initial ray locations in the initial objects' frame.
 
     Note
@@ -102,7 +125,9 @@ def plotobject(geo_params, rays, pltparams = {'color': 'blue'}):
     plt.text(0.65, 1.015, 'RMS=%.2E' % rms, transform=plt.gca().transAxes)
     plt.plot(X, Y, '+', **pltparams)
 
-def rayaberration(geo_params, rays, pltparams= {'color': 'red', 'linewidth': 2}):
+def rayaberration(geo_params: List[Dict],
+                  rays: List[ray],
+                  pltparams: Dict = {'color': 'red', 'linewidth': 2}) -> None:
     """Plots the ray aberration diagrams for an optical system.
 
     Note
diff --git a/tracepy/ray.py b/tracepy/ray.py
index f69a586..d404f40 100644
--- a/tracepy/ray.py
+++ b/tracepy/ray.py
@@ -1,9 +1,9 @@
 import numpy as np
 
 from .transforms import *
-from .exceptions import NormalizationError, NotOnSurfaceError
-
+from .constants import *
 from .geometry import geometry
+from .exceptions import NormalizationError, NotOnSurfaceError
 
 from typing import Union, List, Optional
 
@@ -28,27 +28,28 @@ class ray:
         Index of refraction of current material.
     wvl: float/int
         Wavelength of the ray in microns 550nm --> 0.55.
-    active: int
-        Is the ray still being propagated? 1 if so, else 0 for failed.
+    active: bool
+        Is the ray still being propagated? True if so, else False for failed.
     """
 
-    def __init__(self, params: dict, N_0: Union[float, int] = 1):
-        self.P = np.array(params['P'])
-        self.D = np.array(params['D'])
-        self.P_hist = [self.P]
-        self.D_hist = [self.D]
-        self.N = N_0
-        self.wvl = params.get('wvl',0.55) #Added default wavelength 550nm
-        self.active = 1 #Is the ray still active?
+    def __init__(self, params: dict, N_0: Union[float, int] = 1.):
+        self.P: np.ndarray = np.array(params['P'])
+        self.D: np.ndarray = np.array(params['D'])
+        self.P_hist: List[np.ndarray] = [self.P]
+        self.D_hist: List[np.ndarray]  = [self.D]
+        self.N: float = N_0
+        self.wvl: float = params.get('wvl',0.55) #Added default wavelength 550nm
+        self.active: bool = True #Is the ray still active?
         if abs(np.linalg.norm(self.D)-1.) > .01:
             #Ray direction cosines are not normalized.
             raise NormalizationError()
 
-    def transform(self, surface: geometry):
+    def transform(self, surface: geometry) -> None:
         """ Updates position and direction of a ray to obj coordinate system. """
-        self.P, self.D = transform(surface.R, surface, np.array([self.P]), np.array([self.D]))
+        self.P = transform_points(surface.R, surface, np.array([self.P]))
+        self.D = transform_dir(surface.R, surface, np.array([self.D]))
 
-    def find_intersection(self, surface: geometry):
+    def find_intersection(self, surface: geometry) -> None:
         """Finds the intersection point of a ray with a surface.
 
         Note
@@ -75,28 +76,28 @@ def find_intersection(self, surface: geometry):
         error = 1.
         n_iter = 0
         #Max iterations allowed.
-        n_max = 1e4
-        while error > 1e-6 and n_iter < n_max:
+        n_max = MAX_INTERSECTION_ITERATIONS
+        while error > INTERSECTION_CONVERGENCE_TOLERANCE and n_iter < n_max:
             X, Y, Z = [X_1, Y_1, 0.]+np.dot(self.D, s_j[0])
             try:
                 #'normal' is the surface direction numbers.
-                func, normal= surface.get_surface([X, Y, Z])
+                func, normal= surface.get_surface(np.array([X, Y, Z]))
                 deriv = np.dot(normal, self.D)
                 #Newton-raphson method
                 s_j = [s_j[1], s_j[1]-func/deriv]
             except NotOnSurfaceError:
-                self.active = 0
+                self.active = False
                 return None
             #Error is how far f(X, Y, Z) is from 0.
             error = abs(func)
             n_iter += 1
         if n_iter == n_max or s_0+s_j[0] < 0 or np.dot(([X, Y, Z]-self.P), self.D) < 0.:
-            self.active = 0
+            self.active = False
         else:
             self.normal = normal
             self.P = np.array([X, Y, Z])
 
-    def interact(self, surface: geometry, typeof: str):
+    def interact(self, surface: geometry, typeof: str) -> None:
         """Updates new direction of a ray for a given interaction type.
 
         Note
@@ -129,7 +130,7 @@ def interact(self, surface: geometry, typeof: str):
         elif typeof == 'refraction':
             self.refraction(surface, mu, a, b)
 
-    def reflection(self, surface: geometry, a: Union[float, int]):
+    def reflection(self, surface: geometry, a: Union[float, int]) -> None:
         """Reflects the ray off a surface and updates the ray's direction.
 
         Note
@@ -154,7 +155,7 @@ def refraction(self,
                    surface: geometry,
                    mu: Union[float, int],
                    a: Union[float, int],
-                   b: Union[float, int]):
+                   b: Union[float, int]) -> None:
         """Simulates refraction of a ray into a surface and updates the ray's direction.
 
         Note
@@ -185,16 +186,16 @@ def refraction(self,
         error = 1.
         niter = 0
         #Max iterations allowed.
-        nmax = 1e5
-        while error > 1e-15 and niter < nmax:
+        n_max = MAX_REFRACTION_ITERATIONS
+        while error > REFRACTION_CONVERGENCE_TOLERANCE and niter < n_max:
             #Newton-raphson method
             G = [G[1], (pow(G[1],2)-b)/(2*(G[1]+a))]
             #See Spencer, Murty for where this is inspired by.
             error = abs(pow(G[1],2)+2*a*G[1]+b)
             niter += 1
-        if niter==nmax:
-            self.active = 0
-            return 0.
+        if niter == n_max:
+            self.active = False
+            return None
         #Update direction and index of refraction of the current material.
         self.D = np.array([mu*k+G[1]*K,mu*l+G[1]*L,mu*m+G[1]*M])
         if hasattr(surface,'glass'):
@@ -202,16 +203,17 @@ def refraction(self,
         else:
             self.N = surface.N
 
-    def ray_lab_frame(self, surface: geometry):
+    def ray_lab_frame(self, surface: geometry) -> None:
         """ Updates position and direction of a ray in the lab frame. """
-        self.P, self.D = lab_frame(surface.R, surface, np.array([self.P]), np.array([self.D]))
+        self.P = lab_frame_points(surface.R, surface, np.array([self.P]))
+        self.D = lab_frame_dir(surface.R, surface, np.array([self.D]))
 
-    def update(self):
+    def update(self) -> None:
         """ Updates the P_hist and D_hist arrays from current P and D arrays. """
         self.P_hist.append(self.P)
         self.D_hist.append(self.D)
 
-    def propagate(self, surfaces: List[geometry]):
+    def propagate(self, surfaces: List[geometry]) -> None:
         """Propagates a ray through a given surfaces list.
 
         Note
diff --git a/tracepy/raygroup.py b/tracepy/raygroup.py
index e8cb0c5..4ba22ef 100644
--- a/tracepy/raygroup.py
+++ b/tracepy/raygroup.py
@@ -10,7 +10,7 @@ def ray_plane(geo_params: List[Dict],
               radius: Union[float, int],
               d: List[float],
               nrays: int = 100,
-              wvl: Union[float, int] = 0.55):
+              wvl: Union[float, int] = 0.55) -> List[ray]:
     """Creates a plane of rays and propagates them through geometry.
 
     Note
diff --git a/tracepy/transforms.py b/tracepy/transforms.py
index dbe7084..d52a4d9 100644
--- a/tracepy/transforms.py
+++ b/tracepy/transforms.py
@@ -4,10 +4,9 @@
 
 from typing import Optional, Tuple, Union
 
-def transform(R: np.ndarray,
-              surface: geometry,
-              points: np.ndarray,
-              D: Optional[np.ndarray] = None) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
+def transform_points(R: np.ndarray,
+                     surface: geometry,
+                     points: np.ndarray) -> np.ndarray:
     """Transforms points into the reference frame of surface.
 
     Note
@@ -22,33 +21,52 @@ def transform(R: np.ndarray,
         Surface whos reference frame to transform into.
     points : 2d np.array
         Point for each row that will be transformed.
-    D (optional): 2d np.array
-        Direction for each row that will be transformed.
 
     Returns
     -------
     tran_points : 2d np.array
         Points in the transformed reference frame.
-    tran_D (optional): 2d np.array
-        Directions in the transformed reference frame.
 
     """
 
     tran_points = np.array(np.dot(R, (points-surface.P).T).T)
-    if D is not None:
-        tran_D = np.array(np.dot(R, D.T).T)
-        if len(tran_points) == 1:
-            tran_points = tran_points.flatten()
-            tran_D = tran_D.flatten()
-        return tran_points, tran_D
     if len(tran_points) == 1:
         tran_points = tran_points.flatten()
     return tran_points
 
-def lab_frame(R: np.ndarray,
-              surface: geometry,
-              points: np.ndarray,
-              D: Optional[np.ndarray] = None) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
+def transform_dir(R: np.ndarray,
+                  surface: geometry,
+                  D: np.ndarray) -> np.ndarray:
+    """Transforms directions into the reference frame of surface.
+
+    Note
+    ----
+    Arrays are flattened before return if they only have one row.
+
+    Parameters
+    ----------
+    R : np.array((3,3))
+        Rotation matrix for surface.
+    surface : geometry object
+        Surface whos reference frame to transform into.
+    D : 2d np.array
+        Direction for each row that will be transformed.
+
+    Returns
+    -------
+    tran_D: 2d np.array
+        Directions in the transformed reference frame.
+
+    """
+
+    tran_D = np.array(np.dot(R, D.T).T)
+    if len(tran_D) == 1:
+        tran_D = tran_D.flatten()
+    return tran_D
+
+def lab_frame_points(R: np.ndarray,
+                     surface: geometry,
+                     points: np.ndarray) -> np.ndarray:
     """Transforms points into the lab frame.
 
     Note
@@ -63,25 +81,45 @@ def lab_frame(R: np.ndarray,
         Surface whos reference frame to transform from.
     points : 2d np.array
         Point for each row that will be transformed.
-    D (optional): 2d np.array
-        Direction for each row that will be transformed.
 
     Returns
     -------
     lab_points : 2d np.array
         Points in the lab reference frame.
-    lab_D (optional): 2d np.array
-        Directions in the lab reference frame.
 
     """
 
     lab_points = np.array(np.dot(R.T, points.T).T)+surface.P
-    if D is not None:
-        lab_D = np.array(np.dot(R.T, D.T).T)
-        if len(lab_points) == 1:
-            lab_points = lab_points.flatten()
-            lab_D = lab_D.flatten()
-        return lab_points, lab_D
     if len(lab_points) == 1:
         lab_points = lab_points.flatten()
     return lab_points
+
+def lab_frame_dir(R: np.ndarray,
+                  surface: geometry,
+                  D: np.ndarray) -> np.ndarray:
+    """Transforms directions into the lab frame.
+
+    Note
+    ----
+    Arrays are flattened before return if they only have one row.
+
+    Parameters
+    ----------
+    R : np.array((3,3))
+        Rotation matrix for surface.
+    surface : geometry object
+        Surface whos reference frame to transform from.
+    D : 2d np.array
+        Direction for each row that will be transformed.
+
+    Returns
+    -------
+    lab_D : 2d np.array
+        Directions in the lab reference frame.
+
+    """
+
+    lab_D = np.array(np.dot(R.T, D.T).T)
+    if len(lab_D) == 1:
+        lab_D = lab_D.flatten()
+    return lab_D

From f8af45abf674c7a646b81a738968cd27ed06e8d0 Mon Sep 17 00:00:00 2001
From: GNiendorf <gavinniendorf@gmail.com>
Date: Sat, 7 Sep 2024 17:38:04 -0400
Subject: [PATCH 4/6] decouple spot diagram plot and optimizer rms calc

---
 tracepy/optplot.py | 48 ++++++++++++++++++++++++++++------------------
 1 file changed, 29 insertions(+), 19 deletions(-)

diff --git a/tracepy/optplot.py b/tracepy/optplot.py
index 8323cf8..94923be 100644
--- a/tracepy/optplot.py
+++ b/tracepy/optplot.py
@@ -8,17 +8,15 @@
 
 from typing import List, Dict, Tuple
 
-# TODO: Decouple plotting the spot diagram from the RMS calculation used by optimizer.
-# This involves breaking up the three functions below to decouple the two tasks.
 def _gen_object_points(surface: geometry,
                        surface_idx: int,
-                       rays: List[ray]) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+                       rays: List[ray]) -> np.ndarray:
     """Transform intersection points into a surfaces' reference frame.
 
     Parameters
     ----------
     surface : geometry object
-        Surface whos reference frame the points will be transformed into.
+        Surface whose reference frame the points will be transformed into.
     surface_idx : int
         Integer corresponding to where the surface is in the propagation
         order. For example, 0 means the surface is the first surface rays
@@ -28,8 +26,6 @@ def _gen_object_points(surface: geometry,
 
     Returns
     -------
-    X, Y : np.array of len(rays)
-        X, Y pair in the surface's reference frame.
     points_obj: 2d np.array
         X, Y pair points in 2d array for easy rms calculation.
 
@@ -39,16 +35,29 @@ def _gen_object_points(surface: geometry,
     if points.size == 0:
         #No rays survived
         raise TraceError()
-    #Get X,Y points in obj. reference frame.
+
+    # Get X,Y points in obj. reference frame.
     points_obj = transform_points(surface.R, surface, points)
-    #Round arrays to upper bound on accuracy.
-    points_obj = np.around(points_obj, 14)
-    if points_obj.ndim == 2:
-        X, Y = points_obj[:,0], points_obj[:,1]
-    elif points_obj.ndim == 1:
-        X, Y = points_obj[0], points_obj[1]
-        points_obj = np.array([points_obj])
-    return X, Y, points_obj
+
+    # Round arrays to upper bound on accuracy.
+    return np.around(points_obj, 14)
+
+def calculate_rms(points: np.ndarray) -> float:
+    """Calculates the RMS of the given points.
+
+    Parameters
+    ----------
+    points : np.ndarray
+        Array of points to calculate RMS for.
+
+    Returns
+    -------
+    float
+        The calculated RMS value.
+
+    """
+
+    return np.std(points[:, [0, 1]] - points[:, [0, 1]].mean(axis=0))
 
 def spot_rms(geo_params: List[Dict], rays: List[ray]) -> float:
     """Calculates the RMS of the spot diagram points.
@@ -68,8 +77,8 @@ def spot_rms(geo_params: List[Dict], rays: List[ray]) -> float:
     """
 
     stop = geometry(geo_params[-1])
-    _, _, points_obj = _gen_object_points(stop, -1, rays)
-    return np.std(points_obj[:, [0, 1]] - points_obj[:, [0, 1]].mean(axis=0))
+    points_obj = _gen_object_points(stop, -1, rays)
+    return calculate_rms(points_obj)
 
 def spotdiagram(geo_params: List[Dict],
                 rays: List[ray],
@@ -88,8 +97,9 @@ def spotdiagram(geo_params: List[Dict],
     """
 
     stop = geometry(geo_params[-1])
-    X, Y, points_obj = _gen_object_points(stop, -1, rays)
-    rms = np.std(points_obj[:, [0, 1]] - points_obj[:, [0, 1]].mean(axis=0))
+    points_obj = _gen_object_points(stop, -1, rays)
+    X, Y = points_obj[:, 0], points_obj[:, 1]
+    rms = calculate_rms(points_obj)
 
     plt.subplot(1, 1, 1, aspect='equal')
     plt.locator_params(axis='x', nbins=8)

From 048ef96a37e43a075f2d87ca688178c911e9bc99 Mon Sep 17 00:00:00 2001
From: GNiendorf <gavinniendorf@gmail.com>
Date: Sat, 7 Sep 2024 18:17:10 -0400
Subject: [PATCH 5/6] working static types for geometry.py, move optplot
 constant to constant.py

---
 tracepy/constants.py  |  3 +++
 tracepy/exceptions.py |  3 +++
 tracepy/geometry.py   | 57 ++++++++++++++++++++++++-------------------
 tracepy/optplot.py    |  3 ++-
 4 files changed, 40 insertions(+), 26 deletions(-)

diff --git a/tracepy/constants.py b/tracepy/constants.py
index b430121..a8f47da 100644
--- a/tracepy/constants.py
+++ b/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.
\ No newline at end of file
diff --git a/tracepy/exceptions.py b/tracepy/exceptions.py
index e62bb42..1b029ae 100644
--- a/tracepy/exceptions.py
+++ b/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. """
diff --git a/tracepy/geometry.py b/tracepy/geometry.py
index 40b28b7..a1c0ff9 100644
--- a/tracepy/geometry.py
+++ b/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 = 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
\ No newline at end of file
diff --git a/tracepy/optplot.py b/tracepy/optplot.py
index 94923be..5befcdf 100644
--- a/tracepy/optplot.py
+++ b/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.

From c5afa7c93a70ddc98a30d15798537f116152151c Mon Sep 17 00:00:00 2001
From: GNiendorf <gavinniendorf@gmail.com>
Date: Sat, 7 Sep 2024 18:26:44 -0400
Subject: [PATCH 6/6] Add TODO statement for glass_index function

---
 tracepy/index.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tracepy/index.py b/tracepy/index.py
index e9a1b28..ec877a2 100755
--- a/tracepy/index.py
+++ b/tracepy/index.py
@@ -27,6 +27,7 @@ def cauchy_two_term(x: float, B: float, C: float) -> float:
     """
     return B + (C / (x ** 2))
 
+# TODO: This function, written by @MikeMork, needs to be broken up into several smaller functions with explicit typing.
 def glass_index(glass):
     '''
     Given a glass name this function will return a function that takes wavelength in microns and returns index of refraction.