Commit ff87fabd authored by Sean Leavey's avatar Sean Leavey

Refactor code for clarity

parent 50815510
......@@ -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
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment