Optimize optika.materials.snells_law() using numexpr (#142)

Author: jacobdparker

optika/materials/_snells_law.py

@@ -272,19 +272,21 @@ def snells_law(
         \right) \hat{\mathbf{n}} \right]}
     """
     a = direction
-    n1 = index_refraction
-    n2 = index_refraction_new
+    n1 = index_refraction  # noqa: F841
+    n2 = index_refraction_new  # noqa: F841
 
     if normal is None:
         normal = na.Cartesian3dVectorArray(0, 0, -1)
 
-    r = n1 / n2
-
-    c = -a @ normal
-
-    mirror = np.sign(c) * (2 * is_mirror - 1)
-
-    t = np.sqrt(np.square(1 / r) + np.square(c) - np.square(a.length))
-    b = r * (a + (c + mirror * t) * normal)
-
-    return b
+    a_x = a.x  # noqa: F841
+    a_y = a.y  # noqa: F841
+    a_z = a.z  # noqa: F841
+    u_x = normal.x  # noqa: F841
+    u_y = normal.y  # noqa: F841
+    u_z = normal.z  # noqa: F841
+
+    return na.numexpr.evaluate(
+        "(n1 / n2) * (a + (-(a_x*u_x + a_y*u_y + a_z*u_z) + sign(-(a_x*u_x + a_y*u_y + a_z*u_z)) "
+        "* (2 * is_mirror - 1) * sqrt(1 / (n1 / n2)**2 + (a_x*u_x + a_y*u_y + a_z*u_z)**2"
+        "- (a_x*a_x + a_y*a_y + a_z*a_z))) * normal)"
+    )

optika/materials/_tests/test_snells_law.py

@@ -57,6 +57,7 @@ def test_snells_law_scalar(
     argnames="direction",
     argvalues=[
         na.Cartesian3dVectorArray(0, 0, 1),
+        optika.direction(na.Cartesian2dVectorArray(1 * u.deg, 2 * u.deg)),
     ],
 )
 @pytest.mark.parametrize(
@@ -101,12 +102,28 @@ def test_snells_law(
         normal=normal,
         is_mirror=is_mirror,
     )
+
+    a = direction
+    n1 = index_refraction  # noqa: F841
+    n2 = index_refraction_new  # noqa: F841
+
     if normal is None:
         normal = na.Cartesian3dVectorArray(0, 0, -1)
 
+    r = n1 / n2
+
+    c = -a @ normal
+
+    mirror = np.sign(c) * (2 * is_mirror - 1)
+
+    t = np.sqrt(np.square(1 / r) + np.square(c) - np.square(a.length))
+    result_expected = r * (a + (c + mirror * t) * normal)
+
     assert isinstance(result, na.AbstractCartesian3dVectorArray)
     assert np.allclose(result.length, 1)
     if is_mirror:
         assert not np.allclose(np.sign(direction @ normal), np.sign(result @ normal))
     else:
         assert np.allclose(np.sign(direction @ normal), np.sign(result @ normal))
+
+    assert np.allclose(result, result_expected)