Repository: DCC-Lab/RayTracing

Issue:
ImagingPath and LaserPath methods in MatrixGroup raise exception.
They respectively return a ImagingPath instance and a LaserPath instance, but those classes are in other files not imported. It is also worth mentioning that they both inherit from MatrixGroup, so it is kind of incoherent to use them in MatrixGroup.

Question:
Are those methods really necessary in MatrixGroup?
- Yes: We have to find a way to import them without having a circular import problem.
- No: We should remove them.

Comments:

Patch:
diff --git a/raytracing/laserpath.py b/raytracing/laserpath.py
index edcf1b2..ec4fd0c 100644
--- a/raytracing/laserpath.py
+++ b/raytracing/laserpath.py
@@ -62,7 +62,7 @@ class LaserPath(MatrixGroup):
 
         return q
 
-    def display(self, inputBeam=None, inputBeams=None, comments=None):
+    def display(self, inputBeam=None, inputBeams=None, comments=None):  # pragma: no cover
         """ Display the optical system and trace the laser beam. 
         If comments are included they will be displayed on a
         graph in the bottom half of the plot.
@@ -71,7 +71,7 @@ class LaserPath(MatrixGroup):
 
         if self.isResonator:
             beams = self.laserModes()
-            if len(self.label) == "":
+            if self.label == "":
                 self.label = "Laser modes as calculated"
         elif inputBeam is not None:
             beams = [inputBeam]
@@ -80,7 +80,7 @@ class LaserPath(MatrixGroup):
         else:
             beams = [self.inputBeam]
 
-        if len(self.label) == "":
+        if self.label == "":
             self.label = "User-specified gaussian beams"
 
         if comments is not None:
@@ -95,7 +95,7 @@ class LaserPath(MatrixGroup):
 
         self._showPlot()
 
-    def createBeamTracePlot(self, axes, beams):
+    def createBeamTracePlot(self, axes, beams):  # pragma: no cover
         """ Create a matplotlib plot to draw the laser beam and the elements.
         """
 
@@ -122,7 +122,7 @@ class LaserPath(MatrixGroup):
             y.append(ray.w)
         return (x, y)
 
-    def drawBeamTrace(self, axes, beam):
+    def drawBeamTrace(self, axes, beam):  # pragma: no cover
         """ Draw beam trace corresponding to input beam 
         Because the laser beam diffracts through space, we cannot
         simply propagate the beam over large distances and trace it
@@ -148,7 +148,7 @@ class LaserPath(MatrixGroup):
         axes.plot(x, y, 'r', linewidth=1)
         axes.plot(x, [-v for v in y], 'r', linewidth=1)
 
-    def drawWaists(self, axes, beam):
+    def drawWaists(self, axes, beam):  # pragma: no cover
         """ Draws the expected waist (i.e. the focal spot or the spot where the
         size is minimum) for all positions of the beam. This will show "waists" that
         are virtual if there is an additional lens between the beam and the expceted
diff --git a/raytracing/matrixgroup.py b/raytracing/matrixgroup.py
index fe54cb8..b34038c 100644
--- a/raytracing/matrixgroup.py
+++ b/raytracing/matrixgroup.py
@@ -10,6 +10,7 @@ class MatrixGroup(Matrix):
     """
 
     def __init__(self, elements=None, label=""):
+        self.iteration = 0
         super(MatrixGroup, self).__init__(1, 0, 0, 1, label=label)
 
         self.elements = []
@@ -36,7 +37,7 @@ class MatrixGroup(Matrix):
         if len(self.elements) != 0:
             lastElement = self.elements[-1]
             if lastElement.backIndex != matrix.frontIndex:
-                if isinstance(matrix, Space): # For Space(), we fix it
+                if isinstance(matrix, Space):  # For Space(), we fix it
                     msg = "Fixing mismatched indices between last element and appended Space(). Use Space(d=someDistance, n=someIndex)."
                     warnings.warn(msg, UserWarning)
                     matrix.frontIndex = lastElement.backIndex
@@ -55,12 +56,6 @@ class MatrixGroup(Matrix):
         self.frontVertex = transferMatrix.frontVertex
         self.backVertex = transferMatrix.backVertex
 
