tidy3d.plugins.adjoint.JaxSimulation#

class JaxSimulation[source]#

Bases: Simulation, JaxObject

A Simulation registered with jax.

Parameters:

attrs (dict = {}) – Dictionary storing arbitrary metadata for a Tidy3D object. This dictionary can be freely used by the user for storing data without affecting the operation of Tidy3D as it is not used internally. Note that, unlike regular Tidy3D fields, attrs are mutable. For example, the following is allowed for setting an attr obj.attrs['foo'] = bar. Also note that Tidy3D` will raise a TypeError if attrs contain objects that can not be serialized. One can check if attrs are serializable by calling obj.json().
center (Union[tuple[Union[float, autograd.tracer.Box], Union[float, autograd.tracer.Box], Union[float, autograd.tracer.Box]], Box] = (0.0, 0.0, 0.0)) – [units = um]. Center of object in x, y, and z.
size (Union[tuple[Union[pydantic.v1.types.NonNegativeFloat, autograd.tracer.Box], Union[pydantic.v1.types.NonNegativeFloat, autograd.tracer.Box], Union[pydantic.v1.types.NonNegativeFloat, autograd.tracer.Box]], Box]) – [units = um]. Size in x, y, and z directions.
medium (Union[Medium, AnisotropicMedium, PECMedium, PoleResidue, Sellmeier, Lorentz, Debye, Drude, FullyAnisotropicMedium, CustomMedium, CustomPoleResidue, CustomSellmeier, CustomLorentz, CustomDebye, CustomDrude, CustomAnisotropicMedium, PerturbationMedium, PerturbationPoleResidue] = Medium(attrs={}, name=None, frequency_range=None, allow_gain=False, nonlinear_spec=None, modulation_spec=None, heat_spec=None, type='Medium', permittivity=1.0, conductivity=0.0)) – Background medium of simulation, defaults to vacuum if not specified.
structures (Tuple[Structure, ...] = ()) – Tuple of structures present in simulation. Note: Structures defined later in this list override the simulation material properties in regions of spatial overlap.
symmetry (Tuple[Literal[0, -1, 1], Literal[0, -1, 1], Literal[0, -1, 1]] = (0, 0, 0)) – Tuple of integers defining reflection symmetry across a plane bisecting the simulation domain normal to the x-, y-, and z-axis at the simulation center of each axis, respectively. Each element can be 0 (no symmetry), 1 (even, i.e. ‘PMC’ symmetry) or -1 (odd, i.e. ‘PEC’ symmetry). Note that the vectorial nature of the fields must be taken into account to correctly determine the symmetry value.
sources (Tuple[Annotated[Union[tidy3d.components.source.UniformCurrentSource, tidy3d.components.source.PointDipole, tidy3d.components.source.GaussianBeam, tidy3d.components.source.AstigmaticGaussianBeam, tidy3d.components.source.ModeSource, tidy3d.components.source.PlaneWave, tidy3d.components.source.CustomFieldSource, tidy3d.components.source.CustomCurrentSource, tidy3d.components.source.TFSF], FieldInfo(default=PydanticUndefined, discriminator='type', extra={})], ...] = ()) – Tuple of electric current sources injecting fields into the simulation.
boundary_spec (BoundarySpec = BoundarySpec(attrs={}, x=Boundary(attrs={},, plus=PML(attrs={},, name=None,, type='PML',, num_layers=12,, parameters=PMLParams(attrs={},, sigma_order=3,, sigma_min=0.0,, sigma_max=1.5,, type='PMLParams',, kappa_order=3,, kappa_min=1.0,, kappa_max=3.0,, alpha_order=1,, alpha_min=0.0,, alpha_max=0.0)),, minus=PML(attrs={},, name=None,, type='PML',, num_layers=12,, parameters=PMLParams(attrs={},, sigma_order=3,, sigma_min=0.0,, sigma_max=1.5,, type='PMLParams',, kappa_order=3,, kappa_min=1.0,, kappa_max=3.0,, alpha_order=1,, alpha_min=0.0,, alpha_max=0.0)),, type='Boundary'), y=Boundary(attrs={},, plus=PML(attrs={},, name=None,, type='PML',, num_layers=12,, parameters=PMLParams(attrs={},, sigma_order=3,, sigma_min=0.0,, sigma_max=1.5,, type='PMLParams',, kappa_order=3,, kappa_min=1.0,, kappa_max=3.0,, alpha_order=1,, alpha_min=0.0,, alpha_max=0.0)),, minus=PML(attrs={},, name=None,, type='PML',, num_layers=12,, parameters=PMLParams(attrs={},, sigma_order=3,, sigma_min=0.0,, sigma_max=1.5,, type='PMLParams',, kappa_order=3,, kappa_min=1.0,, kappa_max=3.0,, alpha_order=1,, alpha_min=0.0,, alpha_max=0.0)),, type='Boundary'), z=Boundary(attrs={},, plus=PML(attrs={},, name=None,, type='PML',, num_layers=12,, parameters=PMLParams(attrs={},, sigma_order=3,, sigma_min=0.0,, sigma_max=1.5,, type='PMLParams',, kappa_order=3,, kappa_min=1.0,, kappa_max=3.0,, alpha_order=1,, alpha_min=0.0,, alpha_max=0.0)),, minus=PML(attrs={},, name=None,, type='PML',, num_layers=12,, parameters=PMLParams(attrs={},, sigma_order=3,, sigma_min=0.0,, sigma_max=1.5,, type='PMLParams',, kappa_order=3,, kappa_min=1.0,, kappa_max=3.0,, alpha_order=1,, alpha_min=0.0,, alpha_max=0.0)),, type='Boundary'), type='BoundarySpec')) – Specification of boundary conditions along each dimension. If None, PML boundary conditions are applied on all sides.
monitors (Tuple[Annotated[Union[tidy3d.components.monitor.FieldMonitor, tidy3d.components.monitor.FieldTimeMonitor, tidy3d.components.monitor.PermittivityMonitor, tidy3d.components.monitor.FluxMonitor, tidy3d.components.monitor.FluxTimeMonitor, tidy3d.components.monitor.ModeMonitor, tidy3d.components.monitor.ModeSolverMonitor, tidy3d.components.monitor.FieldProjectionAngleMonitor, tidy3d.components.monitor.FieldProjectionCartesianMonitor, tidy3d.components.monitor.FieldProjectionKSpaceMonitor, tidy3d.components.monitor.DiffractionMonitor], FieldInfo(default=PydanticUndefined, discriminator='type', extra={})], ...] = ()) – Tuple of monitors in the simulation. Note: monitor names are used to access data after simulation is run.
grid_spec (GridSpec = GridSpec(attrs={}, grid_x=AutoGrid(attrs={},, type='AutoGrid',, min_steps_per_wvl=10.0,, max_scale=1.4,, dl_min=0.0,, mesher=GradedMesher(attrs={},, type='GradedMesher')), grid_y=AutoGrid(attrs={},, type='AutoGrid',, min_steps_per_wvl=10.0,, max_scale=1.4,, dl_min=0.0,, mesher=GradedMesher(attrs={},, type='GradedMesher')), grid_z=AutoGrid(attrs={},, type='AutoGrid',, min_steps_per_wvl=10.0,, max_scale=1.4,, dl_min=0.0,, mesher=GradedMesher(attrs={},, type='GradedMesher')), wavelength=None, override_structures=(), snapping_points=(), type='GridSpec')) – Specifications for the simulation grid along each of the three directions.
version (str = 2.7.7) – String specifying the front end version number.
lumped_elements (Tuple[Union[LumpedResistor, CoaxialLumpedResistor], ...] = ()) – Tuple of lumped elements in the simulation. Note: only tidy3d.LumpedResistor is supported currently.
subpixel (Union[bool, SubpixelSpec] = SubpixelSpec(attrs={}, dielectric=PolarizedAveraging(attrs={},, type='PolarizedAveraging'), metal=Staircasing(attrs={},, type='Staircasing'), pec=PECConformal(attrs={},, type='PECConformal',, timestep_reduction=0.3), type='SubpixelSpec')) – Apply subpixel averaging methods of the permittivity on structure interfaces to result in much higher accuracy for a given grid size. Supply a SubpixelSpec to this field to select subpixel averaging methods separately on dielectric, metal, and PEC material interfaces. Alternatively, user may supply a boolean value: True to apply the default subpixel averaging methods corresponding to SubpixelSpec() , or False to apply staircasing.
simulation_type (Optional[Literal['autograd_fwd', 'autograd_bwd', 'tidy3d', None]] = tidy3d) – Tag used internally to distinguish types of simulations for autograd gradient processing.
post_norm (Union[float, FreqDataArray] = 1.0) – Factor to multiply the fields by after running, given the adjoint source pipeline used. Note: this is used internally only.
courant (ConstrainedFloatValue = 0.99) – Normalized Courant stability factor that is no larger than 1 when CFL stability condition is met. It controls time step to spatial step ratio. Lower values lead to more stable simulations for dispersive materials, but result in longer simulation times.
normalize_index (Optional[NonNegativeInt] = 0) – Index of the source in the tuple of sources whose spectrum will be used to normalize the frequency-dependent data. If None, the raw field data is returned unnormalized.
shutoff (NonNegativeFloat = 1e-05) – Ratio of the instantaneous integrated E-field intensity to the maximum value at which the simulation will automatically terminate time stepping. Used to prevent extraneous run time of simulations with fully decayed fields. Set to 0 to disable this feature.
run_time (Union[PositiveFloat, RunTimeSpec]) – [units = sec]. Total electromagnetic evolution time in seconds. Note: If simulation ‘shutoff’ is specified, simulation will terminate early when shutoff condition met. Alternatively, user may supply a RunTimeSpec to this field, which will auto-compute the run_time based on the contents of the spec. If this option is used, the evaluated run_time value is available in the Simulation._run_time property.
input_structures (Tuple[Annotated[Union[tidy3d.plugins.adjoint.components.structure.JaxStructure, tidy3d.plugins.adjoint.components.structure.JaxStructureStaticMedium, tidy3d.plugins.adjoint.components.structure.JaxStructureStaticGeometry], FieldInfo(default=PydanticUndefined, discriminator='type', extra={})], ...] = ()) – Tuple of jax-compatible structures that may depend on differentiable parameters.
output_monitors (Tuple[Annotated[Union[tidy3d.components.monitor.DiffractionMonitor, tidy3d.components.monitor.FieldMonitor, tidy3d.components.monitor.ModeMonitor], FieldInfo(default=PydanticUndefined, discriminator='type', extra={})], ...] = ()) – Tuple of monitors whose data the differentiable output depends on.
grad_monitors (Tuple[FieldMonitor, ...] = ()) – Tuple of monitors used for storing fields, used internally for gradients.
grad_eps_monitors (Tuple[PermittivityMonitor, ...] = ()) – Tuple of monitors used for storing epsilon, used internally for gradients.
fwidth_adjoint (Optional[PositiveFloat] = None) – [units = Hz]. Custom frequency width to use for source_time of adjoint sources. If not supplied or None, uses the average fwidth of the original simulation’s sources.
run_time_adjoint (Optional[PositiveFloat] = None) – [units = sec]. Custom run_time to use for adjoint simulation. If not supplied or None, uses a factor times the adjoint source fwidth.

Attributes

`freqs_adjoint`	Return sorted list of frequencies stripped from the output monitors.
`num_time_steps_adjoint`	Number of time steps in the adjoint simulation.
`tmesh_adjoint`	FDTD time stepping points.
`attrs`

Methods

`epsilon`(box[, coord_key, freq])	Get array of permittivity at volume specified by box and freq.
`from_simulation`(simulation, jax_info)	Convert `Simulation` to `JaxSimulation` with extra info.
`get_freqs_adjoint`(output_monitors)	Return sorted list of unique frequencies stripped from a collection of monitors.
`get_grad_monitors`(input_structures, ...[, ...])	Return dictionary of gradient monitors for simulation.
`make_sim_fwd`(simulation, jax_info)	Make the forward `JaxSimulation` from the supplied `Simulation`.
`plot`([x, y, z, ax, source_alpha, ...])	Wrapper around regular `Simulation` structure plotting.
`plot_eps`([x, y, z, freq, alpha, ...])	Wrapper around regular `Simulation` permittivity plotting.
`plot_structures`([x, y, z, ax, hlim, vlim])	Plot each of simulation's structures on a plane defined by one nonzero x,y,z coordinate.
`plot_structures_eps`([x, y, z, freq, alpha, ...])	Plot each of simulation's structures on a plane defined by one nonzero x,y,z coordinate.
`split_monitors`(monitors, jax_info)	Split monitors into user and adjoint required based on jax info.
`split_structures`(structures, jax_info)	Split structures into regular and input based on jax info.
`store_vjp`(grad_data_fwd, grad_data_adj, ...)	Store the vjp w.r.t.
`store_vjp_parallel`(grad_data_fwd, ...)	Store the vjp w.r.t.
`store_vjp_sequential`(grad_data_fwd, ...)	Store the vjp w.r.t.
`to_gds`(cell[, x, y, z, ...])	Append the simulation structures to a .gds cell. :param cell: Cell object to which the generated polygons are added. :type cell: `gdstk.Cell` or `gdspy.Cell` :param x: Position of plane in x direction, only one of x,y,z can be specified to define plane. :type x: float = None :param y: Position of plane in y direction, only one of x,y,z can be specified to define plane. :type y: float = None :param z: Position of plane in z direction, only one of x,y,z can be specified to define plane. :type z: float = None :param permittivity_threshold: Permittivity value used to define the shape boundaries for structures with custom medim :type permittivity_threshold: float = 1 :param frequency: Frequency for permittivity evaluation in case of custom medium (Hz). :type frequency: float = 0 :param gds_layer_dtype_map: Dictionary mapping mediums to GDSII layer and data type tuples. :type gds_layer_dtype_map: Dict.
`to_gdspy`([x, y, z, gds_layer_dtype_map])	Convert a simulation's planar slice to a .gds type polygon list.
`to_gdstk`([x, y, z, permittivity_threshold, ...])	Convert a simulation's planar slice to a .gds type polygon list. :param x: Position of plane in x direction, only one of x,y,z can be specified to define plane. :type x: float = None :param y: Position of plane in y direction, only one of x,y,z can be specified to define plane. :type y: float = None :param z: Position of plane in z direction, only one of x,y,z can be specified to define plane. :type z: float = None :param permittivity_threshold: Permittivity value used to define the shape boundaries for structures with custom medim :type permittivity_threshold: float = 1 :param frequency: Frequency for permittivity evaluation in case of custom medium (Hz). :type frequency: float = 0 :param gds_layer_dtype_map: Dictionary mapping mediums to GDSII layer and data type tuples. :type gds_layer_dtype_map: Dict.
`to_simulation`()	Convert `JaxSimulation` instance to `Simulation` with an info dict.
`to_simulation_fwd`()	Like `to_simulation()` but the gradient monitors are included.

Inherited Common Usage

input_structures#

output_monitors#

grad_monitors#

grad_eps_monitors#

fwidth_adjoint#

run_time_adjoint#

static get_freqs_adjoint(output_monitors)[source]#: Return sorted list of unique frequencies stripped from a collection of monitors.

property freqs_adjoint#: Return sorted list of frequencies stripped from the output monitors.

property tmesh_adjoint#

FDTD time stepping points.

Returns:: Times (seconds) that the simulation time steps through.
Return type:: np.ndarray

property num_time_steps_adjoint#: Number of time steps in the adjoint simulation.

to_simulation()[source]#: Convert JaxSimulation instance to Simulation with an info dict.

to_gds(cell, x=None, y=None, z=None, permittivity_threshold=1, frequency=0, gds_layer_dtype_map=None)[source]#

Append the simulation structures to a .gds cell. :param cell: Cell object to which the generated polygons are added. :type cell: gdstk.Cell or gdspy.Cell :param x: Position of plane in x direction, only one of x,y,z can be specified to define plane. :type x: float = None :param y: Position of plane in y direction, only one of x,y,z can be specified to define plane. :type y: float = None :param z: Position of plane in z direction, only one of x,y,z can be specified to define plane. :type z: float = None :param permittivity_threshold: Permittivity value used to define the shape boundaries for structures with custom

medim

Parameters:

frequency (float = 0) – Frequency for permittivity evaluation in case of custom medium (Hz).
gds_layer_dtype_map (Dict) – Dictionary mapping mediums to GDSII layer and data type tuples.

to_gdstk(x=None, y=None, z=None, permittivity_threshold=1, frequency=0, gds_layer_dtype_map=None)[source]#

Convert a simulation’s planar slice to a .gds type polygon list. :param x: Position of plane in x direction, only one of x,y,z can be specified to define plane. :type x: float = None :param y: Position of plane in y direction, only one of x,y,z can be specified to define plane. :type y: float = None :param z: Position of plane in z direction, only one of x,y,z can be specified to define plane. :type z: float = None :param permittivity_threshold: Permittivity value used to define the shape boundaries for structures with custom

medim

Parameters:

frequency (float = 0) – Frequency for permittivity evaluation in case of custom medium (Hz).
gds_layer_dtype_map (Dict) – Dictionary mapping mediums to GDSII layer and data type tuples.

Returns:

List of gdstk.Polygon.

Return type:

List

to_gdspy(x=None, y=None, z=None, gds_layer_dtype_map=None)[source]#

Convert a simulation’s planar slice to a .gds type polygon list. :param x: Position of plane in x direction, only one of x,y,z can be specified to define plane. :type x: float = None :param y: Position of plane in y direction, only one of x,y,z can be specified to define plane. :type y: float = None :param z: Position of plane in z direction, only one of x,y,z can be specified to define plane. :type z: float = None :param gds_layer_dtype_map: Dictionary mapping mediums to GDSII layer and data type tuples. :type gds_layer_dtype_map: Dict

Returns:: List of gdspy.Polygon and gdspy.PolygonSet.
Return type:: List

plot(x=None, y=None, z=None, ax=None, source_alpha=None, monitor_alpha=None, hlim=None, vlim=None, **patch_kwargs)[source]#: Wrapper around regular Simulation structure plotting.

plot_eps(x=None, y=None, z=None, freq=None, alpha=None, source_alpha=None, monitor_alpha=None, hlim=None, vlim=None, ax=None)[source]#: Wrapper around regular Simulation permittivity plotting.

plot_structures(x=None, y=None, z=None, ax=None, hlim=None, vlim=None)[source]#

Plot each of simulation’s structures on a plane defined by one nonzero x,y,z coordinate.

Parameters:

x (float = None) – position of plane in x direction, only one of x, y, z must be specified to define plane.
y (float = None) – position of plane in y direction, only one of x, y, z must be specified to define plane.
z (float = None) – position of plane in z direction, only one of x, y, z must be specified to define plane.
ax (matplotlib.axes._subplots.Axes = None) – Matplotlib axes to plot on, if not specified, one is created.
hlim (Tuple[float, float] = None) – The x range if plotting on xy or xz planes, y range if plotting on yz plane.
vlim (Tuple[float, float] = None) – The z range if plotting on xz or yz planes, y plane if plotting on xy plane.

Returns:

The supplied or created matplotlib axes.

Return type:

matplotlib.axes._subplots.Axes

plot_structures_eps(x=None, y=None, z=None, freq=None, alpha=None, cbar=True, reverse=False, ax=None, hlim=None, vlim=None)[source]#

Plot each of simulation’s structures on a plane defined by one nonzero x,y,z coordinate. The permittivity is plotted in grayscale based on its value at the specified frequency.

Parameters:

x (float = None) – position of plane in x direction, only one of x, y, z must be specified to define plane.
y (float = None) – position of plane in y direction, only one of x, y, z must be specified to define plane.
z (float = None) – position of plane in z direction, only one of x, y, z must be specified to define plane.
freq (float = None) – Frequency to evaluate the relative permittivity of all mediums. If not specified, evaluates at infinite frequency.
reverse (bool = False) – If False, the highest permittivity is plotted in black. If True, it is plotteed in white (suitable for black backgrounds).
cbar (bool = True) – Whether to plot a colorbar for the relative permittivity.
alpha (float = None) – Opacity of the structures being plotted. Defaults to the structure default alpha.
ax (matplotlib.axes._subplots.Axes = None) – Matplotlib axes to plot on, if not specified, one is created.
hlim (Tuple[float, float] = None) – The x range if plotting on xy or xz planes, y range if plotting on yz plane.
vlim (Tuple[float, float] = None) – The z range if plotting on xz or yz planes, y plane if plotting on xy plane.

Returns:

The supplied or created matplotlib axes.

Return type:

matplotlib.axes._subplots.Axes

epsilon(box, coord_key='centers', freq=None)[source]#

Get array of permittivity at volume specified by box and freq.

Parameters:

box (Box) – Rectangular geometry specifying where to measure the permittivity.
coord_key (str = 'centers') – Specifies at what part of the grid to return the permittivity at. Accepted values are {'centers', 'boundaries', 'Ex', 'Ey', 'Ez', 'Exy', 'Exz', 'Eyx', 'Eyz', 'Ezx', Ezy'}. The field values (eg. 'Ex') correspond to the corresponding field locations on the yee lattice. If field values are selected, the corresponding diagonal (eg. eps_xx in case of 'Ex') or off-diagonal (eg. eps_xy in case of 'Exy') epsilon component from the epsilon tensor is returned. Otherwise, the average of the main values is returned.
freq (float = None) – The frequency to evaluate the mediums at. If not specified, evaluates at infinite frequency.

Returns:

Datastructure containing the relative permittivity values and location coordinates. For details on xarray DataArray objects, refer to xarray’s Documentation.

Return type:

xarray.DataArray

tidy3d.plugins.adjoint.JaxSimulation

Contents

tidy3d.plugins.adjoint.JaxSimulation#