construct medium from a transfer function describing the differential…

… equation relating polarization current density to electric field

CHANGELOG.md

@@ -21,6 +21,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added optimization methods to the Design plugin. The plugin has been expanded to include Bayesian optimization, genetic algorithms and particle swarm optimization. Explanations of these methods are available in new and updated notebooks.
 - Added new support functions for the Design plugin: automated batching of `Simulation` objects, and summary functions with `DesignSpace.estimate_cost` and `DesignSpace.summarize`.
 - Added validation and repair methods for `TriangleMesh` with inward-facing normals.
+- Added `from_admittance_coeffs` to `PoleResidue`, which allows for directly constructing a `PoleResidue` medium from an admittance function in the Laplace domain.
 
 ### Changed
 - Priority is given to `snapping_points` in `GridSpec` when close to structure boundaries, which reduces the chance of them being skipped.

tests/test_components/test_medium.py

@@ -776,3 +776,71 @@ def create_mediums(n_dataset):
         n_data2 = np.vstack((n_data[0, :, :, :].reshape(1, Ny, Nz, Nf), n_data))
         n_dataset2 = td.ScalarFieldDataArray(n_data2, coords=dict(x=X2, y=Y, z=Z, f=freqs))
         create_mediums(n_dataset=n_dataset2)
+
+
+def test_medium_from_admittance_coeffs():
+    """Test that ``from_admittance_coeffs`` produces same PoleResidue model as
+    conversions from Drude and Lorentz models. Also test some special cases."""
+    freqs = np.linspace(0.01, 1, 1001)
+    twopi = 2 * np.pi
+    m_DR = td.Drude(eps_inf=1.0, coeffs=[(1.5, 3)])
+    # test from transfer function using Drude model
+    f1 = m_DR.coeffs[0][0]
+    d1 = m_DR.coeffs[0][1]
+    # admittance function in Laplace domain (a/b) modeling Drude differential equation is:
+    a = np.array([0, td.EPSILON_0 * (twopi * f1) ** 2, 0])
+    b = np.array([0, twopi * d1, 1])
+
+    m_transfer = td.PoleResidue.from_admittance_coeffs(a, b)
+    assert np.allclose(
+        m_transfer.eps_model(freqs),
+        m_DR.eps_model(freqs),
+    )
+
+    # test from transfer function using Lorentz model
+    m_L = td.Lorentz(eps_inf=1.0, coeffs=[(1.5, 3, 5)])
+    deps = m_L.coeffs[0][0]
+    f1 = m_L.coeffs[0][1]
+    d1 = m_L.coeffs[0][2]
+    # admittance function in Laplace domain (a/b) modeling Lorentz differential equation is:
+    a = np.array([0, td.EPSILON_0 * (twopi * f1) ** 2 * deps, 0])
+    b = np.array([(twopi * f1) ** 2, 2 * twopi * d1, 1])
+
+    m_transfer = td.PoleResidue.from_admittance_coeffs(a, b)
+    assert np.allclose(
+        m_transfer.eps_model(freqs),
+        m_L.eps_model(freqs),
+    )
+
+    # test corner case of an admittance function representing a pure capacitance
+    C = 1e-12  # 1 pF capacitor
+    a = np.array([0, C])
+    b = np.array([1, 0])
+    m_transfer = td.PoleResidue.from_admittance_coeffs(a, b)
+
+    assert len(m_transfer.poles) == 0
+    assert np.isclose(1 + C / td.EPSILON_0, m_transfer.eps_inf)
+
+    #### Test validation
+    # Test improper admittance function resulting in a negative direct polynomial part
+    a = np.array([0, -C])
+    b = np.array([1, 0])
+    with pytest.raises(ValidationError):
+        _ = td.PoleResidue.from_admittance_coeffs(a, b)
+
+    # Test improper admittance function with numerator order too large
+    a = np.array([0, 1, 2])
+    b = np.array([1, 0])
+    with pytest.raises(ValidationError):
+        _ = td.PoleResidue.from_admittance_coeffs(a, b)
+
+    # Test transfer function that will result in a higher order pole
+    a = np.array([1])
+    b = np.array([0, 2])
+    with pytest.raises(ValidationError):
+        _ = td.PoleResidue.from_admittance_coeffs(a, b)
+
+    # Test transfer function that will result in a higher order pole, but is not in simplest form
+    a = np.array([0, 1])
+    b = np.array([0, 2])
+    _ = td.PoleResidue.from_admittance_coeffs(a, b)

