from more_itertools import pairwise
from scipy.sparse import diags
import numpy as np
from ...solutions import BaseSolution
from .base import Action
from . import Operator
from finesse.exceptions import FinesseException
from finesse.components import Mirror, Beamsplitter
[docs]class PseudoLockCavitySolution(BaseSolution):
    """Solution to the :class:`PseudoLockCavity` action.
    Attributes
    ----------
    operator : :class:`OperatorSolution`
        Solution of the operator action that is used to compute the roundtrip
        operator in the cavity
    w : array_like
        2D Eigenvector array
    v : array_like
        1D Eigenvalue array
    v_idx : int
        Index of the eigenmode being locked to
    phase : float
        Phase in degrees of the eigenmode that is being locked to. This is what
        is used to set the operating point of the cavity
    cavity : str
        Which cavity is being locked
    lowest_loss : bool
        If False, the target `mode` attribute is used. Otherwise
    mode : array_like
        [n, m] target mode to try and lock to
    """ 
# IMPORTANT: renaming this class impacts the katscript spec and should be avoided!
[docs]class PseudoLockCavity(Action):
    """An action that locks a cavity defined by a `Cavity` element to a specific mode
    without using any radio-frequency sensing scheme. This will only work on simple
    cavities that are not coupled in any way. You can specify whether to try and lock to
    a particular HG mode with the `mode=[n,m]` keyword argument, or just pick the lowest
    loss mode, `lowest_loss=True`.
    Parameters
    ----------
    cavity : Cavity
        Cavity element describing some Fabry-Perot like optical cavity
    mode : (n, m), optional
        HG mode to try and lock to, default is [0,0]
    lowest_loss : bool, optional
        Select the eigenmode which has the lowest loss, most likely the fundamental
        mode of the cavity. Using lowest loss will override the `mode` selection.
    feedback : Parameter optional
        If `None` the required cavity tuning to lock to the calculated mode will be
        determined from the cavity objects source node element, and the relevant
        `phi` parameter will be used. Alternatively you can specify which tuning
        parameter is used instead. Which should be a `phi` of some mirror in the
        cavity or a `DegreeOfFreedom` which controls the cavity length.
    name : str, optional
        Name of the solution generated by this action
    Examples
    --------
    A Fabry-Perot based on aLIGO cavities:
        >>> import finesse
        >>> model = finesse.Model()
        >>> model.parse('''
        ... l l1 P=1
        ... m ITM T=0.014 L=0 Rc=-1945
        ... s sARM ITM.p2 ETM.p1 L=3994
        ... m ETM R=1 L=0 Rc=2145
        ... link(l1, ITM)
        ... cav arm ETM.p1.o
        ... modes(maxtem=1)
        ... ''')
    The cavity source node is on the ETM so the pseudo-lock will use that node and
    component `phi` parameter to feedback to. In this case the ETM.phi will be
    corrected to match the ITM.phi which will make the HG00 resonant.
        >>> model.ITM.phi = 0
        >>> model.ETM.phi = 10
        >>> model.run("pseudo_lock_cavity(arm)")
        >>> print(model.ITM.phi, model.ETM.phi)
        0.0 degrees 0.0 degrees
        >>> model.ITM.phi = 10
        >>> model.ETM.phi = 0
        >>> model.run("pseudo_lock_cavity(arm)")
        >>> print(model.ITM.phi, model.ETM.phi)
        10.0 degrees 10.0 degrees
    This lock will also handle any misalignments, mismatches, or maps applied to
    the cavity.
        >>> model.ITM.xbeta = 5e-8
        >>> model.ITM.phi = 11
        >>> model.ETM.phi = 0
        >>> model.run("pseudo_lock_cavity(arm)")
        >>> print(model.ITM.phi, model.ETM.phi)
        11.0 degrees 10.984168676865762 degrees
    We can also lock to other HG modes:
        >>> sol = model.run("series(pseudo_lock_cavity(arm, mode=[1,0]), noxaxis())")
        >>> print(model.ITM.phi, model.ETM.phi)
        >>> print(model.homs)
        >>> print(sol['noxaxis']['Ecirc'])
        11.0 degrees 28.663091273523214 degrees
        [[0 0]
        [1 0]
        [0 1]]
        [-0.16465261-0.11468501j  0.41262682-0.27112167j  0.        +0.j        ]
    """
    def __init__(
        self,
        cavity,
        *,
        mode=[0, 0],
        lowest_loss=False,
        feedback=None,
        name="pseudo_lock_cavity",
    ):
        super().__init__(name)
        self.cavity = cavity
        self.mode = mode
        self.lowest_loss = lowest_loss
        self.feedback = feedback
    def _requests(self, model, memo, first=True):
        cav = (
            model.elements[self.cavity] if isinstance(self.cavity, str) else self.cavity
        )
        memo["keep_nodes"].extend((_.full_name for _ in cav.path.nodes))
        if not any(np.all(model.homs == self.mode, axis=1)):
            raise FinesseException(
                f"HG mode {self.mode} was not found in the simulation"
            )
        # If a user provided feedback is given then use that to send locking
        # signal towards
        if self.feedback is not None:
            if isinstance(self.feedback, str):
                memo["changing_parameters"].append(model.get(self.feedback))
            elif hasattr(self.feedback, "parameter"):
                memo["changing_parameters"].append(self.feedback.parameter)
            else:
                memo["changing_parameters"].append(self.feedback)
        else:
            # if not use the cavity source component
            memo["changing_parameters"].append(cav.source.component.phi)
    def _do(self, state):
        model = state.model
        key = id(self)
        sol = PseudoLockCavitySolution(self.name)
        sol.mode = self.mode
        sol.lowest_loss = self.lowest_loss
        sol.cavity = str(self.cavity)
        # Try and get data already computed for this action
        ws = state.action_workspaces.get(key, None)
        if ws is None:
            state.action_workspaces[key] = ws = {}
            cav = ws["cavity"] = (
                model.elements[self.cavity]
                if isinstance(self.cavity, str)
                else self.cavity
            )
            # if the cavity is using a via definition for it's path use that
            # if not use the opposite to ensure we don't select a dumb option
            via = cav.source.opposite if cav.via is None else cav.via
            ws["operator"] = Operator(cav.source, cav.source, via=via)
            ws["hg_idx"] = int(np.argmax(np.all(model.homs == self.mode, axis=1)))
            # Get the phi parameter of the cavity source component
            component = cav.source.component
            if self.feedback is None:
                ws["phi"] = component.phi
            elif isinstance(self.feedback, str):
                ws["phi"] = model.get(self.feedback)
            elif hasattr(self.feedback, "parameter"):
                ws["phi"] = self.feedback.parameter
            else:
                ws["phi"] = self.feedback
            # Decide on the direction of phase application, depending on side
            # of mirror/bs the cav source node is on
            if isinstance(component, Mirror):
                if cav.source in component.p1:
                    ws["scale"] = -1
                else:
                    ws["scale"] = +1
            elif isinstance(component, Beamsplitter):
                # TODO need angle of incidence here
                if cav.source in component.p1 or cav.source in component.p2:
                    ws["scale"] = -1
                else:
                    ws["scale"] = +1
            else:
                raise FinesseException(
                    f"From the cavity node {cav.source}, its owner {cav.source.component} is neither a mirror nor a beamsplitter."
                )
        # Compute eigenmodes and vectors
        sol.operator = state.apply(ws["operator"])
        sol.w, sol.v = np.linalg.eig(sol.operator.operator)
        if not self.lowest_loss:
            # select mode with largest self.mode component
            # which is probably the best guess for making that
            # mode resonant
            sol.v_idx = np.argmax(abs(sol.v[:, ws["hg_idx"]]))
        else:
            # select eigenmode which has the lowest loss, likely the fundamental
            # minimize -> roundtip loss = 1 - abs(sol.w)**2
            sol.v_idx = np.argmax(abs(sol.w))
        sol.phase = np.angle(sol.w[sol.v_idx], deg=True)
        ws["phi"].value += ws["scale"] * sol.phase / 2
        return sol 
