schema-backed postprocess option for invdes

CHANGELOG.md

@@ -5,6 +5,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+- Postprocessing functions in `invdes` plugin can be optionally specified by supplying a `postprocess` instance to the `InverseDesign` object.
+
 ### Fixed
 - Error when loading a previously run `Batch` or `ComponentModeler` containing custom data.
 

tests/test_plugins/test_invdes.py

@@ -140,6 +140,13 @@ def post_process_fn(sim_data: td.SimulationData, **kwargs) -> float:
     return anp.sum(intensity.values)
 
 
+postprocess_obj = tdi.Sum(
+    operations=[
+        tdi.ModePower(monitor_name=MNT_NAME2, direction="+", mode_index=0, weight=0.5),
+    ]
+)
+
+
 def post_process_fn_kwargless(sim_data: td.SimulationData) -> float:
     """Define a post-processing function with no kwargs specified."""
     intensity = sim_data.get_intensity(MNT_NAME1)
@@ -271,14 +278,18 @@ def make_optimizer():
     )
 
 
-def make_result(use_emulated_run_autograd):
+def make_result(use_emulated_run_autograd, fn_postprocess: bool = True):
     """Test running the optimization defined in the ``InverseDesign`` object."""
 
     optimizer = make_optimizer()
 
     PARAMS_0 = np.random.random(optimizer.design.design_region.params_shape)
 
-    return optimizer.run(params0=PARAMS_0, post_process_fn=post_process_fn)
+    if fn_postprocess:
+        return optimizer.run(params0=PARAMS_0, post_process_fn=post_process_fn)
+    else:
+        optimizer = optimizer.updated_copy(postprocess=postprocess_obj, path="design")
+        return optimizer.run(params0=PARAMS_0)
 
 
 def test_default_params(use_emulated_run_autograd):  # noqa: F811
@@ -366,15 +377,56 @@ def test_continue_run_from_file(use_emulated_run_autograd):
     result_full = optimizer.continue_run_from_history(post_process_fn=post_process_fn)
 
 
+def test_postprocess_init(use_emulated_run):  # noqa: F811
+    """Test the intiialization of an ``InverseDesign`` with different ``postprocess`` options."""
+
+    power_obj = tdi.ModePower(monitor_name=MNT_NAME2, direction="+", mode_index=0, weight=0.5)
+    postprocess_obj = tdi.Sum(operations=[power_obj])
+
+    def fn(sim_data):
+        return power_obj.evaluate(sim_data)
+
+    # basline
+    invdes = make_invdes()
+    params = invdes.design_region.params_random
+    sim_data = invdes.to_simulation_data(params=params, task_name="test")
+    obj_fn_val = fn(sim_data)
+
+    # supply postprocess object directly
+    invdes_obj = invdes.updated_copy(postprocess=postprocess_obj)
+
+    # supply invdes with a function from classmethod constructor
+    invdes_from_fn = invdes.from_function(
+        fn,
+        **invdes.dict(
+            exclude={
+                "postprocess",
+            }
+        ),
+    )
+
+    # use the custom postprocess operation.
+    invdes_from_custom = invdes.updated_copy(
+        postprocess=tdi.CustomPostprocessOperation.from_function(fn)
+    )
+
+    # check that they all give the same objective function values, and calling evaluate works
+    for invdes_ in (invdes_obj, invdes_from_fn, invdes_from_custom):
+        obj_fn_val_ = invdes_.postprocess.evaluate(sim_data)
+        assert np.isclose(obj_fn_val, obj_fn_val_)
+
+
+@pytest.mark.parametrize("fn_postprocess", (True, False))
 def test_result(
     use_emulated_run,  # noqa: F811
     use_emulated_run_autograd,
     use_emulated_run_autograd_async,
     tmp_path,
+    fn_postprocess,
 ):
     """Test methods of the ``InverseDesignResult`` object."""
 
-    result = make_result(use_emulated_run_autograd)
+    result = make_result(use_emulated_run_autograd, fn_postprocess=fn_postprocess)
 
     with pytest.raises(KeyError):
         _ = result.get_last("blah")

tidy3d/plugins/invdes/__init__.py

@@ -4,6 +4,7 @@
 from .design import InverseDesign, InverseDesignMulti
 from .optimizer import AdamOptimizer
 from .penalty import ErosionDilationPenalty
