tidy3d.plugins.adjoint.JaxSimulation

class JaxSimulation

Bases: Simulation, JaxObject

A Simulation registered with jax.

Parameters:

• center (Tuple[float, float, float] = (0.0, 0.0, 0.0)) – [units = um]. Center of object in x, y, and z.

• size (Tuple[NonNegativeFloat, NonNegativeFloat, NonNegativeFloat]) – [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(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(x=Boundary(plus=PML(name=None,, type='PML',, num_layers=12,, parameters=PMLParams(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(name=None,, type='PML',, num_layers=12,, parameters=PMLParams(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(plus=PML(name=None,, type='PML',, num_layers=12,, parameters=PMLParams(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(name=None,, type='PML',, num_layers=12,, parameters=PMLParams(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(plus=PML(name=None,, type='PML',, num_layers=12,, parameters=PMLParams(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(name=None,, type='PML',, num_layers=12,, parameters=PMLParams(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(grid_x=AutoGrid(type='AutoGrid',, min_steps_per_wvl=10.0,, max_scale=1.4,, dl_min=0.0,, mesher=GradedMesher(type='GradedMesher')), grid_y=AutoGrid(type='AutoGrid',, min_steps_per_wvl=10.0,, max_scale=1.4,, dl_min=0.0,, mesher=GradedMesher(type='GradedMesher')), grid_z=AutoGrid(type='AutoGrid',, min_steps_per_wvl=10.0,, max_scale=1.4,, dl_min=0.0,, mesher=GradedMesher(type='GradedMesher')), wavelength=None, override_structures=(), type='GridSpec')) – Specifications for the simulation grid along each of the three directions.

• version (str = 2.6.0) – String specifying the front end version number.

• courant (ConstrainedFloatValue = 0.99) – Courant stability factor, controls time step to spatial step ratio. Lower values lead to more stable simulations for dispersive materials, but result in longer simulation times. This factor is normalized to no larger than 1 when CFL stability condition is met in 3D.

• 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.

• subpixel (bool = True) – If True, uses subpixel averaging of the permittivity based on structure definition, resulting in much higher accuracy for a given grid size.

• run_time (PositiveFloat) – [units = sec]. Total electromagnetic evolution time in seconds. Note: If simulation ‘shutoff’ is specified, simulation will terminate early when shutoff condition met.

• 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.
version	Defines the front end version of tidy3d

Methods

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_simulation()	Convert JaxSimulation instance to Simulation with an info dict.
to_simulation_fwd()	Like to_simulation() but the gradient monitors are included.

input_structures

output_monitors

grad_monitors

grad_eps_monitors

fwidth_adjoint

run_time_adjoint

static get_freqs_adjoint(output_monitors)

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.

to_simulation()

Convert JaxSimulation instance to Simulation with an info dict.

plot(x=None, y=None, z=None, ax=None, source_alpha=None, monitor_alpha=None, hlim=None, vlim=None, **patch_kwargs)

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)

Wrapper around regular Simulation permittivity plotting.

plot_structures(x=None, y=None, z=None, ax=None, hlim=None, vlim=None)

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)

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

__eq__(other)

Are two JaxSimulation objects equal?

classmethod split_monitors(monitors, jax_info)

Split monitors into user and adjoint required based on jax info.

classmethod split_structures(structures, jax_info)

Split structures into regular and input based on jax info.

classmethod from_simulation(simulation, jax_info)

Convert Simulation to JaxSimulation with extra info.

classmethod make_sim_fwd(simulation, jax_info)

Make the forward JaxSimulation from the supplied Simulation.

to_simulation_fwd()

Like to_simulation() but the gradient monitors are included.

static get_grad_monitors(input_structures, freqs_adjoint, include_eps_mnts=True)

Return dictionary of gradient monitors for simulation.

store_vjp(grad_data_fwd, grad_data_adj, grad_eps_data, num_proc=1)

Store the vjp w.r.t. each input_structure as a sim using fwd and adj grad_data.

store_vjp_sequential(grad_data_fwd, grad_data_adj, grad_eps_data)

Store the vjp w.r.t. each input_structure without multiprocessing.

store_vjp_parallel(grad_data_fwd, grad_data_adj, grad_eps_data, num_proc)

Store the vjp w.r.t. each input_structure as a sim using fwd and adj grad_data, and parallel processing over num_proc processes.

__hash__()

Hash method.