Refactor code for clarity

finesse/tracing/tracesolution.py

@@ -315,9 +315,12 @@ class UtilityTraceSolution(BaseTraceSolution):
         data = self._data
         gouy = 0
 
+        # Position of last component.
+        llast = sum([data[comp]['L'] for comp in data['components'] if data[comp]['is_space']])
+
         if lstop is None:
             # Use last component position.
-            lstop = sum([data[comp]['L'] for comp in data['components'] if data[comp]['is_space']])
+            lstop = llast
 
         # Distance along optical axis.
         zpath = np.linspace(lstart, lstop, npoints)
@@ -330,39 +333,36 @@ class UtilityTraceSolution(BaseTraceSolution):
         wpath[0] = first_node['q'].beamsize()
         gpath[0] = 0
 
+        def compute(dl, q, z, gref):
+            # Indices of zpath which lie in this space.
+            zind = np.where(np.logical_and(zpath > z, zpath <= (z + dl)))
+            # Distance steps from zero for this space.
+            dz = zpath[zind] - z
+
+            g = np.degrees(q.gouy(q.z + dz))
+            w = q.beamsize(q.z + dz)
+            wpath[zind] = w
+            # want to plot accumulated gouy phase so need to use
+            # a reference from where it started
+            gpath[zind] = g + gref
+
         if lstart < 0:
             raise NotImplementedError("backwards propagation not yet implemented")
 
-        last_space = [comp for comp in data['components'] if data[comp]['is_space']][-1]
+        space_params = []
 
+        # Build parameters from spaces between start and end node.
         for comp, from_node in zip(data['components'], data['nodes']):
-            if data[comp]['is_space']:
-                gouy = data[comp]['gouy_i']
-                gouy_ref = data[comp]['gouy_ref']
-
-                # Space length.
-                L = data[comp]['L']
-                # Parameters at the space's start node.
-                q = data[from_node]['q']
-                # Distance from start node.
-                z = data[from_node]['z']
-
-                # Index selection criteria for this space.
-                isel = zpath > z
-                if comp != last_space:
-                    isel = np.logical_and(isel, zpath <= (z + L))
-                # Indices of zpath which lie in this space.
-                zsel = np.where(isel)
-                # Distance steps from zero for this space.
-                _z = zpath[zsel] - z
-
-                g = np.degrees(q.gouy(_z+q.z))
-                # want to plot accumulated gouy phase so need to use
-                # a reference from where it started
-                _g = gouy + g - gouy_ref
-                w = q.beamsize(_z + q.z)
-                wpath[zsel] = w
-                gpath[zsel] = _g
+            if not data[comp]['is_space']:
+                continue
+            L = data[comp]['L']
+            q = data[from_node]['q']
+            z = data[from_node]['z']
+            gref = data[comp]['gouy_i'] - data[comp]['gouy_ref']
+            space_params.append((L, q, z, gref))
+
+        for params in space_params:
+            compute(*params)
 
         return zpath, wpath, gpath