From 4f6f53c90a19d4d5041cc10104db04dfdbc7431e Mon Sep 17 00:00:00 2001
From: Tyler Hughes <tyler@flexcompute.com>
Date: Thu, 2 Oct 2025 15:19:07 -0400
Subject: [PATCH 1/8] human attempt

---
 docs/faq                                 |  2 +-
 docs/notebooks                           |  2 +-
 tidy3d/plugins/invdes2/__init__.py       |  9 ++++++
 tidy3d/plugins/invdes2/design_region.py  | 15 +++++++++
 tidy3d/plugins/invdes2/device_spec.py    | 41 ++++++++++++++++++++++++
 tidy3d/plugins/invdes2/inverse_design.py | 41 ++++++++++++++++++++++++
 tidy3d/plugins/invdes2/metric.py         | 15 +++++++++
 tidy3d/plugins/invdes2/optimizer_spec.py |  9 ++++++
 tidy3d/plugins/invdes2/test.py           |  0
 9 files changed, 132 insertions(+), 2 deletions(-)
 create mode 100644 tidy3d/plugins/invdes2/__init__.py
 create mode 100644 tidy3d/plugins/invdes2/design_region.py
 create mode 100644 tidy3d/plugins/invdes2/device_spec.py
 create mode 100644 tidy3d/plugins/invdes2/inverse_design.py
 create mode 100644 tidy3d/plugins/invdes2/metric.py
 create mode 100644 tidy3d/plugins/invdes2/optimizer_spec.py
 create mode 100644 tidy3d/plugins/invdes2/test.py

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/tidy3d/plugins/invdes2/__init__.py b/tidy3d/plugins/invdes2/__init__.py
new file mode 100644
index 0000000000..af2720c3ff
--- /dev/null
+++ b/tidy3d/plugins/invdes2/__init__.py
@@ -0,0 +1,9 @@
+from __future__ import annotations
+
+from .design_region import DesignRegion
+from .device_spec import DeviceSpec
+from .inverse_design import InverseDesign
+from .metric import Metric
+from .optimizer_spec import OptimizerSpec
+
+__all__ = ["DesignRegion", "DeviceSpec", "InverseDesign", "Metric", "OptimizerSpec"]
diff --git a/tidy3d/plugins/invdes2/design_region.py b/tidy3d/plugins/invdes2/design_region.py
new file mode 100644
index 0000000000..f9c454f155
--- /dev/null
+++ b/tidy3d/plugins/invdes2/design_region.py
@@ -0,0 +1,15 @@
+from __future__ import annotations
+
+from abc import abstractmethod
+from dataclasses import dataclass
+
+import autograd.numpy as np
+
+import tidy3d as td
+
+
+@dataclass
+class DesignRegion:
+    @abstractmethod
+    def to_structure(self, params: np.ndarray) -> td.Structure:
+        pass
diff --git a/tidy3d/plugins/invdes2/device_spec.py b/tidy3d/plugins/invdes2/device_spec.py
new file mode 100644
index 0000000000..ad9b6ac5e0
--- /dev/null
+++ b/tidy3d/plugins/invdes2/device_spec.py
@@ -0,0 +1,41 @@
+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 DesignRegion
+from .metric import Metric
+
+
+@dataclass
+class DeviceSpec:
+    simulation: td.Simulation
+    design_regions: list[DesignRegion]
+    metrics: list[Metric]
+    name: str
+
+    def get_simulation(self, params: list[np.ndarray]) -> td.Simulation:
+        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:
+        return web.run(simulation, task_name=self.name)
+
+    def get_metric(self, sim_data: web.SimulationData) -> float:
+        """Get metric from running the simulation"""
+        value = 0.0
+        for metric in self.metrics:
+            value = value + metric.evaluate(sim_data)
+        return value
+
+    def get_objective(self, params: list[np.ndarray]) -> float:
+        sim = self.get_simulation(params)
+        sim_data = self.run_simulation(sim)
+        return self.get_metric(sim_data)
diff --git a/tidy3d/plugins/invdes2/inverse_design.py b/tidy3d/plugins/invdes2/inverse_design.py
new file mode 100644
index 0000000000..94503100fb
--- /dev/null
+++ b/tidy3d/plugins/invdes2/inverse_design.py
@@ -0,0 +1,41 @@
+from __future__ import annotations
+
+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]
+
+    # TODO: ensure that all device specs have unique names
+
+    def get_simulations(self, params: list[list[np.ndarray]]) -> dict[str, td.Simulation]:
+        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:
+        return web.run_async(sims)
+
+    def get_metric(self, batch_data: web.BatchData) -> float:
+        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:
+        sims = self.get_simulations(params)
+        batch_data = self.run_simulations(sims)
+        return self.get_metric(batch_data)
diff --git a/tidy3d/plugins/invdes2/metric.py b/tidy3d/plugins/invdes2/metric.py
new file mode 100644
index 0000000000..7212d6dfb1
--- /dev/null
+++ b/tidy3d/plugins/invdes2/metric.py
@@ -0,0 +1,15 @@
+from __future__ import annotations
+
+from abc import abstractmethod
+from dataclasses import dataclass
+
+import tidy3d.web as web
+
+
+@dataclass
+class Metric:
+    weight: float = 1.0
+
+    @abstractmethod
+    def evaluate(self, sim_data: web.SimulationData) -> float:
+        pass
diff --git a/tidy3d/plugins/invdes2/optimizer_spec.py b/tidy3d/plugins/invdes2/optimizer_spec.py
new file mode 100644
index 0000000000..8471e4a663
--- /dev/null
+++ b/tidy3d/plugins/invdes2/optimizer_spec.py
@@ -0,0 +1,9 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass
+class OptimizerSpec:
+    learning_rate: float
+    num_steps: int
diff --git a/tidy3d/plugins/invdes2/test.py b/tidy3d/plugins/invdes2/test.py
new file mode 100644
index 0000000000..e69de29bb2

