From 2d2d5894f317739421c7c5c2b5036462a671fc04 Mon Sep 17 00:00:00 2001
From: momchil <momchil@flexcompute.com>
Date: Thu, 9 Dec 2021 18:19:03 -0800
Subject: [PATCH] Added Medium.eps_diagonal(); fixes #106

---
 tests/test_components.py    | 15 ++++++-
 tidy3d/components/medium.py | 81 +++++++++++++++++++++++++++----------
 2 files changed, 74 insertions(+), 22 deletions(-)

diff --git a/tests/test_components.py b/tests/test_components.py
index fe88414ec..2fe67ff60 100644
--- a/tests/test_components.py
+++ b/tests/test_components.py
@@ -264,7 +264,7 @@ def eps_compare(medium: Medium, expected: Dict, tol: float = 1e-5):
 
 
 def test_epsilon_eval():
-    """Compare epsilon evaluated from a dispersive model to expected."""
+    """Compare epsilon evaluated from a dispersive various models to expected."""
 
     # Dispersive silver model
     poles_silver = [
@@ -310,6 +310,19 @@ def test_epsilon_eval():
     }
     eps_compare(material, expected)
 
+    # Anisotropic material
+    eps = (1.5, 2.0, 2.3)
+    sig = (0.01, 0.03, 0.015)
+    mediums = [Medium(permittivity=eps[i], conductivity=sig[i]) for i in range(3)]
+    material = AnisotropicMedium(xx=mediums[0], yy=mediums[1], zz=mediums[2])
+
+    eps_diag_2 = material.eps_diagonal(2e14)
+    eps_diag_5 = material.eps_diagonal(5e14)
+    assert np.all(np.array(eps_diag_2) == np.array([medium.eps_model(2e14) for medium in mediums]))
+
+    expected = {2e14: np.mean(eps_diag_2), 5e14: np.mean(eps_diag_5)}
+    eps_compare(material, expected)
+
 
 """ modes """
 
diff --git a/tidy3d/components/medium.py b/tidy3d/components/medium.py
index c7b144cc4..78d34ea65 100644
--- a/tidy3d/components/medium.py
+++ b/tidy3d/components/medium.py
@@ -15,6 +15,27 @@
 from ..log import log
 
 
+def ensure_freq_in_range(eps_model: Callable[[float], complex]) -> Callable[[float], complex]:
+    """Decorate ``eps_model`` to log warning if frequency supplied is out of bounds."""
+
+    def _eps_model(self, frequency: float) -> complex:
+        """New eps_model function."""
+
+        # if frequency is none, don't check, return original function
+        if frequency is None or self.frequency_range is None:
+            return eps_model(self, frequency)
+
+        fmin, fmax = self.frequency_range
+        if np.any(frequency < fmin) or np.any(frequency > fmax):
+            log.warning(
+                "frequency passed to `Medium.eps_model()`"
+                f"is outside of `Medium.frequency_range` = {self.frequency_range}"
+            )
+        return eps_model(self, frequency)
+
+    return _eps_model
+
+
 """ Medium Definitions """
 
 
@@ -51,6 +72,25 @@ def eps_model(self, frequency: float) -> complex:
             Complex-valued relative permittivity evaluated at ``frequency``.
         """
 
+    @ensure_freq_in_range
+    def eps_diagonal(self, frequency: float) -> Tuple[complex, complex, complex]:
+        """Main diagonal of the complex-valued permittivity tensor as a function of frequency.
+
+        Parameters
+        ----------
+        frequency : float
+            Frequency to evaluate permittivity at (Hz).
+
+        Returns
+        -------
+        complex
+            The diagonal elements of the relative permittivity tensor evaluated at ``frequency``.
+        """
+
+        # This only needs to be overwritten for anisotropic materials
+        eps = self.eps_model(frequency)
+        return (eps, eps, eps)
+
     @add_ax_if_none
     def plot(self, freqs: float, ax: Ax = None) -> Ax:  # pylint: disable=invalid-name
         """Plot n, k of a :class:`Medium` as a function of frequency.
@@ -170,27 +210,6 @@ def eps_sigma_to_eps_complex(eps_real: float, sigma: float, freq: float) -> comp
         return eps_real + 1j * sigma / omega / EPSILON_0
 
 
-def ensure_freq_in_range(eps_model: Callable[[float], complex]) -> Callable[[float], complex]:
-    """Decorate ``eps_model`` to log warning if frequency supplied is out of bounds."""
-
-    def _eps_model(self, frequency: float) -> complex:
-        """New eps_model function."""
-
-        # if frequency is none, don't check, return original function
-        if frequency is None or self.frequency_range is None:
-            return eps_model(self, frequency)
-
-        fmin, fmax = self.frequency_range
-        if np.any(frequency < fmin) or np.any(frequency > fmax):
-            log.warning(
-                "frequency passed to `Medium.eps_model()`"
-                f"is outside of `Medium.frequency_range` = {self.frequency_range}"
-            )
-        return eps_model(self, frequency)
-
-    return _eps_model
-
-
 """ Dispersionless Medium """
 
 # PEC keyword
@@ -338,6 +357,26 @@ def eps_model(self, frequency: float) -> complex:
         Tuple[complex, complex, complex]
             Complex-valued relative permittivity for each component evaluated at ``frequency``.
         """
+        eps_xx = self.xx.eps_model(frequency)
+        eps_yy = self.yy.eps_model(frequency)
+        eps_zz = self.zz.eps_model(frequency)
+        return np.mean((eps_xx, eps_yy, eps_zz))
+
+    @ensure_freq_in_range
+    def eps_diagonal(self, frequency: float) -> Tuple[complex, complex, complex]:
+        """Main diagonal of the complex-valued permittivity tensor as a function of frequency.
+
+        Parameters
+        ----------
+        frequency : float
+            Frequency to evaluate permittivity at (Hz).
+
+        Returns
+        -------
+        complex
+            The diagonal elements of the relative permittivity tensor evaluated at ``frequency``.
+        """
+
         eps_xx = self.xx.eps_model(frequency)
         eps_yy = self.yy.eps_model(frequency)
         eps_zz = self.zz.eps_model(frequency)