"""Helper classes for performing path integrals with fields on the Yee grid"""
from __future__ import annotations
from abc import ABC, abstractmethod
from typing import Union
import numpy as np
import pydantic.v1 as pd
import shapely as shapely
import xarray as xr
from ...components.base import Tidy3dBaseModel, cached_property
from ...components.data.data_array import (
FreqDataArray,
FreqModeDataArray,
ScalarFieldDataArray,
ScalarFieldTimeDataArray,
ScalarModeFieldDataArray,
TimeDataArray,
)
from ...components.data.monitor_data import FieldData, FieldTimeData, ModeData, ModeSolverData
from ...components.geometry.base import Box, Geometry
from ...components.types import Ax, Axis, Coordinate2D, Direction
from ...components.validators import assert_line, assert_plane
from ...components.viz import add_ax_if_none
from ...constants import AMP, VOLT, fp_eps
from ...exceptions import DataError, Tidy3dError
from ...log import log
from .viz import (
ARROW_CURRENT,
plot_params_current_path,
plot_params_voltage_minus,
plot_params_voltage_path,
plot_params_voltage_plus,
)
MonitorDataTypes = Union[FieldData, FieldTimeData, ModeData, ModeSolverData]
EMScalarFieldType = Union[ScalarFieldDataArray, ScalarFieldTimeDataArray, ScalarModeFieldDataArray]
IntegralResultTypes = Union[FreqDataArray, FreqModeDataArray, TimeDataArray]
class AbstractAxesRH(Tidy3dBaseModel, ABC):
"""Represents an axis-aligned right-handed coordinate system with one axis preferred.
Typically `main_axis` would refer to the normal axis of a plane.
"""
@cached_property
@abstractmethod
def main_axis(self) -> Axis:
"""Get the preferred axis."""
@cached_property
def remaining_axes(self) -> tuple[Axis, Axis]:
"""Get in-plane axes, ordered to maintain a right-handed coordinate system."""
axes: list[Axis] = [0, 1, 2]
axes.pop(self.main_axis)
if self.main_axis == 1:
return (axes[1], axes[0])
else:
return (axes[0], axes[1])
@cached_property
def remaining_dims(self) -> tuple[str, str]:
"""Get in-plane dimensions, ordered to maintain a right-handed coordinate system."""
dim1 = "xyz"[self.remaining_axes[0]]
dim2 = "xyz"[self.remaining_axes[1]]
return (dim1, dim2)
@cached_property
def local_dims(self) -> tuple[str, str, str]:
"""Get in-plane dimensions with in-plane dims first, followed by the `main_axis` dimension."""
dim3 = "xyz"[self.main_axis]
return self.remaining_dims + tuple(dim3)
@pd.root_validator(pre=False)
def _warn_rf_license(cls, values):
log.warning(
"ℹ️ ⚠️ RF simulations are subject to new license requirements in the future. You have instantiated at least one RF-specific component.",
log_once=True,
)
return values
[docs]
class AxisAlignedPathIntegral(AbstractAxesRH, Box):
"""Class for defining the simplest type of path integral, which is aligned with Cartesian axes."""
_line_validator = assert_line()
extrapolate_to_endpoints: bool = pd.Field(
False,
title="Extrapolate to Endpoints",
description="If the endpoints of the path integral terminate at or near a material interface, "
"the field is likely discontinuous. When this field is ``True``, fields that are outside and on the bounds "
"of the integral are ignored. Should be enabled when computing voltage between two conductors.",
)
snap_path_to_grid: bool = pd.Field(
False,
title="Snap Path to Grid",
description="It might be desireable to integrate exactly along the Yee grid associated with "
"a field. When this field is ``True``, the integration path will be snapped to the grid.",
)
[docs]
def compute_integral(self, scalar_field: EMScalarFieldType) -> IntegralResultTypes:
"""Computes the defined integral given the input ``scalar_field``."""
if not scalar_field.does_cover(self.bounds):
raise DataError("Scalar field does not cover the integration domain.")
coord = "xyz"[self.main_axis]
scalar_field = self._get_field_along_path(scalar_field)
# Get the boundaries
min_bound = self.bounds[0][self.main_axis]
max_bound = self.bounds[1][self.main_axis]
if self.extrapolate_to_endpoints:
# Remove field outside the boundaries
scalar_field = scalar_field.sel({coord: slice(min_bound, max_bound)})
# Ignore values on the boundary (sel is inclusive)
scalar_field = scalar_field.drop_sel({coord: (min_bound, max_bound)}, errors="ignore")
coordinates = scalar_field.coords[coord].values
else:
coordinates = scalar_field.coords[coord].sel({coord: slice(min_bound, max_bound)})
# Integration is along the original coordinates plus ensure that
# endpoints corresponding to the precise bounds of the port are included
coords_interp = np.array([min_bound])
coords_interp = np.concatenate((coords_interp, coordinates))
coords_interp = np.concatenate((coords_interp, [max_bound]))
coords_interp = {coord: coords_interp}
# Use extrapolation for the 2 additional endpoints, unless there is only a single sample point
method = "linear"
if len(coordinates) == 1 and self.extrapolate_to_endpoints:
method = "nearest"
scalar_field = scalar_field.interp(
coords_interp, method=method, kwargs={"fill_value": "extrapolate"}
)
result = scalar_field.integrate(coord=coord)
return self._make_result_data_array(result)
def _get_field_along_path(self, scalar_field: EMScalarFieldType) -> EMScalarFieldType:
"""Returns a selection of the input ``scalar_field`` ready for integration."""
(axis1, axis2) = self.remaining_axes
(coord1, coord2) = self.remaining_dims
if self.snap_path_to_grid:
# Coordinates that are not integrated
remaining_coords = {
coord1: self.center[axis1],
coord2: self.center[axis2],
}
# Select field nearest to center of integration line
scalar_field = scalar_field.sel(
remaining_coords,
method="nearest",
drop=False,
)
else:
# Try to interpolate unless there is only a single coordinate
coord1dict = {coord1: self.center[axis1]}
if scalar_field.sizes[coord1] == 1:
scalar_field = scalar_field.sel(coord1dict, method="nearest")
else:
scalar_field = scalar_field.interp(
coord1dict, method="linear", kwargs={"bounds_error": True}
)
coord2dict = {coord2: self.center[axis2]}
if scalar_field.sizes[coord2] == 1:
scalar_field = scalar_field.sel(coord2dict, method="nearest")
else:
scalar_field = scalar_field.interp(
coord2dict, method="linear", kwargs={"bounds_error": True}
)
# Remove unneeded coordinates
scalar_field = scalar_field.reset_coords(drop=True)
return scalar_field
@cached_property
def main_axis(self) -> Axis:
"""Axis for performing integration."""
for index, value in enumerate(self.size):
if value != 0:
return index
raise Tidy3dError("Failed to identify axis.")
def _vertices_2D(self, axis: Axis) -> tuple[Coordinate2D, Coordinate2D]:
"""Returns the two vertices of this path in the plane defined by ``axis``."""
min = self.bounds[0]
max = self.bounds[1]
_, min = Box.pop_axis(min, axis)
_, max = Box.pop_axis(max, axis)
u = [min[0], max[0]]
v = [min[1], max[1]]
return (u, v)
@staticmethod
def _check_monitor_data_supported(em_field: MonitorDataTypes):
"""Helper for validating that monitor data is supported."""
if not isinstance(em_field, (FieldData, FieldTimeData, ModeData, ModeSolverData)):
supported_types = list(MonitorDataTypes.__args__)
raise DataError(
f"'em_field' type {type(em_field)} not supported. Supported types are "
f"{supported_types}"
)
@staticmethod
def _make_result_data_array(result: xr.DataArray) -> IntegralResultTypes:
"""Helper for creating the proper result type."""
if "t" in result.coords:
return TimeDataArray(data=result.data, coords=result.coords)
elif "f" in result.coords and "mode_index" in result.coords:
return FreqModeDataArray(data=result.data, coords=result.coords)
else:
return FreqDataArray(data=result.data, coords=result.coords)
[docs]
class VoltageIntegralAxisAligned(AxisAlignedPathIntegral):
"""Class for computing the voltage between two points defined by an axis-aligned line."""
sign: Direction = pd.Field(
...,
title="Direction of Path Integral",
description="Positive indicates V=Vb-Va where position b has a larger coordinate along the axis of integration.",
)
[docs]
def compute_voltage(self, em_field: MonitorDataTypes) -> IntegralResultTypes:
"""Compute voltage along path defined by a line."""
self._check_monitor_data_supported(em_field=em_field)
e_component = "xyz"[self.main_axis]
field_name = f"E{e_component}"
# Validate that fields are present
em_field._check_fields_stored([field_name])
e_field = em_field.field_components[field_name]
voltage = self.compute_integral(e_field)
if self.sign == "+":
voltage *= -1
voltage = VoltageIntegralAxisAligned._set_data_array_attributes(voltage)
# Return data array of voltage while keeping coordinates of frequency|time|mode index
return voltage
@staticmethod
def _set_data_array_attributes(data_array: IntegralResultTypes) -> IntegralResultTypes:
"""Add explanatory attributes to the data array."""
data_array.name = "V"
return data_array.assign_attrs(units=VOLT, long_name="voltage")
[docs]
@staticmethod
def from_terminal_positions(
plus_terminal: float,
minus_terminal: float,
x: float = None,
y: float = None,
z: float = None,
extrapolate_to_endpoints: bool = True,
snap_path_to_grid: bool = True,
) -> VoltageIntegralAxisAligned:
"""Helper to create a :class:`VoltageIntegralAxisAligned` from two coordinates that
define a line and two positions indicating the endpoints of the path integral.
Parameters
----------
plus_terminal : float
Position along the voltage axis of the positive terminal.
minus_terminal : float
Position along the voltage axis of the negative terminal.
x : float = None
Position in x direction, only two of x,y,z can be specified to define line.
y : float = None
Position in y direction, only two of x,y,z can be specified to define line.
z : float = None
Position in z direction, only two of x,y,z can be specified to define line.
extrapolate_to_endpoints: bool = True
Passed directly to :class:`VoltageIntegralAxisAligned`
snap_path_to_grid: bool = True
Passed directly to :class:`VoltageIntegralAxisAligned`
Returns
-------
VoltageIntegralAxisAligned
The created path integral for computing voltage between the two terminals.
"""
axis_positions = Geometry.parse_two_xyz_kwargs(x=x, y=y, z=z)
# Calculate center and size of the future box
midpoint = (plus_terminal + minus_terminal) / 2
length = np.abs(plus_terminal - minus_terminal)
center = [midpoint, midpoint, midpoint]
size = [length, length, length]
for axis, position in axis_positions:
size[axis] = 0
center[axis] = position
direction = "+"
if plus_terminal < minus_terminal:
direction = "-"
return VoltageIntegralAxisAligned(
center=center,
size=size,
extrapolate_to_endpoints=extrapolate_to_endpoints,
snap_path_to_grid=snap_path_to_grid,
sign=direction,
)
[docs]
@add_ax_if_none
def plot(
self,
x: float = None,
y: float = None,
z: float = None,
ax: Ax = None,
**path_kwargs,
) -> Ax:
"""Plot path integral at single (x,y,z) coordinate.
Parameters
----------
x : float = None
Position of plane in x direction, only one of x,y,z can be specified to define plane.
y : float = None
Position of plane in y direction, only one of x,y,z can be specified to define plane.
z : float = None
Position of plane in z direction, only one of x,y,z can be specified to define plane.
ax : matplotlib.axes._subplots.Axes = None
Matplotlib axes to plot on, if not specified, one is created.
**path_kwargs
Optional keyword arguments passed to the matplotlib plotting of the line.
For details on accepted values, refer to
`Matplotlib's documentation <https://tinyurl.com/36marrat>`_.
Returns
-------
matplotlib.axes._subplots.Axes
The supplied or created matplotlib axes.
"""
axis, position = self.parse_xyz_kwargs(x=x, y=y, z=z)
if axis == self.main_axis or not np.isclose(position, self.center[axis], rtol=fp_eps):
return ax
(xs, ys) = self._vertices_2D(axis)
# Plot the path
plot_params = plot_params_voltage_path.include_kwargs(**path_kwargs)
plot_kwargs = plot_params.to_kwargs()
ax.plot(xs, ys, markevery=[0, -1], **plot_kwargs)
# Plot special end points
end_kwargs = plot_params_voltage_plus.include_kwargs(**path_kwargs).to_kwargs()
start_kwargs = plot_params_voltage_minus.include_kwargs(**path_kwargs).to_kwargs()
if self.sign == "-":
start_kwargs, end_kwargs = end_kwargs, start_kwargs
ax.plot(xs[0], ys[0], **start_kwargs)
ax.plot(xs[1], ys[1], **end_kwargs)
return ax
[docs]
class CurrentIntegralAxisAligned(AbstractAxesRH, Box):
"""Class for computing conduction current via Ampère's circuital law on an axis-aligned loop."""
_plane_validator = assert_plane()
sign: Direction = pd.Field(
...,
title="Direction of Contour Integral",
description="Positive indicates current flowing in the positive normal axis direction.",
)
extrapolate_to_endpoints: bool = pd.Field(
False,
title="Extrapolate to Endpoints",
description="This parameter is passed to :class:`AxisAlignedPathIntegral` objects when computing the contour integral.",
)
snap_contour_to_grid: bool = pd.Field(
False,
title="Snap Contour to Grid",
description="This parameter is passed to :class:`AxisAlignedPathIntegral` objects when computing the contour integral.",
)
[docs]
def compute_current(self, em_field: MonitorDataTypes) -> IntegralResultTypes:
"""Compute current flowing in loop defined by the outer edge of a rectangle."""
AxisAlignedPathIntegral._check_monitor_data_supported(em_field=em_field)
ax1 = self.remaining_axes[0]
ax2 = self.remaining_axes[1]
h_component = "xyz"[ax1]
v_component = "xyz"[ax2]
h_field_name = f"H{h_component}"
v_field_name = f"H{v_component}"
# Validate that fields are present
em_field._check_fields_stored([h_field_name, v_field_name])
h_horizontal = em_field.field_components[h_field_name]
h_vertical = em_field.field_components[v_field_name]
# Decompose contour into path integrals
(bottom, right, top, left) = self._to_path_integrals(h_horizontal, h_vertical)
current = 0
# Compute and add contributions from each part of the contour
current += bottom.compute_integral(h_horizontal)
current += right.compute_integral(h_vertical)
current -= top.compute_integral(h_horizontal)
current -= left.compute_integral(h_vertical)
if self.sign == "-":
current *= -1
current = CurrentIntegralAxisAligned._set_data_array_attributes(current)
return current
@cached_property
def main_axis(self) -> Axis:
"""Axis normal to loop"""
for index, value in enumerate(self.size):
if value == 0:
return index
raise Tidy3dError("Failed to identify axis.")
def _to_path_integrals(
self, h_horizontal=None, h_vertical=None
) -> tuple[AxisAlignedPathIntegral, ...]:
"""Returns four ``AxisAlignedPathIntegral`` instances, which represent a contour
integral around the surface defined by ``self.size``."""
ax1 = self.remaining_axes[0]
ax2 = self.remaining_axes[1]
horizontal_passed = h_horizontal is not None
vertical_passed = h_vertical is not None
if self.snap_contour_to_grid and horizontal_passed and vertical_passed:
(coord1, coord2) = self.remaining_dims
# Locations where horizontal paths will be snapped
v_bounds = [
self.center[ax2] - self.size[ax2] / 2,
self.center[ax2] + self.size[ax2] / 2,
]
h_snaps = h_horizontal.sel({coord2: v_bounds}, method="nearest").coords[coord2].values
# Locations where vertical paths will be snapped
h_bounds = [
self.center[ax1] - self.size[ax1] / 2,
self.center[ax1] + self.size[ax1] / 2,
]
v_snaps = h_vertical.sel({coord1: h_bounds}, method="nearest").coords[coord1].values
bottom_bound = h_snaps[0]
top_bound = h_snaps[1]
left_bound = v_snaps[0]
right_bound = v_snaps[1]
else:
bottom_bound = self.bounds[0][ax2]
top_bound = self.bounds[1][ax2]
left_bound = self.bounds[0][ax1]
right_bound = self.bounds[1][ax1]
# Horizontal paths
path_size = list(self.size)
path_size[ax1] = right_bound - left_bound
path_size[ax2] = 0
path_center = list(self.center)
path_center[ax2] = bottom_bound
bottom = AxisAlignedPathIntegral(
center=path_center,
size=path_size,
extrapolate_to_endpoints=self.extrapolate_to_endpoints,
snap_path_to_grid=self.snap_contour_to_grid,
)
path_center[ax2] = top_bound
top = AxisAlignedPathIntegral(
center=path_center,
size=path_size,
extrapolate_to_endpoints=self.extrapolate_to_endpoints,
snap_path_to_grid=self.snap_contour_to_grid,
)
# Vertical paths
path_size = list(self.size)
path_size[ax1] = 0
path_size[ax2] = top_bound - bottom_bound
path_center = list(self.center)
path_center[ax1] = left_bound
left = AxisAlignedPathIntegral(
center=path_center,
size=path_size,
extrapolate_to_endpoints=self.extrapolate_to_endpoints,
snap_path_to_grid=self.snap_contour_to_grid,
)
path_center[ax1] = right_bound
right = AxisAlignedPathIntegral(
center=path_center,
size=path_size,
extrapolate_to_endpoints=self.extrapolate_to_endpoints,
snap_path_to_grid=self.snap_contour_to_grid,
)
return (bottom, right, top, left)
@staticmethod
def _set_data_array_attributes(data_array: IntegralResultTypes) -> IntegralResultTypes:
"""Add explanatory attributes to the data array."""
data_array.name = "I"
return data_array.assign_attrs(units=AMP, long_name="current")
[docs]
@add_ax_if_none
def plot(
self,
x: float = None,
y: float = None,
z: float = None,
ax: Ax = None,
**path_kwargs,
) -> Ax:
"""Plot path integral at single (x,y,z) coordinate.
Parameters
----------
x : float = None
Position of plane in x direction, only one of x,y,z can be specified to define plane.
y : float = None
Position of plane in y direction, only one of x,y,z can be specified to define plane.
z : float = None
Position of plane in z direction, only one of x,y,z can be specified to define plane.
ax : matplotlib.axes._subplots.Axes = None
Matplotlib axes to plot on, if not specified, one is created.
**path_kwargs
Optional keyword arguments passed to the matplotlib plotting of the line.
For details on accepted values, refer to
`Matplotlib's documentation <https://tinyurl.com/36marrat>`_.
Returns
-------
matplotlib.axes._subplots.Axes
The supplied or created matplotlib axes.
"""
axis, position = self.parse_xyz_kwargs(x=x, y=y, z=z)
if axis != self.main_axis or not np.isclose(position, self.center[axis], rtol=fp_eps):
return ax
plot_params = plot_params_current_path.include_kwargs(**path_kwargs)
plot_kwargs = plot_params.to_kwargs()
path_integrals = self._to_path_integrals()
# Plot the path
for path in path_integrals:
(xs, ys) = path._vertices_2D(axis)
ax.plot(xs, ys, **plot_kwargs)
(ax1, ax2) = self.remaining_axes
# Add arrow to bottom path, unless right path is longer
arrow_path = path_integrals[0]
if self.size[ax2] > self.size[ax1]:
arrow_path = path_integrals[1]
(xs, ys) = arrow_path._vertices_2D(axis)
X = (xs[0] + xs[1]) / 2
Y = (ys[0] + ys[1]) / 2
center = np.array([X, Y])
dx = xs[1] - xs[0]
dy = ys[1] - ys[0]
direction = np.array([dx, dy])
segment_length = np.linalg.norm(direction)
unit_dir = direction / segment_length
# Change direction of arrow depending on sign of current definition
if self.sign == "-":
unit_dir *= -1.0
# Change direction of arrow when the "y" axis is dropped,
# since the plotted coordinate system will be left-handed (x, z)
if self.main_axis == 1:
unit_dir *= -1.0
start = center - unit_dir * segment_length
end = center
ax.annotate(
"",
xytext=(start[0], start[1]),
xy=(end[0], end[1]),
arrowprops=ARROW_CURRENT,
)
return ax