From 9f6885147e8392870cf1b879cdce40b4d6dd0027 Mon Sep 17 00:00:00 2001
From: Tyler Hughes <tyler@flexcompute.com>
Date: Thu, 2 Oct 2025 15:43:55 -0400
Subject: [PATCH 2/8] basic structure

---
 tests/test_plugins/test_invdes2.py       | 112 +++++++++++++++++++++++
 tidy3d/plugins/invdes2/__init__.py       |   2 +
 tidy3d/plugins/invdes2/design_region.py  |   8 +-
 tidy3d/plugins/invdes2/device_spec.py    |  47 +++++++++-
 tidy3d/plugins/invdes2/inverse_design.py |  21 ++++-
 tidy3d/plugins/invdes2/metric.py         |  10 +-
 tidy3d/plugins/invdes2/optimizer_spec.py |   2 +
 tidy3d/plugins/invdes2/test.py           |   0
 8 files changed, 197 insertions(+), 5 deletions(-)
 create mode 100644 tests/test_plugins/test_invdes2.py
 delete mode 100644 tidy3d/plugins/invdes2/test.py

diff --git a/tests/test_plugins/test_invdes2.py b/tests/test_plugins/test_invdes2.py
new file mode 100644
index 0000000000..bdf0b7ddf7
--- /dev/null
+++ b/tests/test_plugins/test_invdes2.py
@@ -0,0 +1,112 @@
+from __future__ import annotations
+
+import autograd.numpy as np
+import pytest
+
+import tidy3d as td
+import tidy3d.web as web
+from tidy3d.plugins.invdes2 import DesignRegion, DeviceSpec, InverseDesign, Metric, OptimizerSpec
+
+
+class DummyRegion(DesignRegion):
+    def to_structure(self, params: np.ndarray) -> td.Structure:
+        # Minimal structure: a box dielectric with constant size, position tuned by params
+        center = (float(params[0]), float(params[1]), float(params[2]))
+        return td.Structure(
+            geometry=td.Box(center=center, size=(1.0, 1.0, 1.0)), medium=td.Medium(permittivity=2.0)
+        )
+
+
+class SumPowerMetric(Metric):
+    def evaluate(self, sim_data: web.SimulationData) -> float:
+        # Return a scalar value recorded on the sim_data mock
+        return float(sim_data._value)
+
+
+def make_base_simulation() -> td.Simulation:
+    return 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=(),
+        sources=(),
+        boundary_spec=td.BoundarySpec.all_sides(boundary=td.PML()),
+        medium=td.Medium(permittivity=1.0),
+    )
+
+
+def test_device_spec_get_simulation_builds_structures():
+    base = make_base_simulation()
+    region1 = DummyRegion()
+    region2 = DummyRegion()
+    spec = DeviceSpec(simulation=base, design_regions=[region1, region2], metrics=[], name="d1")
+    params = [np.array([0.0, 0.0, 0.0]), np.array([1.0, 1.0, 1.0])]
+    sim = spec.get_simulation(params)
+    assert len(sim.structures) == len(base.structures) + 2
+
+
+def test_device_spec_metric_weighting_and_objective(monkeypatch):
+    base = make_base_simulation()
+    region = DummyRegion()
+    m1 = SumPowerMetric(weight=2.0)
+    m2 = SumPowerMetric(weight=3.0)
+    spec = DeviceSpec(simulation=base, design_regions=[region], metrics=[m1, m2], name="d1")
+
+    # Monkeypatch run_simulation to avoid network; provide a dummy value
+    class DummyData:
+        def __init__(self, value: float) -> None:
+            self._value = value
+
+    monkeypatch.setattr(spec, "run_simulation", lambda sim: DummyData(5.0))
+    params = [np.array([0.0, 0.0, 0.0])]
+    val = spec.get_objective(params)
+    assert val == pytest.approx(2.0 * 5.0 + 3.0 * 5.0)
+
+
+def test_inverse_design_unique_names_validation():
+    base = make_base_simulation()
+    region = DummyRegion()
+    spec1 = DeviceSpec(simulation=base, design_regions=[region], metrics=[], name="dup")
+    spec2 = DeviceSpec(simulation=base, design_regions=[region], metrics=[], name="dup")
+    with pytest.raises(ValueError):
+        InverseDesign(
+            optimizer_spec=OptimizerSpec(learning_rate=0.1, num_steps=1),
+            device_specs=[spec1, spec2],
+        )
+
+
+def test_inverse_design_builds_and_aggregates(monkeypatch):
+    base = make_base_simulation()
+    region = DummyRegion()
+    m = SumPowerMetric(weight=1.0)
+    s1 = DeviceSpec(simulation=base, design_regions=[region], metrics=[m], name="a")
+    s2 = DeviceSpec(simulation=base, design_regions=[region], metrics=[m], name="b")
+
+    inv = InverseDesign(
+        optimizer_spec=OptimizerSpec(learning_rate=0.1, num_steps=1), device_specs=[s1, s2]
+    )
+
+    # Monkeypatch run_simulation on each DeviceSpec via batch submission monkeypatch
+    class DummyData:
+        def __init__(self, value: float) -> None:
+            self._value = value
+
+    def fake_run_async(sims_dict):
+        # sims_dict is name->simulation mapping; return name->DummyData with unique values
+        return {name: DummyData(1.0 if name == "a" else 2.0) for name in sims_dict.keys()}
+
+    monkeypatch.setattr(web, "run_async", fake_run_async)
+
+    params = [[np.array([0.0, 0.0, 0.0])], [np.array([1.0, 1.0, 1.0])]]
+    sims = inv.get_simulations(params)
+    assert set(sims.keys()) == {"a", "b"}
+
+    batch = inv.run_simulations(sims)
+    assert isinstance(batch, dict) and set(batch.keys()) == {"a", "b"}
+
+    total = inv.get_metric(batch)
+    assert total == pytest.approx(1.0 + 2.0)
+
+    obj = inv.get_objective(params)
+    assert obj == pytest.approx(1.0 + 2.0)
diff --git a/tidy3d/plugins/invdes2/__init__.py b/tidy3d/plugins/invdes2/__init__.py
index af2720c3ff..508ae6819a 100644
--- a/tidy3d/plugins/invdes2/__init__.py
+++ b/tidy3d/plugins/invdes2/__init__.py
@@ -1,3 +1,5 @@
+"""Public API for the `invdes2` inverse design scaffold."""
+
 from __future__ import annotations
 
 from .design_region import DesignRegion
