Source code for tidy3d.plugins.adjoint.components.simulation

"""Defines a jax-compatible simulation."""

from __future__ import annotations

from typing import Tuple, Union, List, Dict, Literal
from joblib import Parallel, delayed


import pydantic.v1 as pd
import numpy as np
import xarray as xr

from jax.tree_util import register_pytree_node_class

from ....log import log
from ....components.base import cached_property, Tidy3dBaseModel, skip_if_fields_missing
from ....components.monitor import FieldMonitor, PermittivityMonitor
from ....components.monitor import ModeMonitor, DiffractionMonitor, Monitor
from ....components.simulation import Simulation
from ....components.data.monitor_data import FieldData, PermittivityData
from ....components.structure import Structure
from ....components.medium import AbstractMedium
from ....components.types import Ax, annotate_type
from ....components.geometry.base import Box
from ....constants import HERTZ, SECOND
from ....exceptions import AdjointError

from .base import JaxObject, WEB_ADJOINT_MESSAGE
from .structure import (
    JaxStructure,
    JaxStructureType,
    JaxStructureStaticMedium,
    JaxStructureStaticGeometry,
)
from .geometry import JaxPolySlab, JaxGeometryGroup


# bandwidth of adjoint source in units of freq0 if no `fwidth_adjoint`, and one output freq
FWIDTH_FACTOR = 1.0 / 10

# bandwidth of adjoint sources in units of the minimum difference between output frequencies
FWIDTH_FACTOR_MULTIFREQ = 0.1

# the adjoint run time is the forward simulation run time + RUN_TIME_FACTOR / fwidth
RUN_TIME_FACTOR = 10

# how many processors to use for server and client side adjoint
NUM_PROC_LOCAL = 1

# number of input structures before it errors
MAX_NUM_INPUT_STRUCTURES = 400

# generic warning for nonlinearity
NL_WARNING = (
    "The 'adjoint' plugin does not currently support nonlinear materials. "
    "While the gradients might be calculated, they will be inaccurate and the "
    "error will increase as the strength of the nonlinearity is increased. "
    "We strongly recommend using linear simulations only with the adjoint plugin."
)

OutputMonitorTypes = (DiffractionMonitor, FieldMonitor, ModeMonitor)
OutputMonitorType = Tuple[annotate_type(Union[OutputMonitorTypes]), ...]


class JaxInfo(Tidy3dBaseModel):
    """Class to store information when converting between jax and tidy3d."""

    num_input_structures: pd.NonNegativeInt = pd.Field(
        ...,
        title="Number of Input Structures",
        description="Number of input structures in the original JaxSimulation.",
    )

    num_output_monitors: pd.NonNegativeInt = pd.Field(
        ...,
        title="Number of Output Monitors",
        description="Number of output monitors in the original JaxSimulation.",
    )

    num_grad_monitors: pd.NonNegativeInt = pd.Field(
        ...,
        title="Number of Gradient Monitors",
        description="Number of gradient monitors in the original JaxSimulation.",
    )

    num_grad_eps_monitors: pd.NonNegativeInt = pd.Field(
        ...,
        title="Number of Permittivity Monitors",
        description="Number of permittivity monitors in the original JaxSimulation.",
    )

    fwidth_adjoint: float = pd.Field(
        None,
        title="Adjoint Frequency Width",
        description="Custom frequency width of the original JaxSimulation.",
        units=HERTZ,
    )

    run_time_adjoint: float = pd.Field(
        None,
        title="Adjoint Run Time",
        description="Custom run time of the original JaxSimulation.",
        units=SECOND,
    )

    input_structure_types: Tuple[
        Literal["JaxStructure", "JaxStructureStaticMedium", "JaxStructureStaticGeometry"], ...
    ] = pd.Field(
        (),
        title="Input Structure Types",
        description="Type of the original input_structures (as strings).",
    )


