33from __future__ import annotations
44
55from abc import ABC , abstractmethod
6- from typing import Any , Union
6+ from typing import Union
77
88import numpy as np
99import pydantic .v1 as pd
1010import shapely as shapely
11+ import xarray as xr
1112
1213from ...components .base import Tidy3dBaseModel , cached_property
1314from ...components .data .data_array import (
1819 ScalarModeFieldDataArray ,
1920 TimeDataArray ,
2021)
21- from ...components .data .monitor_data import FieldData , FieldTimeData , ModeSolverData
22+ from ...components .data .monitor_data import FieldData , FieldTimeData , ModeData , ModeSolverData
2223from ...components .geometry .base import Box , Geometry
2324from ...components .types import Ax , Axis , Coordinate2D , Direction
2425from ...components .validators import assert_line , assert_plane
3334 plot_params_voltage_plus ,
3435)
3536
36- MonitorDataTypes = Union [FieldData , FieldTimeData , ModeSolverData ]
37+ MonitorDataTypes = Union [FieldData , FieldTimeData , ModeData , ModeSolverData ]
3738EMScalarFieldType = Union [ScalarFieldDataArray , ScalarFieldTimeDataArray , ScalarModeFieldDataArray ]
3839IntegralResultTypes = Union [FreqDataArray , FreqModeDataArray , TimeDataArray ]
3940
4041
41- def _check_em_field_supported (em_field : Any ):
42- """Function for validating correct data arrays."""
43- if not isinstance (em_field , (FieldData , FieldTimeData , ModeSolverData )):
44- raise DataError (
45- "'em_field' type not supported. Supported types are "
46- "'FieldData', 'FieldTimeData', 'ModeSolverData'."
47- )
48-
49-
5042class AbstractAxesRH (Tidy3dBaseModel , ABC ):
5143 """Represents an axis-aligned right-handed coordinate system with one axis preferred.
5244 Typically `main_axis` would refer to the normal axis of a plane.
@@ -137,18 +129,7 @@ def compute_integral(self, scalar_field: EMScalarFieldType) -> IntegralResultTyp
137129 coords_interp , method = method , kwargs = {"fill_value" : "extrapolate" }
138130 )
139131 result = scalar_field .integrate (coord = coord )
140- if isinstance (scalar_field , ScalarFieldDataArray ):
141- return FreqDataArray (data = result .data , coords = result .coords )
142- elif isinstance (scalar_field , ScalarFieldTimeDataArray ):
143- return TimeDataArray (data = result .data , coords = result .coords )
144- else :
145- if not isinstance (scalar_field , ScalarModeFieldDataArray ):
146- raise TypeError (
147- f"Unsupported 'scalar_field' type: { type (scalar_field )}
148- "Expected one of 'ScalarFieldDataArray', 'ScalarFieldTimeDataArray', "
149- "'ScalarModeFieldDataArray'."
150- )
151- return FreqModeDataArray (data = result .data , coords = result .coords )
132+ return self ._make_result_data_array (result )
152133
153134 def _get_field_along_path (self , scalar_field : EMScalarFieldType ) -> EMScalarFieldType :
154135 """Returns a selection of the input ``scalar_field`` ready for integration."""
@@ -206,6 +187,26 @@ def _vertices_2D(self, axis: Axis) -> tuple[Coordinate2D, Coordinate2D]:
206187 v = [min [1 ], max [1 ]]
207188 return (u , v )
208189
190+ @staticmethod
191+ def _check_monitor_data_supported (em_field : MonitorDataTypes ):
192+ """Helper for validating that monitor data is supported."""
193+ if not isinstance (em_field , (FieldData , FieldTimeData , ModeData , ModeSolverData )):
194+ supported_types = list (MonitorDataTypes .__args__ )
195+ raise DataError (
196+ f"'em_field' type { type (em_field )}
197+ f"{ supported_types }
198+ )
199+
200+ @staticmethod
201+ def _make_result_data_array (result : xr .DataArray ) -> IntegralResultTypes :
202+ """Helper for creating the proper result type."""
203+ if "t" in result .coords :
204+ return TimeDataArray (data = result .data , coords = result .coords )
205+ elif "f" in result .coords and "mode_index" in result .coords :
206+ return FreqModeDataArray (data = result .data , coords = result .coords )
207+ else :
208+ return FreqDataArray (data = result .data , coords = result .coords )
209+
209210
210211class VoltageIntegralAxisAligned (AxisAlignedPathIntegral ):
211212 """Class for computing the voltage between two points defined by an axis-aligned line."""
@@ -218,12 +219,11 @@ class VoltageIntegralAxisAligned(AxisAlignedPathIntegral):
218219
219220 def compute_voltage (self , em_field : MonitorDataTypes ) -> IntegralResultTypes :
220221 """Compute voltage along path defined by a line."""
221- _check_em_field_supported (em_field = em_field )
222+ self . _check_monitor_data_supported (em_field = em_field )
222223 e_component = "xyz" [self .main_axis ]
223224 field_name = f"E{ e_component }
224- # Validate that the field is present
225- if field_name not in em_field .field_components :
226- raise DataError (f"'field_name' '{ field_name } )
225+ # Validate that fields are present
226+ em_field ._check_fields_stored ([field_name ])
227227 e_field = em_field .field_components [field_name ]
228228
229229 voltage = self .compute_integral (e_field )
@@ -376,18 +376,15 @@ class CurrentIntegralAxisAligned(AbstractAxesRH, Box):
376376
377377 def compute_current (self , em_field : MonitorDataTypes ) -> IntegralResultTypes :
378378 """Compute current flowing in loop defined by the outer edge of a rectangle."""
379- _check_em_field_supported (em_field = em_field )
379+ AxisAlignedPathIntegral . _check_monitor_data_supported (em_field = em_field )
380380 ax1 = self .remaining_axes [0 ]
381381 ax2 = self .remaining_axes [1 ]
382382 h_component = "xyz" [ax1 ]
383383 v_component = "xyz" [ax2 ]
384384 h_field_name = f"H{ h_component }
385385 v_field_name = f"H{ v_component }
386386 # Validate that fields are present
387- if h_field_name not in em_field .field_components :
388- raise DataError (f"'field_name' '{ h_field_name } )
389- if v_field_name not in em_field .field_components :
390- raise DataError (f"'field_name' '{ v_field_name } )
387+ em_field ._check_fields_stored ([h_field_name , v_field_name ])
391388 h_horizontal = em_field .field_components [h_field_name ]
392389 h_vertical = em_field .field_components [v_field_name ]
393390
0 commit comments