diff --git a/tidy3d/plugins/invdes2/design_region.py b/tidy3d/plugins/invdes2/design_region.py
index f9c454f155..21f816f532 100644
--- a/tidy3d/plugins/invdes2/design_region.py
+++ b/tidy3d/plugins/invdes2/design_region.py
@@ -10,6 +10,12 @@
 
 @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:
-        pass
+        """Return a `td.Structure` built from the provided parameters."""
diff --git a/tidy3d/plugins/invdes2/device_spec.py b/tidy3d/plugins/invdes2/device_spec.py
index ad9b6ac5e0..2a84cfbc4e 100644
--- a/tidy3d/plugins/invdes2/device_spec.py
+++ b/tidy3d/plugins/invdes2/device_spec.py
@@ -13,12 +13,42 @@
 
 @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[DesignRegion]
     metrics: list[Metric]
     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)
@@ -26,16 +56,29 @@ def get_simulation(self, params: list[np.ndarray]) -> td.Simulation:
         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:
-        """Get metric from running the simulation"""
+        """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:
-            value = value + metric.evaluate(sim_data)
+            value = value + getattr(metric, "weight", 1.0) * metric.evaluate(sim_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)
diff --git a/tidy3d/plugins/invdes2/inverse_design.py b/tidy3d/plugins/invdes2/inverse_design.py
index 94503100fb..3aa77e32fa 100644
--- a/tidy3d/plugins/invdes2/inverse_design.py
+++ b/tidy3d/plugins/invdes2/inverse_design.py
@@ -1,3 +1,11 @@
+"""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 dataclasses import dataclass
@@ -16,9 +24,17 @@ class InverseDesign:
     optimizer_spec: OptimizerSpec
     device_specs: list[DeviceSpec]
 
-    # TODO: ensure that all device specs have unique names
+    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)
@@ -26,9 +42,11 @@ def get_simulations(self, params: list[list[np.ndarray]]) -> dict[str, td.Simula
         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: web.BatchData) -> float:
+        """Aggregate metrics across all devices."""
         value = 0.0
         for device_spec in self.device_specs:
             sim_data = batch_data[device_spec.name]
@@ -36,6 +54,7 @@ def get_metric(self, batch_data: web.BatchData) -> float:
         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)
diff --git a/tidy3d/plugins/invdes2/metric.py b/tidy3d/plugins/invdes2/metric.py
index 7212d6dfb1..db25fcefbd 100644
--- a/tidy3d/plugins/invdes2/metric.py
+++ b/tidy3d/plugins/invdes2/metric.py
@@ -8,8 +8,16 @@
 
 @dataclass
 class Metric:
+    """Abstract base class for simulation-derived objective terms.
+
+    Attributes
+    ----------
+    weight:
+        Scalar multiplied with the metric value during aggregation.
+    """
+
     weight: float = 1.0
 
     @abstractmethod
     def evaluate(self, sim_data: web.SimulationData) -> float:
