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.

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

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 = '') → 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.

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 = '', **parse_obj_kwargs) → tidy3d.components.base.Tidy3dBaseModel#

Loads Tidy3dBaseModel instance to .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. Starting / is optional.
**parse_obj_kwargs – Keyword arguments passed to pydantic’s parse_obj method.

Example

>>> simulation.to_hdf5(fname='folder/sim.hdf5') 

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.

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(**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.

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().

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(**sel_kwargs) → tidy3d.plugins.adjoint.components.data.data_array.JaxDataArray#: Select a value from the JaxDataArray by indexing into coordinate values.

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) → None#

Exports Tidy3dBaseModel instance to .hdf5 file.

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

Example

>>> simulation.to_hdf5(fname='folder/sim.hdf5') 

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.