# container for everything defining the inverse design
from __future__ import annotations
import abc
import typing
import autograd.numpy as anp
import numpy as np
import pydantic.v1 as pd
import tidy3d as td
from tidy3d.components.autograd import get_static
from tidy3d.exceptions import ValidationError
from tidy3d.plugins.expressions.metrics import Metric, generate_validation_data
from tidy3d.plugins.expressions.types import ExpressionType
from .base import InvdesBaseModel
from .region import DesignRegionType
from .validators import check_pixel_size
PostProcessFnType = typing.Callable[[td.SimulationData], float]
class AbstractInverseDesign(InvdesBaseModel, abc.ABC):
"""Container for an inverse design problem."""
design_region: DesignRegionType = pd.Field(
...,
title="Design Region",
description="Region within which we will optimize the simulation.",
)
task_name: str = pd.Field(
...,
title="Task Name",
description="Task name to use in the objective function when running the ``JaxSimulation``.",
)
verbose: bool = pd.Field(
False,
title="Task Verbosity",
description="If ``True``, will print the regular output from ``web`` functions.",
)
metric: typing.Optional[ExpressionType] = pd.Field(
None,
title="Objective Metric",
description="Serializable expression defining the objective function.",
)
def make_objective_fn(
self, post_process_fn: typing.Optional[typing.Callable] = None, maximize: bool = True
) -> typing.Callable[[anp.ndarray], tuple[float, dict]]:
"""Construct the objective function for this InverseDesign object."""
if (post_process_fn is None) and (self.metric is None):
raise ValueError("Either 'post_process_fn' or 'metric' must be provided.")
if (post_process_fn is not None) and (self.metric is not None):
raise ValueError("Provide only one of 'post_process_fn' or 'metric', not both.")
direction_multiplier = 1 if maximize else -1
def objective_fn(params: anp.ndarray, aux_data: dict = None) -> float:
"""Full objective function."""
data = self.to_simulation_data(params=params)
if self.metric is None:
post_process_val = post_process_fn(data)
elif isinstance(data, td.SimulationData):
post_process_val = self.metric.evaluate(data)
elif getattr(data, "type", None) == "BatchData":
raise NotImplementedError("Metrics currently do not support 'BatchData'")
else:
raise ValueError(f"Invalid data type: {type(data)}")
penalty_value = self.design_region.penalty_value(params)
objective_fn_val = direction_multiplier * post_process_val - penalty_value
# Store auxiliary data if provided
if aux_data is not None:
aux_data["penalty"] = get_static(penalty_value)
aux_data["post_process_val"] = get_static(post_process_val)
aux_data["objective_fn_val"] = get_static(objective_fn_val) * direction_multiplier
if isinstance(data, td.SimulationData):
aux_data["sim_data"] = data.to_static()
else:
aux_data["sim_data"] = {k: v.to_static() for k, v in data.items()}
aux_data["params"] = params
return objective_fn_val
return objective_fn
@property
def initial_simulation(self) -> td.Simulation:
"""Return a simulation with the initial design region parameters."""
initial_params = self.design_region.initial_parameters
return self.to_simulation(initial_params)
def run(self, simulation, **kwargs) -> td.SimulationData:
"""Run a single tidy3d simulation."""
from tidy3d.web import run
kwargs.setdefault("verbose", self.verbose)
kwargs.setdefault("task_name", self.task_name)
return run(simulation, **kwargs)
def run_async(self, simulations, **kwargs) -> web.BatchData: # noqa: F821
"""Run a batch of tidy3d simulations."""
from tidy3d.web import run_async
kwargs.setdefault("verbose", self.verbose)
return run_async(simulations, **kwargs)
[docs]
class InverseDesign(AbstractInverseDesign):
"""Container for an inverse design problem."""
simulation: td.Simulation = pd.Field(
...,
title="Base Simulation",
description="Simulation without the design regions or monitors used in the objective fn.",
)
output_monitor_names: typing.Tuple[str, ...] = pd.Field(
None,
title="Output Monitor Names",
description="Optional names of monitors whose data the differentiable output depends on."
"If this field is left ``None``, the plugin will try to add all compatible monitors to "
"``JaxSimulation.output_monitors``. While this will work, there may be warnings if the "
"monitors are not compatible with the ``adjoint`` plugin, for example if there are "
"``FieldMonitor`` instances with ``.colocate != False``.",
)
_check_sim_pixel_size = check_pixel_size("simulation")
@pd.root_validator(pre=False)
def _validate_model(cls, values: dict) -> dict:
cls._validate_metric(values)
return values
@staticmethod
def _validate_metric(values: dict) -> dict:
metric_expr = values.get("metric")
if not metric_expr:
return values
simulation = values.get("simulation")
for metric in metric_expr.filter(Metric):
InverseDesign._validate_metric_monitor_name(metric, simulation)
InverseDesign._validate_metric_mode_index(metric, simulation)
InverseDesign._validate_metric_f(metric, simulation)
InverseDesign._validate_metric_data(metric_expr, simulation)
return values
@staticmethod
def _validate_metric_monitor_name(metric: Metric, simulation: td.Simulation) -> None:
"""Validate that the monitor name of the metric exists in the simulation."""
monitor = next((m for m in simulation.monitors if m.name == metric.monitor_name), None)
if monitor is None:
raise ValidationError(
f"Monitor named '{metric.monitor_name}' associated with the metric not found in the simulation monitors."
)
@staticmethod
def _validate_metric_mode_index(metric: Metric, simulation: td.Simulation) -> None:
"""Validate that the mode index of the metric is within the bounds of the monitor's ``ModeSpec.num_modes``."""
monitor = next((m for m in simulation.monitors if m.name == metric.monitor_name), None)
if metric.mode_index >= monitor.mode_spec.num_modes:
raise ValidationError(
f"Mode index '{metric.mode_index}' for metric associated with monitor "
f"'{metric.monitor_name}' is out of bounds. "
f"Maximum allowed mode index is '{monitor.mode_spec.num_modes - 1}'."
)
@staticmethod
def _validate_metric_f(metric: Metric, simulation: td.Simulation) -> None:
"""Validate that the frequencies of the metric are present in the monitor."""
monitor = next((m for m in simulation.monitors if m.name == metric.monitor_name), None)
if metric.f is not None:
metric_f_list = [metric.f] if isinstance(metric.f, float) else metric.f
if len(metric_f_list) != 1:
raise ValidationError("Only a single frequency is supported for the metric.")
for freq in metric_f_list:
if not any(np.isclose(freq, monitor.freqs, atol=1.0)):
raise ValidationError(
f"Frequency '{freq}' for metric associated with monitor "
f"'{metric.monitor_name}' not found in monitor frequencies."
)
else:
if len(monitor.freqs) != 1:
raise ValidationError(
f"Monitor '{metric.monitor_name}' must contain only a single frequency when metric.f is None."
)
@staticmethod
def _validate_metric_data(expr: ExpressionType, simulation: td.Simulation) -> None:
"""Validate that expression can be evaluated and returns a real scalar."""
data = generate_validation_data(expr)
try:
result = expr(data)
except Exception as e:
raise ValidationError(f"Failed to evaluate the metric expression: {str(e)}") from e
if len(np.ravel(result)) > 1:
raise ValidationError(
f"The expression must return a scalar value or an array of length 1 (got {result})."
)
if not np.all(np.isreal(result)):
raise ValidationError(
f"The expression must return a real (not complex) value (got {result})."
)
[docs]
def is_output_monitor(self, monitor: td.Monitor) -> bool:
"""Whether a monitor is added to the ``JaxSimulation`` as an ``output_monitor``."""
output_mnt_types = td.components.simulation.OutputMonitorTypes
if self.output_monitor_names is None:
return any(isinstance(monitor, mnt_type) for mnt_type in output_mnt_types)
return monitor.name in self.output_monitor_names
[docs]
def separate_output_monitors(self, monitors: typing.Tuple[td.Monitor]) -> dict:
"""Separate monitors into output_monitors and regular monitors."""
monitor_fields = dict(monitors=[], output_monitors=[])
for monitor in monitors:
key = "output_monitors" if self.is_output_monitor(monitor) else "monitors"
monitor_fields[key].append(monitor)
return monitor_fields
[docs]
def to_simulation(self, params: anp.ndarray) -> td.Simulation:
"""Convert the ``InverseDesign`` to a corresponding ``td.Simulation`` with traced fields."""
# construct the design region to a regular structure
design_region_structure = self.design_region.to_structure(params)
# construct mesh override structures and a new grid spec, if applicable
grid_spec = self.simulation.grid_spec
mesh_override_structure = self.design_region.mesh_override_structure
if mesh_override_structure:
override_structures = list(self.simulation.grid_spec.override_structures)
override_structures += [mesh_override_structure]
grid_spec = grid_spec.updated_copy(override_structures=override_structures)
return self.simulation.updated_copy(
structures=list(self.simulation.structures) + [design_region_structure],
grid_spec=grid_spec,
)
[docs]
def to_simulation_data(self, params: anp.ndarray, **kwargs) -> td.SimulationData:
"""Convert the ``InverseDesign`` to a ``td.Simulation`` and run it."""
simulation = self.to_simulation(params=params)
return self.run(simulation, **kwargs)
[docs]
class InverseDesignMulti(AbstractInverseDesign):
"""``InverseDesign`` with multiple simulations and corresponding postprocess functions."""
simulations: typing.Tuple[td.Simulation, ...] = pd.Field(
...,
title="Base Simulations",
description="Set of simulation without the design regions or monitors used in the objective fn.",
)
output_monitor_names: typing.Tuple[typing.Union[typing.Tuple[str, ...], None], ...] = pd.Field(
None,
title="Output Monitor Names",
description="Optional names of monitors whose data the differentiable output depends on."
"If this field is left ``None``, the plugin will try to add all compatible monitors to "
"``JaxSimulation.output_monitors``. While this will work, there may be warnings if the "
"monitors are not compatible with the ``adjoint`` plugin, for example if there are "
"``FieldMonitor`` instances with ``.colocate != False``.",
)
_check_sim_pixel_size = check_pixel_size("simulations")
@pd.root_validator()
def _check_lengths(cls, values):
"""Check the lengths of all of the multi fields."""
keys = ("simulations", "post_process_fns", "output_monitor_names", "override_structure_dl")
multi_dict = {key: values.get(key) for key in keys}
sizes = {key: len(val) for key, val in multi_dict.items() if val is not None}
if len(set(sizes.values())) != 1:
raise ValueError(
f"'MultiInverseDesign' requires that the fields {keys} must either "
"have the same length or be left ``None``, if optional. Given fields with "
"corresponding sizes of '{sizes}'."
)
return values
@property
def task_names(self) -> list[str]:
"""Task names associated with each of the simulations."""
return [f"{self.task_name}_{i}" for i in range(len(self.simulations))]
@property
def designs(self) -> typing.List[InverseDesign]:
"""List of individual ``InverseDesign`` objects corresponding to this instance."""
designs_list = []
for i, (task_name, sim) in enumerate(zip(self.task_names, self.simulations)):
des_i = InverseDesign(
design_region=self.design_region,
simulation=sim,
verbose=self.verbose,
task_name=task_name,
)
if self.output_monitor_names is not None:
des_i = des_i.updated_copy(output_monitor_names=self.output_monitor_names[i])
designs_list.append(des_i)
return designs_list
[docs]
def to_simulation(self, params: anp.ndarray) -> dict[str, td.Simulation]:
"""Convert the ``InverseDesign`` to a corresponding dict of ``td.Simulation``s."""
simulation_list = [design.to_simulation(params) for design in self.designs]
return dict(zip(self.task_names, simulation_list))
[docs]
def to_simulation_data(self, params: anp.ndarray, **kwargs) -> web.BatchData: # noqa: F821
"""Convert the ``InverseDesignMulti`` to a set of ``td.Simulation``s and run async."""
simulations = self.to_simulation(params)
return self.run_async(simulations, **kwargs)
InverseDesignType = typing.Union[InverseDesign, InverseDesignMulti]
