tidy3d.plugins.adjoint.JaxCustomMedium#

class tidy3d.plugins.adjoint.JaxCustomMedium#

A Medium registered with jax.

Parameters

name (Optional[str] = None) – Optional unique name for medium.
frequency_range (Optional[Tuple[float, float]] = None) – [units = (Hz, Hz)]. Optional range of validity for the medium.
eps_dataset (JaxPermittivityDataset) – User-supplied dataset containing complex-valued permittivity as a function of space. Permittivity distribution over the Yee-grid will be interpolated based on interp_method.
interp_method (Literal['nearest', 'linear'] = nearest) – Interpolation method to obtain permittivity values that are not supplied at the Yee grids; For grids outside the range of the supplied data, extrapolation will be applied. When the extrapolated value is smaller (greater) than the minimal (maximal) of the supplied data, the extrapolated value will take the minimal (maximal) of the supplied data.

Show JSON schema

{
   "title": "JaxCustomMedium",
   "description": "A :class:`.Medium` registered with jax.\n\nParameters\n----------\nname : Optional[str] = None\n    Optional unique name for medium.\nfrequency_range : Optional[Tuple[float, float]] = None\n    [units = (Hz, Hz)].  Optional range of validity for the medium.\neps_dataset : JaxPermittivityDataset\n    User-supplied dataset containing complex-valued permittivity as a function of space. Permittivity distribution over the Yee-grid will be interpolated based on ``interp_method``.\ninterp_method : Literal['nearest', 'linear'] = nearest\n    Interpolation method to obtain permittivity values that are not supplied at the Yee grids; For grids outside the range of the supplied data, extrapolation will be applied. When the extrapolated value is smaller (greater) than the minimal (maximal) of the supplied data, the extrapolated value will take the minimal (maximal) of the supplied data.",
   "type": "object",
   "properties": {
      "type": {
         "title": "Type",
         "default": "JaxCustomMedium",
         "enum": [
            "JaxCustomMedium"
         ],
         "type": "string"
      },
      "name": {
         "title": "Name",
         "description": "Optional unique name for medium.",
         "type": "string"
      },
      "frequency_range": {
         "title": "Frequency Range",
         "description": "Optional range of validity for the medium.",
         "units": [
            "Hz",
            "Hz"
         ],
         "type": "array",
         "minItems": 2,
         "maxItems": 2,
         "items": [
            {
               "type": "number"
            },
            {
               "type": "number"
            }
         ]
      },
      "eps_dataset": {
         "title": "Permittivity Dataset",
         "description": "User-supplied dataset containing complex-valued permittivity as a function of space. Permittivity distribution over the Yee-grid will be interpolated based on ``interp_method``.",
         "jax_field": true,
         "allOf": [
            {
               "$ref": "#/definitions/JaxPermittivityDataset"
            }
         ]
      },
      "interp_method": {
         "title": "Interpolation method",
         "description": "Interpolation method to obtain permittivity values that are not supplied at the Yee grids; For grids outside the range of the supplied data, extrapolation will be applied. When the extrapolated value is smaller (greater) than the minimal (maximal) of the supplied data, the extrapolated value will take the minimal (maximal) of the supplied data.",
         "default": "nearest",
         "enum": [
            "nearest",
            "linear"
         ],
         "type": "string"
      }
   },
   "required": [
      "eps_dataset"
   ],
   "additionalProperties": false,
   "definitions": {
      "JaxDataArray": {
         "title": "JaxDataArray",
         "description": "A :class:`.DataArray`-like class that only wraps xarray for jax compability.\n\nParameters\n----------\nvalues : Optional[Any]\n    Nested list containing the raw values, which can be tracked by jax.\ncoords : Mapping[str, list]\n    Dictionary storing the coordinates, namely ``(direction, f, mode_index)``.",
         "type": "object",
         "properties": {
            "values": {
               "title": "Values",
               "description": "Nested list containing the raw values, which can be tracked by jax.",
               "jax_field": true
            },
            "coords": {
               "title": "Coords",
               "description": "Dictionary storing the coordinates, namely ``(direction, f, mode_index)``.",
               "type": "object",
               "additionalProperties": {
                  "type": "array",
                  "items": {}
               }
            },
            "type": {
               "title": "Type",
               "default": "JaxDataArray",
               "enum": [
                  "JaxDataArray"
               ],
               "type": "string"
            }
         },
         "required": [
            "values",
            "coords"
         ],
         "additionalProperties": false
      },
      "JaxPermittivityDataset": {
         "title": "JaxPermittivityDataset",
         "description": "A :class:`.PermittivityDataset` registered with jax.\n\nParameters\n----------\neps_xx : JaxDataArray\n    Spatial distribution of the xx-component of the relative permittivity.\neps_yy : JaxDataArray\n    Spatial distribution of the yy-component of the relative permittivity.\neps_zz : JaxDataArray\n    Spatial distribution of the zz-component of the relative permittivity.",
         "type": "object",
         "properties": {
            "type": {
               "title": "Type",
               "default": "JaxPermittivityDataset",
               "enum": [
                  "JaxPermittivityDataset"
               ],
               "type": "string"
            },
            "eps_xx": {
               "title": "Epsilon xx",
               "description": "Spatial distribution of the xx-component of the relative permittivity.",
               "jax_field": true,
               "allOf": [
                  {
                     "$ref": "#/definitions/JaxDataArray"
                  }
               ]
            },
            "eps_yy": {
               "title": "Epsilon yy",
               "description": "Spatial distribution of the yy-component of the relative permittivity.",
               "jax_field": true,
               "allOf": [
                  {
                     "$ref": "#/definitions/JaxDataArray"
                  }
               ]
            },
            "eps_zz": {
               "title": "Epsilon zz",
               "description": "Spatial distribution of the zz-component of the relative permittivity.",
               "jax_field": true,
               "allOf": [
                  {
                     "$ref": "#/definitions/JaxDataArray"
                  }
               ]
            }
         },
         "required": [
            "eps_xx",
            "eps_yy",
            "eps_zz"
         ],
         "additionalProperties": false
      }
   }
}

attribute eps_dataset: JaxPermittivityDataset [Required]#

User-supplied dataset containing complex-valued permittivity as a function of space. Permittivity distribution over the Yee-grid will be interpolated based on interp_method.

Validated by

_eps_inf_greater_no_less_than_one_sigma_positive
_eps_inf_greater_no_less_than_one_sigma_positive
_single_frequency
_single_frequency

eps_dataset_freq(frequency: float) → tidy3d.components.data.dataset.PermittivityDataset#: Override of inherited validator.

classmethod from_tidy3d(tidy3d_obj: tidy3d.components.medium.CustomMedium) → tidy3d.plugins.adjoint.components.medium.JaxCustomMedium#: Convert Tidy3dBaseModel instance to JaxObject.

store_vjp(grad_data_fwd: tidy3d.components.data.monitor_data.FieldData, grad_data_adj: tidy3d.components.data.monitor_data.FieldData, sim_bounds: Tuple[Tuple[float, float, float], Tuple[float, float, float]], wvl_mat: float) → tidy3d.plugins.adjoint.components.medium.JaxMedium#: Returns the gradient of the medium parameters given forward and adjoint field data.

to_medium() → tidy3d.components.medium.CustomMedium#: Convert JaxMedium instance to Medium