add type annotations to the code with new tests

.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_optimizer.py

@@ -0,0 +1,65 @@
+import numpy as np
+import tracepy as tp
+
+lens = {
+    'name': 'lens',
+    'action': 'refraction',
+    'P': 2.,
+    'kappa': -.004,
+    'c': -.5,
+    'Diam': 2.2
+}
+
+back_lens = {
+    'name': 'back_lens',
+    'action': 'refraction',
+    'P': 1.,
+    'c': .5,
+    'kappa': -.004,
+    'N': 1.5,
+    'Diam': 2.2
+}
+
+stop = {
+    'name': 'stop',
+    'action': 'stop',
+    'P': 3.,
+    'Diam': 2.2
+}
+
+# Initial geometry
+geo_1 = [back_lens, lens, stop]
+
+ray_group = tp.ray_plane(geo_1, [0., 0., 0.], 1.1, d=[0., 0., 1.], nrays=100)
+
+# Define variable elements for optimization
+vlens = {
+    'name': 'stop',
+    'vary': ['P']
+}
+
+vlens2 = {
+    'name': 'back_lens',
+    'vary': ['c', 'kappa', 'N']
+}
+
+vlens3 = {
+    'name': 'lens',
+    'vary': ['c', 'kappa']
+}
+
+vary_list = [vlens, vlens2, vlens3]
+
+def test_optimizer():
+    try:
+        geo_opt = tp.optimize(geo_1, vary_list, typeof='least_squares', max_iter=10)
+    except Exception as e:
+        assert False, f"Optimization raised an exception: {e}"
+
+    try:
+        ray_group_3 = tp.ray_plane(geo_opt, [0., 0., 0.], 1.1, d=[0., 0., 1.], nrays=1000)
+        tp.rayaberration(geo_1, ray_group_3)
+    except Exception as e:
+        assert False, f"rayaberration raised an exception: {e}"
+
+

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.

tests/test_plotting.py

@@ -0,0 +1,83 @@
+import numpy as np
+import tracepy as tp
+
+mirror = {
+    'action': 'reflection',
+    'D': np.array([0., np.pi, 0.]),
+    'P': np.array([0., 0., 10.]),
+    'kappa': 0.,
+    'c': 0.05,
+    'Diam': 3.6,
+}
+
+small_mirror = {
+    'action': 'reflection',
+    'D': np.array([0., -np.pi/4., 0.]),
+    'P': np.array([0., 0., 2.]),
+    'Diam': 1.4,
+}
+
+selector = {
+    'action': 'stop',
+    'P': np.array([0., 0., 0.1]),
+    'Diam': 3.6,
+    'diam': 1.4
+}
+
+stop = {
+    'action': 'stop',
+    'P': np.array([0., -4.6, 2.]),
+    'D': np.array([0., np.pi/2., 0.]),
+    'Diam': .8
+}
+
+lensy = -4.
+thickness = .1
+
+lens1 = {
+    'action': 'refraction',
+    'P': np.array([0., lensy, 2.]),
+    'D': np.array([0., 3*np.pi/2., 0.]),
+    'Diam': 0.8,
+    'kappa': 0.,
+    'glass': 'ZF4 cdgm',  # Modified with glass
+    'c': 0.3
+}
+
+lens2 = {
+    'action': 'refraction',
+    'P': np.array([0., lensy - thickness, 2.]),
+    'D': np.array([0., 3*np.pi/2., 0.]),
+    'Diam': 0.8,
+    'kappa': 0.,
+    'c': -0.3
+}
+
+def test_plotting_functions():
+    geo = [selector, mirror, small_mirror, lens1, lens2, stop]
+    ray_group = tp.ray_plane(geo, [0., 0., 0.], 1.8, d=[0., 0., 1.], nrays=50, wvl=.55)
+
+    # Test plotyz
+    try:
+        tp.plotyz(geo, ray_group)
+    except Exception as e:
+        assert False, f"plotyz raised an exception: {e}"
+
+    # Test plotxz
+    try:
+        tp.plotxz(geo, ray_group)
+    except Exception as e:
+        assert False, f"plotxz raised an exception: {e}"
+
+    # Test spotdiagram
+    try:
+        tp.spotdiagram(geo, ray_group)
+    except Exception as e:
+        assert False, f"spotdiagram raised an exception: {e}"
+
+    # Test plotobject
+    try:
+        tp.plotobject(geo, ray_group)
+    except Exception as e:
+        assert False, f"plotobject raised an exception: {e}"
+

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. """