-    def ImagingPath(self):
-        return ImagingPath(elements=self.elements, label=self.label)
-
-    def LaserPath(self):
-        return LaserPath(elements=self.elements, label=self.label)
-
     def transferMatrix(self, upTo=float('+Inf')):
         """ The transfer matrix between front edge and distance=upTo
 
@@ -227,3 +222,16 @@ class MatrixGroup(Matrix):
             axes.annotate(label, xy=(z, 0.0), xytext=(z, -halfHeight * 0.5),
                           xycoords='data', fontsize=12,
                           ha='center', va='bottom')
+
+    def __iter__(self):
+        self.iteration = 0
+        return self
+
+    def __next__(self):
+        if self.elements is None:
+            raise StopIteration
+        if self.iteration < len(self.elements):
+            element = self.elements[self.iteration]
+            self.iteration += 1
+            return element
+        raise StopIteration


Test patch:
diff --git a/raytracing/tests/testsLaserPath.py b/raytracing/tests/testsLaserPath.py
index 6b90b03..38134af 100644
--- a/raytracing/tests/testsLaserPath.py
+++ b/raytracing/tests/testsLaserPath.py
@@ -7,20 +7,75 @@ inf = float("+inf")
 
 class TestLaserPath(unittest.TestCase):
 
-    def testEigenModes(self):
+    def testLaserPathNoElements(self):
+        lasPath = LaserPath()
+        self.assertIsNone(lasPath.inputBeam)
+        self.assertFalse(lasPath.isResonator)
+        self.assertTrue(lasPath.showElementLabels)
+        self.assertTrue(lasPath.showPointsOfInterest)
+        self.assertTrue(lasPath.showPointsOfInterestLabels)
+        self.assertTrue(lasPath.showPlanesAcrossPointsOfInterest)
+        self.assertListEqual(lasPath.elements, [])
+
+    def testLaserPath(self):
+        elements = [Space(5), Lens(5), Space(20), Lens(15), Space(15)]
+        lp = LaserPath(elements, "Laser Path")
+        self.assertListEqual(lp.elements, elements)
+        self.assertEqual(lp.label, "Laser Path")
+
+    def testLaserPathIncorrectElements(self):
+        elements = [Ray(), Lens(10)]
+        with self.assertRaises(TypeError):
+            LaserPath(elements)
+
+    def testEigenModesNoPower(self):
         lp = LaserPath([Space(10)])
         self.assertTupleEqual(lp.eigenModes(), (None, None))
 
-        lp = LaserPath()
-        self.assertTupleEqual(lp.eigenModes(), (None, None))
+    def testEigenModes(self):
+        lp = LaserPath([Space(10), Lens(10)])
+        beam1, beam2 = lp.eigenModes()
+        self.assertEqual(beam1.q.real, -5)
+        self.assertEqual(beam2.q.real, -5)
+        self.assertAlmostEqual(beam1.q.imag, -5 * 3 ** 0.5)
+        self.assertAlmostEqual(beam2.q.imag, 5 * 3 ** 0.5)
 
-        lp = LaserPath([CurvedMirror(-10)])
-        self.assertNotEqual(lp.eigenModes(), (None, None))
+        lp = LaserPath([Lens(10)])
+        beam1, beam2 = lp.eigenModes()
+        self.assertEqual(beam1.q, beam2.q)
+        self.assertEqual(beam1.q, 0)
+
+    def testLaserModesNoPower(self):
+        lp = LaserPath([Space(10)])
+        self.assertListEqual(lp.laserModes(), [])
 
     def testLaserModes(self):
+        lp = LaserPath([Space(10), Lens(10)])
+        laserModes = lp.laserModes()
+        self.assertEqual(len(laserModes), 1)
+        beam = laserModes[0]
+        self.assertEqual(beam.q.real, -5)
+        self.assertAlmostEqual(beam.q.imag, 5 * 3 ** 0.5)
+
+        elements = [Space(1, 1.33), DielectricInterface(1.33, 1, 1), ThickLens(1.33, -10, -5, -20)]
+        lp = LaserPath(elements)
+        laserModes = lp.laserModes()
+        self.assertEqual(len(laserModes), 1)
+        beam = laserModes[0]
+        self.assertAlmostEqual(beam.q.real, -5.90770102)
+        self.assertAlmostEqual(beam.q.imag, 1.52036515)
+
+        lp = LaserPath([Space(10), CurvedMirror(5)])
+        self.assertListEqual(lp.laserModes(), [])
         lp = LaserPath()
         self.assertListEqual(lp.laserModes(), [])
 
+    def testRearrangeBeamTraceForPlotting(self):
+        x = [x for x in range(1, 6)]
+        y = [y for y in range(1, 6)]
+        rayList = [GaussianBeam(w=x_, z=y_) for (x_, y_) in zip(x, y)]
+        lp = LaserPath()
+        self.assertTupleEqual(lp.rearrangeBeamTraceForPlotting(rayList), (x, y))
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/raytracing/tests/testsMatrixGroup.py b/raytracing/tests/testsMatrixGroup.py
index 6d587cb..d473aa7 100644
--- a/raytracing/tests/testsMatrixGroup.py
+++ b/raytracing/tests/testsMatrixGroup.py
@@ -319,6 +319,11 @@ class TestMatrixGroup(unittest.TestCase):
         self.assertEqual(mg.D, supposedMatrix.D)
         self.assertEqual(mg.L, supposedMatrix.L)
 
+    def testInitWithAnotherMatrixGroup(self):
+        mg = MatrixGroup([Lens(5)])
+        mg2 = MatrixGroup(mg)
+        self.assertListEqual(mg.elements, mg2.elements)
+
 
 if __name__ == '__main__':
     unittest.main()