a minimum refactor of optimizer and design region

tidy3d/plugins/invdes/design.py

@@ -25,10 +25,10 @@
 class AbstractInverseDesign(InvdesBaseModel, abc.ABC):
     """Container for an inverse design problem."""
 
-    design_region: DesignRegionType = pd.Field(
+    design_regions: list[DesignRegionType] = pd.Field(
         ...,
-        title="Design Region",
-        description="Region within which we will optimize the simulation.",
+        title="Design Regions",
+        description="Regions within which we will optimize the simulation.",
     )
 
     task_name: str = pd.Field(

tidy3d/plugins/invdes/optimization_spec.py

@@ -0,0 +1,105 @@
+# specification for running the optimizer
+
+import abc
+
+import autograd.numpy as anp
+import numpy as np
+import pydantic.v1 as pd
+
+from .base import InvdesBaseModel
+from .result import InverseDesignResult
+
+
+class AbstractOptimizationSpec(InvdesBaseModel, abc.ABC):
+    """Specification for an optimization."""
+
+    learning_rate: pd.PositiveFloat = pd.Field(
+        ...,
+        title="Learning Rate",
+        description="Step size for the gradient descent optimizer.",
+    )
+
+    maximize: bool = pd.Field(
+        True,
+        title="Direction of Optimization",
+        description="If ``True``, the optimizer will maximize the objective function. If ``False``, the optimizer will minimize the objective function.",
+    )
+
+    num_steps: pd.PositiveInt = pd.Field(
+        ...,
+        title="Number of Steps",
+        description="Number of steps in the gradient descent optimizer.",
+    )
+
+    @abc.abstractmethod
+    def initial_state(self, parameters: np.ndarray) -> dict:
+        """The initial state of the optimizer."""
+
+    def display_fn(self, result: InverseDesignResult, step_index: int) -> None:
+        """Default display function while optimizing."""
+        print(f"step ({step_index + 1}/{self.num_steps})")
+        print(f"\tobjective_fn_val = {result.objective_fn_val[-1]:.3e}")
+        print(f"\tgrad_norm = {anp.linalg.norm(result.grad[-1]):.3e}")
+        print(f"\tpost_process_val = {result.post_process_val[-1]:.3e}")
+        print(f"\tpenalty = {result.penalty[-1]:.3e}")
+
+
+class AdamOptimizationSpec(AbstractOptimizationSpec):
+    """Specification for an optimization."""
+
+    beta1: float = pd.Field(
+        0.9,
+        ge=0.0,
+        le=1.0,
+        title="Beta 1",
+        description="Beta 1 parameter in the Adam optimization method.",
+    )
+
+    beta2: float = pd.Field(
+        0.999,
+        ge=0.0,
+        le=1.0,
+        title="Beta 2",
+        description="Beta 2 parameter in the Adam optimization method.",
+    )
+
+    eps: pd.PositiveFloat = pd.Field(
+        1e-8,
+        title="Epsilon",
+        description="Epsilon parameter in the Adam optimization method.",
+    )
+
+    def initial_state(self, parameters: np.ndarray) -> dict:
+        """initial state of the optimizer"""
+        zeros = np.zeros_like(parameters)
+        return dict(m=zeros, v=zeros, t=0)
+
+    def update(
+        self, parameters: np.ndarray, gradient: np.ndarray, state: dict = None
+    ) -> tuple[np.ndarray, dict]:
+        if state is None:
+            state = self.initial_state(parameters)
+
+        # get state
+        m = np.array(state["m"])
+        v = np.array(state["v"])
+        t = int(state["t"])
+
+        # update time step
+        t = t + 1
+
+        # update moment variables
+        m = self.beta1 * m + (1 - self.beta1) * gradient
+        v = self.beta2 * v + (1 - self.beta2) * (gradient**2)
+
+        # compute bias-corrected moment variables
+        m_ = m / (1 - self.beta1**t)
+        v_ = v / (1 - self.beta2**t)
+
+        # update parameters and state
+        parameters -= self.learning_rate * m_ / (np.sqrt(v_) + self.eps)
+        state = dict(m=m, v=v, t=t)
+        return parameters, state
+
+
+OptimizationSpecType = AdamOptimizationSpec

tidy3d/plugins/invdes/optimizer.py

@@ -1,12 +1,10 @@
 # specification for running the optimizer
 
-import abc
 import typing
 from copy import deepcopy
 
 import autograd as ag
 import autograd.numpy as anp
-import numpy as np
 import pydantic.v1 as pd
 
 import tidy3d as td
@@ -17,7 +15,7 @@
 from .result import InverseDesignResult
 
 
-class AbstractOptimizer(InvdesBaseModel, abc.ABC):
+class Optimizer(InvdesBaseModel):
     """Specification for an optimization."""
 
     design: InverseDesignType = pd.Field(
@@ -27,24 +25,6 @@ class AbstractOptimizer(InvdesBaseModel, abc.ABC):
         discriminator=TYPE_TAG_STR,
     )
 
-    learning_rate: pd.PositiveFloat = pd.Field(
-        ...,
-        title="Learning Rate",
-        description="Step size for the gradient descent optimizer.",
-    )
-
-    maximize: bool = pd.Field(
-        True,
-        title="Direction of Optimization",
-        description="If ``True``, the optimizer will maximize the objective function. If ``False``, the optimizer will minimize the objective function.",
-    )
-
-    num_steps: pd.PositiveInt = pd.Field(
-        ...,
-        title="Number of Steps",
-        description="Number of steps in the gradient descent optimizer.",
-    )
-
     results_cache_fname: str = pd.Field(
         None,
         title="History Storage File",
@@ -67,10 +47,6 @@ class AbstractOptimizer(InvdesBaseModel, abc.ABC):
         "last computed state of these variables.",
     )
 
-    @abc.abstractmethod
-    def initial_state(self, parameters: np.ndarray) -> dict:
-        """The initial state of the optimizer."""
-
     def validate_pre_upload(self) -> None:
         """Validate the fully initialized optimizer is ok for upload to our servers."""
         self.design.simulation.validate_pre_upload()
@@ -197,7 +173,10 @@ def continue_run(
             post_process_val = aux_data["post_process_val"]
 
             # update optimizer and parameters
-            params, opt_state = self.update(parameters=params, state=opt_state, gradient=-grad)
+            # note: would need to update every region in the list here
+            params, opt_state = self.design.region.update(
+                parameters=params, state=opt_state, gradient=-grad
+            )
 
             # cap the parameters
             params = anp.clip(params, a_min=0.0, a_max=1.0)
@@ -255,61 +234,3 @@ def continue_run_from_history(
             post_process_fn=post_process_fn,
             callback=callback,
         )
-
-
-class AdamOptimizer(AbstractOptimizer):
-    """Specification for an optimization."""
-
-    beta1: float = pd.Field(
-        0.9,
-        ge=0.0,
-        le=1.0,
-        title="Beta 1",
-        description="Beta 1 parameter in the Adam optimization method.",
-    )
-
-    beta2: float = pd.Field(
-        0.999,
-        ge=0.0,
-        le=1.0,
-        title="Beta 2",
-        description="Beta 2 parameter in the Adam optimization method.",
-    )
-
-    eps: pd.PositiveFloat = pd.Field(
-        1e-8,
-        title="Epsilon",
-        description="Epsilon parameter in the Adam optimization method.",
-    )
-
-    def initial_state(self, parameters: np.ndarray) -> dict:
-        """initial state of the optimizer"""
-        zeros = np.zeros_like(parameters)
-        return dict(m=zeros, v=zeros, t=0)
-
-    def update(
-        self, parameters: np.ndarray, gradient: np.ndarray, state: dict = None
-    ) -> tuple[np.ndarray, dict]:
-        if state is None:
-            state = self.initial_state(parameters)
-
-        # get state
-        m = np.array(state["m"])
-        v = np.array(state["v"])
-        t = int(state["t"])
-
-        # update time step
-        t = t + 1
-
-        # update moment variables
-        m = self.beta1 * m + (1 - self.beta1) * gradient
-        v = self.beta2 * v + (1 - self.beta2) * (gradient**2)
-
-        # compute bias-corrected moment variables
-        m_ = m / (1 - self.beta1**t)
-        v_ = v / (1 - self.beta2**t)
-
-        # update parameters and state
-        parameters -= self.learning_rate * m_ / (np.sqrt(v_) + self.eps)
-        state = dict(m=m, v=v, t=t)
-        return parameters, state

tidy3d/plugins/invdes/region.py

@@ -15,6 +15,7 @@
 
 from .base import InvdesBaseModel
 from .initialization import InitializationSpecType, UniformInitializationSpec
+from .optimization_spec import OptimizationSpecType
 from .penalty import PenaltyType
 from .transformation import TransformationType
 
@@ -70,6 +71,11 @@ class DesignRegion(InvdesBaseModel, abc.ABC):
         discriminator=TYPE_TAG_STR,
     )
 
+    optimization_spec: OptimizationSpecType = pd.Field(
+        title="Optimization Spec",
+        description="specifices how this design region will be optimized by the Optimizer",
+    )
+
     def _post_init_validators(self):
         """Automatically call any `_validate_XXX` method."""
         for attr_name in dir(self):