-        pass
+        """Return a scalar score computed from `sim_data`."""
diff --git a/tidy3d/plugins/invdes2/optimizer_spec.py b/tidy3d/plugins/invdes2/optimizer_spec.py
index 8471e4a663..dc4c181882 100644
--- a/tidy3d/plugins/invdes2/optimizer_spec.py
+++ b/tidy3d/plugins/invdes2/optimizer_spec.py
@@ -5,5 +5,7 @@
 
 @dataclass
 class OptimizerSpec:
+    """Hyperparameters describing the optimization loop to be used externally."""
+
     learning_rate: float
     num_steps: int
diff --git a/tidy3d/plugins/invdes2/test.py b/tidy3d/plugins/invdes2/test.py
deleted file mode 100644
index e69de29bb2..0000000000

From b897b5652f0e63502e3e135a5e86fb350930a5ef Mon Sep 17 00:00:00 2001
From: Tyler Hughes <tyler@flexcompute.com>
Date: Thu, 2 Oct 2025 17:29:06 -0400
Subject: [PATCH 3/8] more human work

---
 tests/test_plugins/test_invdes2.py       | 216 ++++++++++++++---------
 tidy3d/plugins/invdes2/__init__.py       |   6 +-
 tidy3d/plugins/invdes2/design_region.py  |  41 +++++
 tidy3d/plugins/invdes2/device_spec.py    |  16 +-
 tidy3d/plugins/invdes2/inverse_design.py |  31 ++++
 tidy3d/plugins/invdes2/metric.py         |  22 ++-
 tidy3d/plugins/invdes2/optimizer_spec.py |   2 +-
 7 files changed, 236 insertions(+), 98 deletions(-)

diff --git a/tests/test_plugins/test_invdes2.py b/tests/test_plugins/test_invdes2.py
index bdf0b7ddf7..d0237c09fb 100644
--- a/tests/test_plugins/test_invdes2.py
+++ b/tests/test_plugins/test_invdes2.py
@@ -5,108 +5,152 @@
 
 import tidy3d as td
 import tidy3d.web as web
-from tidy3d.plugins.invdes2 import DesignRegion, DeviceSpec, InverseDesign, Metric, OptimizerSpec
-
-
-class DummyRegion(DesignRegion):
-    def to_structure(self, params: np.ndarray) -> td.Structure:
-        # Minimal structure: a box dielectric with constant size, position tuned by params
-        center = (float(params[0]), float(params[1]), float(params[2]))
-        return td.Structure(
-            geometry=td.Box(center=center, size=(1.0, 1.0, 1.0)), medium=td.Medium(permittivity=2.0)
-        )
-
-
-class SumPowerMetric(Metric):
-    def evaluate(self, sim_data: web.SimulationData) -> float:
-        # Return a scalar value recorded on the sim_data mock
-        return float(sim_data._value)
-
-
-def make_base_simulation() -> td.Simulation:
-    return 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=(),
-        sources=(),
-        boundary_spec=td.BoundarySpec.all_sides(boundary=td.PML()),
-        medium=td.Medium(permittivity=1.0),
+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"
+)
 
-def test_device_spec_get_simulation_builds_structures():
-    base = make_base_simulation()
-    region1 = DummyRegion()
-    region2 = DummyRegion()
-    spec = DeviceSpec(simulation=base, design_regions=[region1, region2], metrics=[], name="d1")
-    params = [np.array([0.0, 0.0, 0.0]), np.array([1.0, 1.0, 1.0])]
-    sim = spec.get_simulation(params)
-    assert len(sim.structures) == len(base.structures) + 2
+device_spec2 = DeviceSpec(
+    simulation=sim_base, design_regions=design_regions, metrics=metrics, name="d2"
+)
 
+device_specs = [device_spec1, device_spec2]
 
-def test_device_spec_metric_weighting_and_objective(monkeypatch):
-    base = make_base_simulation()
-    region = DummyRegion()
-    m1 = SumPowerMetric(weight=2.0)
-    m2 = SumPowerMetric(weight=3.0)
-    spec = DeviceSpec(simulation=base, design_regions=[region], metrics=[m1, m2], name="d1")
+optimizer_spec = OptimizerSpec(learning_rate=0.1, num_steps=1)
 
-    # Monkeypatch run_simulation to avoid network; provide a dummy value
-    class DummyData:
-        def __init__(self, value: float) -> None:
-            self._value = value
+invdes = InverseDesign(optimizer_spec=optimizer_spec, device_specs=device_specs)
 
-    monkeypatch.setattr(spec, "run_simulation", lambda sim: DummyData(5.0))
-    params = [np.array([0.0, 0.0, 0.0])]
-    val = spec.get_objective(params)
-    assert val == pytest.approx(2.0 * 5.0 + 3.0 * 5.0)
 
+def test_parameter_shapes():
+    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_inverse_design_unique_names_validation():
-    base = make_base_simulation()
-    region = DummyRegion()
-    spec1 = DeviceSpec(simulation=base, design_regions=[region], metrics=[], name="dup")
-    spec2 = DeviceSpec(simulation=base, design_regions=[region], metrics=[], name="dup")
-    with pytest.raises(ValueError):
-        InverseDesign(
-            optimizer_spec=OptimizerSpec(learning_rate=0.1, num_steps=1),
-            device_specs=[spec1, spec2],
-        )
 