tidy3d/components/medium.py

@@ -15,6 +15,7 @@
 import numpy as npo
 import pydantic.v1 as pd
 import xarray as xr
+from scipy import signal
 
 from ..constants import (
     C_0,
@@ -3188,6 +3189,180 @@ def compute_derivatives(self, derivative_info: DerivativeInfo) -> AutogradFieldM
 
         return derivative_map
 
+    @classmethod
+    def _real_partial_fraction_decomposition(
+        cls, a: np.ndarray, b: np.ndarray, tol: pd.PositiveFloat = 1e-2
+    ) -> tuple[list[tuple[Complex, Complex]], np.ndarray]:
+        """Computes the complex conjugate pole residue pairs given a rational expression with
+        real coefficients.
+
+        Parameters
+        ----------
+
+        a : np.ndarray
+            Coefficients of the numerator polynomial in increasing monomial order.
+        b : np.ndarray
+            Coefficients of the denominator polynomial in increasing monomial order.
+        tol : pd.PositiveFloat
+            Tolerance for pole finding. Two poles are considered equal, if their spacing is less
+            than ``tol``.
+
+        Returns
+        -------
+        tuple[list[tuple[Complex, Complex]], np.ndarray]
+            The list of complex conjugate poles and their associated residues. The second element of the
+            ``tuple`` is an array of coefficients representing any direct polynomial term.
+
+        """
+
+        if a.ndim != 1 or np.any(np.iscomplex(a)):
+            raise ValidationError(
+                "Numerator coefficients must be a one-dimensional array of real numbers."
+            )
+        if b.ndim != 1 or np.any(np.iscomplex(b)):
+            raise ValidationError(
+                "Denominator coefficients must be a one-dimensional array of real numbers."
+            )
+
+        # Compute residues and poles using scipy
+        (r, p, k) = signal.residue(np.flip(a), np.flip(b), tol=tol, rtype="avg")
+
+        # Assuming real coefficients for the polynomials, the poles should be real or come as
+        # complex conjugate pairs
+        r_filtered = []
+        p_filtered = []
+        for res, (idx, pole) in zip(list(r), enumerate(list(p))):
+            # Residue equal to zero interpreted as rational expression was not
+            # in simplest form. So skip this pole.
+            if res == 0:
+                continue
+            # Causal and stability check
+            if np.real(pole) > 0:
+                raise ValidationError("Transfer function is invalid. It is non-causal.")
+            # Check for higher order pole, which come in consecutive order
+            if idx > 0 and p[idx - 1] == pole:
+                raise ValidationError(
+                    "Transfer function is invalid. A higher order pole was detected. Try reducing ``tol``, "
+                    "or ensure that the rational expression does not have repeated poles. "
+                )
+            if np.imag(pole) == 0:
+                r_filtered.append(res / 2)
+                p_filtered.append(pole)
+            else:
+                pair_found = len(np.argwhere(np.array(p) == np.conj(pole))) == 1
+                if not pair_found:
+                    raise ValueError(
+                        "Failed to find complex-conjugate of pole in poles computed by SciPy."
+                    )
+                previously_added = len(np.argwhere(np.array(p_filtered) == np.conj(pole))) == 1
+                if not previously_added:
+                    r_filtered.append(res)
+                    p_filtered.append(pole)
+
+        poles_residues = list(zip(p_filtered, r_filtered))
+        k_increasing_order = np.flip(k)
+        return (poles_residues, k_increasing_order)
+
+    @classmethod
+    def from_admittance_coeffs(
+        cls,
+        a: np.ndarray,
+        b: np.ndarray,
+        eps_inf: pd.PositiveFloat = 1,
+        pole_tol: pd.PositiveFloat = 1e-2,
+    ) -> PoleResidue:
+        """Construct a :class:`.PoleResidue` model from an admittance function defining the
+        relationship between the electric field and the polarization current density in the
+        Laplace domain.
+
+        Parameters
+        ----------
+        a : np.ndarray
+            Coefficients of the numerator polynomial in increasing monomial order.
+        b : np.ndarray
+            Coefficients of the denominator polynomial in increasing monomial order.
+        eps_inf: pd.PositiveFloat
+            The relative permittivity at infinite frequency.
+        pole_tol: pd.PositiveFloat
+            Tolerance for the pole finding algorithm in Hertz. Two poles are considered equal, if their
+            spacing is closer than ``pole_tol`.
+        Returns
+        -------
+        :class:`.PoleResidue`
+            The pole residue equivalent.
+
+        Notes
+        -----
+
+            The supplied admittance function relates the electric field to the polarization current density
+            in the Laplace domain and is equivalent to a frequency-dependent complex conductivity
+            :math:`\\sigma(\\omega)`.
+
+            .. math::
+                J_p(s) = Y(s)E(s)
+
+            .. math::
+                Y(s) = \\frac{a_0 + a_1 s + \\dots + a_M s^M}{b_0 + b_1 s + \\dots + b_N s^N}
+
+            An equivalent :class:`.PoleResidue` medium is constructed using an equivalent frequency-dependent
+            complex permittivity defined as
+
+            .. math::
+                \\epsilon(s) = \\epsilon_\\infty - \\frac{1}{\\epsilon_0 s}
+                \\frac{a_0 + a_1 s + \\dots + a_M s^M}{b_0 + b_1 s + \\dots + b_N s^N}.
+        """
+
+        if a.ndim != 1 or np.any(np.logical_or(np.iscomplex(a), a < 0)):
+            raise ValidationError(
+                "Numerator coefficients must be a one-dimensional array of non-negative real numbers."
+            )
+        if b.ndim != 1 or np.any(np.logical_or(np.iscomplex(b), b < 0)):
+            raise ValidationError(
+                "Denominator coefficients must be a one-dimensional array of non-negative real numbers."
+            )
+
+        # Trim any trailing zeros, so that length corresponds with polynomial order
+        a = np.trim_zeros(a, "b")
+        b = np.trim_zeros(b, "b")
+
+        # Validate that transfer function will result in a proper transfer function, once converted to
+        # the complex permittivity version
+        # Let q equal the order of the numerator polynomial, and p equal the order
+        # of the denominator polynomal. Then, q < p is strictly proper rational transfer function (RTF)
+        # q <= p is a proper RTF, and q > p is an improper RTF.
+        q = len(a) - 1
+        p = len(b) - 1
+
+        if q > p + 1:
+            raise ValidationError(
+                "Transfer function is improper, the order of the numerator polynomial must be at most "
+                "one greater than the order of the denominator polynomial."
+            )
+
+        # Modify the transfer function defining a complex conductivity to match the complex
+        # frequency-dependent portion of the pole residue model
+        # Meaning divide by -j*omega*epsilon (s*epsilon)
+        b = np.concatenate(([0], b * EPSILON_0))
+
+        poles_and_residues, k = cls._real_partial_fraction_decomposition(
+            a=a, b=b, tol=pole_tol * 2 * np.pi
+        )
+
+        # A direct polynomial term of zeroth order is interpreted as an additional contribution to eps_inf.
+        # So we only handle that special case.
+        if len(k) == 1:
+            if np.iscomplex(k[0]) or k[0] < 0:
+                raise ValidationError(
+                    "Transfer function is invalid. Direct polynomial term must be real and positive for "
+                    "conversion to an equivalent PoleResidue medium."
+                )
+            else:
+                # A pure capacitance will translate to an increased permittivity at infinite frequency.
+                eps_inf = eps_inf + k[0]
+
+        pole_residue_from_transfer = PoleResidue(eps_inf=eps_inf, poles=poles_and_residues)
+        return pole_residue_from_transfer
+
 
 class CustomPoleResidue(CustomDispersiveMedium, PoleResidue):
     """A spatially varying dispersive medium described by the pole-residue pair model.