tidy3d.plugins.adjoint.JaxDataArray#

class tidy3d.plugins.adjoint.JaxDataArray(*, values: Any, coords: Dict[str, list], type: Literal['JaxDataArray'] = '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).

__init__(**kwargs)#: Init method, includes post-init validators.

Methods

`__init__`(**kwargs)	Init method, includes post-init validators.
`add_type_field`()	Automatically place "type" field with model name in the model field dictionary.
`assign_coords`([coords])	Assign new coordinates to this object.
`conj`()	Complex conjugate of self.
`construct`([_fields_set])	Creates a new model setting __dict__ and __fields_set__ from trusted or pre-validated data.
`copy`(**kwargs)	Copy a Tidy3dBaseModel.
`dict`(*[, include, exclude, by_alias, ...])	Generate a dictionary representation of the model, optionally specifying which fields to include or exclude.
`dict_from_file`(fname[, group_path])	Loads a dictionary containing the model from a .yaml, .json, .hdf5, or .hdf5.gz file.
`dict_from_hdf5`(fname[, group_path, ...])	Loads a dictionary containing the model contents from a .hdf5 file.
`dict_from_hdf5_gz`(fname[, group_path, ...])	Loads a dictionary containing the model contents from a .hdf5.gz file.
`dict_from_json`(fname)	Load dictionary of the model from a .json file.
`dict_from_yaml`(fname)	Load dictionary of the model from a .yaml file.
`from_file`(fname[, group_path])	Loads a `Tidy3dBaseModel` from .yaml, .json, .hdf5, or .hdf5.gz file.
`from_hdf5`(fname, group_path)	Load an DataArray from an hdf5 file with a given path to the group.
`from_hdf5_gz`(fname[, group_path, ...])	Loads `Tidy3dBaseModel` instance to .hdf5.gz file.
`from_json`(fname, **parse_obj_kwargs)	Load a `Tidy3dBaseModel` from .json file.
`from_orm`(obj)
`from_yaml`(fname, **parse_obj_kwargs)	Loads `Tidy3dBaseModel` from .yaml file.
`generate_docstring`()	Generates a docstring for a Tidy3D mode and saves it to the __doc__ of the class.
`get_coord_list`(coord_name)	Get a coordinate list by name.
`get_sub_model`(group_path, model_dict)	Get the sub model for a given group path.
`get_submodels_by_hash`()	Return a dictionary of this object's sub-models indexed by their hash values.
`get_tuple_group_name`(index)	Get the group name of a tuple element.
`get_tuple_index`(key_name)	Get the index into the tuple based on its group name.
`help`([methods])	Prints message describing the fields and methods of a `Tidy3dBaseModel`.
`interp`([kwargs, assume_sorted])	Linearly interpolate into the `JaxDataArray` at values into coordinates.
`interp_single`(key, val)	Interpolate into a single dimension of self.
`isel`(**isel_kwargs)	Select a value from the `JaxDataArray` by indexing into coordinates by index.
`isel_single`(coord_name, coord_index)	Select a value cooresponding to a single coordinate from the `JaxDataArray`.
`json`(*[, include, exclude, by_alias, ...])	Generate a JSON representation of the model, include and exclude arguments as per dict().
`multiply_at`(value, coord_name, indices)	Multiply self by value at indices into .
`parse_file`(path, *[, content_type, ...])
`parse_obj`(obj)
`parse_raw`(b, *[, content_type, encoding, ...])
`schema`([by_alias, ref_template])
`schema_json`(*[, by_alias, ref_template])
`sel`([indexers, method])	Select a value from the `JaxDataArray` by indexing into coordinate values.
`squeeze`([dim, drop])	Remove any non-zero dims.
`sum`([dim])	Sum (optionally along a single or multiple dimensions).
`to_file`(fname)	Exports `Tidy3dBaseModel` instance to .yaml, .json, or .hdf5 file
`to_hdf5`(fname, group_path)	Save an xr.DataArray to the hdf5 file with a given path to the group.
`to_hdf5_gz`(fname[, custom_encoders])	Exports `Tidy3dBaseModel` instance to .hdf5.gz file.
`to_json`(fname)	Exports `Tidy3dBaseModel` instance to .json file
`to_yaml`(fname)	Exports `Tidy3dBaseModel` instance to .yaml file.
`tree_flatten`()	Jax works on the values, stash the coords for reconstruction.
`tree_unflatten`(aux_data, children)	How to unflatten the values and coords.
`tuple_to_dict`(tuple_values)	How we generate a dictionary mapping new keys to tuple values for hdf5.
`update_forward_refs`(**localns)	Try to update ForwardRefs on fields based on this Model, globalns and localns.
`updated_copy`(**kwargs)	Make copy of a component instance with `**kwargs` indicating updated field values.
`validate`(value)

Attributes

`as_jnp_array`	`self.values` as a jax array.
`as_list`	`self.values` as a numpy array converted to a list.
`as_ndarray`	`self.values` as a numpy array.
`imag`	Imaginary part of self.
`nonzero_val_coords`	The value and coordinate associated with the only non-zero element of `self.values`.
`real`	Real part of self.
`shape`	Shape of self.values.
`values`
`coords`

class Config#

Bases: object

Sets config for all Tidy3dBaseModel objects.

allow_population_by_field_namebool = True: Allow properties to stand in for fields(?).
arbitrary_types_allowedbool = True: Allow types like numpy arrays.
extrastr = ‘forbid’: Forbid extra kwargs not specified in model.
json_encodersDict[type, Callable]: Defines how to encode type in json file.
validate_allbool = True: Validate default values just to be safe.
validate_assignmentbool: Re-validate after re-assignment of field in model.

__abs__() → tidy3d.plugins.adjoint.components.data.data_array.JaxDataArray#: Absolute value of self’s values.

__add__(other: tidy3d.plugins.adjoint.components.data.data_array.JaxDataArray) → tidy3d.plugins.adjoint.components.data.data_array.JaxDataArray#: Sum self with something else.

__eq__(other) → bool#: Check if two JaxDataArray instances are equal.

__ge__(other)#: define >= for getting unique indices based on hash.

__gt__(other)#: define > for getting unique indices based on hash.

__hash__() → int#: Hash method.

classmethod __init_subclass__() → None#: Things that are done to each of the models.

__iter__() → TupleGenerator#: so dict(model) works

__le__(other)#: define <= for getting unique indices based on hash.

__lt__(other)#: define < for getting unique indices based on hash.

__mul__(other: tidy3d.plugins.adjoint.components.data.data_array.JaxDataArray) → tidy3d.plugins.adjoint.components.data.data_array.JaxDataArray#: Multiply self with something else.

__neg__() → tidy3d.plugins.adjoint.components.data.data_array.JaxDataArray#: Negative of self.

__pow__(power: int) → tidy3d.plugins.adjoint.components.data.data_array.JaxDataArray#: Values raised to a power.

__pretty__(fmt: Callable[[Any], Any], **kwargs: Any) → Generator[Any, None, None]#: Used by devtools (https://python-devtools.helpmanual.io/) to provide a human readable representations of objects

__radd__(other) → tidy3d.plugins.adjoint.components.data.data_array.JaxDataArray#: Sum self with something else.

__repr_name__() → str#: Name of the instance’s class, used in __repr__.

__rich_repr__() → RichReprResult#: Get fields for Rich library

__rmul__(other) → tidy3d.plugins.adjoint.components.data.data_array.JaxDataArray#: Multiply self with something else.

__sub__(other) → tidy3d.plugins.adjoint.components.data.data_array.JaxDataArray#: Subtraction

classmethod __try_update_forward_refs__(**localns: Any) → None#: Same as update_forward_refs but will not raise exception when forward references are not defined.

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

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.

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, .hdf5, or .hdf5.gz file.

Parameters

fname (str) – Full path to the 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_hdf5_gz(fname: str, group_path: str = '', custom_decoders: Optional[List[Callable]] = None) → dict#

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

Parameters

fname (str) – Full path to the .hdf5.gz 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.gz') 

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, .hdf5, or .hdf5.gz file.

Parameters

fname (str) – Full path to the 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_hdf5_gz(fname: str, group_path: str = '', custom_decoders: Optional[List[Callable]] = None, **parse_obj_kwargs) → tidy3d.components.base.Tidy3dBaseModel#

Loads Tidy3dBaseModel instance to .hdf5.gz file.

Parameters

fname (str) – Full path to the .hdf5.gz 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.
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.
**parse_obj_kwargs – Keyword arguments passed to pydantic’s parse_obj method.

Example

>>> simulation = Simulation.from_hdf5_gz(fname='folder/sim.hdf5.gz') 

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_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) 

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

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

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 .

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.

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.

property shape: tuple#: Shape of self.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_hdf5_gz(fname: str, custom_encoders: Optional[List[Callable]] = None) → None#

Exports Tidy3dBaseModel instance to .hdf5.gz file.

Parameters

fname (str) – Full path to the .hdf5.gz file to save the Tidy3dBaseModel to.
custom_encoders (List[Callable]) – List of functions accepting (fname: str, group_path: str, value: Any) that take the value supplied and write it to the hdf5 fname at group_path.

Example

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

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.