tidy3d.plugins.adjoint.JaxDataArray#

class tidy3d.plugins.adjoint.JaxDataArray#

A DataArray-like class that only wraps xarray for jax compability.

Parameters

values (Optional[Any]) – Nested list containing the raw values, which can be tracked by jax.
coords (Mapping[str, list]) – Dictionary storing the coordinates, namely (direction, f, mode_index).

Show JSON schema

{
   "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
}

attribute coords: Dict[str, list] [Required]#

Dictionary storing the coordinates, namely (direction, f, mode_index).

Validated by

_convert_coords_to_list
_coords_match_values

attribute values: Any [Required]#

Nested list containing the raw values, which can be tracked by jax.

Validated by

_convert_values_to_jax_array

get_coord_list(coord_name: str) → list#: Get a coordinate list by name.

interp(**interp_kwargs)#: Interpolate into the JaxDataArray. Not yet supported.

isel(**isel_kwargs) → tidy3d.plugins.adjoint.components.data.data_array.JaxDataArray#: Select a value from the JaxDataArray by indexing into coordinates by index.

isel_single(coord_name: str, coord_index: int) → tidy3d.plugins.adjoint.components.data.data_array.JaxDataArray#: Select a value cooresponding to a single coordinate from the JaxDataArray.

sel(**sel_kwargs) → tidy3d.plugins.adjoint.components.data.data_array.JaxDataArray#: Select a value from the JaxDataArray by indexing into coordinate values.

tree_flatten() → Tuple[list, dict]#: Jax works on the values, stash the coords for reconstruction.

classmethod tree_unflatten(aux_data, children) → tidy3d.plugins.adjoint.components.data.data_array.JaxDataArray#: How to unflatten the values and coords.

property as_list: list#: self.values as a numpy array converted to a list.

property as_ndarray: numpy.ndarray#: self.values as a numpy array.

property imag: numpy.ndarray#: Imaginary part of self.

property nonzero_val_coords: Tuple[List[complex], Dict[str, Any]]#: The value and coordinate associated with the only non-zero element of self.values.

property real: numpy.ndarray#: Real part of self.