"""Collections of DataArrays."""
from __future__ import annotations
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, Any, get_args
import numpy as np
import xarray as xr
from pydantic import Field, model_validator
from tidy3d.components.base import Tidy3dBaseModel
from tidy3d.components.types import xyz
from tidy3d.constants import (
AMP,
C_0,
MICROMETER,
PICOSECOND_PER_NANOMETER_PER_KILOMETER,
VOLT,
UnitScaling,
)
from tidy3d.exceptions import DataError
from tidy3d.log import log
from .data_array import (
DataArray,
EMEScalarFieldDataArray,
EMEScalarModeFieldDataArray,
GroupIndexDataArray,
IndexedFreqDataArray,
ModeDispersionDataArray,
ModeIndexDataArray,
PointDataArray,
ScalarFieldDataArray,
ScalarFieldTimeDataArray,
ScalarModeFieldCylindricalDataArray,
ScalarModeFieldDataArray,
TimeDataArray,
TriangleMeshDataArray,
_TracedDataset,
)
from .em_fields import em_field_symmetry_eigenvalues, point_cloud_field_symmetry_eigenvalues
from .point_cloud import POINT_CLOUD_PERMITTIVITY_COMPONENTS
from .unstructured.surface import TriangularSurfaceDataset
from .zbf import ZBFData
if TYPE_CHECKING:
from collections.abc import Callable
from typing import Literal
from numpy.typing import ArrayLike
from tidy3d.compat import Self
from tidy3d.components.types import Axis, BoundOptional, FreqArray
DEFAULT_MAX_SAMPLES_PER_STEP = 10_000
DEFAULT_MAX_CELLS_PER_STEP = 10_000
DEFAULT_TOLERANCE_CELL_FINDING = 1e-6
ELECTRIC_FIELD_UNITS = f"{VOLT}/{MICROMETER}"
MAGNETIC_FIELD_UNITS = f"{AMP}/{MICROMETER}"
class Dataset(Tidy3dBaseModel, ABC):
"""Abstract base class for objects that store collections of `:class:`.DataArray`s."""
@property
def data_arrs(self) -> dict:
"""Returns a dictionary of all `:class:`.DataArray`s in the dataset."""
data_arrs = {}
for key in self.__class__.model_fields.keys():
data = getattr(self, key)
if isinstance(data, DataArray):
data_arrs[key] = data
return data_arrs
class FreqDataset(Dataset, ABC):
"""Abstract base class for objects that store collections of `:class:`.DataArray`s."""
def _interp_in_freq_update_dict(
self,
freqs: FreqArray,
method: Literal["linear", "cubic", "poly"] = "linear",
assume_sorted: bool = False,
) -> dict[str, DataArray]:
"""Interpolate mode data to new frequency points.
Interpolates all stored mode data (effective indices, field components, group indices,
and dispersion) from the current frequency grid to a new set of frequencies. This is
useful for obtaining mode data at many frequencies from computations at fewer frequencies,
when modes vary smoothly with frequency.
Parameters
----------
freqs : FreqArray
New frequency points to interpolate to. Should generally span a similar range
as the original frequencies to avoid extrapolation.
method : Literal["linear", "cubic", "poly"]
Interpolation method. ``"linear"`` for linear interpolation (requires 2+ source
frequencies), ``"cubic"`` for cubic spline interpolation (requires 4+ source
frequencies), ``"poly"`` for polynomial interpolation using barycentric
formula (requires 3+ source frequencies).
For complex-valued data, real and imaginary parts are interpolated independently.
assume_sorted: bool = False,
Whether to assume the frequency points are sorted.
Returns
-------
ModeSolverData
New :class:`ModeSolverData` object with data interpolated to the requested frequencies.
Note
----
Interpolation assumes modes vary smoothly with frequency. Results may be inaccurate
near mode crossings or regions of rapid mode variation. Use frequency tracking
(``mode_spec.sort_spec.track_freq``) to help maintain mode ordering consistency.
For polynomial interpolation, source frequencies at Chebyshev nodes provide
optimal accuracy within the frequency range.
Example
-------
>>> # Compute modes at 5 frequencies
>>> import numpy as np
>>> freqs_sparse = np.linspace(1e14, 2e14, 5)
>>> # ... create mode_solver and compute modes ...
>>> # mode_data = mode_solver.solve()
>>> # Interpolate to 50 frequencies
>>> freqs_dense = np.linspace(1e14, 2e14, 50)
>>> # mode_data_interp = mode_data.interp(freqs=freqs_dense, method='linear')
"""
freqs = np.array(freqs)
modify_data = {}
for key, data in self.data_arrs.items():
modify_data[key] = self._interp_dataarray_in_freq(data, freqs, method, assume_sorted)
return modify_data
@staticmethod
def _interp_dataarray_in_freq(
data: DataArray,
freqs: FreqArray,
method: Literal["linear", "cubic", "poly", "nearest"],
assume_sorted: bool = False,
) -> DataArray:
"""Interpolate a DataArray along the frequency coordinate.
Parameters
----------
data : DataArray
Data array to interpolate. Must have a frequency coordinate ``"f"``.
freqs : FreqArray
New frequency points.
method : Literal["linear", "cubic", "poly", "nearest"]
Interpolation method (``"linear"``, ``"cubic"``, ``"poly"``, or ``"nearest"``).
For ``"poly"``, uses barycentric formula for polynomial interpolation.
assume_sorted: bool = False,
Whether to assume the frequency points are sorted.
Returns
-------
DataArray
Interpolated data array with the same structure but new frequency points.
"""
# if dataarray is already stored at the correct frequencies, do nothing
if np.array_equal(freqs, data.f):
return data
# Map 'poly' to xarray's 'barycentric' method
xr_method = "barycentric" if method == "poly" else method
# Use xarray's built-in interpolation
# For complex data, this automatically interpolates real and imaginary parts
interp_kwargs = {"method": xr_method}
if method == "nearest":
return data.sel(f=freqs, method="nearest")
else:
if method != "poly":
interp_kwargs["kwargs"] = {"fill_value": "extrapolate"}
return data.interp(f=freqs, assume_sorted=assume_sorted, **interp_kwargs)
class ModeFreqDataset(FreqDataset, ABC):
"""Abstract base class for objects that store collections of `:class:`.DataArray`s."""
def _apply_mode_reorder(self, sort_inds_2d: np.ndarray) -> Self:
"""Apply a mode reordering along mode_index for all frequency indices.
Parameters
----------
sort_inds_2d : np.ndarray
Array of shape (num_freqs, num_modes) where each row is the
permutation to apply to the mode_index for that frequency.
"""
num_freqs, num_modes = sort_inds_2d.shape
modify_data = {}
for key, data in self.data_arrs.items():
if "mode_index" not in data.dims or "f" not in data.dims:
continue
dims_orig = data.dims
f_coord = data.coords["f"]
slices = []
for ifreq in range(num_freqs):
sl = data.isel(f=ifreq, mode_index=sort_inds_2d[ifreq])
slices.append(sl.assign_coords(mode_index=np.arange(num_modes)))
# Concatenate along the 'f' dimension name and then restore original frequency coordinates
data = xr.concat(slices, dim="f").assign_coords(f=f_coord).transpose(*dims_orig)
modify_data[key] = data
return self.updated_copy(**modify_data)
[docs]
class AbstractFieldDataset(Dataset, ABC):
"""Collection of scalar fields with some symmetry properties."""
@property
@abstractmethod
def field_components(self) -> dict[str, DataArray]:
"""Maps the field components to their associated data."""
@property
def solver_field_bounds(self) -> BoundOptional | None:
"""Per-axis bounds where solver field data is physically valid.
Returns ``None`` by default. Subclasses that produce zero-padded
output grids (e.g. mode-solver data) override this to return actual
bounds so that colocation clips to the valid region.
"""
return None
[docs]
def apply_phase(self, phase: float) -> AbstractFieldDataset:
"""Create a copy where all elements are phase-shifted by a value (in radians)."""
if phase == 0.0:
return self
phasor = np.exp(1j * phase)
field_components_shifted = {}
for fld_name, fld_cmp in self.field_components.items():
fld_cmp_shifted = phasor * fld_cmp
field_components_shifted[fld_name] = fld_cmp_shifted
return self.updated_copy(**field_components_shifted)
@property
@abstractmethod
def grid_locations(self) -> dict[str, str]:
"""Maps field components to the string key of their grid locations on the yee lattice."""
@property
@abstractmethod
def symmetry_eigenvalues(self) -> dict[str, Callable[[Axis], float]]:
"""Maps field components to their (positive) symmetry eigenvalues."""
[docs]
def package_colocate_results(
self, centered_fields: dict[str, ScalarFieldDataArray]
) -> xr.Dataset:
"""How to package the dictionary of fields computed via self.colocate()."""
return _TracedDataset(centered_fields)
[docs]
def colocate(self, x: ArrayLike = None, y: ArrayLike = None, z: ArrayLike = None) -> xr.Dataset:
"""Colocate all of the data at a set of x, y, z coordinates.
Parameters
----------
x : Optional[array-like] = None
x coordinates of locations.
If not supplied, does not try to colocate on this dimension.
y : Optional[array-like] = None
y coordinates of locations.
If not supplied, does not try to colocate on this dimension.
z : Optional[array-like] = None
z coordinates of locations.
If not supplied, does not try to colocate on this dimension.
Returns
-------
xr.Dataset
Dataset containing all fields at the same spatial locations.
For more details refer to `xarray's Documentation <https://tinyurl.com/cyca3krz>`_.
Note
----
For many operations (such as flux calculations and plotting),
it is important that the fields are colocated at the same spatial locations.
Be sure to apply this method to your field data in those cases.
"""
if hasattr(self, "monitor") and self.monitor.colocate:
with log as consolidated_logger:
consolidated_logger.warning(
"Colocating data that has already been colocated during the solver "
"run. For most accurate results when colocating to custom coordinates set "
"'Monitor.colocate' to 'False' to use the raw data on the Yee grid "
"and avoid double interpolation. Note: the default value was changed to 'True' "
"in Tidy3D version 2.4.0."
)
# convert supplied coordinates to array and assign string mapping to them
supplied_coord_map = {k: np.array(v) for k, v in zip("xyz", (x, y, z)) if v is not None}
# dict of data arrays to combine in dataset and return
centered_fields = {}
# loop through field components
for field_name, field_data in self.field_components.items():
# loop through x, y, z dimensions and raise an error if only one element along dim
for coord_name, coords_supplied in supplied_coord_map.items():
coord_data = np.array(field_data.coords[coord_name])
if coord_data.size == 1:
raise DataError(
f"colocate given {coord_name}={coords_supplied}, but "
f"data only has one coordinate at {coord_name}={coord_data[0]}. "
"Therefore, can't colocate along this dimension. "
f"supply {coord_name}=None to skip it."
)
if self.solver_field_bounds is not None:
centered_fields[field_name] = field_data.interp_within_domain(
supplied_coord_map, self.solver_field_bounds, assume_sorted=True
)
else:
centered_fields[field_name] = field_data.interp(
**supplied_coord_map, kwargs={"bounds_error": True}
)
# combine all centered fields in a dataset
return self.package_colocate_results(centered_fields)
EMScalarFieldType = (
ScalarFieldDataArray
| ScalarFieldTimeDataArray
| ScalarModeFieldDataArray
| ScalarModeFieldCylindricalDataArray
| EMEScalarModeFieldDataArray
| EMEScalarFieldDataArray
)
[docs]
class ElectromagneticFieldDataset(AbstractFieldDataset, ABC):
"""Stores a collection of E and H fields with x, y, z components."""
Ex: EMScalarFieldType | None = Field(
None,
title="Ex",
description="Spatial distribution of the x-component of the electric field.",
)
Ey: EMScalarFieldType | None = Field(
None,
title="Ey",
description="Spatial distribution of the y-component of the electric field.",
)
Ez: EMScalarFieldType | None = Field(
None,
title="Ez",
description="Spatial distribution of the z-component of the electric field.",
)
Hx: EMScalarFieldType | None = Field(
None,
title="Hx",
description="Spatial distribution of the x-component of the magnetic field.",
)
Hy: EMScalarFieldType | None = Field(
None,
title="Hy",
description="Spatial distribution of the y-component of the magnetic field.",
)
Hz: EMScalarFieldType | None = Field(
None,
title="Hz",
description="Spatial distribution of the z-component of the magnetic field.",
)
@property
def field_components(self) -> dict[str, DataArray]:
"""Maps the field components to their associated data."""
fields = {
"Ex": self.Ex,
"Ey": self.Ey,
"Ez": self.Ez,
"Hx": self.Hx,
"Hy": self.Hy,
"Hz": self.Hz,
}
return {field_name: field for field_name, field in fields.items() if field is not None}
@property
def grid_locations(self) -> dict[str, str]:
"""Maps field components to the string key of their grid locations on the yee lattice."""
return {"Ex": "Ex", "Ey": "Ey", "Ez": "Ez", "Hx": "Hx", "Hy": "Hy", "Hz": "Hz"}
@property
def symmetry_eigenvalues(self) -> dict[str, Callable[[Axis], float]]:
"""Maps field components to their (positive) symmetry eigenvalues."""
return em_field_symmetry_eigenvalues()
[docs]
class FieldDataset(ElectromagneticFieldDataset):
"""Dataset storing a collection of the scalar components of E and H fields in the freq. domain.
Example
-------
>>> x = [-1,1]
>>> y = [-2,0,2]
>>> z = [-3,-1,1,3]
>>> f = [2e14, 3e14]
>>> coords = dict(x=x, y=y, z=z, f=f)
>>> scalar_field = ScalarFieldDataArray((1+1j) * np.random.random((2,3,4,2)), coords=coords)
>>> data = FieldDataset(Ex=scalar_field, Hz=scalar_field)
"""
Ex: ScalarFieldDataArray | None = Field(
None,
title="Ex",
description="Spatial distribution of the x-component of the electric field.",
)
Ey: ScalarFieldDataArray | None = Field(
None,
title="Ey",
description="Spatial distribution of the y-component of the electric field.",
)
Ez: ScalarFieldDataArray | None = Field(
None,
title="Ez",
description="Spatial distribution of the z-component of the electric field.",
)
Hx: ScalarFieldDataArray | None = Field(
None,
title="Hx",
description="Spatial distribution of the x-component of the magnetic field.",
)
Hy: ScalarFieldDataArray | None = Field(
None,
title="Hy",
description="Spatial distribution of the y-component of the magnetic field.",
)
Hz: ScalarFieldDataArray | None = Field(
None,
title="Hz",
description="Spatial distribution of the z-component of the magnetic field.",
)
[docs]
def from_zbf(filename: str, dim1: xyz, dim2: xyz) -> FieldDataset:
"""Creates a :class:`.FieldDataset` from a Zemax Beam File (``.zbf``).
Parameters
----------
filename: str
The file name of the .zbf file to read.
dim1: xyz
Tangential field component to map the x-dimension of the zbf data to.
eg. ``dim1 = "z"`` sets ``FieldDataset.Ez`` to ``Ex`` of the zbf data.
dim2: xyz
Tangential field component to map the y-dimension of the zbf data to.
eg. ``dim2 = "z"`` sets ``FieldDataset.Ez`` to ``Ey`` of the zbf data.
Returns
-------
:class:`.FieldDataset`
A :class:`.FieldDataset` object with two tangential E field components populated
by zbf data.
See Also
--------
:class:`.ZBFData`:
A class containing data read in from a ``.zbf`` file.
"""
log.warning(
"'FieldDataset.from_zbf()' is currently an experimental feature."
" If any issues are encountered, please contact Flexcompute support 'https://www.flexcompute.com/tidy3d/technical-support/'"
)
if dim1 not in get_args(xyz):
raise ValueError(f"'dim1' = '{dim1}' is not allowed, must be one of 'x', 'y', or 'z'.")
if dim2 not in get_args(xyz):
raise ValueError(f"'dim2' = '{dim2}' is not allowed, must be one of 'x', 'y', or 'z'.")
if dim1 == dim2:
raise ValueError("'dim1' and 'dim2' must be different.")
# get the third dimension
dim3 = list(set(get_args(xyz)) - {dim1, dim2})[0]
dims = {"x": 0, "y": 1, "z": 2}
dim2expand = dims[dim3] # this is for expanding E field arrays
# load zbf data
zbfdata = ZBFData.read_zbf(filename)
# Grab E fields, dimensions, wavelength
edim1 = zbfdata.Ex
edim2 = zbfdata.Ey
n1 = zbfdata.nx
n2 = zbfdata.ny
d1 = zbfdata.dx / UnitScaling[zbfdata.unit]
d2 = zbfdata.dy / UnitScaling[zbfdata.unit]
wavelength = zbfdata.wavelength / UnitScaling[zbfdata.unit]
# make scalar field data arrays
len1 = d1 * (n1 - 1)
len2 = d2 * (n2 - 1)
coords1 = np.linspace(-len1 / 2, len1 / 2, n1)
coords2 = np.linspace(-len2 / 2, len2 / 2, n2)
f = [C_0 / wavelength]
Edim1 = ScalarFieldDataArray(
np.expand_dims(edim1, axis=(dim2expand, 3)),
coords={
dim1: coords1,
dim2: coords2,
dim3: [0],
"f": f,
},
)
Edim2 = ScalarFieldDataArray(
np.expand_dims(edim2, axis=(dim2expand, 3)),
coords={
dim1: coords1,
dim2: coords2,
dim3: [0],
"f": f,
},
)
return FieldDataset(
**{
f"E{dim1}": Edim1,
f"E{dim2}": Edim2,
}
)
[docs]
class PointCloudFieldDataset(AbstractFieldDataset):
"""Dataset storing electromagnetic field components at point-cloud coordinates.
Field components are scalar data arrays indexed by ``("index", "f")``. The ``points`` array
maps each ``index`` to its Cartesian coordinate.
"""
points: PointDataArray = Field(
...,
title="Points",
description="Point coordinates associated with the indexed field data.",
)
Ex: IndexedFreqDataArray | None = Field(
None,
title="Ex",
description="Point-cloud x-component of the electric field.",
json_schema_extra={"units": ELECTRIC_FIELD_UNITS},
)
Ey: IndexedFreqDataArray | None = Field(
None,
title="Ey",
description="Point-cloud y-component of the electric field.",
json_schema_extra={"units": ELECTRIC_FIELD_UNITS},
)
Ez: IndexedFreqDataArray | None = Field(
None,
title="Ez",
description="Point-cloud z-component of the electric field.",
json_schema_extra={"units": ELECTRIC_FIELD_UNITS},
)
Hx: IndexedFreqDataArray | None = Field(
None,
title="Hx",
description="Point-cloud x-component of the magnetic field.",
json_schema_extra={"units": MAGNETIC_FIELD_UNITS},
)
Hy: IndexedFreqDataArray | None = Field(
None,
title="Hy",
description="Point-cloud y-component of the magnetic field.",
json_schema_extra={"units": MAGNETIC_FIELD_UNITS},
)
Hz: IndexedFreqDataArray | None = Field(
None,
title="Hz",
description="Point-cloud z-component of the magnetic field.",
json_schema_extra={"units": MAGNETIC_FIELD_UNITS},
)
Dx: IndexedFreqDataArray | None = Field(
None,
title="Dx",
description="Point-cloud x-component of ``D / epsilon_0``, computed from Ex and "
"the local x-direction relative permittivity.",
json_schema_extra={"units": ELECTRIC_FIELD_UNITS},
)
Dy: IndexedFreqDataArray | None = Field(
None,
title="Dy",
description="Point-cloud y-component of ``D / epsilon_0``, computed from Ey and "
"the local y-direction relative permittivity.",
json_schema_extra={"units": ELECTRIC_FIELD_UNITS},
)
Dz: IndexedFreqDataArray | None = Field(
None,
title="Dz",
description="Point-cloud z-component of ``D / epsilon_0``, computed from Ez and "
"the local z-direction relative permittivity.",
json_schema_extra={"units": ELECTRIC_FIELD_UNITS},
)
@property
def field_components(self) -> dict[str, DataArray]:
"""Maps the field components to their associated data."""
fields = {
"Ex": self.Ex,
"Ey": self.Ey,
"Ez": self.Ez,
"Hx": self.Hx,
"Hy": self.Hy,
"Hz": self.Hz,
"Dx": self.Dx,
"Dy": self.Dy,
"Dz": self.Dz,
}
return {field_name: field for field_name, field in fields.items() if field is not None}
@property
def grid_locations(self) -> dict[str, str]:
"""Point-cloud data is not sampled on named Yee-grid locations."""
raise DataError("Point-cloud field data does not have structured Yee-grid locations.")
@property
def symmetry_eigenvalues(self) -> dict[str, Callable[[Axis], float]]:
"""Maps field components to their (positive) symmetry eigenvalues."""
return point_cloud_field_symmetry_eigenvalues()
@model_validator(mode="after")
def _validate_field_indices(self) -> Self:
"""Ensure point-indexed field data is aligned with the point cloud."""
num_points = self.points.sizes["index"]
point_index = np.asarray(self.points.coords["index"])
for field_name, field_data in self.field_components.items():
if field_data.sizes["index"] != num_points:
self._raise_validation_error_at_loc(
f"Field component '{field_name}' has {field_data.sizes['index']} points, "
f"but 'points' contains {num_points} points.",
field_name,
)
if not np.array_equal(np.asarray(field_data.coords["index"]), point_index):
self._raise_validation_error_at_loc(
f"Field component '{field_name}' has index coordinates that do not match "
"the point-cloud index coordinates.",
field_name,
)
return self
[docs]
def colocate(self, x: ArrayLike = None, y: ArrayLike = None, z: ArrayLike = None) -> xr.Dataset:
"""Point-cloud field data is already sampled at requested points and cannot be colocated."""
raise DataError("PointCloudFieldDataset data cannot be colocated on a structured grid.")
[docs]
class PointCloudPermittivityDataset(AbstractFieldDataset):
"""Dataset storing permittivity components for requested point-cloud coordinates.
Components are scalar data arrays indexed by ``("index", "f")``. The ``points`` array maps
each ``index`` row to the requested Cartesian coordinate. Component values are sampled from the
nearest native Yee-grid locations, which may differ from the requested coordinates and from each
other across components.
"""
points: PointDataArray = Field(
...,
title="Points",
description="Requested point coordinates associated with the indexed permittivity data.",
)
eps_xx: IndexedFreqDataArray = Field(
title="Epsilon xx",
description="Point-cloud xx-component of the relative permittivity.",
)
eps_yy: IndexedFreqDataArray = Field(
title="Epsilon yy",
description="Point-cloud yy-component of the relative permittivity.",
)
eps_zz: IndexedFreqDataArray = Field(
title="Epsilon zz",
description="Point-cloud zz-component of the relative permittivity.",
)
@property
def field_components(self) -> dict[str, IndexedFreqDataArray]:
"""Maps the permittivity components to their associated data."""
return {
field_name: getattr(self, field_name)
for field_name in POINT_CLOUD_PERMITTIVITY_COMPONENTS
}
@property
def grid_locations(self) -> dict[str, str]:
"""Maps permittivity components to their native Yee-grid field locations."""
return {"eps_xx": "Ex", "eps_yy": "Ey", "eps_zz": "Ez"}
@property
def symmetry_eigenvalues(self) -> dict[str, None]:
"""Maps permittivity components to their scalar symmetry eigenvalues."""
return {"eps_xx": None, "eps_yy": None, "eps_zz": None}
@model_validator(mode="after")
def _validate_component_indices(self) -> Self:
"""Ensure point-indexed permittivity data is aligned with the point cloud."""
num_points = self.points.sizes["index"]
point_index = np.asarray(self.points.coords["index"])
for component_name, component_data in self.field_components.items():
if component_data.sizes["index"] != num_points:
self._raise_validation_error_at_loc(
f"Permittivity component '{component_name}' has "
f"{component_data.sizes['index']} points, but 'points' contains "
f"{num_points} points.",
component_name,
)
if not np.array_equal(np.asarray(component_data.coords["index"]), point_index):
self._raise_validation_error_at_loc(
f"Permittivity component '{component_name}' has index coordinates that do "
"not match the point-cloud index coordinates.",
component_name,
)
return self
[docs]
def colocate(self, x: ArrayLike = None, y: ArrayLike = None, z: ArrayLike = None) -> xr.Dataset:
"""Point-cloud permittivity data cannot be colocated on a structured grid."""
raise DataError(
"PointCloudPermittivityDataset data cannot be colocated on a structured grid."
)
[docs]
class FieldTimeDataset(ElectromagneticFieldDataset):
"""Dataset storing a collection of the scalar components of E and H fields in the time domain
Example
-------
>>> x = [-1,1]
>>> y = [-2,0,2]
>>> z = [-3,-1,1,3]
>>> t = [0, 1e-12, 2e-12]
>>> coords = dict(x=x, y=y, z=z, t=t)
>>> scalar_field = ScalarFieldTimeDataArray(np.random.random((2,3,4,3)), coords=coords)
>>> data = FieldTimeDataset(Ex=scalar_field, Hz=scalar_field)
"""
Ex: ScalarFieldTimeDataArray | None = Field(
None,
title="Ex",
description="Spatial distribution of the x-component of the electric field.",
)
Ey: ScalarFieldTimeDataArray | None = Field(
None,
title="Ey",
description="Spatial distribution of the y-component of the electric field.",
)
Ez: ScalarFieldTimeDataArray | None = Field(
None,
title="Ez",
description="Spatial distribution of the z-component of the electric field.",
)
Hx: ScalarFieldTimeDataArray | None = Field(
None,
title="Hx",
description="Spatial distribution of the x-component of the magnetic field.",
)
Hy: ScalarFieldTimeDataArray | None = Field(
None,
title="Hy",
description="Spatial distribution of the y-component of the magnetic field.",
)
Hz: ScalarFieldTimeDataArray | None = Field(
None,
title="Hz",
description="Spatial distribution of the z-component of the magnetic field.",
)
[docs]
def apply_phase(self, phase: float) -> AbstractFieldDataset:
"""Create a copy where all elements are phase-shifted by a value (in radians)."""
if phase != 0.0:
raise ValueError("Can't apply phase to time-domain field data, which is real-valued.")
return self
class AuxFieldDataset(AbstractFieldDataset, ABC):
"""Stores a collection of aux fields with x, y, z components."""
Nfx: EMScalarFieldType | None = Field(
None,
title="Nfx",
description="Spatial distribution of the free carrier density for "
"polarization in the x-direction.",
)
Nfy: EMScalarFieldType | None = Field(
None,
title="Nfy",
description="Spatial distribution of the free carrier density for "
"polarization in the y-direction.",
)
Nfz: EMScalarFieldType | None = Field(
None,
title="Nfz",
description="Spatial distribution of the free carrier density for "
"polarization in the z-direction.",
)
@property
def field_components(self) -> dict[str, DataArray]:
"""Maps the field components to their associated data."""
fields = {
"Nfx": self.Nfx,
"Nfy": self.Nfy,
"Nfz": self.Nfz,
}
return {field_name: field for field_name, field in fields.items() if field is not None}
@property
def grid_locations(self) -> dict[str, str]:
"""Maps field components to the string key of their grid locations on the yee lattice."""
return {"Nfx": "Ex", "Nfy": "Ey", "Nfz": "Ez"}
@property
def symmetry_eigenvalues(self) -> dict[str, Callable[[Axis], float]]:
"""Maps field components to their (positive) symmetry eigenvalues."""
return {
"Nfx": lambda dim: +1,
"Nfy": lambda dim: +1,
"Nfz": lambda dim: +1,
}
[docs]
class AuxFieldTimeDataset(AuxFieldDataset):
"""Dataset storing a collection of the scalar components of aux fields in the time domain
Example
-------
>>> x = [-1,1]
>>> y = [-2,0,2]
>>> z = [-3,-1,1,3]
>>> t = [0, 1e-12, 2e-12]
>>> coords = dict(x=x, y=y, z=z, t=t)
>>> scalar_field = ScalarFieldTimeDataArray(np.random.random((2,3,4,3)), coords=coords)
>>> data = AuxFieldTimeDataset(Nfx=scalar_field)
"""
Nfx: ScalarFieldTimeDataArray | None = Field(
None,
title="Nfx",
description="Spatial distribution of the free carrier density for polarization "
"in the x-direction.",
)
Nfy: ScalarFieldTimeDataArray | None = Field(
None,
title="Nfy",
description="Spatial distribution of the free carrier density for polarization "
"in the y-direction.",
)
Nfz: ScalarFieldTimeDataArray | None = Field(
None,
title="Nfz",
description="Spatial distribution of the free carrier density for polarization "
"in the z-direction.",
)
[docs]
class ElectromagneticSurfaceFieldDataset(AbstractFieldDataset, ABC):
"""Stores a collection of E and H fields with x, y, z components on one side of the surface."""
E: TriangularSurfaceDataset | None = Field(
None,
title="E",
description="Spatial distribution of the electric field on the one side of the surface.",
)
H: TriangularSurfaceDataset | None = Field(
None,
title="H",
description="Spatial distribution of the magnetic field on the one side of the surface.",
)
normal: TriangularSurfaceDataset = Field(
...,
title="Normal",
description="Normal direction of the surface oriented outward from the surface.",
)
@property
def field_components(self) -> dict[str, DataArray]:
"""Maps the field components to their associated data."""
fields = {
"E": self.E,
"H": self.H,
}
return {field_name: field for field_name, field in fields.items() if field is not None}
@property
def intensity(self) -> TriangularSurfaceDataset:
"""Return the sum of the squared absolute electric field components."""
if self.E is None:
raise DataError(
"Could not calculate intensity: the dataset does not contain E field information."
)
intensity = self.E.norm(dim="axis") ** 2
return intensity
@property
def current_density(self) -> TriangularSurfaceDataset:
"""Surface current density."""
h_diff = 0
H_inside = None
H_outside = None
if self.H is not None:
# we assume that if data is None it means field is zero on that side (e.g. PEC)
# NOTE: we use self.H.values.sel() (raw xarray) rather than self.H.sel() because the
# latter goes through _non_spatial_sel which wraps scalar selectors in lists, defeating
# drop=True and keeping the 'side' dimension; this would cause NaN on subtraction due to
# xarray coordinate alignment on mismatched 'side' values.
H_inside = (
self.H.values.sel(side="inside", drop=True)
if "inside" in self.H.values.side
else None
)
H_outside = (
self.H.values.sel(side="outside", drop=True)
if "outside" in self.H.values.side
else None
)
if H_inside is not None:
h_diff = h_diff + H_inside
if H_outside is not None:
h_diff = h_diff - H_outside
if H_inside is None and H_outside is None:
raise DataError(
"Could not calculate current density: the dataset does not contain H field information."
)
return self.H.updated_copy(values=xr.cross(h_diff, self.normal.values, dim="axis"))
@property
def grid_locations(self) -> dict[str, str]:
"""Maps field components to the string key of their grid locations on the yee lattice."""
raise RuntimeError("Function 'grid_location' does not apply to surface monitors.")
@property
def symmetry_eigenvalues(self) -> dict[str, Callable[[Axis], float]]:
"""Maps field components to their (positive) symmetry eigenvalues."""
return em_field_symmetry_eigenvalues()
[docs]
class ModeSolverDataset(ElectromagneticFieldDataset, ModeFreqDataset):
"""Dataset storing scalar components of E and H fields as a function of freq. and mode_index.
Example
-------
>>> from tidy3d import ModeSpec
>>> x = [-1,1]
>>> y = [0]
>>> z = [-3,-1,1,3]
>>> f = [2e14, 3e14]
>>> mode_index = np.arange(5)
>>> field_coords = dict(x=x, y=y, z=z, f=f, mode_index=mode_index)
>>> field = ScalarModeFieldDataArray((1+1j)*np.random.random((2,1,4,2,5)), coords=field_coords)
>>> index_coords = dict(f=f, mode_index=mode_index)
>>> index_data = ModeIndexDataArray((1+1j) * np.random.random((2,5)), coords=index_coords)
>>> data = ModeSolverDataset(
... Ex=field,
... Ey=field,
... Ez=field,
... Hx=field,
... Hy=field,
... Hz=field,
... n_complex=index_data
... )
"""
Ex: ScalarModeFieldDataArray | None = Field(
None,
title="Ex",
description="Spatial distribution of the x-component of the electric field of the mode.",
)
Ey: ScalarModeFieldDataArray | None = Field(
None,
title="Ey",
description="Spatial distribution of the y-component of the electric field of the mode.",
)
Ez: ScalarModeFieldDataArray | None = Field(
None,
title="Ez",
description="Spatial distribution of the z-component of the electric field of the mode.",
)
Hx: ScalarModeFieldDataArray | None = Field(
None,
title="Hx",
description="Spatial distribution of the x-component of the magnetic field of the mode.",
)
Hy: ScalarModeFieldDataArray | None = Field(
None,
title="Hy",
description="Spatial distribution of the y-component of the magnetic field of the mode.",
)
Hz: ScalarModeFieldDataArray | None = Field(
None,
title="Hz",
description="Spatial distribution of the z-component of the magnetic field of the mode.",
)
n_complex: ModeIndexDataArray = Field(
title="Propagation Index",
description="Complex-valued effective propagation constants associated with the mode.",
)
n_group_raw: GroupIndexDataArray | None = Field(
None,
alias="n_group", # This is for backwards compatibility only when loading old data
title="Group Index",
description="Index associated with group velocity of the mode.",
)
dispersion_raw: ModeDispersionDataArray | None = Field(
None,
title="Dispersion",
description="Dispersion parameter for the mode.",
json_schema_extra={"units": PICOSECOND_PER_NANOMETER_PER_KILOMETER},
)
@property
def n_eff(self) -> ModeIndexDataArray:
"""Real part of the propagation index."""
return self.n_complex.real
@property
def k_eff(self) -> ModeIndexDataArray:
"""Imaginary part of the propagation index."""
return self.n_complex.imag
@property
def n_group(self) -> GroupIndexDataArray:
"""Group index."""
if self.n_group_raw is None:
log.warning(
"The group index was not computed. To calculate group index, pass "
"'group_index_step = True' in the 'ModeSpec'.",
log_once=True,
)
return self.n_group_raw
@property
def dispersion(self) -> ModeDispersionDataArray:
r"""Dispersion parameter.
.. math::
D = -\frac{\lambda}{c_0} \frac{{\rm d}^2 n_{\text{eff}}}{{\rm d}\lambda^2}
"""
if self.dispersion_raw is None:
log.warning(
"The dispersion was not computed. To calculate dispersion, pass "
"'group_index_step = True' in the 'ModeSpec'.",
log_once=True,
)
return self.dispersion_raw
[docs]
def plot_field(self, *args: Any, **kwargs: Any) -> None:
"""Warn user to use the :class:`.ModeSolver` ``plot_field`` function now."""
raise DeprecationWarning(
"The 'plot_field()' method was moved to the 'ModeSolver' object."
"Once the 'ModeSolver' is constructed, one may call '.plot_field()' on the object and "
"the modes will be computed and displayed with 'Simulation' overlay."
)
class AbstractMediumPropertyDataset(AbstractFieldDataset, ABC):
"""Dataset storing medium property."""
eps_xx: ScalarFieldDataArray = Field(
title="Epsilon xx",
description="Spatial distribution of the xx-component of the relative permittivity.",
)
eps_yy: ScalarFieldDataArray = Field(
title="Epsilon yy",
description="Spatial distribution of the yy-component of the relative permittivity.",
)
eps_zz: ScalarFieldDataArray = Field(
title="Epsilon zz",
description="Spatial distribution of the zz-component of the relative permittivity.",
)
[docs]
class PermittivityDataset(AbstractMediumPropertyDataset):
"""Dataset storing the diagonal components of the permittivity tensor.
Example
-------
>>> x = [-1,1]
>>> y = [-2,0,2]
>>> z = [-3,-1,1,3]
>>> f = [2e14, 3e14]
>>> coords = dict(x=x, y=y, z=z, f=f)
>>> sclr_fld = ScalarFieldDataArray((1+1j) * np.random.random((2,3,4,2)), coords=coords)
>>> data = PermittivityDataset(eps_xx=sclr_fld, eps_yy=sclr_fld, eps_zz=sclr_fld)
"""
@property
def field_components(self) -> dict[str, ScalarFieldDataArray]:
"""Maps the field components to their associated data."""
return {"eps_xx": self.eps_xx, "eps_yy": self.eps_yy, "eps_zz": self.eps_zz}
@property
def grid_locations(self) -> dict[str, str]:
"""Maps field components to the string key of their grid locations on the yee lattice."""
return {"eps_xx": "Ex", "eps_yy": "Ey", "eps_zz": "Ez"}
@property
def symmetry_eigenvalues(self) -> dict[str, None]:
"""Maps field components to their (positive) symmetry eigenvalues."""
return {"eps_xx": None, "eps_yy": None, "eps_zz": None}
class MediumDataset(AbstractMediumPropertyDataset):
"""Dataset storing the diagonal components of the permittivity and permeability tensor.
Example
-------
>>> x = [-1,1]
>>> y = [-2,0,2]
>>> z = [-3,-1,1,3]
>>> f = [2e14, 3e14]
>>> coords = dict(x=x, y=y, z=z, f=f)
>>> sclr_fld = ScalarFieldDataArray((1+1j) * np.random.random((2,3,4,2)), coords=coords)
>>> data = MediumDataset(eps_xx=sclr_fld, eps_yy=sclr_fld, eps_zz=sclr_fld, mu_xx=sclr_fld, mu_yy=sclr_fld, mu_zz=sclr_fld)
"""
mu_xx: ScalarFieldDataArray = Field(
title="Mu xx",
description="Spatial distribution of the xx-component of the relative permeability.",
)
mu_yy: ScalarFieldDataArray = Field(
title="Mu yy",
description="Spatial distribution of the yy-component of the relative permeability.",
)
mu_zz: ScalarFieldDataArray = Field(
title="Mu zz",
description="Spatial distribution of the zz-component of the relative permeability.",
)
@property
def field_components(self) -> dict[str, ScalarFieldDataArray]:
"""Maps the field components to their associated data."""
return {
"eps_xx": self.eps_xx,
"eps_yy": self.eps_yy,
"eps_zz": self.eps_zz,
"mu_xx": self.mu_xx,
"mu_yy": self.mu_yy,
"mu_zz": self.mu_zz,
}
@property
def grid_locations(self) -> dict[str, str]:
"""Maps field components to the string key of their grid locations on the yee lattice."""
return {
"eps_xx": "Ex",
"eps_yy": "Ey",
"eps_zz": "Ez",
"mu_xx": "Hx",
"mu_yy": "Hy",
"mu_zz": "Hz",
}
@property
def symmetry_eigenvalues(self) -> dict[str, None]:
"""Maps field components to their (positive) symmetry eigenvalues."""
return {
"eps_xx": None,
"eps_yy": None,
"eps_zz": None,
"mu_xx": None,
"mu_yy": None,
"mu_zz": None,
}
class TriangleMeshDataset(Dataset):
"""Dataset for storing triangular surface data."""
surface_mesh: TriangleMeshDataArray = Field(
title="Surface mesh data",
description="Dataset containing the surface triangles and corresponding face indices "
"for a surface mesh.",
)
class TimeDataset(Dataset):
"""Dataset for storing a function of time."""
values: TimeDataArray = Field(
title="Values",
description="Values as a function of time.",
)