wip

Author: dmarek-flex

tidy3d/components/microwave/data/data_array.py

@@ -0,0 +1,170 @@
+"""Storing data associated with microwave quantities"""
+
+from __future__ import annotations
+
+from tidy3d.components.data.data_array import (
+    DataArray,
+    FreqDataArray,
+    FreqModeDataArray,
+    TimeDataArray,
+)
+from tidy3d.constants import AMP, OHM, VOLT
+
+
+class VoltageArray(DataArray):
+    # Always set __slots__ = () to avoid xarray warnings
+    __slots__ = ()
+    _data_attrs = {"units": VOLT, "long_name": "voltage"}
+
+
+class CurrentArray(DataArray):
+    # Always set __slots__ = () to avoid xarray warnings
+    __slots__ = ()
+    _data_attrs = {"units": AMP, "long_name": "current"}
+
+
+class ImpedanceArray(DataArray):
+    # Always set __slots__ = () to avoid xarray warnings
+    __slots__ = ()
+    _data_attrs = {"units": OHM, "long_name": "impedance"}
+
+
+# Voltage arrays
+class VoltageFreqDataArray(FreqDataArray):
+    """Voltage data array in frequency domain.
+
+    Example
+    -------
+    >>> import numpy as np
+    >>> f = [2e9, 3e9, 4e9]
+    >>> coords = dict(f=f)
+    >>> data = np.random.random(3) + 1j * np.random.random(3)
+    >>> vfd = VoltageFreqDataArray(data, coords=coords)
+    """
+
+    __slots__ = ()
+
+
+class VoltageTimeDataArray(VoltageArray, TimeDataArray):
+    """Voltage data array in time domain.
+
+    Example
+    -------
+    >>> import numpy as np
+    >>> t = [0, 1e-9, 2e-9, 3e-9]
+    >>> coords = dict(t=t)
+    >>> data = np.sin(2 * np.pi * 1e9 * np.array(t))
+    >>> vtd = VoltageTimeDataArray(data, coords=coords)
+    """
+
+    __slots__ = ()
+
+
+class VoltageFreqModeDataArray(VoltageArray, FreqModeDataArray):
+    """Voltage data array in frequency-mode domain.
+
+    Example
+    -------
+    >>> import numpy as np
+    >>> f = [2e9, 3e9]
+    >>> mode_index = [0, 1]
+    >>> coords = dict(f=f, mode_index=mode_index)
+    >>> data = np.random.random((2, 2)) + 1j * np.random.random((2, 2))
+    >>> vfmd = VoltageFreqModeDataArray(data, coords=coords)
+    """
+
+    __slots__ = ()
+
+
+# Current arrays
+class CurrentFreqDataArray(CurrentArray, FreqDataArray):
+    """Current data array in frequency domain.
+
+    Example
+    -------
+    >>> import numpy as np
+    >>> f = [2e9, 3e9, 4e9]
+    >>> coords = dict(f=f)
+    >>> data = np.random.random(3) + 1j * np.random.random(3)
+    >>> cfd = CurrentFreqDataArray(data, coords=coords)
+    """
+
+    __slots__ = ()
+
+
+class CurrentTimeDataArray(CurrentArray, TimeDataArray):
+    """Current data array in time domain.
+
+    Example
+    -------
+    >>> import numpy as np
+    >>> t = [0, 1e-9, 2e-9, 3e-9]
+    >>> coords = dict(t=t)
+    >>> data = np.cos(2 * np.pi * 1e9 * np.array(t))
+    >>> ctd = CurrentTimeDataArray(data, coords=coords)
+    """
+
+    __slots__ = ()
+
+
+class CurrentFreqModeDataArray(CurrentArray, FreqModeDataArray):
+    """Current data array in frequency-mode domain.
+
+    Example
+    -------
+    >>> import numpy as np
+    >>> f = [2e9, 3e9]
+    >>> mode_index = [0, 1]
+    >>> coords = dict(f=f, mode_index=mode_index)
+    >>> data = np.random.random((2, 2)) + 1j * np.random.random((2, 2))
+    >>> cfmd = CurrentFreqModeDataArray(data, coords=coords)
+    """
+
+    __slots__ = ()
+
+
+# Impedance arrays
+class ImpedanceFreqDataArray(ImpedanceArray, FreqDataArray):
+    """Impedance data array in frequency domain.
+
+    Example
+    -------
+    >>> import numpy as np
+    >>> f = [2e9, 3e9, 4e9]
+    >>> coords = dict(f=f)
+    >>> data = 50.0 + 1j * np.random.random(3)
+    >>> zfd = ImpedanceFreqDataArray(data, coords=coords)
+    """
+
+    __slots__ = ()
+
+
+class ImpedanceTimeDataArray(ImpedanceArray, TimeDataArray):
+    """Impedance data array in time domain.
+
+    Example
+    -------
+    >>> import numpy as np
+    >>> t = [0, 1e-9, 2e-9, 3e-9]
+    >>> coords = dict(t=t)
+    >>> data = 50.0 * np.ones_like(t)
+    >>> ztd = ImpedanceTimeDataArray(data, coords=coords)
+    """
+
+    __slots__ = ()
+
+
+class ImpedanceFreqModeDataArray(ImpedanceArray, FreqModeDataArray):
+    """Impedance data array in frequency-mode domain.
+
+    Example
+    -------
+    >>> import numpy as np
+    >>> f = [2e9, 3e9]
+    >>> mode_index = [0, 1]
+    >>> coords = dict(f=f, mode_index=mode_index)
+    >>> data = 50.0 + 10.0 * np.random.random((2, 2))
+    >>> zfmd = ImpedanceFreqModeDataArray(data, coords=coords)
+    """
+
+    __slots__ = ()

