diff --git a/docs/faq b/docs/faq index 9b8479c00a..802f861885 160000 --- a/docs/faq +++ b/docs/faq @@ -1 +1 @@ -Subproject commit 9b8479c00a7b2d5fb81a90dd21c7b8e92423ef58 +Subproject commit 802f8618852493bbb70e437d02ef6a33594ca6bd diff --git a/docs/notebooks b/docs/notebooks index cde5d9502b..26e33355f0 160000 --- a/docs/notebooks +++ b/docs/notebooks @@ -1 +1 @@ -Subproject commit cde5d9502b6925e118c30dcd2ec24e910eaee181 +Subproject commit 26e33355f099819fc215ffdd12be0c5e519392c9 diff --git a/tidy3d/plugins/invdes/design.py b/tidy3d/plugins/invdes/design.py index f8a4487b07..85e9b56f3a 100644 --- a/tidy3d/plugins/invdes/design.py +++ b/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( diff --git a/tidy3d/plugins/invdes/optimization_spec.py b/tidy3d/plugins/invdes/optimization_spec.py new file mode 100644 index 0000000000..f00fd8e3e6 --- /dev/null +++ b/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 diff --git a/tidy3d/plugins/invdes/optimizer.py b/tidy3d/plugins/invdes/optimizer.py index 08758f12b6..1ea70a807a 100644 --- a/tidy3d/plugins/invdes/optimizer.py +++ b/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 diff --git a/tidy3d/plugins/invdes/region.py b/tidy3d/plugins/invdes/region.py index b94eaf8682..8b9882b323 100644 --- a/tidy3d/plugins/invdes/region.py +++ b/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):