tidy3d.plugins.adjoint.JaxDataArray#

class tidy3d.plugins.adjoint.JaxDataArray#

Bases: tidy3d.components.base.Tidy3dBaseModel

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

attribute values: Any [Required]#

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

Validated by

_convert_values_to_np

classmethod add_type_field() → None#: Automatically place “type” field with model name in the model field dictionary.

assign_coords(coords: Optional[dict] = None, **coords_kwargs) → tidy3d.plugins.adjoint.components.data.data_array.JaxDataArray#: Assign new coordinates to this object.

conj() → tidy3d.plugins.adjoint.components.data.data_array.JaxDataArray#: Complex conjugate of self.

classmethod construct(_fields_set: Optional[SetStr] = None, **values: Any) → Model#: Creates a new model setting __dict__ and __fields_set__ from trusted or pre-validated data. Default values are respected, but no other validation is performed. Behaves as if Config.extra = ‘allow’ was set since it adds all passed values

copy(**kwargs) → tidy3d.components.base.Tidy3dBaseModel#: Copy a Tidy3dBaseModel. With deep=True as default.

dict(*, include: Optional[Union[AbstractSetIntStr, MappingIntStrAny]] = None, exclude: Optional[Union[AbstractSetIntStr, MappingIntStrAny]] = None, by_alias: bool = False, skip_defaults: Optional[bool] = None, exclude_unset: bool = False, exclude_defaults: bool = False, exclude_none: bool = False) → DictStrAny#: Generate a dictionary representation of the model, optionally specifying which fields to include or exclude.

classmethod dict_from_file(fname: str, group_path: Optional[str] = None) → dict#

Loads a dictionary containing the model from a .yaml, .json, or .hdf5 file.

Parameters

fname (str) – Full path to the .yaml or .json file to load the Tidy3dBaseModel from.
group_path (str, optional) – Path to a group inside the file to use as the base level.

Returns

A dictionary containing the model.

Return type

dict

Example

>>> simulation = Simulation.from_file(fname='folder/sim.json') 

classmethod dict_from_hdf5(fname: str, group_path: str = '', custom_decoders: Optional[List[Callable]] = None) → dict#

Loads a dictionary containing the model contents from a .hdf5 file.

Parameters

fname (str) – Full path to the .hdf5 file to load the Tidy3dBaseModel from.
group_path (str, optional) – Path to a group inside the file to selectively load a sub-element of the model only.
custom_decoders (List[Callable]) – List of functions accepting (fname: str, group_path: str, model_dict: dict, key: str, value: Any) that store the value in the model dict after a custom decoding.

Returns

Dictionary containing the model.

Return type

dict

Example

>>> sim_dict = Simulation.dict_from_hdf5(fname='folder/sim.hdf5') 

classmethod dict_from_json(fname: str) → dict#

Load dictionary of the model from a .json file.

Parameters: fname (str) – Full path to the .json file to load the Tidy3dBaseModel from.
Returns: A dictionary containing the model.
Return type: dict

Example

>>> sim_dict = Simulation.dict_from_json(fname='folder/sim.json') 

classmethod dict_from_yaml(fname: str) → dict#

Load dictionary of the model from a .yaml file.

Parameters: fname (str) – Full path to the .yaml file to load the Tidy3dBaseModel from.
Returns: A dictionary containing the model.
Return type: dict

Example

>>> sim_dict = Simulation.dict_from_yaml(fname='folder/sim.yaml') 

classmethod from_file(fname: str, group_path: Optional[str] = None, **parse_obj_kwargs) → tidy3d.components.base.Tidy3dBaseModel#

Loads a Tidy3dBaseModel from .yaml, .json, or .hdf5 file.

Parameters

fname (str) – Full path to the .yaml or .json file to load the Tidy3dBaseModel from.
group_path (str, optional) – Path to a group inside the file to use as the base level. Only for .hdf5 files. Starting / is optional.
**parse_obj_kwargs – Keyword arguments passed to either pydantic’s parse_obj function when loading model.

Returns

An instance of the component class calling load.

Return type

Tidy3dBaseModel

Example

>>> simulation = Simulation.from_file(fname='folder/sim.json') 

classmethod from_hdf5(fname: str, group_path: str) → tidy3d.plugins.adjoint.components.data.data_array.JaxDataArray#: Load an DataArray from an hdf5 file with a given path to the group.

classmethod from_json(fname: str, **parse_obj_kwargs) → tidy3d.components.base.Tidy3dBaseModel#

Load a Tidy3dBaseModel from .json file.

Parameters

fname (str) – Full path to the .json file to load the Tidy3dBaseModel from.

Returns

Tidy3dBaseModel – An instance of the component class calling load.
**parse_obj_kwargs – Keyword arguments passed to pydantic’s parse_obj method.

Example

>>> simulation = Simulation.from_json(fname='folder/sim.json') 

classmethod from_orm(obj: Any) → Model#

classmethod from_yaml(fname: str, **parse_obj_kwargs) → tidy3d.components.base.Tidy3dBaseModel#

Loads Tidy3dBaseModel from .yaml file.

Parameters

fname (str) – Full path to the .yaml file to load the Tidy3dBaseModel from.
**parse_obj_kwargs – Keyword arguments passed to pydantic’s parse_obj method.