+def test_flatten_unflatten_params():
+    params = [
+        [np.ones_like(design_region.shape_3d) for design_region in device_spec.design_regions]
+        for device_spec in invdes.device_specs
+    ]
+    flat = invdes._flatten_params(params)
+    restored = invdes._unflatten_params(flat)
+    flat2 = invdes._flatten_params(restored)
 
-def test_inverse_design_builds_and_aggregates(monkeypatch):
-    base = make_base_simulation()
-    region = DummyRegion()
-    m = SumPowerMetric(weight=1.0)
-    s1 = DeviceSpec(simulation=base, design_regions=[region], metrics=[m], name="a")
-    s2 = DeviceSpec(simulation=base, design_regions=[region], metrics=[m], name="b")
+    assert np.allclose(flat, flat2)
+
+
+def test_design_region_to_structure():
+    for design_region in design_regions:
+        params = np.ones_like(design_region.shape_3d)
+        _ = design_region.to_structure(params)
+
+
+def test_device_spec_get_simulation():
+    for device_spec in device_specs:
+        params = [
+            np.ones_like(design_region.shape_3d) for design_region in device_spec.design_regions
+        ]
+        sim = device_spec.get_simulation(params)
+
+
+def test_invdes_get_simulations():
+    params = [
+        [np.ones_like(design_region.shape_3d) for design_region in device_spec.design_regions]
+        for device_spec in invdes.device_specs
+    ]
+    sims = invdes.get_simulations(params)
+
+    assert len(sims) == len(invdes.device_specs)
+    assert {sims.keys()} == {device_spec.name for device_spec in invdes.device_specs}
 
-    inv = InverseDesign(
-        optimizer_spec=OptimizerSpec(learning_rate=0.1, num_steps=1), device_specs=[s1, s2]
-    )
 
-    # Monkeypatch run_simulation on each DeviceSpec via batch submission monkeypatch
-    class DummyData:
-        def __init__(self, value: float) -> None:
-            self._value = value
+def test_metric_evaluate():
+    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 fake_run_async(sims_dict):
-        # sims_dict is name->simulation mapping; return name->DummyData with unique values
-        return {name: DummyData(1.0 if name == "a" else 2.0) for name in sims_dict.keys()}
 
-    monkeypatch.setattr(web, "run_async", fake_run_async)
+def test_device_spec_get_metric():
+    for device_spec in device_specs:
+        for metric in device_spec.metrics:
+            mnt_data = sim_data_base[metric.monitor_name]
+            val = device_spec.get_metric(mnt_data, metric)
+            assert not np.allclose(val, 0.0)
 
-    params = [[np.array([0.0, 0.0, 0.0])], [np.array([1.0, 1.0, 1.0])]]
-    sims = inv.get_simulations(params)
-    assert set(sims.keys()) == {"a", "b"}
 
-    batch = inv.run_simulations(sims)
-    assert isinstance(batch, dict) and set(batch.keys()) == {"a", "b"}
+def test_invdes_get_metric():
+    batch_data = web.BatchData(
+        {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():
+    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()
+        },
+    )
 
-    total = inv.get_metric(batch)
-    assert total == pytest.approx(1.0 + 2.0)
 
-    obj = inv.get_objective(params)
-    assert obj == pytest.approx(1.0 + 2.0)
+def test_objective_function(use_emulated):
+    params = [
+        [np.ones_like(design_region.shape_3d) for design_region in device_spec.design_regions]
+        for device_spec in invdes.device_specs
+    ]
+    val = invdes.get_objective(params)
+    assert not np.allclose(val, 0.0)
diff --git a/tidy3d/plugins/invdes2/__init__.py b/tidy3d/plugins/invdes2/__init__.py
index 508ae6819a..761aef148f 100644
--- a/tidy3d/plugins/invdes2/__init__.py
+++ b/tidy3d/plugins/invdes2/__init__.py
@@ -2,10 +2,10 @@
 
 from __future__ import annotations
 
-from .design_region import DesignRegion
+from .design_region import TopologyDesignRegion
 from .device_spec import DeviceSpec
 from .inverse_design import InverseDesign
-from .metric import Metric
+from .metric import FluxMetric
 from .optimizer_spec import OptimizerSpec
 
-__all__ = ["DesignRegion", "DeviceSpec", "InverseDesign", "Metric", "OptimizerSpec"]
+__all__ = ["DeviceSpec", "FluxMetric", "InverseDesign", "OptimizerSpec", "TopologyDesignRegion"]
diff --git a/tidy3d/plugins/invdes2/design_region.py b/tidy3d/plugins/invdes2/design_region.py
index 21f816f532..d47fe4ea05 100644
--- a/tidy3d/plugins/invdes2/design_region.py
+++ b/tidy3d/plugins/invdes2/design_region.py
@@ -2,6 +2,7 @@
 
 from abc import abstractmethod
 from dataclasses import dataclass
+from typing import Union
 
 import autograd.numpy as np
 
