Add option to use dispersion fitter without rich.progress

Author: caseyflex

tests/test_plugins/test_dispersion_fitter.py

@@ -4,6 +4,7 @@
 import pytest
 import responses
 import tidy3d as td
+from tidy3d.components.dispersion_fitter import fit
 from tidy3d.exceptions import SetupError, ValidationError
 from tidy3d.plugins.dispersion import (
     AdvancedFastFitterParam,

tidy3d/components/dispersion_fitter.py

@@ -7,11 +7,10 @@
 import numpy as np
 import scipy
 from pydantic.v1 import Field, NonNegativeFloat, PositiveFloat, PositiveInt, validator
-from rich.progress import Progress
 
 from ..constants import fp_eps
 from ..exceptions import ValidationError
-from ..log import get_logging_console, log
+from ..log import Progress, get_logging_console, log
 from .base import Tidy3dBaseModel, cached_property, skip_if_fields_missing
 from .types import ArrayComplex1D, ArrayComplex2D, ArrayFloat1D, ArrayFloat2D
 
@@ -823,7 +822,6 @@ def fit(
         The dispersive medium parameters have the form (resp_inf, poles, residues)
         and are in the original unscaled units.
     """
-
     if max_num_poles < min_num_poles:
         raise ValidationError(
             "Dispersion fitter cannot have 'max_num_poles' less than 'min_num_poles'."
@@ -864,86 +862,82 @@ def make_configs():
 
     with Progress(console=get_logging_console()) as progress:
         task = progress.add_task(
-            f"Fitting to weighted RMS of {tolerance_rms}...",
+            description=f"Fitting to weighted RMS of {tolerance_rms}...",
             total=len(configs),
             visible=init_model.show_progress,
         )
 
-        while not progress.finished:
-            # try different initial pole configurations
-            for num_poles, relaxed, smooth, logspacing, optimize_eps_inf in configs:
-                model = init_model.updated_copy(
-                    num_poles=num_poles,
-                    relaxed=relaxed,
-                    smooth=smooth,
-                    logspacing=logspacing,
-                    optimize_eps_inf=optimize_eps_inf,
+        # try different initial pole configurations
+        for num_poles, relaxed, smooth, logspacing, optimize_eps_inf in configs:
+            model = init_model.updated_copy(
+                num_poles=num_poles,
+                relaxed=relaxed,
+                smooth=smooth,
+                logspacing=logspacing,
+                optimize_eps_inf=optimize_eps_inf,
+            )
+            model = _fit_fixed_parameters((min_num_poles, max_num_poles), model)
+
+            if model.rms_error < best_model.rms_error:
+                log.debug(
+                    f"Fitter: possible improved fit with "
+                    f"rms_error={model.rms_error:.3g} found using "
+                    f"relaxed={model.relaxed}, "
+                    f"smooth={model.smooth}, "
+                    f"logspacing={model.logspacing}, "
+                    f"optimize_eps_inf={model.optimize_eps_inf}, "
+                    f"loss_in_bounds={model.loss_in_bounds}, "
+                    f"passivity_optimized={model.passivity_optimized}, "
+                    f"sellmeier_passivity={model.sellmeier_passivity}."
                 )
-                model = _fit_fixed_parameters((min_num_poles, max_num_poles), model)
-
-                if model.rms_error < best_model.rms_error:
-                    log.debug(
-                        f"Fitter: possible improved fit with "
-                        f"rms_error={model.rms_error:.3g} found using "
-                        f"relaxed={model.relaxed}, "
-                        f"smooth={model.smooth}, "
-                        f"logspacing={model.logspacing}, "
-                        f"optimize_eps_inf={model.optimize_eps_inf}, "
-                        f"loss_in_bounds={model.loss_in_bounds}, "
-                        f"passivity_optimized={model.passivity_optimized}, "
-                        f"sellmeier_passivity={model.sellmeier_passivity}."
-                    )
-                    if model.loss_in_bounds and model.sellmeier_passivity:
-                        best_model = model
-                    else:
-                        if (
-                            not warned_about_passivity_num_iters
-                            and model.passivity_num_iters_too_small
-                        ):
-                            warned_about_passivity_num_iters = True
-                            log.warning(
-                                "Did not finish enforcing passivity in dispersion fitter. "
-                                "If the fit is not good enough, consider increasing "
-                                "'AdvancedFastFitterParam.passivity_num_iters'."
-                            )
-                        if (
-                            not warned_about_slsqp_constraint_scale
-                            and model.slsqp_constraint_scale_too_small
-                        ):
-                            warned_about_slsqp_constraint_scale = True
-                            log.warning(
-                                "SLSQP constraint scale may be too small. "
-                                "If the fit is not good enough, consider increasing "
-                                "'AdvancedFastFitterParam.slsqp_constraint_scale'."
-                            )
+                if model.loss_in_bounds and model.sellmeier_passivity:
+                    best_model = model
+                else:
+                    if not warned_about_passivity_num_iters and model.passivity_num_iters_too_small:
+                        warned_about_passivity_num_iters = True
+                        log.warning(
+                            "Did not finish enforcing passivity in dispersion fitter. "
+                            "If the fit is not good enough, consider increasing "
+                            "'AdvancedFastFitterParam.passivity_num_iters'."
+                        )
+                    if (
+                        not warned_about_slsqp_constraint_scale
+                        and model.slsqp_constraint_scale_too_small
+                    ):
+                        warned_about_slsqp_constraint_scale = True
+                        log.warning(
+                            "SLSQP constraint scale may be too small. "
+                            "If the fit is not good enough, consider increasing "
+                            "'AdvancedFastFitterParam.slsqp_constraint_scale'."
+                        )
+            progress.update(
+                task,
+                advance=1,
+                description=f"Best weighted RMS error so far: {best_model.rms_error:.3g}",
+                refresh=True,
+            )
+
+            # if below tolerance, return
+            if best_model.rms_error < tolerance_rms:
                 progress.update(
                     task,
-                    advance=1,
-                    description=f"Best weighted RMS error so far: {best_model.rms_error:.3g}",
+                    completed=len(configs),
+                    description=f"Best weighted RMS error: {best_model.rms_error:.3g}",
                     refresh=True,
                 )
-
-                # if below tolerance, return
-                if best_model.rms_error < tolerance_rms:
-                    progress.update(
-                        task,
-                        completed=len(configs),
-                        description=f"Best weighted RMS error: {best_model.rms_error:.3g}",
-                        refresh=True,
-                    )
+                log.info(
+                    "Found optimal fit with weighted RMS error %.3g",
+                    best_model.rms_error,
+                )
+                if best_model.show_unweighted_rms:
                     log.info(
-                        "Found optimal fit with weighted RMS error %.3g",
-                        best_model.rms_error,
-                    )
-                    if best_model.show_unweighted_rms:
-                        log.info(
-                            "Unweighted RMS error %.3g",
-                            best_model.unweighted_rms_error,
-                        )
-                    return (
-                        best_model.pole_residue,
-                        best_model.rms_error,
+                        "Unweighted RMS error %.3g",
+                        best_model.unweighted_rms_error,
                     )
+                return (
+                    best_model.pole_residue,
+                    best_model.rms_error,
+                )
 
     # if exited loop, did not reach tolerance (warn)
     progress.update(
@@ -967,3 +961,57 @@ def make_configs():
         best_model.pole_residue,
         best_model.rms_error,
     )
+
+
+def constant_loss_tangent_model(
+    eps_real: float,
+    loss_tangent: float,
+    frequency_range: Tuple[float, float],
+    max_num_poles: PositiveInt = DEFAULT_MAX_POLES,
+    number_sampling_frequency: PositiveInt = 10,
+    tolerance_rms: NonNegativeFloat = DEFAULT_TOLERANCE_RMS,
+    scale_factor: float = 1,
+) -> Tuple[Tuple[float, ArrayComplex1D, ArrayComplex1D], float]:
+    """Fit a constant loss tangent material model.
+
+    Parameters
+    ----------
+    eps_real : float
+        Real part of permittivity
+    loss_tangent : float
+        Loss tangent.
+    frequency_range : Tuple[float, float]
+        Freqquency range for the material to exhibit constant loss tangent response.
+    max_num_poles : PositiveInt, optional
+        Maximum number of poles in the model.
+    number_sampling_frequency : PositiveInt, optional
+        Number of sampling frequencies to compute RMS error for fitting.
+    tolerance_rms : float, optional
+        Weighted RMS error below which the fit is successful and the result is returned.
+    scale_factor : PositiveFloat, optional
+        Factor to rescale frequency by before fitting.
+
+    Returns
+    -------
+    Tuple[Tuple[float, ArrayComplex1D, ArrayComplex1D], float]
+        Best fitting result: (dispersive medium parameters, weighted RMS error).
+        The dispersive medium parameters have the form (resp_inf, poles, residues)
+        and are in the original unscaled units.
+    """
+    if number_sampling_frequency < 2:
+        frequencies = np.array([np.mean(frequency_range)])
+    else:
+        frequencies = np.linspace(frequency_range[0], frequency_range[1], number_sampling_frequency)
+    eps_real_array = np.ones_like(frequencies) * eps_real
+    loss_tangent_array = np.ones_like(frequencies) * loss_tangent
+
+    omega_data = frequencies * 2 * np.pi
+    eps_complex = eps_real_array * (1 + 1j * loss_tangent_array)
+
+    return fit(
+        omega_data=omega_data,
+        resp_data=eps_complex,
+        max_num_poles=max_num_poles,
+        tolerance_rms=tolerance_rms,
+        scale_factor=scale_factor,
+    )

tidy3d/log.py

@@ -1,6 +1,7 @@
 """Logging for Tidy3d."""
 
 import inspect
+from contextlib import contextmanager
 from datetime import datetime
 from typing import Callable, List, Tuple, Union
 
@@ -442,3 +443,29 @@ def get_logging_console() -> Console:
     if "console" not in log.handlers:
         set_logging_console()
     return log.handlers["console"].console
+
+
+class NoOpProgress:
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args, **kwargs):
+        pass
+
+    def add_task(self, *args, **kwargs):
+        pass
+
+    def update(self, *args, **kwargs):
+        pass
+
+
+@contextmanager
+def Progress(console):
+    try:
+        from rich.progress import Progress
+
+        with Progress(console=console) as progress:
+            yield progress
+    except ImportError:
+        with NoOpProgress() as progress:
+            yield progress

tidy3d/plugins/dispersion/fit_fast.py

@@ -7,9 +7,13 @@
 import numpy as np
 from pydantic.v1 import NonNegativeFloat, PositiveInt
 
-from ...components.dispersion_fitter import AdvancedFastFitterParam, fit
+from ...components.dispersion_fitter import (
+    AdvancedFastFitterParam,
+    constant_loss_tangent_model,
+    fit,
+)
 from ...components.medium import PoleResidue
-from ...constants import C_0, HBAR
+from ...constants import HBAR
 from .fit import DispersionFitter
 
 # numerical tolerance for pole relocation for fast fitter
@@ -144,15 +148,18 @@ def constant_loss_tangent_model(
         :class:`.PoleResidue
             Best results of multiple fits.
         """
-        if number_sampling_frequency < 2:
-            frequencies = np.array([np.mean(frequency_range)])
-        else:
-            frequencies = np.linspace(
-                frequency_range[0], frequency_range[1], number_sampling_frequency
-            )
-        wvl_um = C_0 / frequencies
-        eps_real_array = np.ones_like(frequencies) * eps_real
-        loss_tangent_array = np.ones_like(frequencies) * loss_tangent
-        fitter = cls.from_loss_tangent(wvl_um, eps_real_array, loss_tangent_array)
-        material, _ = fitter.fit(max_num_poles=max_num_poles, tolerance_rms=tolerance_rms)
-        return material
+        params, _ = constant_loss_tangent_model(
+            eps_real=eps_real,
+            loss_tangent=loss_tangent,
+            frequency_range=frequency_range,
+            max_num_poles=max_num_poles,
+            number_sampling_frequency=number_sampling_frequency,
+            tolerance_rms=tolerance_rms,
+            scale_factor=HBAR,
+        )
+
+        eps_inf, poles, residues = params
+
+        medium = PoleResidue(eps_inf=eps_inf, poles=list(zip(poles, residues)))
+
+        return medium