Add initial support for expressions/metrics in invdes plugin

CHANGELOG.md

@@ -6,7 +6,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased]
 
 ### Added
-- Added `tidy3d.plugins.expressions` module for constructing and serializing mathematical expressions and simulation metrics like `ModeAmps` and `ModePower`.
+- Added `tidy3d.plugins.expressions` module for constructing and serializing mathematical expressions and simulation metrics like `ModeAmp` and `ModePower`.
+- Support for serializable expressions in the `invdes` plugin (`InverseDesign(metric=ModePower(...))`).
 
 ## [2.7.3] - 2024-09-12
 

tests/test_plugins/test_invdes.py

@@ -5,6 +5,7 @@
 import pytest
 import tidy3d as td
 import tidy3d.plugins.invdes as tdi
+from tidy3d.plugins.expressions import ModePower
 
 # use single threading pipeline
 from ..test_components.test_autograd import use_emulated_run  # noqa: F401
@@ -478,3 +479,26 @@ def test_pixel_size_warn_validator(log_capture):
 
     with AssertLogLevel(log_capture, "WARNING", contains_str="pixel_size"):
         invdes_multi = invdes_multi.updated_copy(design_region=region_too_coarse)
+
+
+def test_invdes_with_metric_objective(use_emulated_run, use_emulated_to_sim_data):  # noqa: F811
+    """Test using a metric as an objective function in InverseDesign."""
+
+    # Create a metric as the objective function
+    metric = 2 * ModePower(monitor_name=MNT_NAME2, freqs=[FREQ0]) ** 2
+
+    invdes = tdi.InverseDesign(
+        simulation=simulation,
+        design_region=make_design_region(),
+        task_name="test_metric",
+        metric=metric,
+    )
+
+    optimizer = tdi.AdamOptimizer(
+        design=invdes,
+        learning_rate=0.2,
+        num_steps=1,
+    )
+
+    params0 = np.random.random(invdes.design_region.params_shape)
+    optimizer.run(params0=params0)

tidy3d/plugins/invdes/design.py

@@ -1,5 +1,7 @@
 # container for everything defining the inverse design
 
+from __future__ import annotations
+
 import abc
 import typing
 
@@ -9,6 +11,7 @@
 import tidy3d as td
 import tidy3d.web as web
 from tidy3d.components.autograd import get_static
+from tidy3d.plugins.expressions.types import ExpressionType
 
 from .base import InvdesBaseModel
 from .region import DesignRegionType
@@ -38,25 +41,42 @@ class AbstractInverseDesign(InvdesBaseModel, abc.ABC):
         description="If ``True``, will print the regular output from ``web`` functions.",
     )
 
+    metric: typing.Optional[ExpressionType] = pd.Field(
+        None,
+        title="Objective Metric",
+        description="Serializable expression defining the objective function.",
+    )
+
     def make_objective_fn(
-        self, post_process_fn: typing.Callable, maximize: bool = True
+        self, post_process_fn: typing.Callable | None = None, maximize: bool = True
     ) -> typing.Callable[[anp.ndarray], tuple[float, dict]]:
-        """construct the objective function for this ``InverseDesignMulti`` object."""
+        """Construct the objective function for this InverseDesign object."""
+
+        if (post_process_fn is None) and (self.metric is None):
+            raise ValueError("Either 'post_process_fn' or 'metric' must be provided.")
+
+        if (post_process_fn is not None) and (self.metric is not None):
+            raise ValueError("Provide only one of 'post_process_fn' or 'metric', not both.")
 
         direction_multiplier = 1 if maximize else -1
 
         def objective_fn(params: anp.ndarray, aux_data: dict = None) -> float:
             """Full objective function."""
-
             data = self.to_simulation_data(params=params)
 
-            # construct objective function values
-            post_process_val = post_process_fn(data)
+            if self.metric is None:
+                post_process_val = post_process_fn(data)
+            elif isinstance(data, td.SimulationData):
+                post_process_val = self.metric.evaluate(data)
+            elif isinstance(data, web.BatchData):
+                raise NotImplementedError("Metrics currently do not support 'BatchData'")
+            else:
+                raise ValueError(f"Invalid data type: {type(data)}")
 
             penalty_value = self.design_region.penalty_value(params)
             objective_fn_val = direction_multiplier * post_process_val - penalty_value
 
-            # store things in ``aux_data`` passed by reference
+            # Store auxiliary data if provided
             if aux_data is not None:
                 aux_data["penalty"] = get_static(penalty_value)
                 aux_data["post_process_val"] = get_static(post_process_val)
@@ -132,8 +152,6 @@ def to_simulation_data(self, params: anp.ndarray, **kwargs) -> td.SimulationData
         simulation = self.to_simulation(params=params)
         kwargs.setdefault("task_name", self.task_name)
         return web.run(simulation, verbose=self.verbose, **kwargs)
-        # sim_data = job.run()
-        # return sim_data
 
 
 class InverseDesignMulti(AbstractInverseDesign):

tidy3d/plugins/invdes/optimizer.py

@@ -91,15 +91,17 @@ def _initialize_result(self, params0: anp.ndarray = None) -> InverseDesignResult
         return InverseDesignResult(design=self.design, opt_state=[state], params=[params0])
 
     def run(
-        self, post_process_fn: typing.Callable, params0: anp.ndarray = None
+        self, post_process_fn: typing.Callable | None = None, params0: anp.ndarray = None
     ) -> InverseDesignResult:
         """Run this inverse design problem from an optional initial set of parameters."""
         self.design.design_region._check_params(params0)
         starting_result = self._initialize_result(params0)
         return self.continue_run(result=starting_result, post_process_fn=post_process_fn)
 
     def continue_run(
-        self, result: InverseDesignResult, post_process_fn: typing.Callable
+        self,
+        result: InverseDesignResult,
+        post_process_fn: typing.Callable | None = None,
     ) -> InverseDesignResult:
         """Run optimizer for a series of steps with an initialized state."""
 
@@ -164,13 +166,15 @@ def continue_run(
         return InverseDesignResult(design=result.design, **history)
 
     def continue_run_from_file(
-        self, fname: str, post_process_fn: typing.Callable
+        self, fname: str, post_process_fn: typing.Callable | None = None
     ) -> InverseDesignResult:
         """Continue the optimization run from a ``.pkl`` file with an ``InverseDesignResult``."""
         result = InverseDesignResult.from_file(fname)
         return self.continue_run(result=result, post_process_fn=post_process_fn)
 
-    def continue_run_from_history(self, post_process_fn: typing.Callable) -> InverseDesignResult:
+    def continue_run_from_history(
+        self, post_process_fn: typing.Callable | None = None
+    ) -> InverseDesignResult:
         """Continue the optimization run from a ``.pkl`` file with an ``InverseDesignResult``."""
         return self.continue_run_from_file(
             fname=self.results_cache_fname, post_process_fn=post_process_fn