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
from finesse.detectors import QuantumNoiseDetector
from finesse.exceptions import ModelMissingAttributeError
from finesse.model import DegreeOfFreedom, Model, SeriesSolution

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

from numpy import linspace


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]
    )