tidy3d/plugins/microwave/custom_path_integrals.py

@@ -20,9 +20,11 @@
     AbstractAxesRH,
     AxisAlignedPathIntegral,
     CurrentIntegralAxisAligned,
+    CurrentIntegralResultTypes,
     IntegralResultTypes,
     MonitorDataTypes,
     VoltageIntegralAxisAligned,
+    VoltageIntegralResultTypes,
 )
 from .viz import (
     ARROW_CURRENT,
@@ -243,7 +245,7 @@ class CustomVoltageIntegral2D(CustomPathIntegral2D):
 
     .. TODO Improve by including extrapolate_to_endpoints field, non-trivial extension."""
 
-    def compute_voltage(self, em_field: MonitorDataTypes) -> IntegralResultTypes:
+    def compute_voltage(self, em_field: MonitorDataTypes) -> VoltageIntegralResultTypes:
         """Compute voltage along path defined by a line.
 
         Parameters
@@ -253,12 +255,12 @@ def compute_voltage(self, em_field: MonitorDataTypes) -> IntegralResultTypes:
 
         Returns
         -------
-        :class:`.IntegralResultTypes`
+        :class:`.VoltageIntegralResultTypes`
             Result of voltage computation over remaining dimensions (frequency, time, mode indices).
         """
         AxisAlignedPathIntegral._check_monitor_data_supported(em_field=em_field)
         voltage = -1.0 * self.compute_integral(field="E", em_field=em_field)
-        voltage = VoltageIntegralAxisAligned._set_data_array_attributes(voltage)
+        voltage = VoltageIntegralAxisAligned._make_result_data_array(voltage)
         return voltage
 
     @add_ax_if_none
@@ -316,7 +318,7 @@ class CustomCurrentIntegral2D(CustomPathIntegral2D):
     To compute the current flowing in the positive ``axis`` direction, the vertices should be
     ordered in a counterclockwise direction."""
 
-    def compute_current(self, em_field: MonitorDataTypes) -> IntegralResultTypes:
+    def compute_current(self, em_field: MonitorDataTypes) -> CurrentIntegralResultTypes:
         """Compute current flowing in a custom loop.
 
         Parameters
@@ -326,12 +328,12 @@ def compute_current(self, em_field: MonitorDataTypes) -> IntegralResultTypes:
 
         Returns
         -------
-        :class:`.IntegralResultTypes`
+        :class:`.CurrentIntegralResultTypes`
             Result of current computation over remaining dimensions (frequency, time, mode indices).
         """
         AxisAlignedPathIntegral._check_monitor_data_supported(em_field=em_field)
         current = self.compute_integral(field="H", em_field=em_field)
-        current = CurrentIntegralAxisAligned._set_data_array_attributes(current)
+        current = CurrentIntegralAxisAligned._make_result_data_array(current)
         return current
 
     @add_ax_if_none

tidy3d/plugins/microwave/impedance_calculator.py

@@ -8,22 +8,28 @@
 import pydantic.v1 as pd
 
 from tidy3d.components.base import Tidy3dBaseModel
-from tidy3d.components.data.data_array import FreqDataArray, FreqModeDataArray, TimeDataArray
+from tidy3d.components.data.data_array import DataArray
 from tidy3d.components.data.monitor_data import FieldTimeData
+from tidy3d.components.microwave.data.data_array import (
+    ImpedanceFreqDataArray,
+    ImpedanceFreqModeDataArray,
+    ImpedanceTimeDataArray,
+)
 from tidy3d.components.monitor import ModeMonitor, ModeSolverMonitor
-from tidy3d.constants import OHM
 from tidy3d.exceptions import ValidationError
 from tidy3d.log import log
 
 from .custom_path_integrals import CustomCurrentIntegral2D, CustomVoltageIntegral2D
 from .path_integrals import (
     AxisAlignedPathIntegral,
     CurrentIntegralAxisAligned,
-    IntegralResultTypes,
     MonitorDataTypes,
     VoltageIntegralAxisAligned,
 )
 
+ImpedanceResultTypes = Union[
+    ImpedanceFreqDataArray, ImpedanceFreqModeDataArray, ImpedanceTimeDataArray
+]
 VoltageIntegralTypes = Union[VoltageIntegralAxisAligned, CustomVoltageIntegral2D]
 CurrentIntegralTypes = Union[CurrentIntegralAxisAligned, CustomCurrentIntegral2D]
 
@@ -43,7 +49,7 @@ class ImpedanceCalculator(Tidy3dBaseModel):
         description="Definition of contour integral for computing current.",
     )
 
