"""Fit PoleResidue Dispersion models to optical NK data"""
from __future__ import annotations
import codecs
import csv
from typing import List, Optional, Tuple
import numpy as np
import requests
import scipy.optimize as opt
from pydantic.v1 import Field, validator
from rich.progress import Progress
from tidy3d.web.core.environment import Env
from ...components.base import Tidy3dBaseModel, cached_property, skip_if_fields_missing
from ...components.medium import AbstractMedium, PoleResidue
from ...components.types import ArrayFloat1D, Ax
from ...components.viz import add_ax_if_none
from ...constants import C_0, HBAR, MICROMETER
from ...exceptions import SetupError, ValidationError, WebError
from ...log import get_logging_console, log
[docs]
class DispersionFitter(Tidy3dBaseModel):
"""Tool for fitting refractive index data to get a
dispersive medium described by :class:`.PoleResidue` model."""
wvl_um: ArrayFloat1D = Field(
...,
title="Wavelength data",
description="Wavelength data in micrometers.",
units=MICROMETER,
)
n_data: ArrayFloat1D = Field(
...,
title="Index of refraction data",
description="Real part of the complex index of refraction.",
)
k_data: ArrayFloat1D = Field(
None,
title="Extinction coefficient data",
description="Imaginary part of the complex index of refraction.",
)
wvl_range: Tuple[Optional[float], Optional[float]] = Field(
(None, None),
title="Wavelength range [wvl_min,wvl_max] for fitting",
description="Truncate the wavelength, n and k data to the wavelength range '[wvl_min, "
"wvl_max]' for fitting.",
units=MICROMETER,
)
@validator("wvl_um", always=True)
def _setup_wvl(cls, val):
"""Convert wvl_um to a numpy array."""
if val.size == 0:
raise ValidationError("Wavelength data cannot be empty.")
return val
@validator("n_data", always=True)
@skip_if_fields_missing(["wvl_um"])
def _ndata_length_match_wvl(cls, val, values):
"""Validate n_data"""
if val.shape != values["wvl_um"].shape:
raise ValidationError("The length of 'n_data' doesn't match 'wvl_um'.")
return val
@validator("k_data", always=True)
@skip_if_fields_missing(["wvl_um"])
def _kdata_setup_and_length_match(cls, val, values):
"""Validate the length of k_data, or setup k if it's None."""
if val is None:
return np.zeros_like(values["wvl_um"])
if val.shape != values["wvl_um"].shape:
raise ValidationError("The length of 'k_data' doesn't match 'wvl_um'.")
return val
@cached_property
def data_in_range(self) -> Tuple[ArrayFloat1D, ArrayFloat1D, ArrayFloat1D]:
"""Filter the wavelength-nk data to wavelength range for fitting.
Returns
-------
Tuple[ArrayFloat1D, ArrayFloat1D, ArrayFloat1D]
Filtered wvl_um, n_data, k_data
"""
ind_select = np.ones(self.wvl_um.shape, dtype=bool)
if self.wvl_range[0] is not None:
ind_select = np.logical_and(self.wvl_um >= self.wvl_range[0], ind_select)
if self.wvl_range[1] is not None:
ind_select = np.logical_and(self.wvl_um <= self.wvl_range[1], ind_select)
if not np.any(ind_select):
raise SetupError("No data within 'wvl_range'")
return self.wvl_um[ind_select], self.n_data[ind_select], self.k_data[ind_select]
@cached_property
def lossy(self) -> bool:
"""Find out if the medium is lossy or lossless
based on the filtered input data.
Returns
-------
bool
True for lossy medium; False for lossless medium
"""
_, _, k_data = self.data_in_range
return k_data is not None and np.any(k_data)
@property
def eps_data(self) -> complex:
"""Convert filtered input n(k) data into complex permittivity.
Returns
-------
complex
Complex-valued relative permittivty.
"""
_, n_data, k_data = self.data_in_range
return AbstractMedium.nk_to_eps_complex(n=n_data, k=k_data)
@property
def freqs(self) -> Tuple[float, ...]:
"""Convert filtered input wavelength data to frequency.
Returns
-------
Tuple[float, ...]
Frequency array converted from filtered input wavelength data
"""
wvl_um, _, _ = self.data_in_range
return C_0 / wvl_um
@property
def frequency_range(self) -> Tuple[float, float]:
"""Frequency range of filtered input data
Returns
-------
Tuple[float, float]
The minimal frequency and the maximal frequency
"""
return self.freqs.min(), self.freqs.max()
@staticmethod
def _unpack_coeffs(coeffs):
"""Unpack coefficient vector into complex pole parameters.
Parameters
----------
coeffs : np.ndarray[real]
Array of real coefficients for the pole residue fit.
Returns
-------
Tuple[np.ndarray[complex], np.ndarray[complex]]
"a" and "c" poles for the PoleResidue model.
"""
if len(coeffs) % 4 != 0:
raise ValueError(f"len(coeffs) must be multiple of 4, got {len(coeffs)=}.")
a_real = coeffs[0::4]
a_imag = coeffs[1::4]
c_real = coeffs[2::4]
c_imag = coeffs[3::4]
poles_a = a_real + 1j * a_imag
poles_c = c_real + 1j * c_imag
return poles_a, poles_c
@staticmethod
def _pack_coeffs(pole_a, pole_c):
"""Pack complex a and c pole parameters into coefficient array.
Parameters
----------
pole_a : np.ndarray[complex]
Array of complex "a" poles for the PoleResidue dispersive model.
pole_c : np.ndarray[complex]
Array of complex "c" poles for the PoleResidue dispersive model.
Returns
-------
np.ndarray[float]
Array of real coefficients for the pole residue fit.
"""
stacked_coeffs = np.stack((pole_a.real, pole_a.imag, pole_c.real, pole_c.imag), axis=1)
return stacked_coeffs.flatten()
@staticmethod
def _coeffs_to_poles(coeffs):
"""Convert model coefficients to poles.
Parameters
----------
coeffs : np.ndarray[float]
Array of real coefficients for the pole residue fit.
Returns
-------
List[Tuple[complex, complex]]
List of complex poles (a, c)
"""
coeffs_scaled = coeffs / HBAR
poles_a, poles_c = DispersionFitter._unpack_coeffs(coeffs_scaled)
return list(zip(poles_a, poles_c))
@staticmethod
def _poles_to_coeffs(poles):
"""Convert poles to model coefficients.
Parameters
----------
poles : List[Tuple[complex, complex]]
List of complex poles (a, c)
Returns
-------
np.ndarray[float]
Array of real coefficients for the pole residue fit.
"""
poles = np.array(poles, dtype=complex)
coeffs = DispersionFitter._pack_coeffs(poles[:, 0], poles[:, 1])
return coeffs * HBAR
@staticmethod
def _eV_to_Hz(f_eV: float):
"""Convert frequency in unit of eV to Hz.
Parameters
----------
f_eV : float
Frequency in unit of eV
"""
return f_eV / (HBAR * 2 * np.pi)
@staticmethod
def _Hz_to_eV(f_Hz: float):
"""Convert frequency in unit of Hz to eV.
Parameters
----------
f_Hz : float
Frequency in unit of Hz
"""
return f_Hz * HBAR * 2 * np.pi
[docs]
def fit(
self,
num_poles: int = 1,
num_tries: int = 50,
tolerance_rms: float = 1e-2,
guess: PoleResidue = None,
) -> Tuple[PoleResidue, float]:
"""Fit data a number of times and returns best results.
Parameters
----------
num_poles : int, optional
Number of poles in the model.
num_tries : int, optional
Number of optimizations to run with different initial guesses.
tolerance_rms : float, optional
RMS error below which the fit is successful and the result is returned.
guess : :class:`.PoleResidue` = None
A :class:`.PoleResidue` medium to use as the initial guess in the first optimization
run.
Returns
-------
Tuple[:class:`.PoleResidue`, float]
Best results of multiple fits: (dispersive medium, RMS error).
"""
# Run it a number of times.
best_medium = None
best_rms = np.inf
with Progress(console=get_logging_console()) as progress:
task = progress.add_task(
f"Fitting with {num_poles} to RMS of {tolerance_rms}...", total=num_tries
)
while not progress.finished:
# if guess is provided use it in the first optimization run
if guess is not None and progress.tasks[0].completed == 0:
medium, rms_error = self._fit_single(num_poles=num_poles, guess=guess)
else:
medium, rms_error = self._fit_single(num_poles=num_poles)
# if improvement, set the best RMS and coeffs
if rms_error < best_rms:
best_rms = rms_error
best_medium = medium
progress.update(
task,
advance=1,
description=f"Best RMS error so far: {best_rms:.3g}",
refresh=True,
)
# if below tolerance, return
if best_rms < tolerance_rms:
progress.update(
task,
completed=num_tries,
description=f"Best RMS error: {best_rms:.3g}",
refresh=True,
)
log.info("Found optimal fit with RMS error %.3g", best_rms)
return best_medium, best_rms
# if exited loop, did not reach tolerance (warn)
log.warning("Unable to fit with RMS error under 'tolerance_rms' of %.3g", tolerance_rms)
log.info("Returning best fit with RMS error %.3g", best_rms)
return best_medium, best_rms
def _make_medium(self, coeffs):
"""Return medium from coeffs from optimizer.
Parameters
----------
coeffs : np.ndarray[float]
Array of real coefficients for the pole residue fit.
Returns
-------
:class:`.PoleResidue`
Dispersive medium corresponding to this set of ``coeffs``.
"""
poles_complex = DispersionFitter._coeffs_to_poles(coeffs)
return PoleResidue(poles=poles_complex, frequency_range=self.frequency_range)
def _fit_single(
self,
num_poles: int = 3,
guess: PoleResidue = None,
) -> Tuple[PoleResidue, float]:
"""Perform a single fit to the data and return optimization result.
Parameters
----------
num_poles : int = 3
Number of poles in the model.
guess : :class:`.PoleResidue` = None
A PoleResidue object to use a guess instead of a random one.
Returns
-------
Tuple[:class:`.PoleResidue`, float]
Results of single fit: (dispersive medium, RMS error).
"""
# NOTE: Not used
def constraint(coeffs, _grad=None):
"""Evaluate the nonlinear stability criterion of Hongjin Choi, Jae-Woo Baek, and
Kyung-Young Jung, "Comprehensive Study on Numerical Aspects of Modified Lorentz Model
Based Dispersive FDTD Formulations," IEEE TAP 2019.
Parameters
----------
coeffs : np.ndarray[float]
Array of real coefficients for the pole residue fit.
_grad : np.ndarray[float]
Gradient of ``constraint`` w.r.t coeffs, not used.
Returns
-------
float
Value of constraint.
"""
poles_a, poles_c = DispersionFitter._unpack_coeffs(coeffs)
a_real = poles_a.real
a_imag = poles_a.imag
c_real = poles_c.real
c_imag = poles_c.imag
prstar = a_real * c_real + a_imag * c_imag
res = 2 * prstar * a_real - c_real * (a_real * a_real + a_imag * a_imag)
res[res >= 0] = 0
return np.sum(res)
def objective(coeffs, _grad=None):
"""Objective function for fit
Parameters
----------
coeffs : np.ndarray[float]
Array of real coefficients for the pole residue fit.
_grad : np.ndarray[float]
Gradient of ``objective`` w.r.t coeffs, not used.
Returns
-------
float
RMS error corresponding to current coeffs.
"""
medium = self._make_medium(coeffs)
eps_model = medium.eps_model(self.freqs)
residual = self.eps_data - eps_model
# cons = constraint(coeffs, _grad)
return np.sqrt(np.sum(np.square(np.abs(residual))) / len(self.eps_data))
# set initial guess
num_coeffs = num_poles * 4
if guess is not None:
if len(guess.poles) != num_poles:
raise ValueError(
f"The number of poles ({len(guess.poles)}) in provided guess 'PoleResidue' "
f"medium does not match argument 'num_poles' = {num_poles})"
)
coeffs0 = self._poles_to_coeffs(guess.poles)
else:
coeffs0 = 2 * (np.random.random(num_coeffs) - 0.5)
# set bounds
bounds_upper = np.zeros(num_coeffs, dtype=float)
bounds_lower = np.zeros(num_coeffs, dtype=float)
if self.lossy:
# if lossy, the real parts can take on values
bounds_lower[0::4] = -np.inf
bounds_upper[2::4] = np.inf
coeffs0[0::4] = -np.abs(coeffs0[0::4])
coeffs0[2::4] = +np.abs(coeffs0[2::4])
else:
# otherwise, they need to be 0
coeffs0[0::2] = 0
bounds_lower[1::2] = -np.inf
bounds_upper[1::2] = np.inf
bounds = list(zip(bounds_lower, bounds_upper))
# TODO: set up constraint properly
scipy_constraint = opt.NonlinearConstraint(constraint, lb=0, ub=np.inf)
# TODO: set options properly
res = opt.minimize(
objective,
coeffs0,
args=(),
method="SLSQP",
bounds=bounds,
constraints=(scipy_constraint,),
tol=1e-7,
callback=None,
options=dict(maxiter=10000),
)
coeffs = res.x
rms_error = objective(coeffs)
# set the latest fit
medium = self._make_medium(coeffs)
return medium, rms_error
[docs]
@add_ax_if_none
def plot(
self,
medium: PoleResidue = None,
wvl_um: ArrayFloat1D = None,
ax: Ax = None,
) -> Ax:
"""Make plot of model vs data, at a set of wavelengths (if supplied).
Parameters
----------
medium : :class:`.PoleResidue` = None
medium containing model to plot against data
wvl_um : ArrayFloat1D = None
Wavelengths to evaluate model at for plot in micrometers.
ax : matplotlib.axes._subplots.Axes = None
Axes to plot the data on, if None, a new one is created.
Returns
-------
matplotlib.axis.Axes
Matplotlib axis corresponding to plot.
"""
ax.plot(self.wvl_um, self.n_data, "x", label="n (data)")
if self.lossy:
ax.plot(self.wvl_um, self.k_data, "+", label="k (data)")
if medium:
if wvl_um is None:
wvl_um = C_0 / self.freqs
eps_model = medium.eps_model(C_0 / wvl_um)
n_model, k_model = AbstractMedium.eps_complex_to_nk(eps_model)
ax.plot(wvl_um, n_model, label="n (model)")
if self.lossy:
ax.plot(wvl_um, k_model, label="k (model)")
ax.set_ylabel("n, k")
ax.set_xlabel("Wavelength ($\\mu m$)")
ax.legend()
return ax
@staticmethod
def _validate_url_load(data_load: List):
"""Validate if the loaded data from URL is valid
The data list should be in this format:
[["wl", "n"],
[float, float],
. .
. .
. .
(if lossy)
["wl", "k"],
[float, float],
. .
. .
. .]]
Parameters
----------
data_load : List
Loaded data from URL
Raises
------
ValidationError
Or other exceptions
"""
has_k = 0
if data_load[0][0] != "wl" or data_load[0][1] != "n":
raise ValidationError(
"Invalid URL. The file should begin with ['wl','n']. "
"Or make sure that you have supplied an appropriate delimiter."
)
for row in data_load[1:]:
if row[0] == "wl":
if row[1] == "k":
has_k += 1
else:
raise ValidationError(
"Invalid URL. The file is not well formatted for ['wl', 'k'] data."
)
else:
# make sure the rest is float type
try:
_ = [float(x) for x in row]
except Exception as e:
raise ValidationError("Invalid URL. Float data cannot be recognized.") from e
if has_k > 1:
raise ValidationError("Invalid URL. Too many k labels.")
[docs]
@classmethod
def from_url(
cls, url_file: str, delimiter: str = ",", ignore_k: bool = False, **kwargs
) -> DispersionFitter:
"""loads :class:`DispersionFitter` from url linked to a csv/txt file that
contains wavelength (micron), n, and optionally k data. Preferred from
refractiveindex.info.
Hint
----
The data file from url should be in this format (delimiter not displayed
here, and note that the strings such as "wl", "n" need to be included
in the file):
* For lossless media::
wl n
[float] [float]
. .
. .
. .
* For lossy media::
wl n
[float] [float]
. .
. .
. .
wl k
[float] [float]
. .
. .
. .
Parameters
----------
url_file : str
Url link to the data file.
e.g. "https://refractiveindex.info/data_csv.php?datafile=database/data-nk/main/Ag/Johnson.yml"
delimiter : str = ","
E.g. in refractiveindex.info, it'll be "," for csv file, and "\\\\t" for txt file.
ignore_k : bool = False
Ignore the k data if they are present, so the fitted material is lossless.
Returns
-------
:class:`DispersionFitter`
A :class:`DispersionFitter` instance.
"""
resp = requests.get(url_file, verify=Env.current.ssl_verify)
try:
resp.raise_for_status()
except Exception as e:
raise WebError("Connection to the website failed. Please provide a valid URL.") from e
data_url = list(
csv.reader(codecs.iterdecode(resp.iter_lines(), "utf-8"), delimiter=delimiter)
)
data_url = list(data_url)
# first validate data
cls._validate_url_load(data_url)
# parsing the data
n_lam = []
k_lam = [] # the two variables contain [wvl_um, n(k)]
has_k = 0 # whether k is in the data
for row in data_url[1:]:
if has_k == 1:
k_lam.append([float(x) for x in row])
elif row[0] == "wl":
has_k += 1
else:
n_lam.append([float(x) for x in row])
n_lam = np.array(n_lam)
k_lam = np.array(k_lam)
if has_k == 1 and not ignore_k:
if n_lam.shape == k_lam.shape and np.allclose(n_lam[:, 0], k_lam[:, 0]):
return cls(wvl_um=n_lam[:, 0], n_data=n_lam[:, 1], k_data=k_lam[:, 1], **kwargs)
raise ValidationError(
"Invalid URL. Both n and k should be provided at each wavelength."
)
return cls(wvl_um=n_lam[:, 0], n_data=n_lam[:, 1], **kwargs)
[docs]
@classmethod
def from_file(cls, fname: str, **loadtxt_kwargs) -> DispersionFitter:
"""Loads :class:`DispersionFitter` from file containing wavelength, n, k data.
Parameters
----------
fname : str
Path to file containing wavelength (um), n, k (optional) data in columns.
**loadtxt_kwargs
Kwargs passed to ``np.loadtxt``, such as ``skiprows``, ``delimiter``.
Hint
----
The data file should be in this format (``delimiter`` and ``skiprows`` can be
customized in ``**loadtxt_kwargs``):
* For lossless media::
wl n
[float] [float]
. .
. .
. .
* For lossy media::
wl n k
[float] [float] [float]
. . .
. . .
. . .
Returns
-------
:class:`DispersionFitter`
A :class:`DispersionFitter` instance.
"""
data = np.loadtxt(fname, **loadtxt_kwargs)
if len(data.shape) != 2:
raise ValueError("data must contain [wavelength, ndata, kdata] in columns")
if data.shape[-1] not in (2, 3):
raise ValueError("data must have either 2 or 3 rows (if k data)")
if data.shape[-1] == 2:
wvl_um, n_data = data.T
k_data = None
else:
wvl_um, n_data, k_data = data.T
return cls(wvl_um=wvl_um, n_data=n_data, k_data=k_data)
[docs]
@classmethod
def from_complex_permittivity(
cls,
wvl_um: ArrayFloat1D,
eps_real: ArrayFloat1D,
eps_imag: ArrayFloat1D = None,
wvl_range: Tuple[Optional[float], Optional[float]] = (None, None),
) -> DispersionFitter:
"""Loads :class:`DispersionFitter` from wavelength and complex relative permittivity data
Parameters
----------
wvl_um : ArrayFloat1D
Wavelength data in micrometers.
eps_real : ArrayFloat1D
Real parts of relative permittivity data
eps_imag : Optional[ArrayFloat1D]
Imaginary parts of relative permittivity data; `None` for lossless medium.
wvg_range : Tuple[Optional[float], Optional[float]]
Wavelength range [wvl_min,wvl_max] for fitting.
Returns
-------
:class:`DispersionFitter`
A :class:`DispersionFitter` instance.
"""
if eps_imag is None:
n, _ = AbstractMedium.eps_complex_to_nk(eps_real + 0j)
return cls(wvl_um=wvl_um, n_data=n, wvl_range=wvl_range)
n, k = AbstractMedium.eps_complex_to_nk(eps_real + eps_imag * 1j)
return cls(wvl_um=wvl_um, n_data=n, k_data=k, wvl_range=wvl_range)
[docs]
@classmethod
def from_loss_tangent(
cls,
wvl_um: ArrayFloat1D,
eps_real: ArrayFloat1D,
loss_tangent: ArrayFloat1D,
wvl_range: Tuple[Optional[float], Optional[float]] = (None, None),
) -> DispersionFitter:
"""Loads :class:`DispersionFitter` from wavelength and loss tangent data.
Parameters
----------
wvl_um : ArrayFloat1D
Wavelength data in micrometers.
eps_real : ArrayFloat1D
Real parts of relative permittivity data
loss_tangent : Optional[ArrayFloat1D]
Loss tangent data, defined as the ratio of imaginary and real parts of permittivity.
wvl_range : Tuple[Optional[float], Optional[float]]
Wavelength range [wvl_min,wvl_max] for fitting.
Returns
-------
:class:`DispersionFitter`
A :class:`DispersionFitter` instance.
"""
eps_complex = AbstractMedium.eps_loss_tangent_to_eps_complex(eps_real, loss_tangent)
n, k = AbstractMedium.eps_complex_to_nk(eps_complex)
return cls(wvl_um=wvl_um, n_data=n, k_data=k, wvl_range=wvl_range)