[docs] @register_pytree_node_class class JaxSimulation(Simulation, JaxObject): """A :class:`.Simulation` registered with jax.""" input_structures: Tuple[annotate_type(JaxStructureType), ...] = pd.Field( (), title="Input Structures", description="Tuple of jax-compatible structures" " that may depend on differentiable parameters.", jax_field=True, ) output_monitors: OutputMonitorType = pd.Field( (), title="Output Monitors", description="Tuple of monitors whose data the differentiable output depends on.", ) grad_monitors: Tuple[FieldMonitor, ...] = pd.Field( (), title="Gradient Field Monitors", description="Tuple of monitors used for storing fields, used internally for gradients.", ) grad_eps_monitors: Tuple[PermittivityMonitor, ...] = pd.Field( (), title="Gradient Permittivity Monitors", description="Tuple of monitors used for storing epsilon, used internally for gradients.", ) fwidth_adjoint: pd.PositiveFloat = pd.Field( None, title="Adjoint Frequency Width", description="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.", units=HERTZ, ) run_time_adjoint: pd.PositiveFloat = pd.Field( None, title="Adjoint Run Time", description="Custom ``run_time`` to use for adjoint simulation. " "If not supplied or ``None``, uses a factor times the adjoint source ``fwidth``.", units=SECOND, ) @pd.validator("output_monitors", always=True) def _output_monitors_colocate_false(cls, val): """Make sure server-side colocation is off.""" new_vals = [] for mnt in val: if mnt.colocate: log.warning( "Output field monitors in the adjoint plugin require explicitly setting " "'colocate=False'. Setting 'colocate=False' in monitor '{mnt.name}'. ", "Use 'SimulationData.at_boundaries' after the solver run to automatically " "colocate the fields to the grid boundaries, or 'MonitorData.colocate' " "if colocating to custom coordinates.", ) mnt = mnt.updated_copy(colocate=False) new_vals.append(mnt) return new_vals @pd.validator("subpixel", always=True) def _subpixel_is_on(cls, val): """Assert subpixel is on.""" if not val: raise AdjointError("'JaxSimulation.subpixel' must be 'True' to use adjoint plugin.") return val @pd.validator("input_structures", always=True) @skip_if_fields_missing(["structures"]) def _warn_overlap(cls, val, values): """Print appropriate warning if structures intersect in ways that cause gradient error.""" input_structures = [s for s in val if "geometry" in s._differentiable_fields] structures = list(values.get("structures")) # if the center and size of all structure geometries do not contain all numbers, skip check for struct in input_structures: geometry = struct.geometry size_all_floats = all(isinstance(s, (float, int)) for s in geometry.bound_size) cent_all_floats = all(isinstance(c, (float, int)) for c in geometry.bound_center) if not (size_all_floats and cent_all_floats): return val with log as consolidated_logger: # check intersections with other input_structures for i, in_struct_i in enumerate(input_structures): geometry_i = in_struct_i.geometry for j in range(i + 1, len(input_structures)): if geometry_i.intersects(input_structures[j].geometry): consolidated_logger.warning( f"'JaxSimulation.input_structures[{i}]' overlaps or touches " f"'JaxSimulation.input_structures[{j}]'. Geometric gradients for " "overlapping input structures may contain errors.", log_once=True, ) # check JaxPolySlab intersections with background structures for i, in_struct_i in enumerate(input_structures): geometry_i = in_struct_i.geometry if not isinstance(geometry_i, JaxPolySlab): continue for j, struct_j in enumerate(structures): if geometry_i.intersects(struct_j.geometry): consolidated_logger.warning( f"'JaxPolySlab'-containing 'JaxSimulation.input_structures[{i}]' " f"intersects with 'JaxSimulation.structures[{j}]'. Note that in this " "version of the adjoint plugin, there may be errors in the gradient " "when 'JaxPolySlab' intersects with background structures." ) return val @pd.validator("output_monitors", always=True) def _warn_if_colocate(cls, val): """warn if any colocate=True in output FieldMonitors.""" for index, mnt in enumerate(val): if isinstance(mnt, FieldMonitor): if mnt.colocate: log.warning( f"'FieldMonitor' at 'JaxSimulation.output_monitors[{index}]' " "has 'colocate=True', " "this may lead to decreased accuracy in adjoint gradient." ) return val return val @pd.validator("medium", always=True) def _warn_nonlinear_medium(cls, val): """warn if the jax simulation medium is nonlinear.""" # hasattr is just an additional check to avoid unnecessary bugs # if a medium is encountered that does not support nonlinear spec, or things change. if hasattr(val, "nonlinear_spec") and val.nonlinear_spec: log.warning( "Nonlinear background medium detected in the 'JaxSimulation'. " + NL_WARNING ) return val @pd.validator("structures", always=True) def _warn_nonlinear_structure(cls, val): """warn if a jax simulation structure.medium is nonlinear.""" for i, struct in enumerate(val): medium = struct.medium # hasattr is just an additional check to avoid unnecessary bugs # if a medium is encountered that does not support nonlinear spec, or things change. if hasattr(medium, "nonlinear_spec") and medium.nonlinear_spec: log.warning(f"Nonlinear medium detected in structures[{i}]. " + NL_WARNING) return val @pd.validator("input_structures", always=True) def _warn_nonlinear_input_structure(cls, val): """warn if a jax simulation input_structure.medium is nonlinear.""" for i, struct in enumerate(val): medium = struct.medium # hasattr is just an additional check to avoid unnecessary bugs # if a medium is encountered that does not support nonlinear spec, or things change. if hasattr(medium, "nonlinear_spec") and medium.nonlinear_spec: log.warning(f"Nonlinear medium detected in input_structures[{i}]. " + NL_WARNING) return val def _restrict_input_structures(self) -> None: """Restrict number of input structures.""" num_input_structures = len(self.input_structures) if num_input_structures > MAX_NUM_INPUT_STRUCTURES: raise AdjointError( "For performance, adjoint plugin restricts the number of input structures to " f"{MAX_NUM_INPUT_STRUCTURES}. Found {num_input_structures}. " + WEB_ADJOINT_MESSAGE ) def _validate_web_adjoint(self) -> None: """Run validators for this component, only if using ``tda.web.run()``.""" self._restrict_input_structures() for structure in self.input_structures: structure._validate_web_adjoint()
[docs] @staticmethod def get_freqs_adjoint(output_monitors: List[Monitor]) -> List[float]: """Return sorted list of unique frequencies stripped from a collection of monitors.""" if len(output_monitors) == 0: raise AdjointError("Can't get adjoint frequency as no output monitors present.") output_freqs = [] for mnt in output_monitors: for freq in mnt.freqs: output_freqs.append(freq) return np.unique(output_freqs).tolist()
@cached_property def freqs_adjoint(self) -> List[float]: """Return sorted list of frequencies stripped from the output monitors.""" return self.get_freqs_adjoint(output_monitors=self.output_monitors) @cached_property def _is_multi_freq(self) -> bool: """Does this simulation have a multi-frequency output?""" return len(self.freqs_adjoint) > 1 @cached_property def _min_delta_freq(self) -> float: """Minimum spacing between output_frequencies (Hz).""" if not self._is_multi_freq: return None delta_freqs = np.abs(np.diff(np.sort(np.array(self.freqs_adjoint)))) return np.min(delta_freqs) @cached_property def _fwidth_adjoint(self) -> float: """Frequency width to use for adjoint source, user-defined or the average of the sources.""" # if user-specified, use that if self.fwidth_adjoint is not None: return self.fwidth_adjoint freqs_adjoint = self.freqs_adjoint # multiple output frequency case if self._is_multi_freq: return FWIDTH_FACTOR_MULTIFREQ * self._min_delta_freq # otherwise, grab from sources and output monitors num_sources = len(self.sources) # should be 0 for adjoint already but worth checking # if no sources, just use a constant factor times the mean adjoint frequency if num_sources == 0: return FWIDTH_FACTOR * np.mean(freqs_adjoint) # if more than one forward source, use their maximum if num_sources > 1: log.warning(f"{num_sources} sources, using their maximum 'fwidth' for adjoint source.") fwidths = [src.source_time.fwidth for src in self.sources] return np.max(fwidths) @cached_property def _run_time_adjoint(self: float) -> float: """Return the run time of the adjoint simulation as a function of its fwidth.""" if self.run_time_adjoint is not None: return self.run_time_adjoint run_time_fwd = self._run_time run_time_adjoint = run_time_fwd + RUN_TIME_FACTOR / self._fwidth_adjoint if self._is_multi_freq: log.warning( f"{len(self.freqs_adjoint)} unique frequencies detected in the output monitors " f"with a minimum spacing of {self._min_delta_freq:.3e} (Hz). " f"Setting the 'fwidth' of the adjoint sources to {FWIDTH_FACTOR_MULTIFREQ} times " f"this value = {self._fwidth_adjoint:.3e} (Hz) to avoid spectral overlap. " "To account for this, the corresponding 'run_time' in the adjoint simulation is " f"will be set to {run_time_adjoint:3e} " f"compared to {self._run_time:3e} in the forward simulation. " "If the adjoint 'run_time' is large due to small frequency spacing, " "it could be better to instead run one simulation per frequency, " "which can be done in parallel using 'tidy3d.plugins.adjoint.web.run_async'." ) return run_time_adjoint @cached_property def tmesh_adjoint(self) -> np.ndarray: """FDTD time stepping points. Returns ------- np.ndarray Times (seconds) that the simulation time steps through. """ dt = self.dt return np.arange(0.0, self._run_time_adjoint + dt, dt) @cached_property def num_time_steps_adjoint(self) -> int: """Number of time steps in the adjoint simulation.""" return len(self.tmesh_adjoint)
[docs] def to_simulation(self) -> Tuple[Simulation, JaxInfo]: """Convert :class:`.JaxSimulation` instance to :class:`.Simulation` with an info dict.""" sim_dict = self.dict( exclude={ "type", "structures", "monitors", "output_monitors", "grad_monitors", "grad_eps_monitors", "input_structures", "fwidth_adjoint", "run_time_adjoint", } ) sim = Simulation.parse_obj(sim_dict) # put all structures and monitors in one list all_structures = list(self.structures) + [js.to_structure() for js in self.input_structures] all_monitors = ( list(self.monitors) + list(self.output_monitors) + list(self.grad_monitors) + list(self.grad_eps_monitors) ) sim = sim.updated_copy(structures=all_structures, monitors=all_monitors) # information about the state of the original JaxSimulation to stash for reconstruction jax_info = JaxInfo( num_input_structures=len(self.input_structures), num_output_monitors=len(self.output_monitors), num_grad_monitors=len(self.grad_monitors), num_grad_eps_monitors=len(self.grad_eps_monitors), fwidth_adjoint=self.fwidth_adjoint, run_time_adjoint=self.run_time_adjoint, input_structure_types=[s.type for s in self.input_structures], ) return sim, jax_info
[docs] def to_gds( self, cell, x: float = None, y: float = None, z: float = None, permittivity_threshold: pd.NonNegativeFloat = 1, frequency: pd.PositiveFloat = 0, gds_layer_dtype_map: Dict[ AbstractMedium, Tuple[pd.NonNegativeInt, pd.NonNegativeInt] ] = None, ) -> None: """Append the simulation structures to a .gds cell. Parameters ---------- cell : ``gdstk.Cell`` or ``gdspy.Cell`` Cell object to which the generated polygons are added. x : float = None Position of plane in x direction, only one of x,y,z can be specified to define plane. y : float = None Position of plane in y direction, only one of x,y,z can be specified to define plane. z : float = None Position of plane in z direction, only one of x,y,z can be specified to define plane. permittivity_threshold : float = 1 Permittivity value used to define the shape boundaries for structures with custom medim 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. """ sim, _ = self.to_simulation() return sim.to_gds( cell=cell, x=x, y=y, z=z, permittivity_threshold=permittivity_threshold, frequency=frequency, gds_layer_dtype_map=gds_layer_dtype_map, )
[docs] def to_gdstk( self, x: float = None, y: float = None, z: float = None, permittivity_threshold: pd.NonNegativeFloat = 1, frequency: pd.PositiveFloat = 0, gds_layer_dtype_map: Dict[ AbstractMedium, Tuple[pd.NonNegativeInt, pd.NonNegativeInt] ] = None, ) -> List: """Convert a simulation's planar slice to a .gds type polygon list. Parameters ---------- x : float = None Position of plane in x direction, only one of x,y,z can be specified to define plane. y : float = None Position of plane in y direction, only one of x,y,z can be specified to define plane. z : float = None Position of plane in z direction, only one of x,y,z can be specified to define plane. permittivity_threshold : float = 1 Permittivity value used to define the shape boundaries for structures with custom medim 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. Return ------ List List of `gdstk.Polygon`. """ sim, _ = self.to_simulation() return sim.to_gdstk( x=x, y=y, z=z, permittivity_threshold=permittivity_threshold, frequency=frequency, gds_layer_dtype_map=gds_layer_dtype_map, )
[docs] def to_gdspy( self, x: float = None, y: float = None, z: float = None, gds_layer_dtype_map: Dict[ AbstractMedium, Tuple[pd.NonNegativeInt, pd.NonNegativeInt] ] = None, ) -> List: """Convert a simulation's planar slice to a .gds type polygon list. Parameters ---------- x : float = None Position of plane in x direction, only one of x,y,z can be specified to define plane. y : float = None Position of plane in y direction, only one of x,y,z can be specified to define plane. z : float = None Position of plane in z direction, only one of x,y,z can be specified to define plane. gds_layer_dtype_map : Dict Dictionary mapping mediums to GDSII layer and data type tuples. Return ------ List List of `gdspy.Polygon` and `gdspy.PolygonSet`. """ sim, _ = self.to_simulation() return sim.to_gdspy(x=x, y=y, z=z, gds_layer_dtype_map=gds_layer_dtype_map)
[docs] def plot( self, x: float = None, y: float = None, z: float = None, ax: Ax = None, source_alpha: float = None, monitor_alpha: float = None, hlim: Tuple[float, float] = None, vlim: Tuple[float, float] = None, **patch_kwargs, ) -> Ax: """Wrapper around regular :class:`.Simulation` structure plotting.""" sim, _ = self.to_simulation() return sim.plot( x=x, y=y, z=z, ax=ax, source_alpha=source_alpha, monitor_alpha=monitor_alpha, hlim=hlim, vlim=vlim, **patch_kwargs, )
[docs] def plot_eps( self, x: float = None, y: float = None, z: float = None, freq: float = None, alpha: float = None, source_alpha: float = None, monitor_alpha: float = None, hlim: Tuple[float, float] = None, vlim: Tuple[float, float] = None, ax: Ax = None, ) -> Ax: """Wrapper around regular :class:`.Simulation` permittivity plotting.""" sim, _ = self.to_simulation() return sim.plot_eps( x=x, y=y, z=z, ax=ax, source_alpha=source_alpha, monitor_alpha=monitor_alpha, hlim=hlim, vlim=vlim, )
[docs] def plot_structures( self, x: float = None, y: float = None, z: float = None, ax: Ax = None, hlim: Tuple[float, float] = None, vlim: Tuple[float, float] = None, ) -> Ax: """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 ------- matplotlib.axes._subplots.Axes The supplied or created matplotlib axes. """ sim, _ = self.to_simulation() return sim.plot_structures( x=x, y=y, z=z, ax=ax, hlim=hlim, vlim=vlim, )
[docs] def plot_structures_eps( self, x: float = None, y: float = None, z: float = None, freq: float = None, alpha: float = None, cbar: bool = True, reverse: bool = False, ax: Ax = None, hlim: Tuple[float, float] = None, vlim: Tuple[float, float] = None, ) -> Ax: """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 ------- matplotlib.axes._subplots.Axes The supplied or created matplotlib axes. """ sim, _ = self.to_simulation() return sim.plot_structures_eps( x=x, y=y, z=z, freq=freq, alpha=alpha, cbar=cbar, reverse=reverse, ax=ax, hlim=hlim, vlim=vlim, )
[docs] def epsilon( self, box: Box, coord_key: str = "centers", freq: float = None, ) -> xr.DataArray: """Get array of permittivity at volume specified by box and freq. Parameters ---------- box : :class:`.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 ------- xarray.DataArray Datastructure containing the relative permittivity values and location coordinates. For details on xarray DataArray objects, refer to `xarray's Documentation <https://tinyurl.com/2zrzsp7b>`_. See Also -------- **Notebooks** * `First walkthrough: permittivity data <../../notebooks/Simulation.html#Permittivity-data>`_ """ sim, _ = self.to_simulation() return sim.epsilon(box=box, coord_key=coord_key, freq=freq)
[docs] def __eq__(self, other: JaxSimulation) -> bool: """Are two JaxSimulation objects equal?""" return self.to_simulation()[0] == other.to_simulation()[0]
[docs] @classmethod def split_monitors(cls, monitors: List[Monitor], jax_info: JaxInfo) -> Dict[str, Monitor]: """Split monitors into user and adjoint required based on jax info.""" all_monitors = list(monitors) # grab or compute the number of type of monitor num_grad_monitors = jax_info.num_grad_monitors num_grad_eps_monitors = jax_info.num_grad_eps_monitors num_output_monitors = jax_info.num_output_monitors num_total_monitors = len(all_monitors) num_mnts = ( num_total_monitors - num_grad_monitors - num_output_monitors - num_grad_eps_monitors ) # split the monitor list based on these numbers monitors = all_monitors[:num_mnts] output_monitors = all_monitors[num_mnts : num_mnts + num_output_monitors] grad_monitors = all_monitors[ num_mnts + num_output_monitors : num_mnts + num_output_monitors + num_grad_monitors ] grad_eps_monitors = all_monitors[num_mnts + num_output_monitors + num_grad_monitors :] # load into a dictionary return dict( monitors=monitors, output_monitors=output_monitors, grad_monitors=grad_monitors, grad_eps_monitors=grad_eps_monitors, )
[docs] @classmethod def split_structures( cls, structures: List[Structure], jax_info: JaxInfo ) -> Dict[str, Structure]: """Split structures into regular and input based on jax info.""" all_structures = list(structures) # get numbers of regular and input structures num_input_structures = jax_info.num_input_structures num_structs = len(structures) - num_input_structures # split the list based on these numbers structures = all_structures[:num_structs] structure_type_map = dict( JaxStructure=JaxStructure, JaxStructureStaticMedium=JaxStructureStaticMedium, JaxStructureStaticGeometry=JaxStructureStaticGeometry, ) input_structures = [] for struct_type_str, struct in zip( jax_info.input_structure_types, all_structures[num_structs:] ): struct_type = structure_type_map[struct_type_str] new_structure = struct_type.from_structure(struct) input_structures.append(new_structure) # return a dictionary containing these split structures return dict(structures=structures, input_structures=input_structures)
[docs] @classmethod def from_simulation(cls, simulation: Simulation, jax_info: JaxInfo) -> JaxSimulation: """Convert :class:`.Simulation` to :class:`.JaxSimulation` with extra info.""" sim_dict = simulation.dict(exclude={"type", "structures", "monitors"}) # split structures and monitors into their respective fields for JaxSimulation structures = cls.split_structures(structures=simulation.structures, jax_info=jax_info) monitors = cls.split_monitors(monitors=simulation.monitors, jax_info=jax_info) # update the dictionary with these and the adjoint fwidth sim_dict.update(**structures) sim_dict.update(**monitors) sim_dict.update( dict( fwidth_adjoint=jax_info.fwidth_adjoint, run_time_adjoint=jax_info.run_time_adjoint, ) ) # load JaxSimulation from the dictionary return cls.parse_obj(sim_dict)
[docs] @classmethod def make_sim_fwd(cls, simulation: Simulation, jax_info: JaxInfo) -> Tuple[Simulation, JaxInfo]: """Make the forward :class:`.JaxSimulation` from the supplied :class:`.Simulation`.""" mnt_dict = JaxSimulation.split_monitors(monitors=simulation.monitors, jax_info=jax_info) structure_dict = JaxSimulation.split_structures( structures=simulation.structures, jax_info=jax_info ) output_monitors = mnt_dict["output_monitors"] input_structures = structure_dict["input_structures"] grad_mnt_dict = cls.get_grad_monitors( input_structures=input_structures, freqs_adjoint=cls.get_freqs_adjoint(output_monitors=output_monitors), ) grad_mnts = grad_mnt_dict["grad_monitors"] grad_eps_mnts = grad_mnt_dict["grad_eps_monitors"] full_monitors = list(simulation.monitors) + grad_mnts + grad_eps_mnts # jax_sim_fwd = jax_sim.updated_copy(**grad_mnts) # sim_fwd, jax_info = jax_sim_fwd.to_simulation() sim_fwd = simulation.updated_copy(monitors=full_monitors) jax_info = jax_info.updated_copy( num_grad_monitors=len(grad_mnts), num_grad_eps_monitors=len(grad_eps_mnts), ) # cls.split_monitors(monitors=simulation.monitors, jax_info=jax_info) # sim_fwd = simulation.updated_copy() return sim_fwd, jax_info
[docs] def to_simulation_fwd(self) -> Tuple[Simulation, JaxInfo, JaxInfo]: """Like ``to_simulation()`` but the gradient monitors are included.""" simulation, jax_info = self.to_simulation() sim_fwd, jax_info_fwd = self.make_sim_fwd(simulation=simulation, jax_info=jax_info) return sim_fwd, jax_info_fwd, jax_info
[docs] @staticmethod def get_grad_monitors( input_structures: List[Structure], freqs_adjoint: List[float], include_eps_mnts: bool = True ) -> dict: """Return dictionary of gradient monitors for simulation.""" grad_mnts = [] grad_eps_mnts = [] for index, structure in enumerate(input_structures): grad_mnt, grad_eps_mnt = structure.make_grad_monitors( freqs=freqs_adjoint, name=f"grad_mnt_{index}" ) grad_mnts.append(grad_mnt) if include_eps_mnts: grad_eps_mnts.append(grad_eps_mnt) return dict(grad_monitors=grad_mnts, grad_eps_monitors=grad_eps_mnts)
def _store_vjp_structure( self, structure: JaxStructure, fld_fwd: FieldData, fld_adj: FieldData, eps_data: PermittivityData, num_proc: int = NUM_PROC_LOCAL, ) -> JaxStructure: """Store the vjp for a single structure.""" freq_max = float(max(eps_data.eps_xx.coords["f"])) eps_out = self.medium.eps_model(frequency=freq_max) return structure.store_vjp( grad_data_fwd=fld_fwd, grad_data_adj=fld_adj, grad_data_eps=eps_data, sim_bounds=self.bounds, eps_out=eps_out, num_proc=num_proc, )
[docs] def store_vjp( self, grad_data_fwd: Tuple[FieldData], grad_data_adj: Tuple[FieldData], grad_eps_data: Tuple[PermittivityData], num_proc: int = NUM_PROC_LOCAL, ) -> JaxSimulation: """Store the vjp w.r.t. each input_structure as a sim using fwd and adj grad_data.""" # if num_proc supplied and greater than 1, run parallel if num_proc is not None and num_proc > 1: return self.store_vjp_parallel( grad_data_fwd=grad_data_fwd, grad_data_adj=grad_data_adj, grad_eps_data=grad_eps_data, num_proc=num_proc, ) # otherwise, call regular sequential one return self.store_vjp_sequential( grad_data_fwd=grad_data_fwd, grad_data_adj=grad_data_adj, grad_eps_data=grad_eps_data )
[docs] def store_vjp_sequential( self, grad_data_fwd: Tuple[FieldData], grad_data_adj: Tuple[FieldData], grad_eps_data: Tuple[PermittivityData], ) -> JaxSimulation: """Store the vjp w.r.t. each input_structure without multiprocessing.""" map_args = [self.input_structures, grad_data_fwd, grad_data_adj, grad_eps_data] input_structures_vjp = list(map(self._store_vjp_structure, *map_args)) return self.copy( update=dict( input_structures=input_structures_vjp, grad_monitors=(), grad_eps_monitors=() ) )
[docs] def store_vjp_parallel( self, grad_data_fwd: Tuple[FieldData], grad_data_adj: Tuple[FieldData], grad_eps_data: Tuple[PermittivityData], num_proc: int, ) -> JaxSimulation: """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.""" # Indexing into structures which use internal parallelization, and those which don't. # For the latter, simple parallelization over the list will be used. internal_par_structs = [JaxGeometryGroup] # Parallelize polyslabs internally or externally depending on total number polyslabs = [struct for struct in self.input_structures if isinstance(struct, JaxPolySlab)] if len(polyslabs) < num_proc: internal_par_structs += [JaxPolySlab] inds_par_internal, inds_par_external = [], [] for index, structure in enumerate(self.input_structures): if isinstance(structure.geometry, tuple(internal_par_structs)): inds_par_internal.append(index) else: inds_par_external.append(index) def make_args(indexes, num_proc_internal) -> list: """Make the arguments to map over selecting over a set of structure ``indexes``.""" args_list = [] for index in indexes: args_i = [ self.input_structures[index], grad_data_fwd[index], grad_data_adj[index], grad_eps_data[index], num_proc_internal, ] args_list.append(args_i) return args_list # Get vjps for structures that parallelize internally using simple map args_list_internal = make_args(inds_par_internal, num_proc_internal=num_proc) vjps_par_internal = [self._store_vjp_structure(*args) for args in args_list_internal] # Get vjps for structures where we parallelize directly here args_list_external = make_args(inds_par_external, num_proc_internal=NUM_PROC_LOCAL) vjps_par_external = list( Parallel(n_jobs=num_proc)( delayed(self._store_vjp_structure)(*args) for args in args_list_external ) ) # Reshuffle the two lists back in the correct order vjps_all = list(vjps_par_internal) + list(vjps_par_external) input_structures_vjp = [None] * len(self.input_structures) for index, vjp in zip(inds_par_internal + inds_par_external, vjps_all): input_structures_vjp[index] = vjp return self.copy( update=dict( input_structures=input_structures_vjp, grad_monitors=(), grad_eps_monitors=() ) )