Add aux field monitors

Author: caseyflex

CHANGELOG.md

@@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 - New `LobeMeasurer` tool in the `microwave` plugin that locates lobes in antenna patterns and calculates lobe measures like half-power beamwidth and sidelobe level.
 - Validation step that raises a `ValueError` when no frequency-domain monitors are present, preventing invalid adjoint runs.
+- Monitor `AuxFieldTimeMonitor` for aux fields like the free carrier density in `TwoPhotonAbsorption`.
 
 ### Changed
 

tests/test_components/test_medium.py

@@ -769,6 +769,29 @@ def test_nonlinear_medium():
         )
         _ = sim.updated_copy(medium=med, path="structures/0")
 
+    grid_spec = td.GridSpec.auto(min_steps_per_wvl=10, wavelength=1)
+    sim = sim.updated_copy(grid_spec=grid_spec)
+    aux_fields = ("Nfz",)
+    with AssertLogLevel(None):
+        med = td.Medium(
+            nonlinear_spec=td.NonlinearSpec(models=[td.TwoPhotonAbsorption(beta=1, tau=1)])
+        )
+        monitor = td.AuxFieldTimeMonitor(
+            interval=1, size=(0, 0, 0), name="aux_field_time", fields=aux_fields
+        )
+        sim = sim.updated_copy(medium=med, path="structures/0")
+        sim = sim.updated_copy(monitors=[monitor])
+
+    with AssertLogLevel("WARNING", contains_str="stores field"):
+        med = td.Medium(
+            nonlinear_spec=td.NonlinearSpec(models=[td.TwoPhotonAbsorption(beta=1, tau=0)])
+        )
+        _ = sim.updated_copy(medium=med, path="structures/0")
+
+    with AssertLogLevel("WARNING", contains_str="stores field"):
+        med = td.Medium(nonlinear_spec=td.NonlinearSpec(models=[td.KerrNonlinearity(n2=1)]))
+        _ = sim.updated_copy(medium=med, path="structures/0")
+
 
 def test_custom_medium():
     Nx, Ny, Nz, Nf = 4, 3, 1, 1

tests/test_components/test_monitor.py

@@ -376,6 +376,7 @@ def test_monitor():
         name="directivity",
     )
     m10 = td.PermittivityMonitor(size=size, center=center, freqs=FREQS, name="perm")
+    m11 = td.AuxFieldTimeMonitor(size=size, center=center, name="aux_field_time", fields=("Nfx",))
 
     tmesh = np.linspace(0, 1, 10)
 

tests/test_data/test_data_arrays.py

@@ -30,6 +30,7 @@
     ),
 ]
 FIELDS = ("Ex", "Ey", "Ez", "Hx", "Hz")
+AUX_FIELDS = ("Nfz",)
 INTERVAL = 2
 ORDERS_X = list(range(-1, 2))
 ORDERS_Y = list(range(-2, 3))
@@ -51,6 +52,9 @@
 FIELD_TIME_MONITOR = td.FieldTimeMonitor(
     size=SIZE_3D, fields=FIELDS, name="field_time", interval=INTERVAL
 )
