SuPyMode 2.4.3__cp313-cp313-win_amd64.whl

SuPyMode/superset_plots.py

@@ -0,0 +1,497 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+from typing import TYPE_CHECKING
+
+import matplotlib.pyplot as plt
+import numpy
+from matplotlib.backends.backend_pdf import PdfPages
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+from MPSPlots import helper
+
+from SuPyMode.taper import AlphaProfile
+from SuPyMode.supermode import SUPERMODE
+from SuPyMode.utils import (
+    parse_filename,
+    interpret_slice_number_and_itr,
+)
+
+if TYPE_CHECKING:
+    from SuPyMode.superset import SuperSet
+
+
+class SuperSetPlots(object):
+
+    @staticmethod
+    def test_valid_input(
+        user_input, valid_inputs: list, variable_name: str = ""
+    ) -> None:
+        if user_input not in valid_inputs:
+            raise ValueError(
+                f"[{variable_name}] user_input: {user_input} argument not valid. Valid choices are: {valid_inputs}"
+            )
+
+    @staticmethod
+    def interpret_mode_of_interest(
+        superset: "SuperSet", mode_of_interest: str | SUPERMODE | list[SUPERMODE]
+    ) -> list[SUPERMODE]:
+        """
+        Resolves the mode of interest from user input to the appropriate list of SUPERMODE instances
+        based on the specified criteria or direct references.
+
+        Parameters
+        ----------
+        superset : SuperSet
+            Superset containing the available modes.
+        mode_of_interest : str or SUPERMODE or list of SUPERMODE
+            A category ("fundamental", "non-fundamental", or "all"), one mode,
+            or an explicit list of modes.
+
+        Returns
+        -------
+        list of SUPERMODE
+            Modes selected by the request.
+
+        Raises
+        ------
+        ValueError
+            If the category or input type is unsupported.
+        """
+        from SuPyMode.supermode import SUPERMODE
+
+        if isinstance(mode_of_interest, str):
+            match mode_of_interest:
+                case "fundamental":
+                    return superset.fundamental_supermodes
+                case "non-fundamental":
+                    return superset.non_fundamental_supermodes
+                case "all":
+                    return superset.supermodes
+                case _:
+                    raise ValueError(
+                        f"Unrecognized mode category '{mode_of_interest}'. Expected 'fundamental', 'non-fundamental', or 'all'."
+                    )
+
+        if isinstance(mode_of_interest, SUPERMODE):
+            return [mode_of_interest]
+
+        if isinstance(mode_of_interest, list) and all(
+            isinstance(item, SUPERMODE) for item in mode_of_interest
+        ):
+            return mode_of_interest
+
+        raise ValueError(
+            "mode_of_interest must be either 'fundamental', 'non-fundamental', 'all', a SUPERMODE instance, or a list of SUPERMODE instances."
+        )
+
+    @helper.pre_plot(nrows=1, ncols=1)
+    def plot_index(
+        self,
+        axes: plt.Axes,
+        mode_of_interest: list[SUPERMODE] = "all",
+    ) -> plt.Figure:
+        """
+        Plot the effective index for each mode as a function of inverse taper ratio (ITR).
+
+        Parameters
+        ----------
+        axes : matplotlib.axes.Axes
+            The axes object on which to plot.
+        mode_of_interest : list of SUPERMODE
+            List of modes to be plotted.
+
+        Returns
+        -------
+        plt.Figure
+            The figure object containing the generated plots. This can be customized further or saved to a file.
+        """
+        mode_of_interest = self.interpret_mode_of_interest(
+            superset=self, mode_of_interest=mode_of_interest
+        )
+
+        for mode in mode_of_interest:
+            mode.index.plot(ax=axes, show=False)
+
+        axes.legend()
+
+    @helper.pre_plot(nrows=1, ncols=1)
+    def plot_beta(
+        self, axes: plt.Axes, mode_of_interest: list[SUPERMODE] | str = "all"
+    ) -> plt.Figure:
+        """
+        Plot the effective propagation constant for each mode as a function of inverse taper ratio (ITR).
+
+        Parameters
+        ----------
+        axes : matplotlib.axes.Axes
+            The axes object on which to plot.
+        mode_of_interest : list of SUPERMODE
+            List of modes to be plotted.
+
+        Examples
+        --------
+        >>> fig, ax = plt.subplots()
+        >>> superset_plots.plot_index(ax=ax, mode_of_interest=[mode1, mode2])
+        >>> plt.show()
+
+        Returns
+        -------
+        plt.Figure
+            The figure object containing the generated plots. This can be customized further or saved to a file.
+        """
+        mode_of_interest = self.interpret_mode_of_interest(
+            superset=self, mode_of_interest=mode_of_interest
+        )
+
+        for mode in mode_of_interest:
+            mode.beta.plot(ax=axes, show=False)
+
+        axes.legend()
+
+    @helper.pre_plot(nrows=1, ncols=1)
+    def plot_eigen_value(
+        self, axes: plt.Axes, mode_of_interest: list[SUPERMODE] | str = "all"
+    ) -> plt.Figure:
+        """
+        Plot the computed eigen-values for each mode as a function of inverse taper ratio (ITR).
+
+        Parameters
+        ----------
+        axes : matplotlib.axes.Axes
+            The axes object on which to plot.
+        mode_of_interest : list of SUPERMODE
+            List of modes to be plotted.
+
+        Examples
+        --------
+        >>> fig, ax = plt.subplots()
+        >>> superset_plots.plot_index(ax=ax, mode_of_interest=[mode1, mode2])
+        >>> plt.show()
+
+        Returns
+        -------
+        plt.Figure
+            The figure object containing the generated plots. This can be customized further or saved to a file.
+        """
+        mode_of_interest = self.interpret_mode_of_interest(
+            superset=self, mode_of_interest=mode_of_interest
+        )
+
+        for mode in mode_of_interest:
+            mode.eigenvalue.plot(ax=axes, show=False)
+
+        axes.legend()
+
+    @helper.pre_plot(nrows=1, ncols=1)
+    def plot_beating_length(
+        self,
+        axes: plt.Axes,
+        mode_of_interest: list[SUPERMODE] = "all",
+        combination: list = "pairs",
+    ) -> plt.Figure:
+        """
+        Render a figure representing the beating length for each mode combination as a function of inverse taper ratio (ITR).
+
+        Parameters
+        ----------
+        axes : matplotlib.axes.Axes
+            The axes object on which to plot.
+        mode_of_interest : list of SUPERMODE
+            List of modes to be plotted.
+        combination : list
+            List of mode combinations.
+
+        Returns
+        -------
+        plt.Figure
+            The figure object containing the generated plots. This can be customized further or saved to a file.
+        """
+        mode_of_interest = self.interpret_mode_of_interest(
+            superset=self, mode_of_interest=mode_of_interest
+        )
+
+        combination = self.interpret_combination(
+            mode_of_interest=mode_of_interest, combination=combination
+        )
+
+        for mode_0, mode_1 in combination:
+            mode_0.beating_length.plot(ax=axes, other_supermode=mode_1)
+
+        axes.legend()
+
+    @helper.pre_plot(nrows=1, ncols=1)
+    def plot_normalized_coupling(
+        self,
+        axes: plt.Axes,
+        mode_of_interest: list[SUPERMODE] | str = "all",
+        combination: list | str = "pairs",
+    ) -> plt.Figure:
+        """
+        Render a figure representing the normalized coupling for each mode combination as a function of inverse taper ratio (ITR).
+
+        Parameters
+        ----------
+        ax : matplotlib.axes.Axes
+            The axes object on which to plot.
+        mode_of_interest : list of SUPERMODE
+            List of modes to be plotted.
+        combination : list
+            List of mode combinations.
+
+        Returns
+        -------
+        plt.Figure
+            The figure object containing the generated plots. This can be customized further or saved to a file.
+        """
+        mode_of_interest = self.interpret_mode_of_interest(
+            superset=self, mode_of_interest=mode_of_interest
+        )
+
+        combination = self.interpret_combination(
+            mode_of_interest=mode_of_interest, combination=combination
+        )
+
+        for mode_0, mode_1 in combination:
+            mode_0.normalized_coupling.plot(ax=axes, other_supermode=mode_1, show=False)
+
+    @helper.pre_plot(nrows=1, ncols=1)
+    def plot_adiabatic(
+        self,
+        axes: plt.Axes,
+        mode_of_interest: list[SUPERMODE] | str = "all",
+        combination: list | str = "pairs",
+        add_profile: list[AlphaProfile] = None,
+    ) -> plt.Figure:
+        """
+        Render a figure representing the adiabatic criterion for each mode combination as a function of inverse taper ratio (ITR).
+
+        Parameters
+        ----------
+        ax : matplotlib.axes.Axes
+            The axes object on which to plot.
+        mode_of_interest : list of SUPERMODE
+            List of modes to be plotted.
+        combination : list
+            List of mode combinations.
+        add_profile : list of AlphaProfile, optional
+            List of profiles to add to the plot (default is an empty list).
+
+        Returns
+        -------
+        plt.Figure
+            The figure object containing the generated plots. This can be customized further or saved to a file.
+        """
+        mode_of_interest = self.interpret_mode_of_interest(
+            superset=self, mode_of_interest=mode_of_interest
+        )
+
+        combination = self.interpret_combination(
+            mode_of_interest=mode_of_interest, combination=combination
+        )
+
+        for mode_0, mode_1 in combination:
+            mode_0.adiabatic.plot(ax=axes, other_supermode=mode_1, show=False)
+
+        if add_profile is not None:
+            for profile in numpy.atleast_1d(add_profile):
+                profile.render_adiabatic_factor_vs_itr_on_ax(ax=axes, line_style="--")
+
+    @helper.post_mpl_plot
+    def plot_field(
+        self,
+        mode_of_interest: list = "all",
+        itr_list: list[float] = None,
+        slice_list: list[int] = None,
+        show_mode_label: bool = True,
+        show_itr: bool = True,
+        show_slice: bool = True,
+    ) -> plt.Figure:
+        """
+        Render the mode field for different ITR values or slice numbers.
+
+        Parameters
+        ----------
+        mode_of_interest : list, optional
+            List of modes to be plotted (default is 'all').
+        itr_list : list of float, optional
+            List of ITR values for plotting (default is None).
+        slice_list : list of int, optional
+            List of slice numbers for plotting (default is None).
+        show_mode_label : bool, optional
+            Flag to display mode labels (default is True).
+        show_itr : bool, optional
+            Flag to display ITR values (default is True).
+        show_slice : bool, optional
+            Flag to display slice numbers (default is True).
+
+        Returns
+        -------
+        plt.Figure
+            The figure object containing the generated plots. This can be customized further or saved to a file.
+        """
+        mode_of_interest = self.interpret_mode_of_interest(
+            superset=self, mode_of_interest=mode_of_interest
+        )
+
+        # Interpret input lists
+        slice_list, itr_list = interpret_slice_number_and_itr(
+            itr_baseline=self.model_parameters.itr_list,
+            itr_list=itr_list,
+            slice_list=slice_list,
+        )
+
+        mode_of_interest = self.interpret_mode_of_interest(
+            superset=self, mode_of_interest=mode_of_interest
+        )
+
+        n_mode = len(mode_of_interest)
+        n_slice = len(slice_list)
+        grid_size = numpy.array([n_slice, n_mode])
+
+        figure, axes = plt.subplots(
+            *grid_size,
+            figsize=3 * numpy.flip(grid_size),
+            squeeze=False,
+            sharex=False,
+            sharey=False,
+        )
+
+        # Plot each mode field on the grid
+        for m, mode in enumerate(mode_of_interest):
+            for n, slice_number in enumerate(slice_list):
+                mode.field.render_on_ax(
+                    ax=axes[n, m],
+                    slice_number=slice_number,
+                    show_mode_label=show_mode_label,
+                    show_itr=show_itr,
+                    show_slice=show_slice,
+                )
+
+        return figure
+
+    @helper.post_mpl_plot
+    def plot_geometry(self) -> plt.Figure:
+        """
+        Render the geometry of the taper on the provided axes.
+
+        Parameters
+        ----------
+        ax : matplotlib.axes.Axes
+            The axes object on which to plot the geometry.
+
+        Returns
+        -------
+        plt.Figure
+            The figure object containing the generated plot. This can be customized further or saved to a file.
+        """
+        figure, ax = plt.subplots(figsize=(5, 5))
+
+        pcm = ax.pcolormesh(
+            self.model_parameters.x_vector,
+            self.model_parameters.y_vector,
+            self.model_parameters.mesh,
+            shading="auto",
+            cmap="Blues",
+        )
+
+        ax.set_aspect("equal")
+
+        divider = make_axes_locatable(ax)
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+
+        figure.colorbar(pcm, cax=cax)
+
+        return figure
+
+    def plot(self, plot_type: str, **kwargs) -> None:
+        """
+        General plotting function to handle different types of supermode plots.
+
+        Parameters
+        ----------
+        plot_type : str
+            The type of plot to generate. Options include 'index', 'beta', 'eigen-value', etc.
+        **kwargs : dict
+            Additional keyword arguments for specific plot configurations.
+
+        Raises
+        ------
+        ValueError
+            If an unrecognized plot type is specified.
+        """
+        match plot_type.lower().replace("_", "-"):
+            case "geometry":
+                return self.plot_geometry(**kwargs)
+            case "index":
+                return self.plot_index(**kwargs)
+            case "beta":
+                return self.plot_beta(**kwargs)
+            case "eigen-value":
+                return self.plot_eigen_value(**kwargs)
+            case "normalized-coupling":
+                return self.plot_normalized_coupling(**kwargs)
+            case "overlap":
+                return self.plot_overlap(**kwargs)
+            case "adiabatic":
+                return self.plot_adiabatic(**kwargs)
+            case "field":
+                return self.plot_field(**kwargs)
+            case "beating-length":
+                return self.plot_beating_length(**kwargs)
+            case "normalized-adiabatic":
+                return self.plot_normalized_adiabatic(**kwargs)
+            case _:
+                raise ValueError(
+                    f"Invalid plot type: {plot_type}. Options are: geometry, index, beta, eigen-value, adiabatic, normalized-adiabatic, normalized-coupling, field, beating-length"
+                )
+
+    @parse_filename
+    def generate_pdf_report(
+        self,
+        filename: str = "auto",
+        itr_list: list[float] | None = None,
+        slice_list: list[int] | None = None,
+        mode_of_interest: list[SUPERMODE] | str = "all",
+        combination: str = "pairs",
+    ) -> None:
+        """
+        Generate a full report of the coupler properties as a PDF file.
+
+        Parameters
+        ----------
+        filename : str, optional, default="auto"
+            Name of the report file to be output. If "auto", a default name will be generated based on the current timestamp.
+        directory : str, optional, default='.'
+            Directory to save the report.
+        itr_list : list of float, optional, default=None
+            List of ITR values to evaluate the mode field. If None, all available ITR values will be used.
+        slice_list : list of int, optional, default=None
+            List of slice values to evaluate the mode field. If None, all available slices will be used.
+        mode_of_interest : list, optional, default='all'
+            List of modes to consider in the report. If 'all', all available modes will be included.
+        combination : str, optional, default='specific'
+            Method for selecting mode combinations ('specific' or 'pairs').
+
+        Examples
+        --------
+        >>> superset_plots.generate_pdf_report(filename="coupler_report", itr_list=[0.1, 0.2, 0.5], mode_of_interest=['LP01', 'LP11'])
+        This will generate a PDF report named 'coupler_report.pdf' containing plots for the specified modes at given ITR values.
+        """
+        kwargs = dict(show=False, mode_of_interest=mode_of_interest)
+
+        figure_list = [
+            self.plot_geometry(show=False),
+            self.plot_field(itr_list=itr_list, slice_list=slice_list, **kwargs),
+            self.plot_index(**kwargs),
+            self.plot_beta(**kwargs),
+            self.plot_normalized_coupling(**kwargs, combination=combination),
+            self.plot_adiabatic(show=False, combination=combination),
+        ]
+
+        pp = PdfPages(filename.with_suffix(".pdf"))
+
+        for fig in figure_list:
+            fig.savefig(pp, format="pdf")
+
+        pp.close()
+
+        plt.close()