@@ -19,3 +20,43 @@ class DesignRegion:
     @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."""
+
+
+@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, eps_bounds=self.eps_bounds
+        )
+
+
+DesignRegionType = Union[TopologyDesignRegion]
diff --git a/tidy3d/plugins/invdes2/device_spec.py b/tidy3d/plugins/invdes2/device_spec.py
index 2a84cfbc4e..5729e7201b 100644
--- a/tidy3d/plugins/invdes2/device_spec.py
+++ b/tidy3d/plugins/invdes2/device_spec.py
@@ -7,8 +7,8 @@
 import tidy3d as td
 import tidy3d.web as web
 
-from .design_region import DesignRegion
-from .metric import Metric
+from .design_region import DesignRegionType
+from .metric import MetricType
 
 
 @dataclass
@@ -31,8 +31,8 @@ class DeviceSpec:
     """
 
     simulation: td.Simulation
-    design_regions: list[DesignRegion]
-    metrics: list[Metric]
+    design_regions: list[DesignRegionType]
+    metrics: list[MetricType]
     name: str
 
     def get_simulation(self, params: list[np.ndarray]) -> td.Simulation:
@@ -74,7 +74,8 @@ def get_metric(self, sim_data: web.SimulationData) -> float:
         """
         value = 0.0
         for metric in self.metrics:
-            value = value + getattr(metric, "weight", 1.0) * metric.evaluate(sim_data)
+            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:
@@ -82,3 +83,8 @@ def get_objective(self, params: list[np.ndarray]) -> float:
         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]
diff --git a/tidy3d/plugins/invdes2/inverse_design.py b/tidy3d/plugins/invdes2/inverse_design.py
index 3aa77e32fa..220c254ab0 100644
--- a/tidy3d/plugins/invdes2/inverse_design.py
+++ b/tidy3d/plugins/invdes2/inverse_design.py
@@ -58,3 +58,34 @@ def get_objective(self, params: list[list[np.ndarray]]) -> float:
         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 _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
index db25fcefbd..570a0b1c2b 100644
--- a/tidy3d/plugins/invdes2/metric.py
+++ b/tidy3d/plugins/invdes2/metric.py
@@ -2,8 +2,11 @@
 
 from abc import abstractmethod
 from dataclasses import dataclass
+from typing import Union
 
-import tidy3d.web as web
+import autograd.numpy as np
+
+import tidy3d as td
 
 
 @dataclass
@@ -16,8 +19,21 @@ class Metric:
         Scalar multiplied with the metric value during aggregation.
     """
 
+    monitor_name: str
     weight: float = 1.0
 
     @abstractmethod
-    def evaluate(self, sim_data: web.SimulationData) -> float:
-        """Return a scalar score computed from `sim_data`."""
+    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
index dc4c181882..db27d4ca79 100644
--- a/tidy3d/plugins/invdes2/optimizer_spec.py
+++ b/tidy3d/plugins/invdes2/optimizer_spec.py
@@ -7,5 +7,5 @@
 class OptimizerSpec:
     """Hyperparameters describing the optimization loop to be used externally."""
 
-    learning_rate: float
     num_steps: int
+    learning_rate: float

From e51070166aa770be55924805eb9df8ce16d6ce79 Mon Sep 17 00:00:00 2001
From: Tyler Hughes <tyler@flexcompute.com>
Date: Fri, 3 Oct 2025 16:04:37 -0400
Subject: [PATCH 4/8] pass ma tests

---
 tests/test_plugins/test_invdes2.py      | 27 ++++++++++++-------------
 tidy3d/plugins/invdes2/design_region.py |  4 +---
 tidy3d/plugins/invdes2/device_spec.py   |  2 +-
 3 files changed, 15 insertions(+), 18 deletions(-)

diff --git a/tests/test_plugins/test_invdes2.py b/tests/test_plugins/test_invdes2.py
index d0237c09fb..c7c50f1721 100644
--- a/tests/test_plugins/test_invdes2.py
+++ b/tests/test_plugins/test_invdes2.py
@@ -4,7 +4,6 @@
 import pytest
 
 import tidy3d as td