Returns

An instance of the component class calling from_yaml.

Return type

Tidy3dBaseModel

Example

>>> simulation = Simulation.from_yaml(fname='folder/sim.yaml') 

classmethod generate_docstring() → str#: Generates a docstring for a Tidy3D mode and saves it to the __doc__ of the class.

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

classmethod get_sub_model(group_path: str, model_dict: dict | list) → dict#: Get the sub model for a given group path.

get_submodels_by_hash() → Dict[int, List[Union[str, Tuple[str, int]]]]#: Return a dictionary of this object’s sub-models indexed by their hash values.

static get_tuple_group_name(index: int) → str#: Get the group name of a tuple element.

static get_tuple_index(key_name: str) → int#: Get the index into the tuple based on its group name.

help(methods: bool = False) → None#

Prints message describing the fields and methods of a Tidy3dBaseModel.

Parameters: methods (bool = False) – Whether to also print out information about object’s methods.

Example

>>> simulation.help(methods=True) 

interp(kwargs=None, assume_sorted=None, **interp_kwargs) → tidy3d.plugins.adjoint.components.data.data_array.JaxDataArray#: Linearly interpolate into the JaxDataArray at values into coordinates.

interp_single(key: str, val: float) → tidy3d.plugins.adjoint.components.data.data_array.JaxDataArray#

Interpolate into a single dimension of self.

Note: this interpolation works by finding the index of the value into the coords list. Instead of an integer value, we use interpolation to get a floating point index. The floor() of this value is the ‘minus’ index and the ceil() gives the ‘plus’ index. We then apply coefficients linearly based on how close to plus or minus we are. This is a workaround to jnp.interp not allowing multi-dimensional interpolation.

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.

json(*, include: Optional[Union[AbstractSetIntStr, MappingIntStrAny]] = None, exclude: Optional[Union[AbstractSetIntStr, MappingIntStrAny]] = None, by_alias: bool = False, skip_defaults: Optional[bool] = None, exclude_unset: bool = False, exclude_defaults: bool = False, exclude_none: bool = False, encoder: Optional[Callable[[Any], Any]] = None, models_as_dict: bool = True, **dumps_kwargs: Any) → unicode#

Generate a JSON representation of the model, include and exclude arguments as per dict().

encoder is an optional function to supply as default to json.dumps(), other arguments as per json.dumps().

multiply_at(value: complex, coord_name: str, indices: List[int]) → tidy3d.plugins.adjoint.components.data.data_array.JaxDataArray#: Multiply self by value at indices into .

classmethod parse_file(path: Union[str, pathlib.Path], *, content_type: unicode = None, encoding: unicode = 'utf8', proto: pydantic.parse.Protocol = None, allow_pickle: bool = False) → Model#

classmethod parse_obj(obj: Any) → Model#

classmethod parse_raw(b: Union[str, bytes], *, content_type: unicode = None, encoding: unicode = 'utf8', proto: pydantic.parse.Protocol = None, allow_pickle: bool = False) → Model#

classmethod schema(by_alias: bool = True, ref_template: unicode = '#/definitions/{model}') → DictStrAny#

classmethod schema_json(*, by_alias: bool = True, ref_template: unicode = '#/definitions/{model}', **dumps_kwargs: Any) → unicode#

sel(indexers: Optional[dict] = None, method: str = 'nearest', **sel_kwargs) → tidy3d.plugins.adjoint.components.data.data_array.JaxDataArray#: Select a value from the JaxDataArray by indexing into coordinate values.

squeeze(dim: Optional[str] = None, drop: bool = True) → tidy3d.plugins.adjoint.components.data.data_array.JaxDataArray#: Remove any non-zero dims.

sum(dim: Optional[str] = None)#: Sum (optionally along a single or multiple dimensions).

to_file(fname: str) → None#

Exports Tidy3dBaseModel instance to .yaml, .json, or .hdf5 file

Parameters: fname (str) – Full path to the .yaml or .json file to save the Tidy3dBaseModel to.

Example

>>> simulation.to_file(fname='folder/sim.json') 

to_hdf5(fname: str, group_path: str) → None#: Save an xr.DataArray to the hdf5 file with a given path to the group.

to_json(fname: str) → None#

Exports Tidy3dBaseModel instance to .json file

Parameters: fname (str) – Full path to the .json file to save the Tidy3dBaseModel to.

Example

>>> simulation.to_json(fname='folder/sim.json') 

to_yaml(fname: str) → None#

Exports Tidy3dBaseModel instance to .yaml file.

Parameters: fname (str) – Full path to the .yaml file to save the Tidy3dBaseModel to.

Example

>>> simulation.to_yaml(fname='folder/sim.yaml') 

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.

classmethod tuple_to_dict(tuple_values: tuple) → dict#: How we generate a dictionary mapping new keys to tuple values for hdf5.

classmethod update_forward_refs(**localns: Any) → None#: Try to update ForwardRefs on fields based on this Model, globalns and localns.

updated_copy(**kwargs) → tidy3d.components.base.Tidy3dBaseModel#: Make copy of a component instance with **kwargs indicating updated field values.

classmethod validate(value: Any) → Model#

property as_jnp_array: jax.Array#: self.values as a jax array.

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.

property shape: tuple#: Shape of self.values.