SuPyMode/utils.py

@@ -0,0 +1,196 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+from typing import Iterable, Callable, List, Union
+import numpy
+import pickle
+from pathlib import Path
+from functools import wraps
+from pathvalidate import sanitize_filepath
+import logging
+
+from SuPyMode.directories import user_data_directory
+
+
+def interpret_slice_number_and_itr(
+    itr_baseline: numpy.ndarray,
+    slice_list: Union[int, List[int], None] = None,
+    itr_list: Union[float, List[float], None] = None,
+    sort_slice_number: bool = True,
+) -> tuple:
+    """
+    Interprets slice numbers and corresponding inverse taper ratios (ITRs), returning arrays of slice numbers
+    and their respective ITRs based on the provided lists of slices and ITRs.
+
+    Parameters:
+        - itr_baseline (numpy.ndarray): Array of baseline ITR values, typically equidistant, representing the full range.
+        - slice_list (Union[int, List[int]], optional): A single slice index or a list of slice indices. If None and itr_list is also None, defaults to [0, -1].
+        - itr_list (Union[float, List[float]], optional): A single ITR value or a list of ITR values to be converted to slice indices. If provided, slice_list defaults to an empty list unless specified.
+        - sort_slice_number (bool, optional): Whether to sort the resulting slice numbers and corresponding ITRs in descending order. Default is False.
+
+    Returns:
+        - Tuple[numpy.ndarray, numpy.ndarray]: Two numpy arrays, the first containing slice indices and the second containing the corresponding ITR values.
+
+    Raises:
+        - ValueError: If the provided ITRs or slice indices are outside the bounds of the baseline ITR array.
+    """
+
+    if slice_list is None and itr_list is None:
+        slice_list = [0, -1]
+    elif itr_list is not None and slice_list is None:
+        slice_list = []
+
+    slice_list = numpy.atleast_1d(slice_list)
+    itr_list = numpy.atleast_1d(itr_list) if itr_list is not None else numpy.array([])
+
+    # Logic to convert ITRs to slice indices, and combine lists
+    slice_from_itr = (
+        itr_to_slice(itr_baseline, itr=itr_list)
+        if itr_list.size > 0
+        else numpy.array([])
+    )
+    total_slice_list = numpy.concatenate([slice_list, slice_from_itr]).astype(int)
+
+    total_itr_list = slice_to_itr(itr_baseline, slice_number=total_slice_list)
+
+    # Sort if required
+    if sort_slice_number:
+        indices = numpy.argsort(total_itr_list)
+        total_slice_list = total_slice_list[indices][::-1]  # Descending order
+        total_itr_list = total_itr_list[indices][::-1]
+
+    return total_slice_list, total_itr_list
+
+
+def parse_filename(save_function: Callable) -> Callable:
+    @wraps(save_function)
+    def wrapper(superset, filename: str, directory: str = "auto", **kwargs):
+        directory = user_data_directory if directory == "auto" else directory
+
+        filename = Path(filename)
+
+        filename = sanitize_filepath(filename)
+
+        filename = Path(directory).joinpath(filename)
+
+        save_function(superset, filename=filename, **kwargs)
+
+        logging.info(f"Saving data into: {filename}")
+
+        return filename
+
+    return wrapper
+
+
+def load_superset(filename: str, directory: str = "default"):
+    """
+    Saves the superset instance as a serialized pickle file.
+
+    :param      filename:  The filename
+    :type       filename:  str
+    """
+    if directory == "default":
+        directory = user_data_directory
+
+    filename = Path(directory).joinpath(filename).with_suffix(".pickle")
+
+    with open(filename, "rb") as input_file:
+        superset = pickle.load(input_file)
+
+    return superset
+
+
+def itr_to_slice(itr_list: numpy.ndarray, itr: float | list) -> int | list:
+    """
+    Converts an inverse taper ratio (itr) or a list of itrs to their closest corresponding slice numbers
+    in the provided itr_list. If a single itr is provided, returns a single slice number. If a list of
+    itrs is provided, returns a list of slice numbers.
+
+    Parameters:
+        itr_list (np.ndarray): Array of itrs, typically representing a continuous range of values.
+        itr (Union[float, List[float]]): A single itr value or a list of itr values.
+
+    Returns:
+        Union[int, List[int]]: The closest slice number or list of slice numbers corresponding to the provided itr or itrs.
+
+    """
+    if numpy.isscalar(itr):
+        return numpy.argmin(abs(itr_list - itr))
+
+    if isinstance(itr_list, Iterable):
+        return [numpy.argmin(abs(itr_list - value)) for value in itr]
+
+    raise TypeError("itr must be a float or a list of floats.")
+
+
+def slice_to_itr(itr_list: numpy.ndarray, slice_number: int | list) -> float | list:
+    """
+    Converts a slice number or a list of slice numbers to their corresponding inverse taper ratios (itrs)
+    from the provided itr_list. If a single slice number is provided, returns a single itr value.
+    If a list of slice numbers is provided, returns a list of itr values.
+
+    Parameters:
+        - itr_list (np.ndarray): Array of itrs, typically representing a continuous range of values.
+        - slice_number (Union[int, List[int]]): A single slice index or a list of slice indices.
+
+    Returns:
+         - Union[float, List[float]]: The itr or list of itrs corresponding to the given slice numbers.
+
+    """
+    if numpy.isscalar(slice_number):
+        return itr_list[slice_number]
+
+    if isinstance(slice_number, Iterable):
+        return numpy.take(itr_list, slice_number)
+
+    raise TypeError("slice_number must be a scalar or a list of scalars.")
+
+
+def get_symmetrized_vector(
+    vector: numpy.ndarray, symmetry_type: str = "last"
+) -> numpy.ndarray:
+    """
+    Generate a symmetric version of the input vector based on the specified symmetry type.
+
+    Parameters:
+    -----------
+    vector : numpy.ndarray
+        A one-dimensional array for which the symmetric version is to be calculated.
+    symmetry_type : str, optional
+        Type of symmetry to apply. Supported types:
+        - 'last': Symmetrize using the last element as reference.
+        - 'first': Symmetrize using the first element as reference.
+        Default is 'last'.
+
+    Returns
+    -------
+    numpy.ndarray
+        A new vector that is the symmetrized version of the input vector.
+
+    Raises:
+    -------
+    ValueError
+        If the input vector is not one-dimensional or the symmetry type is unsupported.
+    """
+
+    if vector.ndim != 1:
+        raise ValueError(
+            f"Expected a 1-dimensional vector, but got a {vector.ndim}-dimensional vector instead."
+        )
+
+    if symmetry_type.lower() not in ["last", "first"]:
+        raise ValueError("Symmetry type must be 'last' or 'first'.")
+
+    size = len(vector)
+    dx = numpy.diff(vector)[0]  # More robust than assuming vector[1] - vector[0]
+
+    if symmetry_type.lower() == "last":
+        start_value = vector[-1]
+        expanded = numpy.arange(0, 2 * size - 1) * dx
+        return start_value - expanded[::-1] if dx > 0 else start_value + expanded[::-1]
+    else:  # 'first'
+        start_value = vector[0]
+        expanded = numpy.arange(0, 2 * size - 1) * dx
+        return start_value + expanded
+
+
+# -