-import tidy3d.web as web
 from tidy3d.plugins.invdes2 import (
     DeviceSpec,
     FluxMetric,
@@ -66,7 +65,10 @@ def test_parameter_shapes():
 
 def test_flatten_unflatten_params():
     params = [
-        [np.ones_like(design_region.shape_3d) for design_region in device_spec.design_regions]
+        [
+            np.ones_like(design_region.parameter_shape)
+            for design_region in device_spec.design_regions
+        ]
         for device_spec in invdes.device_specs
     ]
     flat = invdes._flatten_params(params)
@@ -78,27 +80,28 @@ def test_flatten_unflatten_params():
 
 def test_design_region_to_structure():
     for design_region in design_regions:
-        params = np.ones_like(design_region.shape_3d)
+        params = np.ones(design_region.parameter_shape)
         _ = design_region.to_structure(params)
 
 
 def test_device_spec_get_simulation():
     for device_spec in device_specs:
         params = [
-            np.ones_like(design_region.shape_3d) for design_region in device_spec.design_regions
+            np.ones(design_region.parameter_shape) for design_region in device_spec.design_regions
         ]
         sim = device_spec.get_simulation(params)
+        assert len(sim.structures) == len(sim_base.structures) + len(device_spec.design_regions)
 
 
 def test_invdes_get_simulations():
     params = [
-        [np.ones_like(design_region.shape_3d) for design_region in device_spec.design_regions]
+        [np.ones(design_region.parameter_shape) for design_region in device_spec.design_regions]
         for device_spec in invdes.device_specs
     ]
     sims = invdes.get_simulations(params)
 
     assert len(sims) == len(invdes.device_specs)
-    assert {sims.keys()} == {device_spec.name for device_spec in invdes.device_specs}
+    assert set(sims.keys()) == {device_spec.name for device_spec in invdes.device_specs}
 
 
 def test_metric_evaluate():
@@ -110,16 +113,12 @@ def test_metric_evaluate():
 
 def test_device_spec_get_metric():
     for device_spec in device_specs:
-        for metric in device_spec.metrics:
-            mnt_data = sim_data_base[metric.monitor_name]
-            val = device_spec.get_metric(mnt_data, metric)
-            assert not np.allclose(val, 0.0)
+        val = device_spec.get_metric(sim_data_base)
+        assert not np.allclose(val, 0.0)
 
 
 def test_invdes_get_metric():
-    batch_data = web.BatchData(
-        {device_spec.name: sim_data_base for device_spec in invdes.device_specs}
-    )
+    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)
 
@@ -149,7 +148,7 @@ def use_emulated(monkeypatch):
 
 def test_objective_function(use_emulated):
     params = [
-        [np.ones_like(design_region.shape_3d) for design_region in device_spec.design_regions]
+        [np.ones(design_region.parameter_shape) for design_region in device_spec.design_regions]
         for device_spec in invdes.device_specs
     ]
     val = invdes.get_objective(params)
diff --git a/tidy3d/plugins/invdes2/design_region.py b/tidy3d/plugins/invdes2/design_region.py
index d47fe4ea05..3bc30fdf86 100644
--- a/tidy3d/plugins/invdes2/design_region.py
+++ b/tidy3d/plugins/invdes2/design_region.py
@@ -54,9 +54,7 @@ def to_structure(self, params: np.ndarray) -> td.Structure:
         # TODO: add transformations
         eps_data = params.reshape(self.shape_3d)
 
-        return td.Structure.from_permittivity_array(
-            geometry=geometry, eps_data=eps_data, eps_bounds=self.eps_bounds
-        )
+        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
index 5729e7201b..520597a0e3 100644
--- a/tidy3d/plugins/invdes2/device_spec.py
+++ b/tidy3d/plugins/invdes2/device_spec.py
@@ -74,7 +74,7 @@ def get_metric(self, sim_data: web.SimulationData) -> float:
         """
         value = 0.0
         for metric in self.metrics:
-            mnt_data = sim_data[metric.monitor_name]
+            mnt_data = sim_data.monitor_data[metric.monitor_name]
             value = value + metric.weight * metric.evaluate(mnt_data)
         return value
 

From 4c4dd050d39f364f499d1e76ab4037d52641dd15 Mon Sep 17 00:00:00 2001
From: Tyler Hughes <tyler@flexcompute.com>
Date: Fri, 3 Oct 2025 16:21:05 -0400
Subject: [PATCH 5/8] improvements

---
 tests/test_plugins/test_invdes2.py       | 24 +++++-------------------
 tidy3d/plugins/invdes2/README.md         |  0
 tidy3d/plugins/invdes2/design_region.py  |  4 ++++
 tidy3d/plugins/invdes2/device_spec.py    |  6 +++++-
 tidy3d/plugins/invdes2/inverse_design.py |  7 ++++++-
 5 files changed, 20 insertions(+), 21 deletions(-)
 create mode 100644 tidy3d/plugins/invdes2/README.md

diff --git a/tests/test_plugins/test_invdes2.py b/tests/test_plugins/test_invdes2.py
index c7c50f1721..9918ccbead 100644
--- a/tests/test_plugins/test_invdes2.py
+++ b/tests/test_plugins/test_invdes2.py
@@ -64,13 +64,7 @@ def test_parameter_shapes():
 
 
 def test_flatten_unflatten_params():
-    params = [
-        [
-            np.ones_like(design_region.parameter_shape)
-            for design_region in device_spec.design_regions
-        ]
-        for device_spec in invdes.device_specs
-    ]
+    params = invdes.ones()
     flat = invdes._flatten_params(params)
     restored = invdes._unflatten_params(flat)
     flat2 = invdes._flatten_params(restored)
@@ -80,24 +74,19 @@ def test_flatten_unflatten_params():
 
 def test_design_region_to_structure():
     for design_region in design_regions:
-        params = np.ones(design_region.parameter_shape)
+        params = design_region.ones()
         _ = design_region.to_structure(params)
 
 
 def test_device_spec_get_simulation():
     for device_spec in device_specs:
-        params = [
-            np.ones(design_region.parameter_shape) for design_region in device_spec.design_regions
-        ]
+        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():
-    params = [
-        [np.ones(design_region.parameter_shape) for design_region in device_spec.design_regions]
-        for device_spec in invdes.device_specs
-    ]
+    params = invdes.ones()
     sims = invdes.get_simulations(params)
 
     assert len(sims) == len(invdes.device_specs)
@@ -147,9 +136,6 @@ def use_emulated(monkeypatch):
 
 
 def test_objective_function(use_emulated):
-    params = [
-        [np.ones(design_region.parameter_shape) for design_region in device_spec.design_regions]
-        for device_spec in invdes.device_specs
-    ]
+    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..e69de29bb2
diff --git a/tidy3d/plugins/invdes2/design_region.py b/tidy3d/plugins/invdes2/design_region.py
index 3bc30fdf86..765c64d07a 100644
--- a/tidy3d/plugins/invdes2/design_region.py
+++ b/tidy3d/plugins/invdes2/design_region.py
@@ -26,6 +26,10 @@ def to_structure(self, params: np.ndarray) -> td.Structure:
     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):
diff --git a/tidy3d/plugins/invdes2/device_spec.py b/tidy3d/plugins/invdes2/device_spec.py
index 520597a0e3..79766d9285 100644
--- a/tidy3d/plugins/invdes2/device_spec.py
+++ b/tidy3d/plugins/invdes2/device_spec.py
@@ -74,7 +74,7 @@ def get_metric(self, sim_data: web.SimulationData) -> float:
         """
         value = 0.0
         for metric in self.metrics:
