From 4f8e89e4fa6497d7a5dc44cd54fe67c15c0fece3 Mon Sep 17 00:00:00 2001
From: Casey Wojcik <charles@flexcompute.com>
Date: Wed, 8 Nov 2023 16:34:10 +0000
Subject: [PATCH] Added LO-TO form of PoleResidue model

---
 CHANGELOG.md                         |  1 +
 tests/test_components/test_medium.py |  8 +++
 tidy3d/components/medium.py          | 86 ++++++++++++++++++++++++++++
 3 files changed, 95 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index ea1e44e846..152c0e10e0 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 - Added support for two-photon absorption via `TwoPhotonAbsorption` class. Added `KerrNonlinearity` that implements Kerr effect without third-harmonic generation.
+- Can create `PoleResidue` from LO-TO form via `PoleResidue.from_lo_to`.
 
 ### Changed
 - Indent for the json string of Tidy3D models has been changed to `None` when used internally; kept as `indent=4` for writing to `json` and `yaml` files.
diff --git a/tests/test_components/test_medium.py b/tests/test_components/test_medium.py
index b0d561bb03..29d8f1966c 100644
--- a/tests/test_components/test_medium.py
+++ b/tests/test_components/test_medium.py
@@ -116,6 +116,14 @@ def test_medium_dispersion():
     # test eps_model for int arguments
     m_SM.eps_model(np.array([1, 2]))
 
+    # test LO-TO form
+    poles = [(1, 0.1, 2, 5), (3, 0.4, 1, 0.4)]
+    m_LO_TO = td.PoleResidue.from_lo_to(poles=poles, eps_inf=2)
+    assert np.allclose(
+        m_LO_TO.eps_model(freqs),
+        td.PoleResidue.lo_to_eps_model(poles=poles, eps_inf=2, frequency=freqs),
+    )
+
 
 def test_medium_dispersion_conversion():
 
diff --git a/tidy3d/components/medium.py b/tidy3d/components/medium.py
index 250dd3c283..7095572cc9 100644
--- a/tidy3d/components/medium.py
+++ b/tidy3d/components/medium.py
@@ -2045,6 +2045,92 @@ def to_medium(self) -> Medium:
             frequency_range=self.frequency_range,
         )
 
+    @classmethod
+    def lo_to_eps_model(
+        cls,
+        poles: Tuple[Tuple[float, float, float, float], ...],
+        eps_inf: pd.PositiveFloat,
+        frequency: float,
+    ) -> complex:
+        """Evaluate eps model for a given set of LO-TO coefficients."""
+        omega = 2 * np.pi * frequency
+        eps = eps_inf
+        for (omega_lo, gamma_lo, omega_to, gamma_to) in poles:
+            eps *= omega_lo**2 - omega**2 - 1j * omega * gamma_lo
+            eps /= omega_to**2 - omega**2 - 1j * omega * gamma_to
+        return eps
+
+    @classmethod
+    def from_lo_to(
+        cls, poles: Tuple[Tuple[float, float, float, float], ...], eps_inf: pd.PositiveFloat = 1
+    ) -> PoleResidue:
+        """Construct a pole residue model from the LO-TO form
+        (longitudinal and transverse optical modes).
+        The LO-TO form is :math:`\\epsilon_\\infty \\prod_{i=1}^l \\frac{\\omega_{LO, i}^2 - \\omega^2 - i \\omega \\gamma_{LO, i}}{\\omega_{TO, i}^2 - \\omega^2 - i \\omega \\gamma_{TO, i}}` as given in the paper:
+
+            M. Schubert, T. E. Tiwald, and C. M. Herzinger,
+            "Infrared dielectric anisotropy and phonon modes of sapphire,"
+            Phys. Rev. B 61, 8187 (2000).
+
+        Parameters
+        ----------
+        poles : Tuple[Tuple[float, float, float, float], ...]
+            The LO-TO poles, given as list of tuples of the form
+            (omega_LO, gamma_LO, omega_TO, gamma_TO).
+
+        Returns
+        -------
+        :class:`.PoleResidue`
+            The pole residue equivalent of the LO-TO form provided.
+        """
+
+        omegas_lo, gammas_lo, omegas_to, gammas_to = (np.array(param) for param in zip(*poles))
+
+        # discriminants of quadratic factors of denominator
+        discs = 2 * np.emath.sqrt((gammas_to / 2) ** 2 - omegas_to**2)
+
+        # require nondegenerate TO poles
+        if len({(omega_to, gamma_to) for (_, _, omega_to, gamma_to) in poles}) != len(poles) or any(
+            disc == 0 for disc in discs
+        ):
+            raise ValidationError(
+                "Unable to construct a pole residue model "
+                "from an LO-TO form with degenerate TO poles. Consider adding a "
+                "perturbation to split the poles, or using "
+                "'PoleResidue.lo_to_eps_model' and fitting with the 'FastDispersionFitter'."
+            )
+
+        # roots of denominator, in pairs
+        roots = []
+        for gamma_to, disc in zip(gammas_to, discs):
+            roots.append(-gamma_to / 2 + disc / 2)
+            roots.append(-gamma_to / 2 - disc / 2)
+
+        # interpolants
+        interpolants = eps_inf * np.ones(len(roots), dtype=complex)
+        for i, a in enumerate(roots):
+            for omega_lo, gamma_lo in zip(omegas_lo, gammas_lo):
+                interpolants[i] *= omega_lo**2 + a**2 + a * gamma_lo
+            for j, a2 in enumerate(roots):
+                if j != i:
+                    interpolants[i] /= a - a2
+
+        a_coeffs = []
+        c_coeffs = []
+
+        for i in range(0, len(roots), 2):
+            if not np.isreal(roots[i]):
+                a_coeffs.append(roots[i])
+                c_coeffs.append(interpolants[i])
+            else:
+                a_coeffs.append(roots[i])
+                a_coeffs.append(roots[i + 1])
+                # factor of two from adding conjugate pole of real pole
+                c_coeffs.append(interpolants[i] / 2)
+                c_coeffs.append(interpolants[i + 1] / 2)
+
+        return PoleResidue(eps_inf=eps_inf, poles=list(zip(a_coeffs, c_coeffs)))
+
     @staticmethod
     def eV_to_angular_freq(f_eV: float):
         """Convert frequency in unit of eV to rad/s.