+from .postprocess import CustomPostProcessOperation, ModePower, Sum
 from .region import TopologyDesignRegion
 from .result import InverseDesignResult
 from .transformation import FilterProject
@@ -16,5 +17,8 @@
     "TopologyDesignRegion",
     "AdamOptimizer",
     "InverseDesignResult",
+    "CustomPostProcessOperation",
+    "Sum",
+    "ModePower",
     "utils",
 )

tidy3d/plugins/invdes/design.py

@@ -1,4 +1,5 @@
 # container for everything defining the inverse design
+from __future__ import annotations
 
 import abc
 import typing
@@ -11,11 +12,10 @@
 from tidy3d.components.autograd import get_static
 
 from .base import InvdesBaseModel
+from .postprocess import CustomPostProcessOperation, PostProcessFnType, PostProcessOperationType
 from .region import DesignRegionType
 from .validators import check_pixel_size
 
-PostProcessFnType = typing.Callable[[td.SimulationData], float]
-
 
 class AbstractInverseDesign(InvdesBaseModel, abc.ABC):
     """Container for an inverse design problem."""
@@ -32,24 +32,63 @@ class AbstractInverseDesign(InvdesBaseModel, abc.ABC):
         description="Task name to use in the objective function when running the ``JaxSimulation``.",
     )
 
+    postprocess: PostProcessOperationType = pd.Field(
+        None,
+        title="PostProcess Object",
+        description="Optional object specifying how to perform weighted sum of ``ModeData`` "
+        "outputs. Can be used instead of passing a generic postprocess function "
+        "to ``Optimizer.run()``.",
+    )
+
     verbose: bool = pd.Field(
         False,
         title="Task Verbosity",
         description="If ``True``, will print the regular output from ``web`` functions.",
     )
 
+    @classmethod
+    def from_function(cls, fn: PostProcessFnType, **kwargs) -> AbstractInverseDesign:
+        """Create an ``InverseDesign`` object from a user-supplied postprocessing function."""
+
+        postprocess = CustomPostProcessOperation.from_function(fn)
+        return cls(postprocess=postprocess, **kwargs)
+
+    def _get_postprocess_fn(
+        self, post_process_fn: typing.Callable = None
+    ) -> typing.Callable[[td.SimulationData], float]:
+        """Returns postprocessing function used internally given some passed function."""
+
+        if post_process_fn is None and self.postprocess is None:
+            raise ValueError(
+                "If ``InverseDesign.postprocess`` is not specified, must pass a "
+                "``post_process_fn`` to ``Optimizer.run``."
+            )
+        if post_process_fn is not None and self.postprocess is not None:
+            raise ValueError(
+                "Can't specify both an ``InverseDesign.postprocess`` and pass a ``post_process_fn``."
+                "to ``Optimizer.run``."
+            )
+
+        if self.postprocess is not None:
+            return self.postprocess.evaluate
+
+        return post_process_fn
+
     def make_objective_fn(
-        self, post_process_fn: typing.Callable
+        self,
+        post_process_fn: typing.Callable = None,
     ) -> typing.Callable[[anp.ndarray], tuple[float, dict]]:
         """construct the objective function for this ``InverseDesignMulti`` object."""
 
+        _postprocess_fn = self._get_postprocess_fn(post_process_fn)
+
         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)
+            post_process_val = _postprocess_fn(data)
 
             penalty_value = self.design_region.penalty_value(params)
             objective_fn_val = post_process_val - penalty_value

tidy3d/plugins/invdes/optimizer.py

@@ -85,15 +85,15 @@ 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, 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
     ) -> InverseDesignResult:
         """Run optimizer for a series of steps with an initialized state."""
 
@@ -158,13 +158,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
     ) -> 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
+    ) -> 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

tidy3d/plugins/invdes/postprocess.py