-            mnt_data = sim_data.monitor_data[metric.monitor_name]
+            mnt_data = sim_data[metric.monitor_name]
             value = value + metric.weight * metric.evaluate(mnt_data)
         return value
 
@@ -88,3 +88,7 @@ def get_objective(self, params: list[np.ndarray]) -> float:
     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
index 220c254ab0..5a88208ade 100644
--- a/tidy3d/plugins/invdes2/inverse_design.py
+++ b/tidy3d/plugins/invdes2/inverse_design.py
@@ -8,6 +8,7 @@
 
 from __future__ import annotations
 
+from collections.abc import Mapping
 from dataclasses import dataclass
 
 import autograd.numpy as np
@@ -45,7 +46,7 @@ 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: web.BatchData) -> float:
+    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:
@@ -64,6 +65,10 @@ 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 = []

From c9a54d6a70e4d352af2d6a6a124cda73d90a3385 Mon Sep 17 00:00:00 2001
From: Tyler Hughes <tyler@flexcompute.com>
Date: Fri, 3 Oct 2025 16:23:42 -0400
Subject: [PATCH 6/8] new branch

---
 tests/test_plugins/test_invdes2.py | 30 ++++++++++++++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/tests/test_plugins/test_invdes2.py b/tests/test_plugins/test_invdes2.py
index 9918ccbead..95fb6a96eb 100644
--- a/tests/test_plugins/test_invdes2.py
+++ b/tests/test_plugins/test_invdes2.py
@@ -56,6 +56,11 @@
 
 
 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 == [
@@ -64,6 +69,11 @@ def test_parameter_shapes():
 
 
 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)
@@ -73,12 +83,22 @@ def test_flatten_unflatten_params():
 
 
 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)
@@ -86,6 +106,11 @@ def test_device_spec_get_simulation():
 
 
 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)
 
@@ -94,6 +119,7 @@ def test_invdes_get_simulations():
 
 
 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)
@@ -101,18 +127,21 @@ def test_metric_evaluate():
 
 
 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)
@@ -136,6 +165,7 @@ def use_emulated(monkeypatch):
 
 
 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)

From d60c2deeebb35f2953df300009bdfefe49fde1c8 Mon Sep 17 00:00:00 2001
From: Tyler Hughes <tyler@flexcompute.com>
Date: Fri, 3 Oct 2025 16:55:47 -0400
Subject: [PATCH 7/8] readme v1

---
 tidy3d/plugins/invdes2/README.md | 263 +++++++++++++++++++++++++++++++
 1 file changed, 263 insertions(+)

diff --git a/tidy3d/plugins/invdes2/README.md b/tidy3d/plugins/invdes2/README.md
index e69de29bb2..bd28c1f4ea 100644
--- a/tidy3d/plugins/invdes2/README.md
+++ b/tidy3d/plugins/invdes2/README.md
@@ -0,0 +1,263 @@
+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)
+```
+
+

From c4acc341d7f529bc5b4c1354a1a2b7e036078fbb Mon Sep 17 00:00:00 2001
From: Tyler Hughes <tyler@flexcompute.com>
Date: Fri, 3 Oct 2025 16:58:40 -0400
Subject: [PATCH 8/8] readme v2

---
 tidy3d/plugins/invdes2/README.md | 57 ++++++++++++++++++++++++++++++++
 1 file changed, 57 insertions(+)

diff --git a/tidy3d/plugins/invdes2/README.md b/tidy3d/plugins/invdes2/README.md
index bd28c1f4ea..0b06b3e4ae 100644
--- a/tidy3d/plugins/invdes2/README.md
+++ b/tidy3d/plugins/invdes2/README.md
@@ -261,3 +261,60 @@ 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.
+
+