finesse-simulation-O4/utils.py

from logging import getLogger
from typing import Any, NamedTuple
from finesse.analysis.actions import Noxaxis
from finesse.analysis.actions.axes import Xaxis
from finesse.components import SignalGenerator, Node
from finesse.detectors import PowerDetectorDemod1, QuantumNoiseDetector
from finesse.exceptions import ModelMissingAttributeError
from finesse.model import DegreeOfFreedom, Model, Port, SeriesSolution

from matplotlib.axes import Axes
from matplotlib.text import Text
from matplotlib.pyplot import figure, show

from numpy import float64, linspace
from numpy.typing import NDArray


def compute_solutions(
    model: Model, DOF: str, padding: float, nb: int = 250
) -> SeriesSolution:
    """
    Run a simulation with a given degree of freedom varying
    """
    if type(model.get(DOF)) is DegreeOfFreedom:
        return model.run(
            Xaxis(
                model.get(DOF).DC,
                "lin",
                model.get(DOF).DC - padding,
                model.get(DOF).DC + padding,
                nb,
            )
        )
    raise TypeError(
        "{} should be a degree of freedom, not {}".format(
            DOF, type(model.get(DOF))
        )
    )


def process_solution(
    model: Model,
    DOF: str,
    padding: float,
    nb: int = 250,
    ax: Axes | None = None,
) -> SeriesSolution:
    """
    Compute and, if needed, display a given simulation with a given degree of freedom varying
    """
    solution = compute_solutions(
        model,
        DOF,
        padding,
        nb,
    )
    if ax is not None:
        x = linspace(
            model.get(DOF).DC - padding,
            model.get(DOF).DC + padding,
            nb + 1,
        )
        ylabel: Text = (
            ax.semilogy(x, solution["B1_DC"], label="dark fringe")[0]
            .axes.semilogy(x, solution["NE_p1"], label="north cavity")[
                0
            ]  # pyright: ignore[reportAttributeAccessIssue, reportOptionalMemberAccess]
            .axes.semilogy(x, solution["WE_p1"], label="west cavity")[0]
            .axes.vlines(
                [model.get(DOF).DC],
                min(solution["B1_DC"]),
                max(solution["NE_p1"]),
                colors="red",
            )
            .axes.legend()
            .axes.set_ylabel("power (W)")
        )
        ax.grid()
    return solution


class DisplayData(NamedTuple):
    """
    Allows to avoid typing issue
    """

    DOF: str
    padding: float


def display_displaydata(model: Model, data: list[DisplayData]):
    Figure = figure(figsize=(13, 10))
    for i in range(len(data)):
        process_solution(
            model,
            data[i].DOF,
            data[i].padding,
            ax=Figure.add_subplot(3, 2, i + 1),
        )
    show()


def typingfix_seriessolution(
    solution: None | SeriesSolution | tuple[Any],
) -> SeriesSolution:
    if solution is None:
        raise ValueError("Result of the simulation is None")
    if isinstance(solution, tuple):
        raise ValueError("Result is not in the expected format")
    return solution


def fix_dark_fringe(model: Model, power: float) -> tuple[int, float]:
    solution: SeriesSolution = typingfix_seriessolution(
        model.run(Noxaxis())
    )
    if "B1_DC" not in solution.outputs:
        raise Exception("Model does not contain B1 readout")

    if type(model.get("DARM")) is DegreeOfFreedom:
        result = solution["B1_DC"]
        start, stop, nb = 0, 1, 0
        while (abs(result - power) > 1e-4) and (nb < 100):
            nb += 1
            temp = start + (stop - start) / 2

            model.get("DARM").DC = temp
            solution: SeriesSolution = typingfix_seriessolution(
                model.run(Noxaxis())
            )
            if "B1_DC" not in solution.outputs:
                raise Exception("Model do not contains B1 readout")
            result = solution["B1_DC"]
            if result > power:
                stop = temp
            else:
                start = temp
        return nb, solution["B1_DC"]
    raise Exception("DARM is not a degree of freedom in the model")


def get_QNLS(
    model: Model, start: int = 5, stop: int = 1000, nb: int = 100
) -> SeriesSolution:
    new_model = model.deepcopy()
    new_model.fsig.f = 1  # pyright: ignore[reportAttributeAccessIssue]
    try:
        new_model.add(
            SignalGenerator("darmx", new_model.space_NI_NE.h, 1, 0)  # pyright: ignore[reportAttributeAccessIssue]
        )
    except ModelMissingAttributeError as exception:
        getLogger().error("no space named space_NI_NE in the model")
        raise exception
    try:
        new_model.add(
            SignalGenerator("darmy", new_model.space_WI_WE.h, 1, 180)  # pyright: ignore[reportAttributeAccessIssue]
        )
    except ModelMissingAttributeError as exception:
        getLogger().error("no space named space_WI_WE in the model")
        raise exception
    try:
        new_model.add(
            QuantumNoiseDetector("NSR_with_RP", new_model.SR.p2.o, True)  # pyright: ignore[reportAttributeAccessIssue]
        )
    except ModelMissingAttributeError as exception:
        getLogger().error("no port named SR.p2.o in the model")
        raise exception
    return new_model.run(
        Xaxis(new_model.fsig.f, "log", start, stop, nb)  # pyright: ignore[reportAttributeAccessIssue]
    )


def compute_TF(
    model: Model,
    frequencies: NDArray[float64],
    inputs: list[Port | str],
) -> SeriesSolution:
    model.fsig.f = 1  # pyright: ignore[reportAttributeAccessIssue]

    for input in inputs:
        model.add(SignalGenerator("sig1", model.get(input), 1, 0))
    model.add(
        PowerDetectorDemod1(
            "TF", model.get("SDB1").p2.o, f=model.get("fsig").f
        )
    )
    model.get("TF").select_mask(exclude=(0, 0))
    return model.run(
        Xaxis(
            model.fsig.f,  # pyright: ignore[reportAttributeAccessIssue]
            "log",
            min(frequencies),
            max(frequencies),
            frequencies.shape[0] - 1,
        )
    )


class TF:
    def __init__(
        self,
        model: Model,
        inputs: list[Port | str],
        outputs: list[Port | str],
        frequencies: NDArray[float64],
    ):
        self.model = model.deepcopy()
        self.inputs = inputs
        self.outputs = outputs
        self.frequencies = frequencies
        self.result: SeriesSolution | None = None
        self.compute()

        # alias
        self.f = self.frequencies

    def compute(self):
        solution = compute_TF(self.model, self.frequencies, self.inputs)
        self.result = solution
        return self

    def get(self, output: Port | str):
        if self.result is None:
            raise Exception(
                "TF need to be computed before getting result"
            )
        return self.result[output]