immutable data arrays for rf

fix setting data attributes

Author: dmarek-flex

tidy3d/components/data/data_array.py

@@ -23,11 +23,14 @@
 from tidy3d.components.geometry.bound_ops import bounds_contains
 from tidy3d.components.types import Axis, Bound
 from tidy3d.constants import (
+    AMP,
     HERTZ,
     MICROMETER,
+    OHM,
     PICOSECOND_PER_NANOMETER_PER_KILOMETER,
     RADIAN,
     SECOND,
+    VOLT,
     WATT,
 )
 from tidy3d.exceptions import DataError, FileError
@@ -1309,6 +1312,165 @@ class PerturbationCoefficientDataArray(DataArray):
     _dims = ("wvl", "coeff")
 
 
+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(VoltageArray, 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__ = ()
+
+
 DATA_ARRAY_TYPES = [
     SpatialDataArray,
     ScalarFieldDataArray,
@@ -1346,6 +1508,15 @@ class PerturbationCoefficientDataArray(DataArray):
     SpatialVoltageDataArray,
     PerturbationCoefficientDataArray,
     IndexedTimeDataArray,
+    VoltageFreqDataArray,
+    VoltageTimeDataArray,
+    VoltageFreqModeDataArray,
+    CurrentFreqDataArray,
+    CurrentTimeDataArray,
+    CurrentFreqModeDataArray,
+    ImpedanceFreqDataArray,
+    ImpedanceTimeDataArray,
+    ImpedanceFreqModeDataArray,
 ]
 DATA_ARRAY_MAP = {data_array.__name__: data_array for data_array in DATA_ARRAY_TYPES}
 
@@ -1356,3 +1527,14 @@ class PerturbationCoefficientDataArray(DataArray):
     IndexedFieldVoltageDataArray,
     PointDataArray,
 ]
+
+IntegralResultTypes = Union[FreqDataArray, FreqModeDataArray, TimeDataArray]
+VoltageIntegralResultTypes = Union[
+    VoltageFreqDataArray, VoltageFreqModeDataArray, VoltageTimeDataArray
+]
+CurrentIntegralResultTypes = Union[
+    CurrentFreqDataArray, CurrentFreqModeDataArray, CurrentTimeDataArray
+]
+ImpedanceResultTypes = Union[
+    ImpedanceFreqDataArray, ImpedanceFreqModeDataArray, ImpedanceTimeDataArray
+]

tidy3d/plugins/microwave/custom_path_integrals.py

@@ -19,10 +19,10 @@
 from .path_integrals import (
     AbstractAxesRH,
     AxisAlignedPathIntegral,
-    CurrentIntegralAxisAligned,
+    CurrentIntegralResultTypes,
     IntegralResultTypes,
     MonitorDataTypes,
-    VoltageIntegralAxisAligned,
+    VoltageIntegralResultTypes,
 )
 from .viz import (
     ARROW_CURRENT,
@@ -89,6 +89,10 @@ def compute_integral(
             Result of integral over remaining dimensions (frequency, time, mode indices).
         """
 
+        from tidy3d.plugins.smatrix.utils import (
+            _make_base_result_data_array,
+        )
+
         (dim1, dim2, dim3) = self.local_dims
 
         h_field_name = f"{field}{dim1}"
@@ -130,7 +134,7 @@ def compute_integral(
         # Integrate along the path
         result = integrand.integrate(coord="s")
         result = result.reset_coords(drop=True)
-        return AxisAlignedPathIntegral._make_result_data_array(result)
+        return _make_base_result_data_array(result)
 
     @staticmethod
     def _compute_dl_component(coord_array: xr.DataArray, closed_contour=False) -> np.array:
@@ -243,7 +247,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,13 +257,16 @@ def compute_voltage(self, em_field: MonitorDataTypes) -> IntegralResultTypes:
 
         Returns
         -------
-        :class:`.IntegralResultTypes`
+        :class:`.VoltageIntegralResultTypes`
             Result of voltage computation over remaining dimensions (frequency, time, mode indices).
         """
+        from tidy3d.plugins.smatrix.utils import (
+            _make_voltage_data_array,
+        )
+
         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)
-        return voltage
+        return _make_voltage_data_array(voltage)
 
     @add_ax_if_none
     def plot(
@@ -316,7 +323,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,13 +333,16 @@ def compute_current(self, em_field: MonitorDataTypes) -> IntegralResultTypes:
 
         Returns
         -------
-        :class:`.IntegralResultTypes`
+        :class:`.CurrentIntegralResultTypes`
             Result of current computation over remaining dimensions (frequency, time, mode indices).
         """
+        from tidy3d.plugins.smatrix.utils import (
+            _make_current_data_array,
+        )
+
         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)
-        return current
+        return _make_current_data_array(current)
 
     @add_ax_if_none
     def plot(

tidy3d/plugins/microwave/impedance_calculator.py

@@ -8,18 +8,16 @@
 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 ImpedanceResultTypes
 from tidy3d.components.data.monitor_data import FieldTimeData
 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,
 )
@@ -43,7 +41,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,9 +54,11 @@ def compute_impedance(self, em_field: MonitorDataTypes) -> IntegralResultTypes:
 
         Returns
         -------
-        :class:`.IntegralResultTypes`
+        :class:`.ImpedanceResultTypes`
             Result of impedance computation over remaining dimensions (frequency, time, mode indices).
         """
+        from tidy3d.plugins.smatrix.utils import _make_impedance_data_array
+
         AxisAlignedPathIntegral._check_monitor_data_supported(em_field=em_field)
 
         # If both voltage and current integrals have been defined then impedance is computed directly
@@ -98,7 +98,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 = _make_impedance_data_array(impedance)
         return impedance
 
     @pd.validator("current_integral", always=True)
@@ -111,19 +111,6 @@ 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")
-
     @pd.root_validator(pre=False)
     def _warn_rf_license(cls, values):
         log.warning(