Source code for finesse.solutions.beamtrace

"""Solution objects for beam propagations."""

import logging
import typing
from functools import partial

import numpy as np

from ..components import Cavity, Space
from ..components.general import Connector
from ..components.node import OpticalNode, Port
from ..env import is_interactive, warn
from ..exceptions import FinesseException
from ..gaussian import BeamParam
from ..solutions.base import BaseSolution
from ..symbols import Symbol, evaluate
from ..utilities.misc import is_iterable, pairwise
from ..utilities.tables import NumberTable
from ..utilities.text import scale_si

LOGGER = logging.getLogger(__name__)


[docs]class ABCDSolution(BaseSolution): """Solution for a composite ABCD calculation. Parameters ---------- M : array-like or two-element tuple of array-like The ABCD matrix / matrices. direction : str Direction / plane of computation. ``"both"`` indicates that `M` is a tuple of the ABCD matrices computed over both the tangential and sagittal planes. ``"x"`` implies M is the ABCD matrix computed over the tangential plane. ``"y"`` implies M is the ABCD matrix computed over the sagittal plane. symbolic : bool Flag indicating whether the calculations are symbolic. """
[docs] def __init__(self, name, M, direction, symbolic): super().__init__(name) self.empty = False # tree drawing fill circle self._M = M self.__direction = direction self.__symbolic = symbolic
def __str__(self): return str(self._M) @property def direction(self): """The plane in which this ABCD matrix was computed - 'x' for tangential, 'y' for sagittal. :`getter`: Returns the ABCD matrix plane. """ return self.__direction @property def symbolic(self): """Indicates whether this ABCD solution is symbolic. :`getter`: Returns `True` if this stores symbolic expressions, `False` otherwise. """ return self.__symbolic @property def M(self): """A copy of the underlying ABCD matrix as a :class:`numpy.ndarray`. :`getter`: Returns a copy of underlying ABCD matrix. """ return self._M.copy()
[docs] def eval(self): """Evaluate the symbolic ABCD matrix. Computes the numeric form of the ABCD matrix using the ``eval`` method of each parameter reference. Returns ------- out : :class:`numpy.ndarray` A numeric matrix for the evaluated ABCD. """ if not self.symbolic: return self.M return evaluate(self.M)
[docs]class PropagationSolution(BaseSolution): """Solution representation of a call to :func:`~finesse.tracing.tools.propagate_beam`. This class contains useful attributes and methods for accessing properties of the beam that was propagated through the path specified by the above function call. If this propagation call was symbolic then each property returned from this class will also be symbolic - evaluate these using the ``eval`` method of the symbolic expression. Note that PropagationSolution objects are returned via :func:`~finesse.tracing.tools.propagate_beam` (or :meth:`.Model.propagate_beam`), they should never need to be created manually. See :ref:`propagating_beams` for details and examples on using this class. """
[docs] def __init__(self, name, node_info, comp_info, symbolic): super().__init__(name) self.node_info = node_info self.comp_info = comp_info self.__symbolic = symbolic self.total_acc_gouy = 0 for info in comp_info.values(): self.total_acc_gouy += info.get("acc_gouy", 0) self.empty = False # NOTE (sjr) For non-symbolic propagations, we need to store the fixed # space length as it was at the time of the propagate_beam # call to avoid any issues (when plotting the solution) if a # space length is changed in the model afterwards if not self.symbolic: self.__frozen_space_lengths = { space: evaluate(space.L.value) for space in self.spaces } else: self.__frozen_space_lengths = {}
def __get_model(self): nodes = list(self.node_info.keys()) if not nodes: return None return nodes[0]._model def __get_component(self, x): if isinstance(x, str): name = x model = self.__get_model() if model is None: raise RuntimeError("Empty PropagationSolution - no path traversed!") x = model.elements.get(name) if x is None: raise ValueError(f"No component of name {name} exists in the model") if not isinstance(x, Connector): raise TypeError( "Invalid type for component. Expected Connector " f"but got {type(x)}" ) if x not in self.comp_info: raise ValueError( f"No component of name {x.name} exists in the propagated path" ) return x def __get_node(self, x: str): if len(x.split(".")) != 3: raise ValueError(f"Unexpected node name format: {x}") model = self.__get_model() if model is None: raise RuntimeError("Empty PropagationSolution - no path traversed!") if not model.network.has_node(x): raise ValueError(f"No node of name {x} exists in the model.") return model.network.nodes[x]["weakref"]() @property def symbolic(self): """Whether the propagation solution is symbolic.""" return self.__symbolic @property def start_node(self): """The starting node of the propagation.""" nodes = list(self.node_info.keys()) if not nodes: return None return nodes[0] @property def end_node(self): """The final node of the propagation.""" nodes = list(self.node_info.keys()) if not nodes: return None return nodes[-1] @property def nodes(self): """A list of all the nodes traversed, in order.""" return list(self.node_info.keys()) @property def ports(self): """A list of all the ports traversed, in order.""" return list(dict.fromkeys(node.port for node in self.nodes)) @property def spaces(self): """A list of all spaces traversed, in order.""" space_entries = dict( filter(lambda x: "acc_gouy" in x[1], self.comp_info.items()) ) return list(space_entries.keys()) @property def components(self): """A list of all components (excluding spaces) traversed, in order.""" return list(filter(lambda x: not isinstance(x, Space), self.comp_info.keys())) @property def positions(self): """A dictionary of the :class:`.Connector` instances to their positions (relative to the start node).""" pos = {} for node, info in self.node_info.items(): if node.component in pos: continue pos[node.component] = info["z"] return pos @property def path_length(self): """The geometric path length of the traversed path. Equal to the sum of each space length in the path. """ nodes = list(self.node_info.keys()) if not nodes: return 0 return self.node_info[nodes[-1]]["z"] @property def optical_path_length(self): """The optical path length of the traversed path. Equal to the sum of the product of each space length and refractive index in the path. """ nodes = self.nodes if not nodes: return 0 return self.node_info[nodes[-1]]["z_optical"] @property def beamsizes(self): """Dictionary of node to beam size mappings.""" return {node: info["q"].w for node, info in self.node_info.items()} @property def ws(self): """Identical to :attr:`PropagationSolution.beamsizes`""" return self.beamsizes @property def waistpositions(self): """Dictionary of node to waist position (as measured from node) mappings.""" return {node: info["q"].z for node, info in self.node_info.items()} @property def z0s(self): """Identical to :attr:`PropagationSolution.waistpositions`""" return self.waistpositions @property def qs(self): """Dictionary of node to beam parameter mappings.""" return {node: info["q"] for node, info in self.node_info.items()} @property def full_ABCD(self): """The full, composite ABCD matrix from the start to the end of the path.""" last_node_info = list(self.node_info.values())[-1] return last_node_info["ABCD"]
[docs] def acc_gouy(self, *args): """Accumulated Gouy phase over a sequence of spaces. Parameters ---------- args : sequence of args Space components or names of spaces. Returns ------- agouy : float or :class:`.Function` The accumulated Gouy phase over the given spaces. """ agouy = 0 model = self.__get_model() if model is None: return agouy for space in args: if isinstance(space, str): space = model.elements.get(space) agouy += self.comp_info[space]["acc_gouy"] return agouy
[docs] def acc_gouy_up_to(self, point): """Accumulated Gouy phase up to a `point` in the traversed path. This computes the cumulative Gouy phase from the :attr:`.PropagationSolution.start_node` to the specified `point`. Parameters ---------- point : :class:`.OpticalNode` or :class:`.Port` or :class:`.Connector` or str A node, port, component or name of component up to which to compute the accumulated Gouy phase. Returns ------- agouy : float or :class:`.Function` The accumulated Gouy phase up to the given `point`. """ if isinstance(point, (OpticalNode, Port)): point = point.component point = self.__get_component(point) agouy = 0 for comp, info in self.comp_info.items(): if comp == point: break agouy += info.get("acc_gouy", 0) return agouy
[docs] def abcd(self, up_to=None): """Composite ABCD matrix up to a specific point in the path. Parameters ---------- up_to : :class:`.OpticalNode` or :class:`.Port`, str, optional The location in the path at which to get the composite ABCD matrix. This can be an optical node or a port. When `None` it will be the `to_node` of the propgagtion. Returns ------- M : :class:`numpy.ndarray` The ABCD matrix, as a NumPy array, computed up to the specified location. Raises ------ ke : KeyError If `up_to` is a node or port which does not exist within the solution. te : TypeError If `up_to` is not an :class:`.OpticalNode` or :class:`.Port`. """ if up_to is None: up_to = self.end_node if isinstance(up_to, str): try: up_to = next(filter(lambda node: node.full_name == up_to, self.nodes)) except StopIteration: raise FinesseException( f"Could not find a node in this propagation with the name `{up_to}`" ) if isinstance(up_to, OpticalNode): if up_to not in self.node_info: raise KeyError( f"No optical node of name {up_to.full_name} in the solution." ) M = self.node_info[up_to]["ABCD"] elif isinstance(up_to, Port): ientry = self.node_info.get(up_to.i, {}) oentry = self.node_info.get(up_to.o, {}) if not ientry and not oentry: raise KeyError(f"No port of name {up_to.full_name} in the solution.") in_M = ientry.get("ABCD", None) out_M = oentry.get("ABCD", None) if in_M is not None: M = in_M elif out_M is not None: M = out_M else: raise RuntimeError( "Bug encountered! No composite ABCD matrix entry found " f"at the port {up_to.full_name}." ) ## TODO (sjr) Add support for up_to being a Connector (or name of Connector) else: raise TypeError(f"Invalid type for up_to: {up_to}") return M
[docs] def q(self, at) -> BeamParam: """Beam parameter at a given node of the path. Parameters ---------- at : :class:`.OpticalNode` or str The location in the path at which to get the beam parameter. This can be an optical node or the name of the optical node. Returns ------- q : :class:`.BeamParam` The beam parameter corresponding to the specified node. Raises ------ ke : KeyError If `at` is a node which does not exist within the solution. ve : ValueError If `at` is a string corresponding to a node which doesn't exist in the associated model. te : TypeError If `at` is not an :class:`.OpticalNode` or a string. """ if isinstance(at, str): at = self.__get_node(at) if isinstance(at, OpticalNode): if at not in self.node_info: if at.opposite in self.node_info: # If the opposite node direction is in the solution # then return the reverse of this beam parameter at # this node (i.e. -q*) q = self.node_info[at.opposite]["q"].reverse() else: raise KeyError( f"No optical node of name {at.full_name} (nor its opposite) " f"in the solution." ) else: q = self.node_info[at]["q"] else: raise TypeError(f"Invalid type for at: {at}") return q
[docs] def beamsize(self, at): """Beam radius at a given location of the path. Parameters ---------- See :meth:`.PropagationSolution.q`. Returns ------- w : float or :class:`.Function` Beam size corresponding to the specified node. Raises ------ See :meth:`.PropagationSolution.q`. """ return self.q(at).w
[docs] def w(self, at): """Identical to :meth:`PropagationSolution.beamsize`.""" return self.beamsize(at)
[docs] def waistsize(self, at): """Waist radius as measured from a given location of the path. Parameters ---------- See :meth:`.PropagationSolution.q`. Returns ------- w : float or :class:`.Function` Waist size using the beam parameter basis at the specified node. Raises ------ See :meth:`.PropagationSolution.q`. """ return self.q(at).w0
[docs] def w0(self, at): """Identical to :meth:`PropagationSolution.waistsize`.""" return self.waistsize(at)
[docs] def waistpos(self, from_point): """Waist position as measured at `from_point`. Parameters ---------- See :meth:`.PropagationSolution.q`. Returns ------- w : float or :class:`.Function` Distance to the waist from the specified node. Raises ------ See :meth:`.PropagationSolution.q`. """ return self.q(from_point).z
[docs] def z0(self, from_point): """Identical to :meth:`PropagationSolution.waistpos`.""" return self.waistpos(from_point)
def __getitem__(self, key): if isinstance(key, OpticalNode): return self.node_info[key] if isinstance(key, Connector) or isinstance(key, str): key = self.__get_component(key) return self.comp_info[key] raise TypeError(f"Invalid type for key: {key}")
[docs] def compute_distances_matrix(self, ztype="geometric", subs=None): """Compute the distances between each optic, relative to each other. Returns a dict of dicts for each "delta z". Note that each distance value is in metres. Use :meth:`PropagationSolution.distances_matrix_table` to create a tabulated representation of this dict. Parameters ---------- ztype : str, optional; default: "geometric" Type of distance, can be either 'geometric' or 'optical'. In the former case the values are the distances between each optic in terms of sums of space lengths of each space between them. In the latter case, each value is instead the optical path length between each component, i.e. the sum of the product of the space length and refractive index of each space between them. subs : dict, optional A dictionary of model parameter to value substitutions to pass to the ``eval`` methods of symbolic expressions. If this solution object is not symbolic then this argument is ignored. Returns ------- deltas : dict Dict of dicts for each dz between components. """ if not self.symbolic and subs is not None: warn(f"Ignoring {subs=} kwarg as PropagationSolution is non-symbolic.") deltas = {} outer_done = set() for node1, info1 in self.node_info.items(): comp1 = node1.component if comp1 in outer_done: continue inner_done = set() deltas[comp1] = {} for node2, info2 in self.node_info.items(): comp2 = node2.component if comp2 in inner_done: continue z2 = self._eval_z_for_display(info2, ztype=ztype, subs=subs) z1 = self._eval_z_for_display(info1, ztype=ztype, subs=subs) deltas[comp1][comp2] = z2 - z1 inner_done.add(comp2) outer_done.add(comp1) return deltas
[docs] def distances_matrix_table( self, ztype="geometric", subs=None, numfmt=None, **kwargs ): """Returns the distances between each optic, relative to each other, in a table. Parameters ---------- ztype : str, optional; default: "geometric": See :meth:`.PropagationSolution.compute_distances_matrix`. subs : dict, optional A dictionary of model parameter to value substitutions to pass to the ``eval`` methods of symbolic expressions. If this solution object is not symbolic then this argument is ignored. numfmt : str or func, optional Either a function to format numbers or a formatting string. The function must return a string. Defaults to using an SI scale function. **kwargs : dict, optional Additional arguments to create the table. See `finesse.utilities.tables.NumberTable` for further documentation. The arguments `table`, `colnames` and `rownames` are already used and should not be passed. """ deltas = self.compute_distances_matrix(ztype=ztype, subs=subs) table = [list(deltas[comp].values()) for comp in deltas] colnames = [comp.name for comp in self.components] rownames = [comp.name for comp in deltas] if numfmt is None: numfmt = partial(scale_si, units="m") return NumberTable( table, colnames=colnames, rownames=rownames, numfmt=numfmt, **kwargs, )
[docs] def position(self, point, ptype="geometric"): """Gets the position of the specified `point` relative to the start node. Parameters ---------- point : :class:`.OpticalNode` or :class:`.Port` or :class:`.Connector` or str The location in the path from which to obtain the relative position. This can be an optical node, a port, a connector or the name of a connector. ptype : str Type of distance, can be either 'geometric' or 'optical'. In the former case the value returned is the distance from the start node to `point` as a sum of each space length. In the latter case the value returned will be the optical path length from the start node to `point`, i.e. a sum of each space length multiplied by its refractive index. Returns ------- z : float or symbol The relative distance from the start node to the measured `point`. """ if ptype.casefold() == "geometric": zentry = "z" elif ptype.casefold() == "optical": zentry = "z_optical" else: raise ValueError( f"Unrecognised value for ptype: {ptype}. This must be either " "'geometric' or 'optical'." ) if isinstance(point, OpticalNode): if point not in self.node_info: raise KeyError( f"No optical node of name {point.full_name} in the solution." ) else: z = self.node_info[point][zentry] elif isinstance(point, Port): ientry = self.node_info.get(point.i, {}) oentry = self.node_info.get(point.o, {}) if not ientry and not oentry: raise KeyError(f"No port of name {point.full_name} in the solution.") z1 = ientry.get(zentry, None) z2 = oentry.get(zentry, None) if z1 is not None and z2 is not None: if z1 != z2: raise RuntimeError( f"Bug encountered! Positions of nodes {point.i.full_name} " f"and {point.o.full_name} not equal ({z1} != {z2})" ) z = z1 else: z = z1 or z2 elif isinstance(point, Connector) or isinstance(point, str): point = self.__get_component(point) z = None for node in point.optical_nodes: if node in self.node_info: z = self.node_info[node][zentry] break if z is None: raise KeyError(f"No component of name {point.name} in the solution.") else: raise TypeError("Unrecognised type for argument point.") return z
[docs] def table(self, subs=None, numfmt=None, **kwargs): """Construct a table showing the beam properties at each node. Parameters ---------- subs : dict, optional A dictionary of model parameter to value substitutions to pass to the ``eval`` methods of symbolic expressions. If this solution object is not symbolic then this argument is ignored. numfmt : str or func, optional Either a function to format numbers or a formatting string. The function must return a string. Defaults to "{:.3f}" for the first 7 columns and `lambda q: f"{q.real:.3f} + {q.imag:.3f}j"` for the last one. **kwargs : dict, optional Additional arguments to create the table. See `finesse.utilities.tables.NumberTable` for further documentation. The arguments `table`, `colnames` and `rownames` are already used and should not be passed. Returns ------- table : finesse.utilities.NumberTable """ if not self.symbolic and subs is not None: warn(f"Ignoring {subs=} kwarg as PropagationSolution is non-symbolic.") table = [] names = [] headers = ["z", "w0", "zr", "w", "RoC", "S", "Acc. Gouy", "q"] if numfmt is None: numfmt = ( [partial(scale_si, units="m")] * 5 + [partial(scale_si, units="D")] + [partial(scale_si, units="°")] + [lambda q: f"{q.real:.3f} + {q.imag:.3f}j"] ) agouy = 0 for node, info in self.node_info.items(): z = self._eval_z_for_display(info, subs=subs) q = self._eval_q_for_display(info, subs=subs) if node.is_input: # accumulate the Gouy phase space_info = self.comp_info.get(node.space) if space_info: agouy += space_info.get("acc_gouy", 0) if isinstance(agouy, Symbol): agouy_val = agouy.eval(subs=subs) else: agouy_val = agouy table.append([z, q.w0, q.zr, q.w, q.Rc, q.S, agouy_val, q]) names.append(node.full_name) return NumberTable(table, headers, names, numfmt=numfmt)
def __str__(self): return str(self.table())
[docs] def print(self, subs=None, numfmt=None, **kwargs): """Print the propagated beam properties at each node in a table format. This internally calls :meth:`.PropagationSolution.table` and prints its return string. Parameters ---------- subs : dict, optional A dictionary of model parameter to value substitutions to pass to the ``eval`` methods of symbolic expressions. If this solution object is not symbolic then this argument is ignored. numfmt : str or func, optional Either a function to format numbers or a formatting string. The function must return a string. Defaults to "{:.3f}" for the first 6 columns and `lambda q: f"{q.real:.3f} + {q.imag:.3f}j"` for the last one. **kwargs : dict, optional Additional arguments to create the table. See `finesse.utilities.tables.NumberTable` for further documentation. The arguments `table`, `colnames` and `rownames` are already used and should not be passed. """ self.table(subs=subs, numfmt=numfmt, **kwargs).print()
def _eval_z_for_display(self, info, ztype="geometric", subs=None): if ztype.casefold() == "geometric": z = info["z"] elif ztype.casefold() == "optical": z = info["z_optical"] else: raise ValueError(f"Unrecognised ztype: {ztype}") if self.symbolic and isinstance(z, Symbol): if subs is not None: if not isinstance(subs, dict): raise TypeError("Expected argument subs to be of type dict") if any(is_iterable(arg) for arg in subs.values()): raise ValueError( "Parameter substitutions for plotting / printing must " "not use arrays of values." ) z = z.eval(subs=subs) return z def _eval_q_for_display(self, info, subs=None): q = info["q"] if self.symbolic and q.symbolic: if subs is not None: if not isinstance(subs, dict): raise TypeError("Expected argument subs to be of type dict") if any(is_iterable(arg) for arg in subs.values()): raise ValueError( "Parameter substitutions for plotting / printing must " "not use arrays of values." ) q = q.eval(subs=subs) return q def _eval_L_for_display(self, space, subs=None): if self.symbolic: if subs is not None: if not isinstance(subs, dict): raise TypeError("Expected argument subs to be of type dict") if any(is_iterable(arg) for arg in subs.values()): raise ValueError( "Parameter substitutions for plotting / printing must " "not use arrays of values." ) L = space.L.ref.eval(subs=subs) else: L = self.__frozen_space_lengths[space] return L
[docs] def segment(self, node, *args, normalise_z=True, w_scale=1, npts=400, subs=None): """Obtain data for a segment of the beam over the space attached to the specified `node`. The expected, valid positional `args` are any combination of: * "beamsize", * "gouy", * "curvature", where all three will be used by default if none of these are given. Use :meth:`.PropagationSolution.all_segments` to obtain the beam data over all spaces of the solution. Parameters ---------- node : :class:`.OpticalNode` The starting node of the segment. normalise_z : bool, optional; default: True Whether to normalise returned ``data["z"]`` array such that first value of this is zero. w_scale : scalar, optional; default: 1 Quantity to scale beam size values by if calculating these. For example, specify `w_scale = 1e3` to get beam sizes in mm. By default the units of the beam size will be in metres. npts : int, optional; default: 400 Number of points to use for computing data values. subs : dict, optional A dictionary of model parameter to value substitutions to pass to the ``eval`` methods of symbolic expressions. If this solution object is not symbolic then this argument is ignored. Returns ------- zs : :class:`numpy.ndarray` Array of z-axis values corresponding to the position of the node up to the length of the attached space. If `normalise_z` is True then the position of `node` will be subtracted from all values, such that the first value in this array will be zero. data : dict Dictionary of data mapping ``args : values``, where `args` are those specified (see above) and `values` are the arrays of values corresponding to each of these args as a function of the z-axis values. """ if not args: args = ("beamsize", "gouy", "curvature") if not isinstance(node, OpticalNode): raise NotImplementedError( "Obtaining beam segments from locations other than OpticalNodes " "is not yet supported." ) if node.is_input or node is self.end_node: raise ValueError( "Cannot obtain a beam segment corresponding to " f"the node {node.full_name}. This node must be an " "output node and cannot be the end node of the solution." ) info = self.node_info.get(node) if info is None: raise ValueError(f"Node {node.full_name} not present in the solution!") # Store the z data and each arg data array here data = {} z = self._eval_z_for_display(info, subs=subs) q = self._eval_q_for_display(info, subs=subs) # Make array of values from [0, space_length] L = self._eval_L_for_display(node.space, subs=subs) intra_zs = np.linspace(0, L, npts) zs = z + intra_zs if normalise_z: zs -= zs[0] for arg in args: values = getattr(q, arg)(q.z + intra_zs) if arg == "beamsize": values *= w_scale # Computing Gouy phase so convert to degrees and # subtract first value to get Gouy phase accumulated # over the segment rather than absolute phases if arg == "gouy": values = np.degrees(values) values -= values[0] data[arg] = values return zs, data
def __adaptive_npts_per_segment(self, N, subs=None): dS_per_space = {} via_waists = {} end_node = self.end_node for node1, node2 in pairwise(self.nodes): if node1.is_input or node1 is end_node: continue space = node1.space q1 = self._eval_q_for_display(self.node_info[node1], subs=subs) q2 = self._eval_q_for_display(self.node_info[node2], subs=subs) # Compute change in defocus over the space dS_per_space[space] = abs(q2.S - q1.S) # And determine whether we go through a waist via_waists[space] = np.sign(q1.S) != np.sign(q2.S) # Get the maximum defocus change max_dS = max(dS_per_space.values()) for space, via_waist in zip(dS_per_space, via_waists.values()): # If a waist exists in this space then want to automatically # increase the weighting by resetting dS to twice the maximum # so that resolution increases around the waist if via_waist: dS_per_space[space] = max_dS * 2 dS_sum = sum(dS_per_space.values()) npts_per_space = {} for space, dS in dS_per_space.items(): # Now determine the number of points to use for this segment # by the size of the defocus change n = int(N * dS / dS_sum) # If the defocus change was zero, or small enough such that n is # less than two, let's just set the number of points for this # segment to two, this means that total number of points can # be a few larger than N but that's not really a problem if n < 2: n = 2 npts_per_space[space] = n return npts_per_space def __equal_npts_per_segment(self, N, **_): nspaces = len(self.spaces) npts_per_space = {} for space in self.spaces: npts_per_space[space] = int(N / nspaces) return npts_per_space def __all_npts_per_segment(self, N, **_): npts_per_space = {} for space in self.spaces: npts_per_space[space] = N return npts_per_space
[docs] def all_segments( self, *args, add_gouy=True, w_scale=1, npts=1000, resolution="adaptive", subs=None, ): """Construct a dictionary containing beam data for all segments of the solution. The expected, valid positional arguments are any combination of: * "beamsize", * "gouy", * "curvature", where all three will be used by default if no args are given. Use :meth:`.PropagationSolution.segment` to obtain the beam data over a single space of the solution. Parameters ---------- add_gouy : bool, optional; default: True Whether to add the last Gouy phase from previous segment to all values of current segment, thereby constructing each "gouy" array entry as accumulated Gouy phases over all segments. w_scale : scalar, optional; default: 1 Quantity to scale beam size values by if calculating these. For example, specify `w_scale = 1e3` to get beam sizes in mm. By default the units of the beam size will be in metres. npts : int, optional; default: 1000 Number of points to use for computing data values. The actual number of data points used per segment depends upon the `resolution` argument. resolution : str, optional; default: "adaptive" The method of segment resolution setting to use. This can be one of three arguments: * "adaptive": Sets the number of points per segment in such a way as to attempt to increase the resolution near the waist. Each segment will have a number of points allocated to it accordingly, with the total number of points across all segments then approximately equal to `npts`. * "equal": Allocates an equal number of points to each segment, i.e. each segment has ``int(npts / len(self.spaces))`` points. * "all": Gives ``npts`` to all segments, such that the total number of data points across all segments is ``len(self.spaces) * npts``. subs : dict, optional A dictionary of model parameter to value substitutions to pass to the ``eval`` methods of symbolic expressions. If this solution object is not symbolic then this argument is ignored. Returns ------- data : dict Dictionary of data mapping ``space : zs, segdata``, where `space` is each space (i.e. segment) in the solution path, `zs` are the z-axis values and `segdata` is the dict of data values for the targeted beam properties over the space. """ data = {} if not args: args = ("beamsize", "gouy", "curvature") resolution_map = { "adaptive": self.__adaptive_npts_per_segment, "equal": self.__equal_npts_per_segment, "all": self.__all_npts_per_segment, } resolution_method = resolution_map.get(resolution.casefold()) if resolution_method is None: raise ValueError( f"Unexpected value for resolution argument: {resolution}. Expected " "one of: 'adaptive', 'equal' or 'all'." ) npts_per_space = resolution_method(npts, subs=subs) prev_gouy = 0 end_node = self.end_node for node in self.nodes: if node.is_input or node is end_node: continue zs, segdata = self.segment( node, *args, # Don't want to normalise each start z to zero now as # we want each to start from actual node position normalise_z=False, w_scale=w_scale, npts=npts_per_space[node.space], subs=subs, ) # Accumulating Gouy over all segments so make adjustments if "gouy" in args and add_gouy: gouys = segdata["gouy"] # Add on the last value of Gouy from prev segment so # that phase is accumulated over full path gouys += prev_gouy prev_gouy = gouys[-1] data[node.space] = zs, segdata return data
[docs] def plot_beamsizes(self, **kwargs): """Plot the beam sizes over the propagated path. This is just a convenience wrapper which is identical to calling :meth:`PropagationSolution.plot` with ``"beamsize"`` as the arg. Returns ------- fig : Figure Handle to the figure. ax : Axis Handle to the axis. """ return self.plot("beamsize", **kwargs)
[docs] def plot_acc_gouy(self, **kwargs): """Plot the accumulated Gouy phases over the propagated path. This is just a convenience wrapper which is identical to calling :meth:`PropagationSolution.plot` with ``"gouy"`` as the arg. Returns ------- fig : Figure Handle to the figure. ax : Axis Handle to the axis. """ return self.plot("gouy", **kwargs)
[docs] def plot_curvatures(self, **kwargs): """Plot the wavefront curvatures over the propagated path. This is just a convenience wrapper which is identical to calling :meth:`PropagationSolution.plot` with ``"curvature"`` as the arg. Returns ------- fig : Figure Handle to the figure. ax : Axis Handle to the axis. """ return self.plot("curvature", **kwargs)
[docs] def plot( self, *args, filename=None, show=False, ignore=None, name_xoffsets=None, name_yoffsets=None, ylims=None, npts=1000, resolution="adaptive", single_sided=True, subs=None, ): """Plot any combination of the beam sizes, accumulated Gouy phases and / or wavefront curvatures over the propagated path. The expected, valid positional arguments are any combination of: - "beamsize", - "gouy", - "curvature", or "all" to plot all of the above. If no positional args are given then the beamsize (first axis) and accumulated Gouy phase (second axis) will be plotted by default. The locations of each component will be marked on the figure, unless the component is in `ignore` or the component has "AR" or "HR" in its name. Parameters ---------- filename : str or file-like, optional Name of a file or existing file object to save the figure to. show : bool, optional; default: True Whether to show the figure. ignore : component, sequence of, optional A component or sequence of components to ignore when making markers. name_xoffsets : dict, optional Dictionary of component names to x-axis offsets for shifting where the component name text is placed. The offset value is interpreted in terms of data co-ordinates. name_yoffsets : dict, optional Dictionary of component names to y-axis offsets for shifting where the component name text is placed. The offset value is interpreted in terms of data co-ordinates. ylims : dict, optional Dictionary of target names (i.e. "beamsize", "gouy" or "curvature") to manual axis y-limits. npts : int, optional; default: 1000 See equivalent argument in :meth:`.PropagationSolution.all_segments`. resolution : str, optional; default: "adaptive" See equivalent argument in :meth:`.PropagationSolution.all_segments`. single_sided : bool If False the beamsize plot is a single positive line. subs : dict, optional A dictionary of model parameter to value substitutions to pass to the ``eval`` methods of symbolic expressions. If this solution object is not symbolic then this argument is ignored. Returns ------- fig : Figure Handle to the figure. axs : axes The axis handles. """ import matplotlib.pyplot as plt if not args: args = ("beamsize", "gouy") if "all" in args: args = ("beamsize", "gouy", "curvature") valid_args = ("beamsize", "gouy", "curvature") if any(arg not in valid_args for arg in args): raise ValueError( "Invalid target argument in args, expected any " f"combination of {valid_args} or 'all'" ) if not self.symbolic and subs is not None: warn(f"Ignoring {subs=} kwarg as PropagationSolution is non-symbolic.") N = len(args) fig, axs = plt.subplots(N, 1, sharex=True) if N == 1: axs = [axs] maximums = {k: 0 for k in args} minimums = {k: float("inf") for k in args} data = self.all_segments( *args, w_scale=1e3, npts=npts, resolution=resolution, subs=subs ) for zs, segd in data.values(): for i, arg in enumerate(args): y = segd[arg] if arg == "beamsize": if single_sided: axs[i].plot(zs, y, color="r") else: axs[i].fill_between(zs, y, -y, color="r") else: axs[i].plot(zs, y) maximums[arg] = max(maximums[arg], y.max()) minimums[arg] = min(minimums[arg], y.min()) if ignore is None: ignore = [] if not is_iterable(ignore): ignore = [ignore] if name_xoffsets is None: name_xoffsets = {} if name_yoffsets is None: name_yoffsets = {} for node, info in self.node_info.items(): if node.is_input: z = self._eval_z_for_display(info, subs=subs) comp = node.component if "AR" not in comp.name and comp not in ignore: name = comp.name x_offset = name_xoffsets.get(name, 0) y_offset = name_yoffsets.get(name, 0) # display the name in a nicer way name = name.replace("_", "\n") n_newlines = name.count("\n") for i, arg in enumerate(args): if not i: axs[i].axvline( z, 0.12 + 0.1 * n_newlines, color="k", linestyle="--" ) if arg == "beamsize" and not single_sided: ytext_pos = -1 * maximums[arg] else: ytext_pos = minimums[arg] ytext_pos += y_offset axs[i].text( z + x_offset, ytext_pos, name, ha="center", va="center" ) else: axs[i].axvline(z, color="k", linestyle="--") for ax in axs: ax.set_xlim(0, None) axs[-1].set_xlabel("Distance [m]") ylabel_mappings = { "beamsize": "Beam size [mm]", "gouy": "Gouy phase\naccumulation [deg]", "curvature": "Wavefront curvature [1/m]", } if ylims is None: ylims = {} for i, arg in enumerate(args): if arg != "beamsize": ylim = ylims.get(arg, None) if ylim is None: axs[i].set_ylim(0 if arg == "gouy" else None, maximums[arg]) else: axs[i].set_ylim(ylim[0], ylim[1]) ylabel = ylabel_mappings.get(arg, arg) axs[i].set_ylabel(ylabel) if filename is not None: fig.savefig(filename) if show: plt.show() if N == 1: return fig, axs[0] return fig, axs
[docs] def animate_beamsizes(self, subs, **kwargs): """Animate the beam sizes over the propagated path. This is just a convenience wrapper which is identical to calling :meth:`PropagationSolution.animate` with ``"beamsize"`` as the arg. Returns ------- fig : Figure Handle to the figure. ax : Axis Handle to the axis. an : FuncAnimation Handle to the animation. """ return self.animate(subs, "beamsize", **kwargs)
[docs] def animate_acc_gouy(self, subs, **kwargs): """Animate the accumulated Gouy phases over the propagated path. This is just a convenience wrapper which is identical to calling :meth:`PropagationSolution.animate` with ``"gouy"`` as the arg. Returns ------- fig : Figure Handle to the figure. ax : Axis Handle to the axis. an : FuncAnimation Handle to the animation. """ return self.animate(subs, "gouy", **kwargs)
[docs] def animate_curvatures(self, subs, **kwargs): """Animate the wavefront curvatures over the propagated path. This is just a convenience wrapper which is identical to calling :meth:`PropagationSolution.animate` with ``"curvature"`` as the arg. Returns ------- fig : Figure Handle to the figure. ax : Axis Handle to the axis. an : FuncAnimation Handle to the animation. """ return self.animate(subs, "curvature", **kwargs)
[docs] def animate( self, subs, *args, filename=None, show=True, ignore=None, name_xoffsets=None, name_yoffsets=None, ylims=None, npts=200, blit=True, interval=200, ): """Animate any combination of the beam sizes, accumulated Gouy phases and / or wavefront curvatures over the propagated path using the substitution parameters in `subs`. The expected, valid positional arguments (i.e. `*args`) are any combination of: - "beamsize", - "gouy", - "curvature", or "all" to animate all of the above. If no positional args are given then the beamsize (first axis) and accumulated Gouy phase (second axis) will be animated by default. At least one model parameter substitution in `subs` must be array-like - this will then be the animation axis. If more than one are array-like then each array must be the same size - the substitutions will then be carried out simulatenously. Any scalar value entry in `subs` will be applied before the animation axis. Parameters ---------- subs : dict Dictionary of model parameter substitutions. At least one entry must be array-like such than animation can be performed over this axis. If multiple substitutions are arrays then they must all be the same size. filename : str, optional Name of a file to save the animation to. show : bool, optional; default: True Whether to show the resulting animation. ignore : component, sequence of, optional A component or sequence of components to ignore when making markers. name_xoffsets : dict, optional Dictionary of component names to x-axis offsets for shifting where the component name text is placed. The offset value is interpreted in terms of data co-ordinates. name_yoffsets : dict, optional Dictionary of component names to y-axis offsets for shifting where the component name text is placed. The offset value is interpreted in terms of data co-ordinates. ylims : dict, optional Dictionary of target names (i.e. "beamsize", "gouy" or "curvature") to manual axis y-limits. npts : int, optional; default: 200 Number of points to use for computing beam sizes and Gouy phases over spaces. blit : bool, optional; default: True Whether blitting is used to optimize drawing. interval : int, optional; default: 200 Delay between frames in milliseconds. Returns ------- fig : Figure Handle to the figure. axs : axes The axis handles. an : FuncAnimation Handle to the animation. """ import matplotlib.animation as animation import matplotlib.pyplot as plt if not self.symbolic: raise RuntimeError( "Animation is only applicable to symbolic PropagationSolution objects" ) if not args: args = ("beamsize", "gouy") if "all" in args: args = ("beamsize", "gouy", "curvature") valid_args = ("beamsize", "gouy", "curvature") if any(arg not in valid_args for arg in args): raise ValueError( "Invalid target argument in args, expected any " f"combination of {valid_args} or 'all'" ) anim_params = dict(filter(lambda x: is_iterable(x[1]), subs.items())) if not anim_params: raise ValueError("Expected at least one array-like substitution in subs.") N = len(list(anim_params.values())[0]) if any(len(arr) != N for arr in anim_params.values()): raise ValueError("Lengths of all array-like substitutions must be equal.") # Dictionary of singular value substitutions to use at each animation point isubs = {} static_params = dict(filter(lambda x: not is_iterable(x[1]), subs.items())) # Set fixed parameter substitutions for param, value in static_params.items(): isubs[param] = value if ignore is None: ignore = [] if not is_iterable(ignore): ignore = [ignore] # Data dictionaries storing arrays of propagated distances, beamsizes # accumulated Gouy phases and component positions at each node # for each animation point node_xdata = {} node_ydata = {k: {} for k in args} comp_positions = {} for node in self.node_info: if not node.is_input: node_xdata[node] = np.zeros((N, npts)) for arg in args: node_ydata[arg][node] = np.zeros((N, npts)) else: comp = node.component if "AR" not in comp.name and comp not in ignore: comp_positions[node] = np.ma.zeros(N) maximums = {k: 0 for k in args} minimums = {k: float("inf") for k in args} end_node = self.end_node for i in range(N): # Update animation parameters for param, array in anim_params.items(): isubs[param] = array[i] # TODO (sjr) Replace this loop with an all_segments call, will require # some careful rejigging of code in this method so I'll leave # this for a later date prev_acc_gouy = 0 for node, info in self.node_info.items(): if not node.is_input and node != end_node: z = self._eval_z_for_display(info, subs=isubs) q = self._eval_q_for_display(info, subs=isubs) # Make array of values from [0, space_length] L = self._eval_L_for_display(node.space, subs=isubs) intra_zs = np.linspace(0, L, npts) node_xdata[node][i][:] = z + intra_zs for arg in args: data = getattr(q, arg)(q.z + intra_zs) if arg == "beamsize": data /= 1e-3 # scale to mm elif arg == "gouy": # Plotting accumulated Gouy so need to modify data as follows: # - convert to degrees # - subtract the Gouy phase directly at the output node to get # Gouy phase accumulated over this space # - then add previous accumulated Gouy to keep cumulative over full path data = np.degrees(data) data -= data[0] data += prev_acc_gouy prev_acc_gouy = data[-1] maximums[arg] = max(maximums[arg], data.max()) minimums[arg] = min(minimums[arg], data.min()) node_ydata[arg][node][i][:] = data else: if node in comp_positions: comp_positions[node][i] = self._eval_z_for_display( info, subs=isubs ) Np = len(args) fig, axs = plt.subplots(Np, 1, sharex=True) if Np == 1: axs = [axs] # Lists of the various artists to animate collections = [] # beam size fill shapes lines = {k: [] for k in args if k != "beamsize"} comp_pos_lines = [] # component position lines comp_name_texts = [] # names of components for i, arg in enumerate(args): for xdata, ydata in zip(node_xdata.values(), node_ydata[arg].values()): if arg == "beamsize": collections.append( axs[i].fill_between(xdata[0], ydata[0], -ydata[0], color="r") ) else: (line,) = axs[i].plot([], []) lines[arg].append(line) if name_xoffsets is None: name_xoffsets = {} if name_yoffsets is None: name_yoffsets = {} comp_name_offsets = [] for node, zs in comp_positions.items(): name = node.component.name x_offset = name_xoffsets.get(name, 0) y_offset = name_yoffsets.get(name, 0) name = name.replace("_", "\n") n_newlines = name.count("\n") for i, arg in enumerate(args): z0 = zs[0] if not i: vline = axs[i].axvline( z0, 0.12 + 0.1 * n_newlines, color="k", linestyle="--" ) if arg == "beamsize": ytext_pos = -1 * maximums[arg] else: ytext_pos = minimums[arg] ytext_pos += y_offset ct = axs[i].text( z0 + x_offset, ytext_pos, name, ha="center", va="center" ) comp_name_texts.append(ct) comp_name_offsets.append(x_offset) else: vline = axs[i].axvline(z0, color="k", linestyle="--") comp_pos_lines.append(vline) anim_txt_former = lambda k: "\n".join( f"{p.full_name} = {v[k]:.2f} {p.units}" for p, v in anim_params.items() ) anim_text = axs[-1].text( 0.15, 0.9, "", ha="center", va="center", transform=axs[-1].transAxes, ) axs[-1].set_xlabel("Distance [m]") ylabel_mappings = { "beamsize": "Beam size [mm]", "gouy": "Gouy phase\naccumulation [deg]", "curvature": "Wavefront curvature [1/m]", } if ylims is None: ylims = {} for i, arg in enumerate(args): if arg != "beamsize": ylim = ylims.get(arg, None) if ylim is None: axs[i].set_ylim(0 if arg == "gouy" else None, maximums[arg]) else: axs[i].set_ylim(ylim[0], ylim[1]) ylabel = ylabel_mappings.get(arg, arg) axs[i].set_ylabel(ylabel) # Set maximum propagated distance max_x = 0 for xdata in node_xdata.values(): max_x = max(max_x, xdata.max()) for ax in axs: ax.set_xlim(0, max_x) try: w_ax_index = args.index("beamsize") except ValueError: w_ax_index = -1 all_artists = collections + comp_pos_lines + comp_name_texts + [anim_text] for arg in args: if arg == "beamsize": continue all_artists += lines[arg] # Need to create a dummy fig and ax for making a fill_between collection # in each segment -> the actual collections then get modified using the # vertices of this dummy collection during the animation dummy_fig, dummy_ax = plt.subplots() def animate(frame_index): if w_ax_index != -1: for coll, xdata, ydata_w in zip( collections, node_xdata.values(), node_ydata["beamsize"].values(), ): dummy_coll = dummy_ax.fill_between( xdata[frame_index], ydata_w[frame_index], -ydata_w[frame_index], color="r", ) dummy_path = dummy_coll.get_paths()[0] path = coll.get_paths()[0] path.vertices[:, 0] = dummy_path.vertices[:, 0] path.vertices[:, 1] = dummy_path.vertices[:, 1] for arg in args: if arg == "beamsize": continue for line, xdata, ydata in zip( lines[arg], node_xdata.values(), node_ydata[arg].values(), ): line.set_xdata(xdata[frame_index]) line.set_ydata(ydata[frame_index]) for j, zs in zip( range(0, len(comp_pos_lines), Np), comp_positions.values() ): for i in range(Np): vline = comp_pos_lines[j + i] vline.set_xdata([zs[frame_index]]) for ctname_txt, offset, zs in zip( comp_name_texts, comp_name_offsets, comp_positions.values() ): ctname_txt.set_x(zs[frame_index] + offset) anim_text.set_text(anim_txt_former(frame_index)) return all_artists an = animation.FuncAnimation( fig, animate, frames=N, blit=blit, interval=interval ) # Get rid of the dummy figure now so that it doesn't get shown / saved plt.close(dummy_fig) if filename is not None: an.save(filename) if show: plt.show() if Np == 1: return fig, axs[0], an return fig, axs, an
[docs]class AstigmaticPropagationSolution(BaseSolution): """Solution representation of a call to :func:`finesse.tracing.tools.propagate_beam_astig`. Internally this stores two :class:`.PropagationSolution` instances which are used to access the per-plane beam parameters. These propagation solutions can be accessed via :attr:`.AstigmaticPropagationSolution.ps_x` and :attr:`.AstigmaticPropagationSolution.ps_y` for the tangential and sagittal planes, respectively. """
[docs] def __init__(self, name, ps_x: PropagationSolution, ps_y: PropagationSolution): super().__init__(name) self.__ps_x = ps_x self.__ps_y = ps_y self.__symbolic = ps_x.symbolic self.empty = False
@property def symbolic(self): """Whether the astigmatism solution is symbolic.""" return self.__symbolic @property def ps_x(self): """The internal :class:`.PropagationSolution` for the tangential plane.""" return self.__ps_x @property def ps_y(self): """The internal :class:`.PropagationSolution` for the sagittal plane.""" return self.__ps_y @property def start_node(self): """The starting node of the propagation.""" return self.ps_x.start_node @property def end_node(self): """The final node of the propagation.""" return self.ps_x.end_node @property def nodes(self): """A list of all the nodes traversed, in order.""" return self.ps_x.nodes @property def ports(self): """A list of all the ports traversed, in order.""" return self.ps_x.ports @property def spaces(self): """A list of all spaces traversed, in order.""" return self.ps_x.spaces @property def components(self): """A list of all components (excluding spaces) traversed, in order.""" return self.ps_x.components @property def path_length(self): """The geometric path length of the traversed path.""" return self.ps_x.path_length @property def overlaps(self): """A dict of nodes to the qx, qy overlaps at these nodes.""" return {node: self.overlap(node) for node in self.nodes}
[docs] def qx(self, at): """Tangential beam parameter at a given location of the path. See :meth:`.PropagationSolution.q` for parameters and return object descriptions. """ return self.ps_x.q(at)
[docs] def qy(self, at): """Sagittal beam parameter at a given location of the path. See :meth:`.PropagationSolution.q` for parameters and return object descriptions. """ return self.ps_y.q(at)
[docs] def overlap(self, at): """Overlap between tangential and sagittal beam parameters at a given location of the path. See :meth:`.PropagationSolution.q` for parameters description. Returns ------- O : float or :class:`.Function` Overlap between the beam parameters in each plane, at the specified node. """ return BeamParam.overlap(self.qx(at), self.qy(at))
[docs] def plot( self, *args, filename=None, show=True, ignore=None, name_xoffsets=None, name_yoffsets=None, ylims=None, npts=1000, resolution="equal", subs=None, ): """Plot any combination of the beam sizes, accumulated Gouy phases and / or wavefront curvatures over the propagated path, showing the values for both planes. The expected, valid positional arguments are any combination of: - "beamsize", - "gouy", - "curvature", or "all" to plot all of the above. If no positional args are given then the beamsize (first axis) and accumulated Gouy phase (second axis) will be plotted by default. .. note:: The resulting figure will be divided into two columns, the first giving the absolute quantity values for both planes and the second giving the difference between the quantities in the two planes (tangential plane minus sagittal plane). For beam size plots, the tangential plane values are shown in red whilst sagittal are shown in blue. Whilst for any other quantity, the tangential plane values are solid lines and sagittal are dashed lines. The locations of each component will be marked on the figure, unless the component is in `ignore` or the component has "AR" or "HR" in its name. Parameters ---------- filename : str or file-like, optional Name of a file or existing file object to save the figure to. show : bool, optional; default: True Whether to show the figure. ignore : component, sequence of, optional A component or sequence of components to ignore when making markers. name_xoffsets : dict, optional Dictionary of component names to x-axis offsets for shifting where the component name text is placed. The offset value is interpreted in terms of data co-ordinates. name_yoffsets : dict, optional Dictionary of component names to y-axis offsets for shifting where the component name text is placed. The offset value is interpreted in terms of data co-ordinates. ylims : dict, optional Dictionary of target names (i.e. "beamsize", "gouy" or "curvature") to manual axis y-limits. npts : int, optional; default: 1000 See equivalent argument in :meth:`.PropagationSolution.all_segments`. resolution : str, optional; default: "equal" See equivalent argument in :meth:`.PropagationSolution.all_segments`. subs : dict, optional A dictionary of model parameter to value substitutions to pass to the ``eval`` methods of symbolic expressions. If this solution object is not symbolic then this argument is ignored. Returns ------- fig : Figure Handle to the figure. axs : axes The axis handles. """ import matplotlib.pyplot as plt if not args: args = ("beamsize", "gouy") if "all" in args: args = ("beamsize", "gouy", "curvature") valid_args = ("beamsize", "gouy", "curvature") if any(arg not in valid_args for arg in args): raise ValueError( "Invalid target argument in args, expected any " f"combination of {valid_args} or 'all'" ) if not self.symbolic and subs is not None: warn( f"Ignoring {subs=} kwarg as AstigmaticPropagationSolution is " f"non-symbolic." ) N = len(args) fig, axs = plt.subplots(N, 2, sharex=True) if N == 1: axs = np.array([axs]) maximums = {k: 0 for k in args} diff_mins = {k: 0 for k in args} if resolution.casefold() == "adaptive": raise NotImplementedError( "Adaptive resolution method for astigmatic propagation solution " "plotting has an outstanding issue and so is not yet supported." ) data_x = self.ps_x.all_segments( w_scale=1e3, npts=npts, resolution=resolution, subs=subs ) data_y = self.ps_y.all_segments( w_scale=1e3, npts=npts, resolution=resolution, subs=subs ) for (zs_x, segd_x), (zs_y, segd_y) in zip(data_x.values(), data_y.values()): for i, arg in enumerate(args): vx = segd_x[arg] vy = segd_y[arg] if arg == "beamsize": axs[i][0].fill_between(zs_x, vx, -vx, color="r", alpha=0.5) axs[i][0].fill_between(zs_y, vy, -vy, color="b", alpha=0.5) else: axs[i][0].plot(zs_x, vx) axs[i][0].plot(zs_y, vy, linestyle="--") # TODO (sjr) This currently won't work with adaptive resolution # method as vx and vy may be different size arrays. # Need to find a way to resolve this. Maybe just force # each segment in data_x, data_y to have same points # between the two (bit annoying to handle though). axs[i][1].plot(zs_x, vx - vy, color="k") maximums[arg] = max(maximums[arg], vx.max(), vy.max()) diff_mins[arg] = min(diff_mins[arg], (vx - vy).min()) if ignore is None: ignore = [] if not is_iterable(ignore): ignore = [ignore] if name_xoffsets is None: name_xoffsets = {} if name_yoffsets is None: name_yoffsets = {} for node, info in self.ps_x.node_info.items(): if node.is_input: z = self.ps_x._eval_z_for_display(info, subs=subs) comp = node.component if "AR" not in comp.name and comp not in ignore: name = comp.name x_offset = name_xoffsets.get(name, 0) y_offset = name_yoffsets.get(name, 0) # display the name in a nicer way name = name.replace("_", "\n") n_newlines = name.count("\n") for i, arg in enumerate(args): if not i: for ax in axs[i]: ax.axvline( z, 0.12 + 0.1 * n_newlines, color="k", linestyle="--", ) if arg == "beamsize": ytext_pos = [-1 * maximums[arg], diff_mins[arg]] else: ytext_pos = [0, 0] for ax, ytp in zip(axs[i], ytext_pos): ax.text( z + x_offset, ytp + y_offset, name, ha="center", va="center", ) else: for ax in axs[i]: ax.axvline(z, color="k", linestyle="--") for ax in axs.flatten(): ax.set_xlim(0, None) for ax in axs[-1]: ax.set_xlabel("Distance [m]") ylabel_mappings = { "beamsize": "Beam size [mm]", "gouy": "Gouy phase\naccumulation [deg]", "curvature": "Wavefront curvature [1/m]", } ylabel_diff_mappings = { "beamsize": r"$\mathrm{w}_\mathrm{x} - \mathrm{w}_\mathrm{y}$ [mm]", "gouy": r"$\psi_\mathrm{x} - \psi_\mathrm{y}$ [deg]", "curvature": r"$\mathrm{S}_\mathrm{x} - \mathrm{S}_\mathrm{y}$ [1/m]", } if ylims is None: ylims = {} for i, arg in enumerate(args): if arg != "beamsize": ylim = ylims.get(arg, None) if ylim is None: axs[i][0].set_ylim(0 if arg == "gouy" else None, maximums[arg]) else: axs[i][0].set_ylim(ylim[0], ylim[1]) ylabel = ylabel_mappings.get(arg, arg) axs[i][0].set_ylabel(ylabel) dylabel = ylabel_diff_mappings.get(arg) axs[i][1].set_ylabel(dylabel) if filename is not None: fig.savefig(filename) if show: plt.show() if N == 1: return fig, axs[0] return fig, axs
[docs]class BeamTraceSolution(BaseSolution): """Trace solution corresponding to calls to :meth:`.Model.beam_trace`. Note that BeamTraceSolution objects are returned via :meth:`.Model.beam_trace` calls, they should never need to be created manually. This class provides a dict-like interface to beam trace solution data. If ``trace`` is an instance of this class then one can access the beam parameters at both planes of a node via:: # Using the look-up key notation qx, qy = trace[node] # Or the get method qx, qy = trace.get(node) One can also access individual plane beam parameters with:: # Tangential plane qx = trace[node].qx # Sagittal plane qy = trace[node].qy A copy of the Python dictionary which stores all the underlying ``node : (qx, qy)`` mappings can be obtained with the :attr:`.BeamTraceSolution.data` property. To draw a forest representation of the trace solution data, one can simply do:: # Prints the forest of beam parameters print(trace) # Stores the forest as a string trace_str = str(trace) This forest will be ordered by the trace order used for the associated :meth:`.Model.beam_trace` call which constructed this solution. """ NodeData = typing.NamedTuple("NodeData", [("qx", BeamParam), ("qy", BeamParam)])
[docs] def __init__(self, name, data, forest=None): super().__init__(name) self._data = {node: self.NodeData(qx, qy) for node, (qx, qy) in data.items()} self.__forest = forest self.empty = False # tree drawing fill circle # TODO (sjr) Temporary workaround for #328, need better output colouring interactive = is_interactive() self.__start_highlight = ": \x1b[0;32;40m[" if not interactive else ": [" self.__end_highlight = "]\x1b[0m" if not interactive else "]"
[docs] def items(self): """A view on the underlying dict items.""" return self._data.items()
[docs] def values(self): """A view on the underlying dict values.""" return self._data.values()
[docs] def keys(self): """A view on the underlying dict keys.""" return self._data.keys()
@property def data(self): """A copy of the underlying dictionary of beam trace solution data. :`getter`: Returns the beam trace solution data. Read-only. """ return self._data.copy() @property def data_qx(self): """A copy of the underlying `data` dictionary but with only the tangential plane beam parameters selected. :`getter`: Returns a dictionary of traced nodes with corresponding tangential plane beam parameters (read-only). """ return {node: qx for node, (qx, _) in self._data.items()} @property def data_qy(self): """A copy of the underlying `data` dictionary but with only the sagittal plane beam parameters selected. :`getter`: Returns a dictionary of traced nodes with corresponding sagittal plane beam parameters (read-only). """ return {node: qy for node, (_, qy) in self._data.items()} def __getitem__(self, node): if isinstance(node, OpticalNode): return self._data[node] if isinstance(node, Port): in_q = self._data.get(node.i, None) out_q = self._data.get(node.o, None) if in_q is None and out_q is None: raise KeyError(f"Port {node} not in the trace data.") return {node.i: in_q, node.o: out_q} if isinstance(node, Connector): return {n: self._data[n] for n in node.optical_nodes} if isinstance(node, str): search = [n for n in self._data.keys() if n.full_name == node] if len(search) == 1: return self.__getitem__(search[0]) elif len(search) > 1: raise RuntimeError(f"Unexpected number of nodes named {node} in data") else: raise KeyError(f"No node named {node} in this trace") raise TypeError( "Expected key to be of type string, OpticalNode, Port, or Connector " f"but got object of type: {type(node)}" )
[docs] def get(self, node, default=None): """Gets the beam parameter(s) at the specified node / port. Note that `node` can be an instance of :class:`.OpticalNode`, resulting in this method returning qx and qy for that specific node, *or* it can be an object of type :class:`.Port` - in which case a dictionary of both the input and output node beam parameters are returned. Parameters ---------- node : :class:`.OpticalNode` or :class:`.Port` The node or port to access. default : any, optional; default: None The value to return if `node` does not exist within the trace data. """ try: return self[node] except KeyError: return default except TypeError: raise
[docs] def q(self, node): """A convenience method for getting the non-astigmatic beam parameter at a node. .. warning:: This is only intended to be used on nodes which do not exhibit astigmatism. If qx != qy at the node then this method will raise a ValueError. To get both qx and qy at a node use either:: qx, qy = trace[node] or:: qx, qy = trace.get(node) where ``trace`` is an instance of this class. Parameters ---------- node : :class:`.OpticalNode` The node at which to obtain q. Returns ------- q : :class:`.BeamParam` The beam parameter (which is the same in both planes) at the node. Raises ------ ex : ValueError If the beam parameters qx != qy at the node. """ qx, qy = self._data[node] if qx != qy: raise ValueError( f"Beam parameters qx != qy at node {node.full_name}. Use " "sub-script operator or BeamTraceSolution.get to obtain both " "qx and qy in general." ) return qx
def __str_q(self, q=None, node=None): if node is not None: qx, qy = self._data[node] if q is not None: if isinstance(q, tuple): qx, qy = q else: qx = qy = q format_q = lambda q: f"{q.real:.2f} + {q.imag:.2f}j" if qx != qy: return f"qx = {format_q(qx)}, qy = {format_q(qy)}" return f"q = {format_q(qx)}" def __draw_tree(self, tree, lpad, lines): branch = "├─" pipe = "│" end = "╰─" dash = "─" ltree = tree.left rtree = tree.right if ltree is not None: s = branch + dash if rtree is None: s = end + dash pad = " " else: s = branch + dash pad = pipe + " " lines.append( lpad + s + "o" + " " + ltree.node.full_name + self.__start_highlight + self.__str_q(node=ltree.node) + self.__end_highlight ) self.__draw_tree(ltree, lpad + pad, lines) if rtree is not None: s = end + dash pad = " " lines.append( lpad + s + "o" + " " + rtree.node.full_name + self.__start_highlight + self.__str_q(node=rtree.node) + self.__end_highlight ) self.__draw_tree(rtree, lpad + pad, lines) def __str__(self): all_trees = "" format_g = lambda gx, gy: f"gx = {gx}, gy = {gy}" if gx != gy else f"g = {gx}" cav_info = ( lambda cav: f"({self.__str_q(q=(cav.qx, cav.qy))}, {format_g(cav.gx, cav.gy)})" ) nodes = list(self._data.keys()) if not nodes: return all_trees if self.__forest is None: # Drawing a full model beam trace forest = nodes[0]._model.trace_forest.forest else: # Drawing a custom forest, e.g. from a single cavity trace forest = self.__forest for tree in forest: lines = [ " " + "o" + " " + tree.node.full_name + self.__start_highlight + self.__str_q(node=tree.node) + self.__end_highlight ] self.__draw_tree(tree, "", lines) tree_s = "\n ".join(lines) dep = tree.dependency if tree.is_source and isinstance(dep, Cavity): all_trees += ( f"\nInternal trace of cavity: {dep.name} {cav_info(dep)}\n\n" ) else: all_trees += f"\nDependency: {dep.name}\n\n" all_trees += tree_s + "\n" return all_trees
[docs] def print(self): """Draws the trace solution as a forest of beam parameters. This uses the :class:`.TraceForest` structure associated with the model that constructed this solution, where each tree is ordered by the trace order of the relevant call to the beam tracing method. """ print(str(self))