[docs]class PseudoLockDRFPMISolution(BaseSolution):
    """TODO: needs updating for all the various attributes this stores."""
[docs]    def plot_PRC_SRC_eigenvectors(self):
        import matplotlib.pyplot as plt
        """Plots eigenvector maxtrix for PRC and SRC"""
        plt.imshow(abs(self.M_evec), origin="upper")
        plt.text(0, 0, "PRC", c="b", horizontalalignment="center")
        plt.text(0, 3, "SRC", c="b", horizontalalignment="center")
        plt.hlines(2.5, -0.5, 5.5)
        plt.colorbar()
        plt.title("Eigenvectors")  
# IMPORTANT: renaming this class impacts the katscript spec and should be avoided!
[docs]class PseudoLockDRFPMI(Action):
    """Pseudo-locking is attempting to find an operating point for a LIGO like model
    without needing to use RF sidebands and readouts. Although it is not physically
    accurate it does provide a useful tool for analysing detectors from a more
    theoretical basis. This generates a :class:`PseudoLockDRFPMISolution` solution
    containing various operators and results.
    This action is hardcoded to work with a LIGO like model. Mirrors should be named
    with ITMX, ETMX, PRM, etc.
    This currently action only really works for finding the lock points for the PRC,
    SRC, XARM, and YARM – which all have defined cavity roundtrips which allow
    eigendecomposition of roundtrip operators. The eigenvectors describe the HOM mix
    for each resonant mode in a cavity and the eigenvalues the roundtrip phase and
    loss of the mode. This code looks for eigenvectors with the largest HG00 content
    and then uses the eigenvalues to compute what cavity tunings need to make this
    mode resonant.
    The corner is the most complicated here as the PRC and SRC are coupled via the
    beamsplitter. The eigendecomposition is performed on the 2x2 operator matrix for
    the PRC, SRC, and the coupling matrices between them. When the coupling is small
    the results are the same as performing the decomposition on each SRC and PRC
    separately.
    Note that the only degree of freedom that this does not handle currently is MICH.
    MICH is awkward because it is not a cavity. It is essentially the beamsplitter
    position that makes the anti-symmetric port dark.
    Parameters
    ----------
    frequency : float, optional
        Frequency to use for calculating the operators
    apply_tunings : bool, optional
        When True the action will modify the model tunings
    name : str, optional
        Name of the solution generated by this action
    """
    def __init__(self, frequency=0, *, apply_tunings=True, name="operator_lock"):
        super().__init__(name)
        self.frequency = frequency
        self.apply_tunings = apply_tunings
        # connections to compute operators for in a LIGO like model
        # name : (from, to, via)
        self.operators = {
            "BS13": ("BS.p1.i", "BS.p3.o", None),
            "BS31": ("BS.p3.i", "BS.p1.o", None),
            "BS12": ("BS.p1.i", "BS.p2.o", None),
            "BS21": ("BS.p2.i", "BS.p1.o", None),
            "BS24": ("BS.p2.i", "BS.p4.o", None),
            "BS42": ("BS.p4.i", "BS.p2.o", None),
            "BS34": ("BS.p3.i", "BS.p4.o", None),
            "BS43": ("BS.p4.i", "BS.p3.o", None),
            "MICHX1": ("BS.p3.o", "ITMX.p1.i", None),
            "MICHY1": ("BS.p2.o", "ITMY.p1.i", None),
            "MICHX2": ("ITMX.p1.o", "BS.p3.i", None),
            "MICHY2": ("ITMY.p1.o", "BS.p2.i", None),
            "R1X": ("ITMX.p1.i", "ITMX.p1.o", None),
            "R2X": ("ITMX.p2.i", "ITMX.p2.o", None),
            "T1X": ("ITMX.p1.i", "ITMX.p2.o", None),
            "T2X": ("ITMX.p2.i", "ITMX.p1.o", None),
            "GX": ("ITMX.p2.o", "ITMX.p2.i", "ETMX.p1.i"),
            "GY": ("ITMY.p2.o", "ITMY.p2.i", "ETMY.p1.i"),
            "R1Y": ("ITMY.p1.i", "ITMY.p1.o", None),
            "R2Y": ("ITMY.p2.i", "ITMY.p2.o", None),
            "T1Y": ("ITMY.p1.i", "ITMY.p2.o", None),
            "T2Y": ("ITMY.p2.i", "ITMY.p1.o", None),
            "SRC": ("BS.p4.o", "BS.p4.i", "SRM.p1.o"),
            "PRC": ("BS.p1.o", "BS.p1.i", "PRM.p2.o"),
            "SQZ": ("sqz.p1.o", "SRM.p2.o", None),
        }
    def _requests(self, model, memo, first=True):
        for name, (_from, _to, via) in self.operators.items():
            memo["keep_nodes"].append(_from)
            memo["keep_nodes"].append(_to)
            if via:
                memo["keep_nodes"].append(via)
        if self.apply_tunings:
            memo["changing_parameters"].append("SRCL.DC")
            memo["changing_parameters"].append("PRCL.DC")
    def _do(self, state):
        sim = state.sim
        model = state.model
        key = id(self)
        f_idx = None
        N = sim.model_settings.num_HOMs
        sol = PseudoLockDRFPMISolution(self.name)
        # Try and get data already computed for this action
        ws = state.action_workspaces.get(key, None)
        if ws is None:
            try:
                frequency = float(self.frequency)
            except TypeError:
                frequency = float(self.frequency.value)
            # find the right frequency index
            for freq in state.sim.carrier.optical_frequencies.frequencies:
                if freq.f == frequency:
                    f_idx = freq.index
                    break
            if f_idx is None:
                raise RuntimeError(
                    f"Could not find an optical carrier frequency with a value of {self.frequency}Hz"
                )
            state.action_workspaces[key] = ws = {}
            # loop through each operator we need to get
            for name, (start, end, via) in self.operators.items():
                nodes = model.path(start, end, via_node=via).nodes
                connections = tuple(
                    (n1.full_name, n2.full_name) for n1, n2 in pairwise(nodes)
                )
                operators = []
                # Get all the matrix views to eventually multiply together later
                for _ in pairwise(nodes):
                    # if we have a 1D array it's just a diagonal matrix
                    # so convert it to a sparse array for easy multiplying
                    # later
                    # TODO can have some option to select between different
                    # frequencies or signal/carrier at some point
                    Mview = sim.carrier.connections[_]
                    if Mview.ndim == 2:
                        # We have a frequency scattering matrix
                        # TODO : not sure on user interface for getting
                        # different frequency couplings yet, for now it's
                        # just same freq in and out
                        M = Mview[f_idx, f_idx]
                    else:
                        # no frequency coupling
                        M = Mview[f_idx]
                    operators.append(M)
                ws[name] = (nodes, connections, operators)
        # Loop through all the opeators are multiply them together to get
        # each reduced operator
        sol.operators = {}
        for name, (nodes, connections, operators) in ws.items():
            M = np.eye(sim.model_settings.num_HOMs)
            for op in operators:
                if op.view.ndim == 1:
                    M = diags(-op.view) @ M
                else:
                    M = -op.view @ M
            sol.operators[name] = M
        # Give all the opertors variable names
        GX = sol.operators["GX"]
        GY = sol.operators["GY"]
        SRC = sol.operators["SRC"]
        PRC = sol.operators["PRC"]
        R1X = sol.operators["R1X"]
        R1Y = sol.operators["R1Y"]
        R2X = sol.operators["R2X"]
        R2Y = sol.operators["R2Y"]
        # T1X = sol.operators["T1X"]
        # T1Y = sol.operators["T1Y"]
        # T2X = sol.operators["T2X"]
        # T2Y = sol.operators["T2Y"]
        BS12 = sol.operators["BS12"]
        BS21 = sol.operators["BS21"]
        BS13 = sol.operators["BS13"]
        BS31 = sol.operators["BS31"]
        BS24 = sol.operators["BS24"]
        BS42 = sol.operators["BS42"]
        BS34 = sol.operators["BS34"]
        BS43 = sol.operators["BS43"]
        MICHX1 = sol.operators["MICHX1"]
        MICHX2 = sol.operators["MICHX2"]
        MICHY1 = sol.operators["MICHY1"]
        MICHY2 = sol.operators["MICHY2"]
        # Now finally multiply the various operators together to compute
        # the different eigendecompositions
        # I = np.eye(N, dtype=np.complex128)
        GRTX = GX @ R2X
        GRTY = GY @ R2Y
        RX = R1X  # + T2X @ np.linalg.inv(I - GRTX) @ GX @ T1X
        RY = R1Y  # + T2Y @ np.linalg.inv(I - GRTY) @ GY @ T1Y
        MX = MICHX2 @ RX @ MICHX1
        MY = MICHY2 @ RY @ MICHY1
        MICHSS = BS34 @ MX @ BS43 + BS24 @ MY @ BS42
        MICHSP = BS31 @ MX @ BS43 + BS21 @ MY @ BS42
        MICHPP = BS31 @ MX @ BS13 + BS21 @ MY @ BS12
        MICHPS = BS34 @ MX @ BS13 + BS24 @ MY @ BS12
        sol.SRC = MICHSS @ SRC
        sol.PRC = MICHPP @ PRC
        sol.SRC_PRC = MICHSP @ SRC
        sol.PRC_SRC = MICHPS @ PRC
        # eigenvalues and vectors of the SRC roundtrip operator gives
        # the various resonant mode structure regardless of HOM content
        # and defects
        sol.X_eval, sol.X_evec = np.linalg.eig(GRTX)
        sol.Y_eval, sol.Y_evec = np.linalg.eig(GRTY)
        sol.X_HG00_idx, sol.Y_HG00_idx = (
            abs(sol.X_evec[0, :]).argmax(),
            abs(sol.Y_evec[0, :]).argmax(),
        )
        sol.xarm_phi = -np.angle(sol.X_eval[sol.X_HG00_idx], deg=True) / 2
        sol.yarm_phi = -np.angle(sol.Y_eval[sol.Y_HG00_idx], deg=True) / 2
        # Make a 2x2 block matrix for corner coupling between SRC and PRC
        sol.M = np.block([[sol.PRC, sol.PRC_SRC], [sol.SRC_PRC, sol.SRC]])
        sol.M_eval, sol.M_evec = np.linalg.eig(sol.M)
        # Compute which eigenvectors are the most HG00 like
        # Grab the HG00 row in and select the column with the
        # max value, will be the mode with the most HG00 content
        sol.PRC_HG00_idx, sol.SRC_HG00_idx = (
            abs(sol.M_evec[0, :]).argmax(),
            abs(sol.M_evec[N, :]).argmax(),
        )
        # Compute phases for each cavity
        sol.corner_phases = -np.angle(sol.M_eval, deg=True) / 2
        # Then select the cavity detuning to pick the most HG00 mode
        # Resonant sideband extraction
        sol.src_phi = sol.corner_phases[sol.SRC_HG00_idx]
        sol.prc_phi = sol.corner_phases[sol.PRC_HG00_idx]
        sol.PRC_evec = sol.M_evec[:N, sol.PRC_HG00_idx]
        sol.SRC_evec = sol.M_evec[N:, sol.SRC_HG00_idx]
        if self.apply_tunings:
            state.model.SRCL.DC += sol.src_phi
            state.model.PRCL.DC += 90 - sol.prc_phi
        return sol