fix setting data attributes

Author: dmarek-flex

tidy3d/plugins/microwave/impedance_calculator.py

@@ -120,11 +120,12 @@ def check_voltage_or_current(cls, val, values):
     @staticmethod
     def _make_result_data_array(result: DataArray) -> ImpedanceResultTypes:
         """Helper for creating the proper result type."""
+        cls = ImpedanceFreqDataArray
         if "t" in result.coords:
-            return ImpedanceTimeDataArray(data=result.data, coords=result.coords)
+            cls = ImpedanceTimeDataArray
         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)
+            cls = ImpedanceFreqModeDataArray
+        return cls.assign_data_attrs(cls(data=result.data, coords=result.coords))
 
     @pd.root_validator(pre=False)
     def _warn_rf_license(cls, values):

tidy3d/plugins/microwave/path_integrals.py

@@ -7,7 +7,6 @@
 
 import numpy as np
 import pydantic.v1 as pd
-import xarray as xr
 
 from tidy3d.components.base import Tidy3dBaseModel, cached_property
 from tidy3d.components.data.data_array import (
@@ -24,6 +23,7 @@
     CurrentFreqDataArray,
     CurrentFreqModeDataArray,
     CurrentTimeDataArray,
+    DataArray,
     VoltageFreqDataArray,
     VoltageFreqModeDataArray,
     VoltageTimeDataArray,
@@ -220,13 +220,14 @@ def _check_monitor_data_supported(em_field: MonitorDataTypes):
             )
 
     @staticmethod
-    def _make_result_data_array(result: xr.DataArray) -> IntegralResultTypes:
+    def _make_result_data_array(result: DataArray) -> IntegralResultTypes:
         """Helper for creating the proper result type."""
+        cls = FreqDataArray
         if "t" in result.coords:
-            return TimeDataArray(data=result.data, coords=result.coords)
+            cls = TimeDataArray
         if "f" in result.coords and "mode_index" in result.coords:
-            return FreqModeDataArray(data=result.data, coords=result.coords)
-        return FreqDataArray(data=result.data, coords=result.coords)
+            cls = FreqModeDataArray
+        return cls.assign_data_attrs(cls(data=result.data, coords=result.coords))
 
 
 class VoltageIntegralAxisAligned(AxisAlignedPathIntegral):
@@ -257,13 +258,14 @@ def compute_voltage(self, em_field: MonitorDataTypes) -> VoltageIntegralResultTy
         return voltage
 
     @staticmethod
-    def _make_result_data_array(result: xr.DataArray) -> VoltageIntegralResultTypes:
+    def _make_result_data_array(result: DataArray) -> VoltageIntegralResultTypes:
         """Helper for creating the proper result type."""
+        cls = VoltageFreqDataArray
         if "t" in result.coords:
-            return VoltageTimeDataArray(data=result.data, coords=result.coords)
+            cls = VoltageTimeDataArray
         if "f" in result.coords and "mode_index" in result.coords:
-            return VoltageFreqModeDataArray(data=result.data, coords=result.coords)
-        return VoltageFreqDataArray(data=result.data, coords=result.coords)
+            cls = VoltageFreqModeDataArray
+        return cls.assign_data_attrs(cls(data=result.data, coords=result.coords))
 
     @staticmethod
     def from_terminal_positions(
@@ -424,7 +426,7 @@ def compute_current(self, em_field: MonitorDataTypes) -> CurrentIntegralResultTy
 
         if self.sign == "-":
             current *= -1
-        current = CurrentIntegralAxisAligned._set_data_array_attributes(current)
+        current = CurrentIntegralAxisAligned._make_result_data_array(current)
         return current
 
     @cached_property
@@ -516,13 +518,14 @@ def _to_path_integrals(
         return (bottom, right, top, left)
 
     @staticmethod
-    def _make_result_data_array(result: xr.DataArray) -> CurrentIntegralResultTypes:
+    def _make_result_data_array(result: DataArray) -> CurrentIntegralResultTypes:
         """Helper for creating the proper result type."""
+        cls = CurrentFreqDataArray
         if "t" in result.coords:
-            return CurrentTimeDataArray(data=result.data, coords=result.coords)
+            cls = CurrentTimeDataArray
         if "f" in result.coords and "mode_index" in result.coords:
-            return CurrentFreqModeDataArray(data=result.data, coords=result.coords)
-        return CurrentFreqDataArray(data=result.data, coords=result.coords)
+            cls = CurrentFreqModeDataArray
+        return cls.assign_data_attrs(cls(data=result.data, coords=result.coords))
 
     @add_ax_if_none
     def plot(