diff --git a/docs/faq b/docs/faq index ed2256c606..f311e85976 160000 --- a/docs/faq +++ b/docs/faq @@ -1 +1 @@ -Subproject commit ed2256c6061a584f86b97051df6b806b02d8d355 +Subproject commit f311e85976cef312f8dabe13c840e340556b2b36 diff --git a/docs/notebooks b/docs/notebooks index 5aba14f236..5120baf869 160000 --- a/docs/notebooks +++ b/docs/notebooks @@ -1 +1 @@ -Subproject commit 5aba14f2361666baec39e30d4f0620fc7557b6b9 +Subproject commit 5120baf8690cfe722b2d3c8f67fb44362d65fcb6 diff --git a/tests/test_plugins/test_invdes2.py b/tests/test_plugins/test_invdes2.py new file mode 100644 index 0000000000..95fb6a96eb --- /dev/null +++ b/tests/test_plugins/test_invdes2.py @@ -0,0 +1,171 @@ +from __future__ import annotations + +import autograd.numpy as np +import pytest + +import tidy3d as td +from tidy3d.plugins.invdes2 import ( + DeviceSpec, + FluxMetric, + InverseDesign, + OptimizerSpec, + TopologyDesignRegion, +) + +from ..test_components.autograd.test_autograd import use_emulated_run # noqa: F401 +from ..utils import run_emulated + +sim_base = td.Simulation( + size=(10.0, 10.0, 10.0), + grid_spec=td.GridSpec.auto(wavelength=1.0, min_steps_per_wvl=10), + run_time=1.0, + structures=(), + monitors=[ + td.FluxMonitor(center=(0.0, 0.0, 0.0), size=(1.0, 1.0, 1.0), freqs=[2e14], name="flux") + ], + sources=(), + boundary_spec=td.BoundarySpec.all_sides(boundary=td.PML()), + medium=td.Medium(permittivity=1.0), +) + +sim_data_base = run_emulated(sim_base, task_name="sim_base") + +# TODO: Add more metrics here +metrics = [FluxMetric(monitor_name="flux", weight=0.5)] + +# TODO: Add more design regions here +design_regions = [ + TopologyDesignRegion( + size=(1.0, 1.0, 1.0), center=(0.0, 0.0, 0.0), eps_bounds=(1.0, 4.0), pixel_size=0.02 + ) +] + +device_spec1 = DeviceSpec( + simulation=sim_base, design_regions=design_regions, metrics=metrics, name="d1" +) + +device_spec2 = DeviceSpec( + simulation=sim_base, design_regions=design_regions, metrics=metrics, name="d2" +) + +device_specs = [device_spec1, device_spec2] + +optimizer_spec = OptimizerSpec(learning_rate=0.1, num_steps=1) + +invdes = InverseDesign(optimizer_spec=optimizer_spec, device_specs=device_specs) + + +def test_parameter_shapes(): + """Ensure parameter shape metadata aligns across devices and regions. + + - `InverseDesign.parameter_shape` should equal the list of each `DeviceSpec.parameter_shape`. + - Each `DeviceSpec.parameter_shape` should equal the list of each region's `parameter_shape`. + """ + assert invdes.parameter_shape == [d.parameter_shape for d in invdes.device_specs] + for device_spec in invdes.device_specs: + assert device_spec.parameter_shape == [ + d.parameter_shape for d in device_spec.design_regions + ] + + +def test_flatten_unflatten_params(): + """Round-trip flatten/unflatten preserves the parameter vector. + + Uses helper constructors to build correctly sized parameter arrays, then verifies that + flatten → unflatten → flatten yields an identical 1D vector. + """ + params = invdes.ones() + flat = invdes._flatten_params(params) + restored = invdes._unflatten_params(flat) + flat2 = invdes._flatten_params(restored) + + assert np.allclose(flat, flat2) + + +def test_design_region_to_structure(): + """Each design region can map its parameter vector to a `td.Structure`. + + Builds per-region parameter arrays with the provided helper and ensures `to_structure` + returns a structure without error. + """ + for design_region in design_regions: + params = design_region.ones() + _ = design_region.to_structure(params) + + +def test_device_spec_get_simulation(): + """`DeviceSpec.get_simulation` appends one structure per design region. + + The resulting simulation should contain the original structures plus the number of + design regions in the spec. + """ + for device_spec in device_specs: + params = device_spec.ones() + sim = device_spec.get_simulation(params) + assert len(sim.structures) == len(sim_base.structures) + len(device_spec.design_regions) + + +def test_invdes_get_simulations(): + """`InverseDesign.get_simulations` returns a batch keyed by device names. + + Confirms the number of simulations equals the number of device specs and that keys are + exactly the device names. + """ + params = invdes.ones() + sims = invdes.get_simulations(params) + + assert len(sims) == len(invdes.device_specs) + assert set(sims.keys()) == {device_spec.name for device_spec in invdes.device_specs} + + +def test_metric_evaluate(): + """`Metric.evaluate` produces a non-zero scalar from emulated monitor data.""" + for metric in metrics: + mnt_data = sim_data_base[metric.monitor_name] + val = metric.evaluate(mnt_data) + assert not np.allclose(val, 0.0) + + +def test_device_spec_get_metric(): + """`DeviceSpec.get_metric` aggregates weighted metric values into a scalar.""" + for device_spec in device_specs: + val = device_spec.get_metric(sim_data_base) + assert not np.allclose(val, 0.0) + + +def test_invdes_get_metric(): + """`InverseDesign.get_metric` sums device metrics from batch results.""" + batch_data = {device_spec.name: sim_data_base for device_spec in invdes.device_specs} + val = invdes.get_metric(batch_data) + assert not np.allclose(val, 0.0) + + +def test_inverse_design_unique_names_validation(): + """Constructing `InverseDesign` with duplicate device names raises `ValueError`.""" + device_specs_fail = [device_spec1, device_spec1] + with pytest.raises(ValueError): + InverseDesign(optimizer_spec=optimizer_spec, device_specs=device_specs_fail) + + +@pytest.fixture +def use_emulated(monkeypatch): + """Emulate the InverseDesign.to_simulation_data to call emulated run.""" + monkeypatch.setattr( + DeviceSpec, + "run_simulation", + lambda self, simulation: run_emulated(simulation, task_name="test"), + ) + monkeypatch.setattr( + InverseDesign, + "run_simulations", + lambda self, sims: { + task_name: run_emulated(sim, task_name=task_name) for task_name, sim in sims.items() + }, + ) + + +def test_objective_function(use_emulated): + """`InverseDesign.get_objective` returns a non-zero scalar using emulated runs.""" + params = invdes.ones() + val = invdes.get_objective(params) + assert not np.allclose(val, 0.0) diff --git a/tidy3d/plugins/invdes2/README.md b/tidy3d/plugins/invdes2/README.md new file mode 100644 index 0000000000..0b06b3e4ae --- /dev/null +++ b/tidy3d/plugins/invdes2/README.md @@ -0,0 +1,320 @@ +Inverse Design Scaffold (invdes2) +================================= + +Overview +-------- +This module provides a lightweight, composable scaffold for photonic inverse design built on Tidy3D. It separates concerns into: + +- Design regions: turn parameter vectors into `td.Structure`s. +- Device specs: combine a base `td.Simulation` with parameterized structures and score it with metrics. +- Inverse design orchestrator: build, run, and aggregate multiple device scenarios. +- Metrics: compute scalar objectives from monitor data. +- Optimizer spec: a hyperparameter container to integrate with your optimizer of choice. + +High-level flow +--------------- +``` + parameters (list[list[np.ndarray]]) + │ + ▼ + InverseDesign.get_simulations(params) + │ ┌─────────────────────────────────────────────────────┐ + ├─ for each device_spec: │ + │ DeviceSpec.get_simulation(device_params) │ + │ ├─ for each region: │ + │ │ DesignRegion.to_structure(region_params) │ + │ └─ append structures to base simulation │ + └────────────────────────────────────────────────────────────────┘ + │ + ▼ + InverseDesign.run_simulations({name: Simulation}) ──> Tidy3D Web (async) + │ + ▼ + InverseDesign.get_metric(batch_data) ── sum over devices + │ + ▼ + scalar objective (float) +``` + +Key concepts and APIs +--------------------- +Design regions (`design_region.py`) +- `DesignRegion` (ABC): contract for parameterized geometry providers. + - `parameter_shape: int`: total number of parameters (flattened). + - `to_structure(params: np.ndarray) -> td.Structure`: maps a 1D parameter vector to a `td.Structure`. + - `ones(**kwargs) -> np.ndarray`: helper to create a correctly sized parameter vector. +- `TopologyDesignRegion(DesignRegion)`: pixellated permittivity grid inside a `td.Box`. + - `size, center, pixel_size`: define a voxelized grid; `shape_3d` is computed from `(size / pixel_size)`. + - `eps_bounds`: available for future projection/constraints (currently not applied). + +Device specs (`device_spec.py`) +- `DeviceSpec`: describes a single device scenario. + - `simulation: td.Simulation`: base sim. + - `design_regions: list[DesignRegion]`: ordered list; each consumes a parameter array. + - `metrics: list[Metric]`: weighted scalar contributions. + - `name: str`: unique task name for batch runs. + - `get_simulation(params)`: appends one `td.Structure` per region to the base sim, returns a new `td.Simulation`. + - `run_simulation(sim)`: executes the simulation via `tidy3d.web.run(sim, task_name=name)`. + - `get_metric(sim_data)`: for each metric, pulls `sim_data[metric.monitor_name]`, applies `metric.evaluate`, multiplies by `metric.weight`, and sums. + - `parameter_shape`: list of per-region sizes. + - `ones(**kwargs)`: list of correctly sized 1D arrays for each region. + +Inverse design orchestrator (`inverse_design.py`) +- `InverseDesign` coordinates multiple `DeviceSpec`s. + - Validates unique device names in `__post_init__`. + - `get_simulations(params) -> dict[str, td.Simulation]`: builds one sim per device. + - `run_simulations(sims) -> BatchData`: runs all sims concurrently via `tidy3d.web.run_async`. + - `get_metric(batch_data) -> float`: sums `DeviceSpec.get_metric(batch_data[name])` across devices. + - `get_objective(params) -> float`: convenience: build → run → aggregate. + - `parameter_shape`: list of per-device lists of per-region sizes. + - `_flatten_params(params) -> np.ndarray` and `_unflatten_params(vec) -> params`: helpers for optimizer integrations. + - `ones(**kwargs)`: list of list of 1D arrays, matching devices→regions. + +Metrics (`metric.py`) +- `Metric` (ABC): contract for computing scalars from monitor data. + - `monitor_name: str`, `weight: float = 1.0`. + - `evaluate(mnt_data) -> float`: returns scalar. +- `FluxMetric(Metric)`: sums `mnt_data.flux.values`. + +Optimizer spec (`optimizer_spec.py`) +- `OptimizerSpec`: hyperparameters for an external optimization loop. + - `num_steps: int`, `learning_rate: float`. + - Not used directly by this scaffold; provided for consistency and integration. + +Parameter shapes and nesting +--------------------------- +- Per-region parameter is a 1D vector with length `region.parameter_shape`. +- Per-device parameters are `list[np.ndarray]`, aligned with `DeviceSpec.design_regions`. +- All-devices parameters are `list[list[np.ndarray]]`, aligned with `InverseDesign.device_specs`. +- `_unflatten_params` returns 1D segments for each region; `TopologyDesignRegion.to_structure` internally reshapes to 3D via `shape_3d`. + +Minimal example +--------------- +```python +import autograd.numpy as np +import tidy3d as td +from tidy3d.plugins.invdes2 import ( + DeviceSpec, + FluxMetric, + InverseDesign, + OptimizerSpec, + TopologyDesignRegion, +) + +# 1) Base simulation with a Flux monitor +sim_base = td.Simulation( + size=(10.0, 10.0, 10.0), + grid_spec=td.GridSpec.auto(wavelength=1.0, min_steps_per_wvl=10), + run_time=1.0, + structures=(), + monitors=[ + td.FluxMonitor(center=(0.0, 0.0, 0.0), size=(1.0, 1.0, 1.0), freqs=[2e14], name="flux") + ], + sources=(), + boundary_spec=td.BoundarySpec.all_sides(boundary=td.PML()), + medium=td.Medium(permittivity=1.0), +) + +# 2) Design regions and metrics +regions = [ + TopologyDesignRegion( + size=(1.0, 1.0, 1.0), center=(0.0, 0.0, 0.0), eps_bounds=(1.0, 4.0), pixel_size=0.02 + ) +] +metrics = [FluxMetric(monitor_name="flux", weight=0.5)] + +# 3) Two device scenarios (could vary monitors, sources, or regions per device) +dev1 = DeviceSpec(simulation=sim_base, design_regions=regions, metrics=metrics, name="d1") +dev2 = DeviceSpec(simulation=sim_base, design_regions=regions, metrics=metrics, name="d2") + +inv = InverseDesign(optimizer_spec=OptimizerSpec(learning_rate=0.1, num_steps=10), device_specs=[dev1, dev2]) + +# 4) Build params, run, and score +params = inv.ones() # nested list: devices → regions → 1D arrays +objective_value = inv.get_objective(params) +print("Objective:", objective_value) +``` + +Batch build/run/aggregate explicitly +------------------------------------ +```python +params = inv.ones() +sims = inv.get_simulations(params) # {"d1": Simulation, "d2": Simulation} +batch_data = inv.run_simulations(sims) # runs async via Tidy3D Web +value = inv.get_metric(batch_data) # aggregates across devices +``` + +Flatten/unflatten for optimizers +-------------------------------- +```python +flat = inv._flatten_params(params) +restored = inv._unflatten_params(flat) +assert np.allclose(flat, inv._flatten_params(restored)) + +# Example: manual gradient loop (requires autograd-enabled backend for Tidy3D execution) +from autograd import grad + +def objective_from_flat(vec): + p = inv._unflatten_params(vec) + return inv.get_objective(p) + +g = grad(objective_from_flat) +vec = flat +for step in range(inv.optimizer_spec.num_steps): + vec = vec - inv.optimizer_spec.learning_rate * g(vec) +final_params = inv._unflatten_params(vec) +``` + +Testing and emulation seam +-------------------------- +The tests demonstrate two seam points for swapping execution backends via monkeypatching: + +```python +# Swap device run for emulation +monkeypatch.setattr( + DeviceSpec, + "run_simulation", + lambda self, simulation: run_emulated(simulation, task_name="test"), +) + +# Swap batch run for emulation +monkeypatch.setattr( + InverseDesign, + "run_simulations", + lambda self, sims: {name: run_emulated(sim, task_name=name) for name, sim in sims.items()}, +) +``` + +Extending the system +-------------------- +Add a new design region +```python +from dataclasses import dataclass +import autograd.numpy as np +import tidy3d as td +from tidy3d.plugins.invdes2 import DeviceSpec, InverseDesign, TopologyDesignRegion + +@dataclass +class ParamBox(td.typing.NoAttrs): # or just a plain dataclass + center: tuple[float, float, float] + size: tuple[float, float, float] + +class CustomRegion(DesignRegion): + @property + def parameter_shape(self) -> int: + return 3 # example: 3 parameters + + def to_structure(self, params: np.ndarray) -> td.Structure: + # use params to drive geometry/material + geometry = td.Box(center=(0,0,0), size=(1,1,1)) + return td.Structure(geometry=geometry, medium=td.Medium(permittivity=1.0)) +``` + +Add a new metric +```python +from dataclasses import dataclass +import autograd.numpy as np +import tidy3d as td + +@dataclass +class PowerAtFreq(Metric): + monitor_name: str + target_idx: int + weight: float = 1.0 + + def evaluate(self, mnt_data: td.FluxData) -> float: + return float(mnt_data.flux.values[self.target_idx]) +``` + +Design decisions and rationale +------------------------------ +- Separation of concerns: geometry (regions), simulation assembly (device), orchestration (inverse design), scoring (metrics). +- Immutability: `DeviceSpec.get_simulation` uses `updated_copy` to avoid mutating the base. +- Batched orchestration: device names index batch results; uniqueness enforced at construction. +- Numeric compatibility: all math uses `autograd.numpy` for easy integration with gradient methods when supported by the backend. +- Simple helpers: `ones()` prevent shape mistakes when creating parameter vectors. + +Limitations and future work +--------------------------- +- `eps_bounds` not applied yet. Consider projections or relaxed binarization during optimization. +- `_unflatten_params` returns 1D segments; reshape happens inside region implementations. This is intentional but should be kept in mind when writing new regions. +- Optimizer loop is intentionally not baked-in; add a small `optimize()` convenience wrapper if a standard optimizer is adopted. +- Consider Protocol-based structural typing for plugin ergonomics once the contract stabilizes. + +Quick reference +--------------- +```python +# Create params +region_params = region.ones() +device_params = device_spec.ones() +all_params = inv.ones() + +# Build and run +sim = device_spec.get_simulation(device_params) +sim_data = device_spec.run_simulation(sim) +score = device_spec.get_metric(sim_data) + +# Multi-device +sims = inv.get_simulations(all_params) +batch = inv.run_simulations(sims) +objective = inv.get_metric(batch) +``` + + +How invdes2 differs from invdes (design philosophy) +--------------------------------------------------- + +At a glance +- **invdes2 (this module)**: minimal, composable, and backend-agnostic scaffold. Focuses on clear seams: regions → device build → batch run → metric aggregation. No optimizer baked-in by design. +- **invdes (existing plugin)**: feature-rich, opinionated, and end-to-end. Includes parameter transforms, penalties, initialization specs, optimizer with history/caching, and analysis helpers. + +Side‑by‑side comparison + +- **Core abstraction** + - invdes2: `DesignRegion` + `DeviceSpec` + `Metric` + `InverseDesign` (orchestration only). + - invdes: `DesignRegion` with transformations/penalties, `InverseDesign`/`InverseDesignMulti`, plus an `Optimizer` producing an `InverseDesignResult`. + +- **Parameter mapping to permittivity** + - invdes2: `TopologyDesignRegion.to_structure(params)` reshapes a 1D vector into a voxel grid; `eps_bounds` reserved for future use; no built-in transforms/penalties. + - invdes: rich pipeline: transformations (e.g., filter + projection) → material density → penalties (e.g., erosion/dilation) → permittivity mapping; validated gradients. + +- **Objective composition** + - invdes2: explicit `Metric` objects per monitor; `DeviceSpec.get_metric` does weighted sum; `InverseDesign.get_metric` aggregates across devices. + - invdes: user supplies a `post_process_fn(sim_data)` (or batch dict) that defines the scalar objective; penalties subtract from the objective; helpers/utilities available. + +- **Execution model** + - invdes2: `web.run` for single and `web.run_async` for batches; easy monkeypatch seams in tests for emulation. + - invdes: same web backends, but wrapped inside an optimizer loop with result persistence. + +- **Optimization** + - invdes2: no optimizer included; provides `_flatten_params/_unflatten_params` and simple helpers like `ones()` to integrate with external optimizers. + - invdes: built‑in Adam optimizer, gradient computation via autograd, progress display, callbacks, result caching to disk, continue/resume, plotting, and export helpers. + +- **State and results** + - invdes2: stateless orchestration; returns scalars and leaves bookkeeping to the caller. + - invdes: stateful `InverseDesignResult` with history of params, grads, states, and convenience accessors. + +- **Validation and typing** + - invdes2: plain dataclasses + simple ABCs; minimal validation; lightweight surface for rapid iteration. + - invdes: heavy validation with Pydantic, explicit error messaging, gradient checks for transforms/penalties. + +Why invdes2 exists +- **Iteration speed**: smaller conceptual surface for experimentation and research prototypes. +- **Composability**: metrics are small objects; devices are thin wrappers around base simulations; orchestration is explicit. +- **Testability**: seams for swapping execution (`run_simulation`, `run_simulations`) enable local/emulated unit tests without network calls. + +Trade‑offs +- Fewer batteries included: no transformations, penalties, optimizer history, or built‑in continuation. +- Less guardrail validation: places more responsibility on the optimization script for bounds, projections, and schedules. + +Migration guidance (invdes → invdes2) +- Map `post_process_fn` to one or more `Metric` implementations. If your post‑process touched multiple monitors, create a metric per monitor and combine via weights, or implement a custom metric that reads multiple monitors from `sim_data`. +- Convert transformations/penalties into your optimization loop: pre‑process the parameter vectors (e.g., filtering + projection) before passing to `to_structure`, and subtract penalty values inside your objective (or define a negative‑weight metric analogue). +- Replace `InverseDesignMulti` with multiple `DeviceSpec`s (one per simulation scenario); `InverseDesign` will build a batch and sum device objectives. +- Replace `Optimizer.run(...)` with an external loop that calls `get_objective(params)` or wraps it with autograd for gradients; use `_flatten_params/_unflatten_params` to interoperate with libraries. + +When to choose which +- Choose invdes2 when you want a lean, hackable kernel for orchestration and prefer to own the optimization and parameter processing. +- Choose invdes when you want a turnkey, validated pipeline with transformations, penalties, built‑in optimizers, and analysis utilities. + + diff --git a/tidy3d/plugins/invdes2/__init__.py b/tidy3d/plugins/invdes2/__init__.py new file mode 100644 index 0000000000..761aef148f --- /dev/null +++ b/tidy3d/plugins/invdes2/__init__.py @@ -0,0 +1,11 @@ +"""Public API for the `invdes2` inverse design scaffold.""" + +from __future__ import annotations + +from .design_region import TopologyDesignRegion +from .device_spec import DeviceSpec +from .inverse_design import InverseDesign +from .metric import FluxMetric +from .optimizer_spec import OptimizerSpec + +__all__ = ["DeviceSpec", "FluxMetric", "InverseDesign", "OptimizerSpec", "TopologyDesignRegion"] diff --git a/tidy3d/plugins/invdes2/design_region.py b/tidy3d/plugins/invdes2/design_region.py new file mode 100644 index 0000000000..765c64d07a --- /dev/null +++ b/tidy3d/plugins/invdes2/design_region.py @@ -0,0 +1,64 @@ +from __future__ import annotations + +from abc import abstractmethod +from dataclasses import dataclass +from typing import Union + +import autograd.numpy as np + +import tidy3d as td + + +@dataclass +class DesignRegion: + """Abstract parameterized geometry provider for inverse design. + + Implementations transform parameter arrays into concrete `td.Structure` + instances that will be appended to a base simulation. + """ + + @abstractmethod + def to_structure(self, params: np.ndarray) -> td.Structure: + """Return a `td.Structure` built from the provided parameters.""" + + @property + @abstractmethod + def parameter_shape(self) -> int: + """Return the (flattened) shape of the parameters for this design region.""" + + def ones(self, **kwargs) -> np.ndarray: + """Return an array of ones with the shape of the parameters for this design region.""" + return np.ones(self.parameter_shape, **kwargs) + + +@dataclass +class TopologyDesignRegion(DesignRegion): + """Design region as a pixellated permittivity grid.""" + + size: tuple[float, float, float] + center: tuple[float, float, float] + eps_bounds: tuple[float, float] + pixel_size: float + + @property + def shape_3d(self) -> tuple[int, int, int]: + """Return the shape of the parameters for this design region.""" + return tuple(int(np.ceil(size / self.pixel_size)) for size in self.size) + + @property + def parameter_shape(self) -> int: + """Return the shape of the parameters for this design region.""" + return int(np.prod(self.shape_3d)) + + def to_structure(self, params: np.ndarray) -> td.Structure: + """Return a `td.Structure` built from the provided parameters.""" + + geometry = td.Box(center=self.center, size=self.size) + + # TODO: add transformations + eps_data = params.reshape(self.shape_3d) + + return td.Structure.from_permittivity_array(geometry=geometry, eps_data=eps_data) + + +DesignRegionType = Union[TopologyDesignRegion] diff --git a/tidy3d/plugins/invdes2/device_spec.py b/tidy3d/plugins/invdes2/device_spec.py new file mode 100644 index 0000000000..79766d9285 --- /dev/null +++ b/tidy3d/plugins/invdes2/device_spec.py @@ -0,0 +1,94 @@ +from __future__ import annotations + +from dataclasses import dataclass + +import autograd.numpy as np + +import tidy3d as td +import tidy3d.web as web + +from .design_region import DesignRegionType +from .metric import MetricType + + +@dataclass +class DeviceSpec: + """Specification of a single device scenario for inverse design. + + Attributes + ---------- + simulation: + Base `td.Simulation` onto which parameterized structures are appended. + design_regions: + Ordered list of `DesignRegion` instances. Each must consume a + corresponding parameter array in `get_simulation`. + metrics: + List of `Metric` instances whose weighted sum forms the device's + objective contribution. + name: + Unique identifier for this device, used as the `task_name` and as the + key in batch results. + """ + + simulation: td.Simulation + design_regions: list[DesignRegionType] + metrics: list[MetricType] + name: str + + def get_simulation(self, params: list[np.ndarray]) -> td.Simulation: + """Construct a simulation by appending parameterized structures. + + Parameters + ---------- + params: + List of arrays matching `design_regions` order; each array is passed + to the corresponding region's `to_structure` to produce a structure. + + Returns + ------- + td.Simulation + A new simulation with appended structures. + """ + structures = list(self.simulation.structures) + for param, design_region in zip(params, self.design_regions): + structure = design_region.to_structure(param) + structures.append(structure) + return self.simulation.updated_copy(structures=structures) + + def run_simulation(self, simulation: td.Simulation) -> web.SimulationData: + """Run the simulation via Tidy3D Web and return results.""" + return web.run(simulation, task_name=self.name) + + def get_metric(self, sim_data: web.SimulationData) -> float: + """Compute the weighted sum of metrics for this device. + + Parameters + ---------- + sim_data: + Simulation results to be consumed by each metric. + + Returns + ------- + float + Weighted sum of metric values. + """ + value = 0.0 + for metric in self.metrics: + mnt_data = sim_data[metric.monitor_name] + value = value + metric.weight * metric.evaluate(mnt_data) + return value + + def get_objective(self, params: list[np.ndarray]) -> float: + """Build, run, and score this device for the given parameters.""" + sim = self.get_simulation(params) + sim_data = self.run_simulation(sim) + return self.get_metric(sim_data) + + @property + def parameter_shape(self) -> list[int]: + """Return the shape of the parameters for each design region.""" + return [design_region.parameter_shape for design_region in self.design_regions] + + def ones(self, **kwargs) -> list[np.ndarray]: + """Return a list of arrays of ones with the shape of the parameters for each design region.""" + return [design_region.ones(**kwargs) for design_region in self.design_regions] diff --git a/tidy3d/plugins/invdes2/inverse_design.py b/tidy3d/plugins/invdes2/inverse_design.py new file mode 100644 index 0000000000..5a88208ade --- /dev/null +++ b/tidy3d/plugins/invdes2/inverse_design.py @@ -0,0 +1,96 @@ +"""Inverse design orchestration utilities. + +Coordinates multiple device scenarios to evaluate a scalar objective by +constructing simulations from parameter vectors, submitting them to Tidy3D Web +in batch, and aggregating metric values across devices. Compatible with +autograd-enabled Tidy3D Web. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass + +import autograd.numpy as np + +import tidy3d as td +import tidy3d.web as web + +from .device_spec import DeviceSpec +from .optimizer_spec import OptimizerSpec + + +@dataclass +class InverseDesign: + optimizer_spec: OptimizerSpec + device_specs: list[DeviceSpec] + + def __post_init__(self) -> None: + """Validate that all `DeviceSpec` names are unique.""" + names = [device_spec.name for device_spec in self.device_specs] + if len(set(names)) != len(names): + duplicates = sorted({name for name in names if names.count(name) > 1}) + raise ValueError( + "All DeviceSpec names must be unique; duplicates found: " + ", ".join(duplicates) + ) + + def get_simulations(self, params: list[list[np.ndarray]]) -> dict[str, td.Simulation]: + """Build simulations for each device from parameter vectors.""" + simulations = {} + for param, device_spec in zip(params, self.device_specs): + simulation = device_spec.get_simulation(param) + simulations[device_spec.name] = simulation + return simulations + + def run_simulations(self, sims: dict[str, td.Simulation]) -> web.BatchData: + """Execute multiple simulations concurrently via Tidy3D Web.""" + return web.run_async(sims) + + def get_metric(self, batch_data: Mapping[str, td.SimulationData]) -> float: + """Aggregate metrics across all devices.""" + value = 0.0 + for device_spec in self.device_specs: + sim_data = batch_data[device_spec.name] + value = value + device_spec.get_metric(sim_data) + return value + + def get_objective(self, params: list[list[np.ndarray]]) -> float: + """Evaluate the scalar objective across all devices given parameters.""" + sims = self.get_simulations(params) + batch_data = self.run_simulations(sims) + return self.get_metric(batch_data) + + @property + def parameter_shape(self) -> list[list[int]]: + """Return the shape of the parameters for each device.""" + return [device_spec.parameter_shape for device_spec in self.device_specs] + + def ones(self, **kwargs) -> list[list[np.ndarray]]: + """Return a list of arrays of ones with the shape of the parameters for each design region.""" + return [device_spec.ones(**kwargs) for device_spec in self.device_specs] + + def _flatten_params(self, params: list[list[np.ndarray]]) -> np.ndarray: + """Flatten the parameters for each device.""" + flattened_params = [] + for device_params in params: + for design_region_params in device_params: + flattened_params.append(design_region_params.flatten()) + return np.concatenate(flattened_params) + + def _unflatten_params(self, params: np.ndarray) -> list[list[np.ndarray]]: + """Unflatten the parameters by slicing using per-region sizes. + + The method reverses `_flatten_params`, reconstructing a nested list + aligned with `device_specs[i].design_regions[j]`. + """ + result: list[list[np.ndarray]] = [] + cursor = 0 + for device_spec in self.device_specs: + device_params: list[np.ndarray] = [] + for region in device_spec.design_regions: + size = int(region.parameter_shape) + segment = params[cursor : cursor + size] + device_params.append(segment) + cursor += size + result.append(device_params) + return result diff --git a/tidy3d/plugins/invdes2/metric.py b/tidy3d/plugins/invdes2/metric.py new file mode 100644 index 0000000000..570a0b1c2b --- /dev/null +++ b/tidy3d/plugins/invdes2/metric.py @@ -0,0 +1,39 @@ +from __future__ import annotations + +from abc import abstractmethod +from dataclasses import dataclass +from typing import Union + +import autograd.numpy as np + +import tidy3d as td + + +@dataclass +class Metric: + """Abstract base class for simulation-derived objective terms. + + Attributes + ---------- + weight: + Scalar multiplied with the metric value during aggregation. + """ + + monitor_name: str + weight: float = 1.0 + + @abstractmethod + def evaluate(self, mnt_data: td.MonitorData) -> float: + """Return a scalar score computed from `monitor_data`.""" + + +@dataclass +class FluxMetric(Metric): + """Metric based on the flux of a monitor.""" + + def evaluate(self, mnt_data: td.FluxData) -> float: + """Return the flux of the monitor.""" + return np.sum(mnt_data.flux.values) + + +MetricType = Union[FluxMetric] diff --git a/tidy3d/plugins/invdes2/optimizer_spec.py b/tidy3d/plugins/invdes2/optimizer_spec.py new file mode 100644 index 0000000000..db27d4ca79 --- /dev/null +++ b/tidy3d/plugins/invdes2/optimizer_spec.py @@ -0,0 +1,11 @@ +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass +class OptimizerSpec: + """Hyperparameters describing the optimization loop to be used externally.""" + + num_steps: int + learning_rate: float