+AUX_FIELD_TIME_MONITOR = td.AuxFieldTimeMonitor(
+    size=SIZE_3D, fields=AUX_FIELDS, name="aux_field_time", interval=INTERVAL
+)
 FIELD_MONITOR_2D = td.FieldMonitor(size=SIZE_2D, fields=FIELDS, name="field_2d", freqs=FREQS)
 FIELD_TIME_MONITOR_2D = td.FieldTimeMonitor(
     size=SIZE_2D, fields=FIELDS, name="field_time_2d", interval=INTERVAL
@@ -80,6 +84,7 @@
 MONITORS = [
     FIELD_MONITOR,
     FIELD_TIME_MONITOR,
+    AUX_FIELD_TIME_MONITOR,
     MODE_MONITOR_WITH_FIELDS,
     PERMITTIVITY_MONITOR,
     MODE_MONITOR,

tests/test_data/test_monitor_data.py

@@ -11,6 +11,7 @@
     FreqModeDataArray,
 )
 from tidy3d.components.data.monitor_data import (
+    AuxFieldTimeData,
     DiffractionData,
     DirectivityData,
     FieldData,
@@ -24,6 +25,7 @@
 
 from ..utils import AssertLogLevel
 from .test_data_arrays import (
+    AUX_FIELD_TIME_MONITOR,
     DIFFRACTION_MONITOR,
     DIRECTIVITY_MONITOR,
     FIELD_MONITOR,
@@ -120,6 +122,17 @@ def make_field_time_data_2d(symmetry: bool = True):
     )
 
 
+def make_aux_field_time_data(symmetry: bool = True):
+    sim = SIM_SYM if symmetry else SIM
+    return AuxFieldTimeData(
+        monitor=AUX_FIELD_TIME_MONITOR,
+        Nfz=make_scalar_field_time_data_array("Ez", symmetry),
+        symmetry=sim.symmetry,
+        symmetry_center=sim.center,
+        grid_expanded=sim.discretize_monitor(AUX_FIELD_TIME_MONITOR),
+    )
+
+
 def make_mode_solver_data():
     mode_data = ModeData(
         monitor=MODE_MONITOR_WITH_FIELDS,

tests/test_data/test_sim_data.py

@@ -15,6 +15,7 @@
 from ..utils import get_nested_shape
 from .test_data_arrays import FIELD_MONITOR, SIM, SIM_SYM
 from .test_monitor_data import (
+    make_aux_field_time_data,
     make_diffraction_data,
     make_directivity_data,
     make_field_data,
@@ -32,6 +33,8 @@
 FIELD = make_field_data(symmetry=False)
 FIELD_TIME_SYM = make_field_time_data()
 FIELD_TIME = make_field_time_data(symmetry=False)
+AUX_FIELD_TIME_SYM = make_aux_field_time_data()
+AUX_FIELD_TIME = make_aux_field_time_data(symmetry=False)
 PERMITTIVITY_SYM = make_permittivity_data()
 PERMITTIVITY = make_permittivity_data(symmetry=False)
 MODE = make_mode_data()
@@ -45,6 +48,7 @@
 MONITOR_DATA = (
     FIELD,
     FIELD_TIME,
+    AUX_FIELD_TIME,
     MODE_SOLVER,
     PERMITTIVITY,
     MODE,
@@ -56,6 +60,7 @@
 MONITOR_DATA_SYM = (
     FIELD_SYM,
     FIELD_TIME_SYM,
+    AUX_FIELD_TIME_SYM,
     MODE_SOLVER,
     PERMITTIVITY_SYM,
     MODE,

tests/utils.py

@@ -780,6 +780,9 @@ def make_custom_data(lims, unstructured):
             size=(0, 0, 0), center=(0, 0, 0), fields=["Ex"], freqs=[1.5e14, 2e14], name="field"
         ),
         td.FieldTimeMonitor(size=(0, 0, 0), center=(0, 0, 0), name="field_time", interval=100),
+        td.AuxFieldTimeMonitor(
+            size=(0, 0, 0), center=(0, 0, 0), fields=("Nfx",), name="aux_field_time", interval=100
+        ),
         td.FluxMonitor(size=(1, 1, 0), center=(0, 0, 0), freqs=[2e14, 2.5e14], name="flux"),
         td.FluxTimeMonitor(size=(1, 1, 0), center=(0, 0, 0), name="flux_time"),
         td.PermittivityMonitor(size=(1, 1, 0.1), name="eps", freqs=[1e14]),

tidy3d/__init__.py

@@ -267,6 +267,7 @@
 
 # monitors
 from .components.monitor import (
+    AuxFieldTimeMonitor,
     DiffractionMonitor,
     DirectivityMonitor,
     FieldMonitor,
@@ -460,6 +461,7 @@ def set_logging_level(level: str) -> None:
     "PlaneWaveBeamProfile",
     "FieldMonitor",
     "FieldTimeMonitor",
+    "AuxFieldTimeMonitor",
     "FluxMonitor",
     "FluxTimeMonitor",
     "ModeMonitor",
@@ -496,6 +498,7 @@ def set_logging_level(level: str) -> None:
     "ModeSolverDataset",
     "FieldData",
     "FieldTimeData",
+    "AuxFieldTimeData",
     "PermittivityData",
     "FluxData",
     "FluxTimeData",

tidy3d/components/data/dataset.py

@@ -312,6 +312,88 @@ def apply_phase(self, phase: float) -> AbstractFieldDataset:
         return self
 
 
+class AuxFieldDataset(AbstractFieldDataset, ABC):
+    """Stores a collection of aux fields with x, y, z components."""
+
+    Nfx: EMScalarFieldType = pd.Field(
+        None,
+        title="Nfx",
+        description="Spatial distribution of the free carrier density for "
+        "polarization in the x-direction.",
+    )
+    Nfy: EMScalarFieldType = pd.Field(
+        None,
+        title="Nfy",
+        description="Spatial distribution of the free carrier density for "
+        "polarization in the y-direction.",
+    )
+    Nfz: EMScalarFieldType = pd.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 dict(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 dict(
+            Nfx=lambda dim: +1,
+            Nfy=lambda dim: +1,
+            Nfz=lambda dim: +1,
+        )
+
+
+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 = pd.Field(
+        None,
+        title="Nfx",
+        description="Spatial distribution of the free carrier density for polarization "
+        "in the x-direction.",
+    )
+    Nfy: ScalarFieldTimeDataArray = pd.Field(
+        None,
+        title="Nfy",
+        description="Spatial distribution of the free carrier density for polarization "
+        "in the y-direction.",
+    )
+    Nfz: ScalarFieldTimeDataArray = pd.Field(
+        None,
+        title="Nfz",
+        description="Spatial distribution of the free carrier density for polarization "
+        "in the z-direction.",
+    )
+
+
 class ModeSolverDataset(ElectromagneticFieldDataset):
     """Dataset storing scalar components of E and H fields as a function of freq. and mode_index.
 

tidy3d/components/data/monitor_data.py

@@ -21,6 +21,7 @@
 from ..grid.grid import Coords, Grid
 from ..medium import Medium, MediumType
 from ..monitor import (
+    AuxFieldTimeMonitor,
     DiffractionMonitor,
     DirectivityMonitor,
     FieldMonitor,
@@ -85,6 +86,7 @@
 )
 from .dataset import (
     AbstractFieldDataset,
+    AuxFieldTimeDataset,
     Dataset,
     ElectromagneticFieldDataset,
     FieldDataset,
@@ -196,7 +198,9 @@ def get_amplitude(x) -> complex:
 class AbstractFieldData(MonitorData, AbstractFieldDataset, ABC):
     """Collection of scalar fields with some symmetry properties."""
 
-    monitor: Union[FieldMonitor, FieldTimeMonitor, PermittivityMonitor, ModeMonitor]
+    monitor: Union[
+        FieldMonitor, FieldTimeMonitor, AuxFieldTimeMonitor, PermittivityMonitor, ModeMonitor
+    ]
 
     symmetry: Tuple[Symmetry, Symmetry, Symmetry] = pd.Field(
         (0, 0, 0),
@@ -1328,6 +1332,41 @@ def time_reversed_copy(self) -> FieldTimeData:
         return self.copy(update=new_data)
 
 
+class AuxFieldTimeData(AuxFieldTimeDataset, AbstractFieldData):
+    """
+    Data associated with a :class:`.AuxFieldTimeMonitor`: scalar components of aux fields.
+
+    Notes
+    -----
+
+        The data is stored as a `DataArray <https://docs.xarray.dev/en/stable/generated/xarray.DataArray.html>`_
+        object using the `xarray <https://docs.xarray.dev/en/stable/index.html>`_ package.
+
+    Example
+    -------
+    >>> from tidy3d import ScalarFieldTimeDataArray
+    >>> x = [-1,1,3]
+    >>> y = [-2,0,2,4]
+    >>> z = [-3,-1,1,3,5]
+    >>> t = [0, 1e-12, 2e-12]
+    >>> coords = dict(x=x[:-1], y=y[:-1], z=z[:-1], t=t)
+    >>> grid = Grid(boundaries=Coords(x=x, y=y, z=z))
+    >>> scalar_field = ScalarFieldTimeDataArray(np.random.random((2,3,4,3)), coords=coords)
+    >>> monitor = AuxFieldTimeMonitor(
+    ...     size=(2,4,6), interval=100, name='field', fields=['Nfx'], colocate=True
+    ... )
+    >>> data = AuxFieldTimeData(monitor=monitor, Nfx=scalar_field, grid_expanded=grid)
+    """
+
+    monitor: AuxFieldTimeMonitor = pd.Field(
+        ...,
+        title="Monitor",
+        description="Time-domain auxiliary field monitor associated with the data.",
+    )
+
+    _contains_monitor_fields = enforce_monitor_fields_present()
+
+
 class PermittivityData(PermittivityDataset, AbstractFieldData):
     """Data for a :class:`.PermittivityMonitor`: diagonal components of the permittivity tensor.
 
@@ -3684,6 +3723,7 @@ def fields_circular_polarization(self) -> xr.Dataset:
     ModeData,
     FluxData,
     FluxTimeData,
+    AuxFieldTimeData,
     FieldProjectionKSpaceData,
     FieldProjectionCartesianData,
     FieldProjectionAngleData,

tidy3d/components/medium.py

@@ -260,6 +260,11 @@ def complex_fields(self) -> bool:
         """Whether the model uses complex fields."""
         return False
 
+    @property
+    def aux_fields(self) -> List[str]:
+        """List of available aux fields in this model."""
+        return []
+
 
 class NonlinearSusceptibility(NonlinearModel):
     """Model for an instantaneous nonlinear chi3 susceptibility.
@@ -496,6 +501,13 @@ def complex_fields(self) -> bool:
         """Whether the model uses complex fields."""
         return self.use_complex_fields
 
+    @property
+    def aux_fields(self) -> List[str]:
+        """List of available aux fields in this model."""
+        if self.tau == 0:
+            return []
+        return ["Nfx", "Nfy", "Nfz"]
+
 
 class KerrNonlinearity(NonlinearModel):
     """Model for Kerr nonlinearity which gives an intensity-dependent refractive index
@@ -713,6 +725,14 @@ def _hardcode_medium_freqs(
             new_models.append(new_model)
         return self.updated_copy(models=new_models)
 
+    @property
+    def aux_fields(self) -> List[str]:
+        """List of available aux fields in all present models."""
+        fields = []
+        for model in self.models:
+            fields += model.aux_fields
+        return fields
+
 
 class AbstractMedium(ABC, Tidy3dBaseModel):
     """A medium within which electromagnetic waves propagate."""