-    def compute_impedance(self, em_field: MonitorDataTypes) -> IntegralResultTypes:
+    def compute_impedance(self, em_field: MonitorDataTypes) -> ImpedanceResultTypes:
         """Compute impedance for the supplied ``em_field`` using ``voltage_integral`` and
         ``current_integral``. If only a single integral has been defined, impedance is
         computed using the total flux in ``em_field``.
@@ -56,7 +62,7 @@ def compute_impedance(self, em_field: MonitorDataTypes) -> IntegralResultTypes:
 
         Returns
         -------
-        :class:`.IntegralResultTypes`
+        :class:`.ImpedanceResultTypes`
             Result of impedance computation over remaining dimensions (frequency, time, mode indices).
         """
         AxisAlignedPathIntegral._check_monitor_data_supported(em_field=em_field)
@@ -98,7 +104,7 @@ def compute_impedance(self, em_field: MonitorDataTypes) -> IntegralResultTypes:
                 impedance = np.real(voltage) / np.real(current)
             else:
                 impedance = voltage / current
-        impedance = ImpedanceCalculator._set_data_array_attributes(impedance)
+        impedance = ImpedanceCalculator._make_result_data_array(impedance)
         return impedance
 
     @pd.validator("current_integral", always=True)
@@ -112,17 +118,13 @@ def check_voltage_or_current(cls, val, values):
         return val
 
     @staticmethod
-    def _set_data_array_attributes(data_array: IntegralResultTypes) -> IntegralResultTypes:
-        """Helper to set additional metadata for ``IntegralResultTypes``."""
-        # Determine type based on coords present
-        if "mode_index" in data_array.coords:
-            data_array = FreqModeDataArray(data_array)
-        elif "f" in data_array.coords:
-            data_array = FreqDataArray(data_array)
-        else:
-            data_array = TimeDataArray(data_array)
-        data_array.name = "Z0"
-        return data_array.assign_attrs(units=OHM, long_name="characteristic impedance")
+    def _make_result_data_array(result: DataArray) -> ImpedanceResultTypes:
+        """Helper for creating the proper result type."""
+        if "t" in result.coords:
+            return ImpedanceTimeDataArray(data=result.data, coords=result.coords)
+        if "f" in result.coords and "mode_index" in result.coords:
+            return ImpedanceFreqModeDataArray(data=result.data, coords=result.coords)
+        return ImpedanceFreqDataArray(data=result.data, coords=result.coords)
 
     @pd.root_validator(pre=False)
     def _warn_rf_license(cls, values):