@@ -0,0 +1,130 @@
+# defines postprocessing classes for `InverseDesign` objects.
+from __future__ import annotations
+
+import abc
+import typing
+
+import autograd.numpy as anp
+import pydantic.v1 as pd
+
+import tidy3d as td
+
+from .base import InvdesBaseModel
+
+PostProcessFnType = typing.Callable[[td.SimulationData], float]
+
+
+class AbstractPostProcessOperation(InvdesBaseModel, abc.ABC):
+    """Abstract base class defining components that make up postprocessing classes."""
+
+    @abc.abstractmethod
+    def evaluate(self, sim_data: td.SimulationData) -> float:
+        """How to evaluate this operation on a ``SimulationData`` object."""
+
+
+class CustomPostProcessOperation(AbstractPostProcessOperation):
+    """A postprocessing operation to subclass and implement ones own ``evalute`` method for."""
+
+    def evaluate(self, sim_data: td.SimulationData) -> float:
+        """How to evaluate this operation on a ``SimulationData`` object."""
+        raise NotImplementedError("Must define the 'self.evaluate(sim_data) -> float' method.")
+
+    @classmethod
+    def from_function(cls, fn: PostProcessFnType, **kwargs) -> CustomPostProcessOperation:
+        """Create a ``CustomPostProcessOperation`` from a function of ``SimulationData``."""
+        cls.evaluate = lambda self, sim_data: fn(sim_data)
+        obj = cls(**kwargs)
+        return obj
+
+
+class ElementaryPostProcessOperation(AbstractPostProcessOperation, abc.ABC):
+    """A postprocess operation that can work on its own."""
+
+
+class ModePower(ElementaryPostProcessOperation):
+    """Grab the power from a ``ModeMonitor`` and apply an optional weight."""
+
+    monitor_name: str = pd.Field(
+        ...,
+        title="Monitor Name",
+        description="Name of the ``ModeMonitor`` corresponding to the ``ModeData`` "
+        "that we want to compute power for.",
+    )
+
+    direction: td.components.types.Direction = pd.Field(
+        None,
+        title="Direction",
+        description="If specified, selects a specific direction ``'-'`` or ``'+'`` from the "
+        "``ModeData`` to include in power. Otherwise, sums over all directions.",
+    )
+
+    mode_index: pd.NonNegativeInt = pd.Field(
+        None,
+        title="Mode Index",
+        description="If specified, selects a specific mode index from the ``ModeData`` "
+        "to include in power. Otherwise, sums over all mode indices.",
+    )
+
+    f: float = pd.Field(
+        None,
+        title="Frequency",
+        description="If specified, selects a specific frequency from the ``ModeData`` "
+        "to include in power. Otherwise, sums over all frequencies.",
+    )
+
+    weight: float = pd.Field(
+        1.0,
+        title="Weight",
+        description="Weight specifying the contribution of this power to the objective function. ",
+    )
+
+    @property
+    def sel_kwargs(self) -> dict[str, typing.Any]:
+        """Selection kwargs corresponding to the fields."""
+        sel_kwargs_all = dict(direction=self.direction, mode_index=self.mode_index, f=self.f)
+        return {key: sel for key, sel in sel_kwargs_all.items() if sel is not None}
+
+    def evaluate(self, sim_data: td.SimulationData) -> float:
+        """Evaluate this instance when passed a simulation dataset."""
+
+        if self.monitor_name not in sim_data:
+            raise KeyError(f"No data found for monitor with name '{self.monitor_name}'.")
+
+        mnt_data = sim_data[self.monitor_name]
+
+        if not isinstance(mnt_data, td.ModeData):
+            raise ValueError(
+                "'ModePower' only works with 'ModeData' corresponding to 'ModeMonitor'. "
+                f"Monitor name of '{self.monitor_name}' returned data of type {type(mnt_data)}."
+            )
+
+        amps = mnt_data.amps
+        amp = amps.sel(**self.sel_kwargs)
+        powers = abs(amp.values) ** 2
+        power = anp.sum(powers)
+        return self.weight * power
+
+
+ElementaryPostProcessOperationType = typing.Union[ModePower]
+
+
+class CombinedPostProcessOperation(AbstractPostProcessOperation, abc.ABC):
+    """A postprocess operation that combines elementary operations."""
+
+    operations: tuple[ElementaryPostProcessOperationType, ...] = pd.Field(
+        (),
+        title="PostProcessing Operations",
+        description="Set of objects specifying the operations combined in this operation.",
+    )
+
+
+class Sum(CombinedPostProcessOperation):
+    """Sum of the evaluate outputs of elementary operation objects."""
+
+    def evaluate(self, sim_data: td.SimulationData) -> float:
+        """sum the evaluation of each operation."""
+
+        return sum(op.evaluate(sim_data) for op in self.operations)
+
+
+PostProcessOperationType = typing.Union[CustomPostProcessOperation, ModePower, Sum]