From dc465e298f97900389977f72aad1c3a51ec7ebda Mon Sep 17 00:00:00 2001 From: Yannick Augenstein Date: Tue, 9 Jul 2024 18:49:31 +0200 Subject: [PATCH] Add class-based penalties and parametrizations --- tidy3d/plugins/autograd/__init__.py | 8 + tidy3d/plugins/autograd/functions.py | 14 +- tidy3d/plugins/autograd/invdes/__init__.py | 16 +- tidy3d/plugins/autograd/invdes/filters.py | 32 ++-- .../autograd/invdes/parametrizations.py | 135 ++++++++++++---- tidy3d/plugins/autograd/invdes/penalties.py | 148 ++++++++++++------ tidy3d/plugins/autograd/utilities.py | 2 +- 7 files changed, 245 insertions(+), 110 deletions(-) diff --git a/tidy3d/plugins/autograd/__init__.py b/tidy3d/plugins/autograd/__init__.py index 394e945f6..adf9686d5 100644 --- a/tidy3d/plugins/autograd/__init__.py +++ b/tidy3d/plugins/autograd/__init__.py @@ -13,6 +13,10 @@ threshold, ) from .invdes import ( + CircularFilter, + ConicFilter, + ErosionDilationPenalty, + FilterAndProject, grey_indicator, make_circular_filter, make_conic_filter, @@ -27,6 +31,10 @@ from .utilities import chain, get_kernel_size_px, make_kernel __all__ = [ + "CircularFilter", + "ConicFilter", + "ErosionDilationPenalty", + "FilterAndProject", "make_filter", "make_conic_filter", "make_circular_filter", diff --git a/tidy3d/plugins/autograd/functions.py b/tidy3d/plugins/autograd/functions.py index 4d642ed82..9bf4f8d1b 100644 --- a/tidy3d/plugins/autograd/functions.py +++ b/tidy3d/plugins/autograd/functions.py @@ -342,7 +342,7 @@ def grey_dilation( The input array to perform grey dilation on. size : Union[Union[int, Tuple[int, int]], None] = None The size of the structuring element. If None, `structure` must be provided. - structure : Union[np.ndarray, None] = None + structure : Union[NDArray, None] = None The structuring element. If None, `size` must be provided. mode : PaddingType = "reflect" The padding mode to use. @@ -396,7 +396,7 @@ def grey_erosion( The input array to perform grey dilation on. size : Union[Union[int, Tuple[int, int]], None] = None The size of the structuring element. If None, `structure` must be provided. - structure : Union[np.ndarray, None] = None + structure : Union[NDArray, None] = None The structuring element. If None, `size` must be provided. mode : PaddingType = "reflect" The padding mode to use. @@ -450,7 +450,7 @@ def grey_opening( The input array to perform grey opening on. size : Union[Union[int, Tuple[int, int]], None] = None The size of the structuring element. If None, `structure` must be provided. - structure : Union[np.ndarray, None] = None + structure : Union[NDArray, None] = None The structuring element. If None, `size` must be provided. mode : PaddingType = "reflect" The padding mode to use. @@ -483,7 +483,7 @@ def grey_closing( The input array to perform grey closing on. size : Union[Union[int, Tuple[int, int]], None] = None The size of the structuring element. If None, `structure` must be provided. - structure : Union[np.ndarray, None] = None + structure : Union[NDArray, None] = None The structuring element. If None, `size` must be provided. mode : PaddingType = "reflect" The padding mode to use. @@ -516,7 +516,7 @@ def morphological_gradient( The input array to compute the morphological gradient of. size : Union[Union[int, Tuple[int, int]], None] = None The size of the structuring element. If None, `structure` must be provided. - structure : Union[np.ndarray, None] = None + structure : Union[NDArray, None] = None The structuring element. If None, `size` must be provided. mode : PaddingType = "reflect" The padding mode to use. @@ -549,7 +549,7 @@ def morphological_gradient_internal( The input array to compute the internal morphological gradient of. size : Union[Union[int, Tuple[int, int]], None] = None The size of the structuring element. If None, `structure` must be provided. - structure : Union[np.ndarray, None] = None + structure : Union[NDArray, None] = None The structuring element. If None, `size` must be provided. mode : PaddingType = "reflect" The padding mode to use. @@ -580,7 +580,7 @@ def morphological_gradient_external( The input array to compute the external morphological gradient of. size : Union[Union[int, Tuple[int, int]], None] = None The size of the structuring element. If None, `structure` must be provided. - structure : Union[np.ndarray, None] = None + structure : Union[NDArray, None] = None The structuring element. If None, `size` must be provided. mode : PaddingType = "reflect" The padding mode to use. diff --git a/tidy3d/plugins/autograd/invdes/__init__.py b/tidy3d/plugins/autograd/invdes/__init__.py index 3ce8255dd..080f7dba8 100644 --- a/tidy3d/plugins/autograd/invdes/__init__.py +++ b/tidy3d/plugins/autograd/invdes/__init__.py @@ -1,17 +1,27 @@ -from .filters import make_circular_filter, make_conic_filter, make_filter +from .filters import ( + CircularFilter, + ConicFilter, + make_circular_filter, + make_conic_filter, + make_filter, +) from .misc import grey_indicator -from .parametrizations import make_filter_and_project -from .penalties import make_curvature_penalty, make_erosion_dilation_penalty +from .parametrizations import FilterAndProject, make_filter_and_project +from .penalties import ErosionDilationPenalty, make_curvature_penalty, make_erosion_dilation_penalty from .projections import ramp_projection, tanh_projection __all__ = [ "grey_indicator", + "CircularFilter", + "ConicFilter", "make_circular_filter", "make_conic_filter", "make_curvature_penalty", "make_erosion_dilation_penalty", + "ErosionDilationPenalty", "make_filter", "make_filter_and_project", + "FilterAndProject", "ramp_projection", "tanh_projection", ] diff --git a/tidy3d/plugins/autograd/invdes/filters.py b/tidy3d/plugins/autograd/invdes/filters.py index 03cc0a82a..9118a5747 100644 --- a/tidy3d/plugins/autograd/invdes/filters.py +++ b/tidy3d/plugins/autograd/invdes/filters.py @@ -3,7 +3,8 @@ from typing import Callable, Iterable, Tuple, Union import numpy as np -import pydantic as pd +import pydantic.v1 as pd +from numpy.typing import NDArray from tidy3d.components.base import Tidy3dBaseModel @@ -25,7 +26,7 @@ class AbstractFilter(Tidy3dBaseModel, abc.ABC): The padding mode to use. """ - kernel_size: Tuple[pd.PositiveInt, ...] = pd.Field( + kernel_size: Union[pd.PositiveInt, Tuple[pd.PositiveInt, ...]] = pd.Field( ..., description="Size of the kernel in pixels for each dimension." ) normalize: bool = pd.Field( @@ -54,11 +55,11 @@ def from_radius_dl( An instance of the filter. """ kernel_size = get_kernel_size_px(radius=radius, dl=dl) - return cls(kernel_size, **kwargs) + return cls(kernel_size=kernel_size, **kwargs) @staticmethod @abc.abstractmethod - def get_kernel(size_px: Iterable[int], normalize: bool) -> np.ndarray: + def get_kernel(size_px: Iterable[int], normalize: bool) -> NDArray: """Get the kernel for the filter. Parameters @@ -70,27 +71,28 @@ def get_kernel(size_px: Iterable[int], normalize: bool) -> np.ndarray: Returns ------- - np.ndarray + NDArray The kernel. """ ... - def __call__(self, array: np.ndarray) -> np.ndarray: + def __call__(self, array: NDArray) -> NDArray: """Apply the filter to an input array. Parameters ---------- - array : np.ndarray + array : NDArray The input array to filter. Returns ------- - np.ndarray + NDArray The filtered array. """ original_shape = array.shape squeezed_array = np.squeeze(array) - size_px = self.kernel_size + size_px = tuple(np.atleast_1d(self.kernel_size)) + print(size_px) if len(size_px) != squeezed_array.ndim: size_px *= squeezed_array.ndim kernel = self.get_kernel(size_px, self.normalize) @@ -103,12 +105,12 @@ class ConicFilter(AbstractFilter): @staticmethod @lru_cache(maxsize=1) - def get_kernel(size_px: Iterable[int], normalize: bool) -> np.ndarray: + def get_kernel(size_px: Iterable[int], normalize: bool) -> NDArray: """Get the conic kernel. See Also -------- - :func:`~filters.AbstractFilter.get_kernel` For full method documentation. + :func:`~filters.AbstractFilter.get_kernel` for full method documentation. """ return make_kernel(kernel_type="conic", size=size_px, normalize=normalize) @@ -118,12 +120,12 @@ class CircularFilter(AbstractFilter): @staticmethod @lru_cache(maxsize=1) - def get_kernel(size_px: Iterable[int], normalize: bool) -> np.ndarray: + def get_kernel(size_px: Iterable[int], normalize: bool) -> NDArray: """Get the circular kernel. See Also -------- - :func:`~filters.AbstractFilter.get_kernel` For full method documentation. + :func:`~filters.AbstractFilter.get_kernel` for full method documentation. """ return make_kernel(kernel_type="circular", size=size_px, normalize=normalize) @@ -171,7 +173,7 @@ def make_filter( normalize: bool = True, padding: PaddingType = "reflect", filter_type: KernelType, -) -> Callable[[np.ndarray], np.ndarray]: +) -> Callable[[NDArray], NDArray]: """Create a filter function based on the specified kernel type and size. Parameters @@ -191,7 +193,7 @@ def make_filter( Returns ------- - Callable[[np.ndarray], np.ndarray] + Callable[[NDArray], NDArray] A function that applies the created filter to an input array. """ kernel_size = _get_kernel_size(radius, dl, size_px) diff --git a/tidy3d/plugins/autograd/invdes/parametrizations.py b/tidy3d/plugins/autograd/invdes/parametrizations.py index 5dd2eac1b..c09b8ec98 100644 --- a/tidy3d/plugins/autograd/invdes/parametrizations.py +++ b/tidy3d/plugins/autograd/invdes/parametrizations.py @@ -1,12 +1,108 @@ +from __future__ import annotations + from typing import Callable, Tuple, Union +import pydantic.v1 as pd from numpy.typing import NDArray +from tidy3d.components.base import Tidy3dBaseModel + from ..types import KernelType, PaddingType -from .filters import make_filter +from .filters import AbstractFilter, CircularFilter, ConicFilter from .projections import tanh_projection +class FilterAndProject(Tidy3dBaseModel): + """A class that combines filtering and projection operations. + + Parameters + ---------- + filter : AbstractFilter + The filter to apply. + beta : float = 1.0 + The beta parameter for the tanh projection. + eta : float = 0.5 + The eta parameter for the tanh projection. + """ + + filter: AbstractFilter + beta: pd.NonNegativeFloat = 1.0 + eta: pd.NonNegativeFloat = 0.5 + + def __call__(self, array: NDArray, beta: float = None, eta: float = None) -> NDArray: + """Apply the filter and projection to an input array. + + Parameters + ---------- + array : NDArray + The input array to filter and project. + beta : float = None + The beta parameter for the tanh projection. If None, uses the instance's beta. + eta : float = None + The eta parameter for the tanh projection. If None, uses the instance's eta. + + Returns + ------- + NDArray + The filtered and projected array. + """ + filtered = self.filter(array) + beta = beta if beta is not None else self.beta + eta = eta if eta is not None else self.eta + projected = tanh_projection(filtered, beta, eta) + return projected + + @classmethod + def from_params( + cls, + radius: Union[float, Tuple[float, ...]] = None, + dl: Union[float, Tuple[float, ...]] = None, + *, + size_px: Union[int, Tuple[int, ...]] = None, + beta: float = 1.0, + eta: float = 0.5, + filter_type: KernelType = "conic", + padding: PaddingType = "reflect", + ) -> FilterAndProject: + """Create a FilterAndProject instance from parameters. + + Parameters + ---------- + radius : Union[float, Tuple[float, ...]] = None + The radius of the kernel. Can be a scalar or a tuple. + dl : Union[float, Tuple[float, ...]] = None + The grid spacing. Can be a scalar or a tuple. + size_px : Union[int, Tuple[int, ...]] = None + The size of the kernel in pixels for each dimension. Can be a scalar or a tuple. + beta : float = 1.0 + The steepness of the tanh projection. Higher values result in a sharper transition. + eta : float = 0.5 + The midpoint of the tanh projection. + filter_type : KernelType = "conic" + The type of filter kernel to use. + padding : PaddingType = "reflect" + The padding type to use for the filter. + + Returns + ------- + FilterAndProject + An instance of FilterAndProject. + """ + if filter_type == "conic": + filter_class = ConicFilter + elif filter_type == "circular": + filter_class = CircularFilter + else: + raise ValueError(f"Unsupported filter_type: {filter_type}") + + if size_px is not None: + filter_instance = filter_class(kernel_size=size_px, padding=padding) + else: + filter_instance = filter_class.from_radius_dl(radius=radius, dl=dl, padding=padding) + + return cls(filter=filter_instance, beta=beta, eta=eta) + + def make_filter_and_project( radius: Union[float, Tuple[float, ...]] = None, dl: Union[float, Tuple[float, ...]] = None, @@ -19,35 +115,10 @@ def make_filter_and_project( ) -> Callable: """Create a function that filters and projects an array. - This is the standard filter-and-project scheme used in topology optimization. - - Parameters - ---------- - radius : Union[float, Tuple[float, ...]] = None - The radius of the kernel. Can be a scalar or a tuple. - dl : Union[float, Tuple[float, ...]] = None - The grid spacing. Can be a scalar or a tuple. - size_px : Union[int, Tuple[int, ...]] = None - The size of the kernel in pixels for each dimension. Can be a scalar or a tuple. - beta : float = 1.0 - The beta parameter for the tanh projection. - eta : float = 0.5 - The eta parameter for the tanh projection. - filter_type : KernelType = "conic" - The type of filter kernel to use. - padding : PaddingType = "reflect" - The padding type to use for the filter. - - Returns - ------- - function - A function that takes an array and applies the filter and projection. + See Also + -------- + :func:`~parametrizations.FilterAndProject.from_params` for full method documentation. """ - _filter = make_filter(radius, dl, size_px=size_px, filter_type=filter_type, padding=padding) - - def _filter_and_project(array: NDArray, beta: float = beta, eta: float = eta) -> NDArray: - array = _filter(array) - array = tanh_projection(array, beta, eta) - return array - - return _filter_and_project + return FilterAndProject.from_params( + radius, dl, size_px=size_px, beta=beta, eta=eta, filter_type=filter_type, padding=padding + ) diff --git a/tidy3d/plugins/autograd/invdes/penalties.py b/tidy3d/plugins/autograd/invdes/penalties.py index 90bd40889..b83781e89 100644 --- a/tidy3d/plugins/autograd/invdes/penalties.py +++ b/tidy3d/plugins/autograd/invdes/penalties.py @@ -3,79 +3,123 @@ import autograd.numpy as np from numpy.typing import NDArray +from tidy3d.components.base import Tidy3dBaseModel from tidy3d.components.types import ArrayFloat2D from ..types import PaddingType -from .parametrizations import make_filter_and_project +from .parametrizations import FilterAndProject -def make_erosion_dilation_penalty( - radius: Union[float, Tuple[float, ...]] = None, - dl: Union[float, Tuple[float, ...]] = None, - *, - size_px: Union[int, Tuple[int, ...]] = None, - beta: float = 100.0, - eta: float = 0.5, - delta_eta: float = 0.01, - padding: PaddingType = "reflect", -) -> Callable: - """Computes a penalty for erosion/dilation of a parameter map not being unity. - - Accepts a parameter array normalized between 0 and 1. Uses filtering and projection methods - to erode and dilate the features within this array. Measures the change in the array after - eroding and dilating (and also dilating and eroding). Returns a penalty proportional to the - magnitude of this change. The amount of change under dilation and erosion is minimized if - the structure has large feature sizes and large radius of curvature relative to the length scale. +class ErosionDilationPenalty(Tidy3dBaseModel): + """A class that computes a penalty for erosion/dilation of a parameter map not being unity. Parameters ---------- - radius : Union[float, Tuple[float, ...]] = None - The radius of the kernel. Can be a scalar or a tuple. - dl : Union[float, Tuple[float, ...]] = None - The grid spacing. Can be a scalar or a tuple. - size_px : Union[int, Tuple[int, ...]] = None - The size of the kernel in pixels for each dimension. Can be a scalar or a tuple. - beta : float = 1.0 - Strength of the tanh projection. - eta : float = 0.5 - Midpoint of the tanh projection. + filtproj : FilterAndProject + The filter and project instance to use for erosion and dilation operations. delta_eta : float = 0.01 The binarization threshold for erosion and dilation operations. - padding : PaddingType = "reflect" - The padding type to use for the filter. - - Returns - ------- - Callable - A function that computes the erosion/dilation penalty for a given array. """ - filtproj = make_filter_and_project( - radius, dl, size_px=size_px, beta=beta, eta=eta, padding=padding - ) - eta_dilate = 0.0 + delta_eta - eta_eroded = 1.0 - delta_eta - def _dilate(array: NDArray, beta: float) -> NDArray: - return filtproj(array, beta=beta, eta=eta_dilate) + filtproj: FilterAndProject + delta_eta: float = 0.01 + + def __call__(self, array: NDArray) -> float: + """Compute the erosion/dilation penalty for a given array. + + Parameters + ---------- + array : NDArray + The input array to compute the penalty for. - def _erode(array: NDArray, beta: float) -> NDArray: - return filtproj(array, beta=beta, eta=eta_eroded) + Returns + ------- + float + The computed erosion/dilation penalty. + """ + eta_dilate = 0.0 + self.delta_eta + eta_eroded = 1.0 - self.delta_eta - def _open(array: NDArray, beta: float) -> NDArray: - return _dilate(_erode(array, beta=beta), beta=beta) + def _dilate(arr: NDArray): + return self.filtproj(arr, eta=eta_dilate) - def _close(array: NDArray, beta: float) -> NDArray: - return _erode(_dilate(array, beta=beta), beta=beta) + def _erode(arr: NDArray): + return self.filtproj(arr, eta=eta_eroded) - def _erosion_dilation_penalty(array: NDArray, beta: float = beta) -> float: - diff = _close(array, beta) - _open(array, beta) + def _open(arr: NDArray): + return _dilate(_erode(arr)) + + def _close(arr: NDArray): + return _erode(_dilate(arr)) + + diff = _close(array) - _open(array) if not np.any(diff): return 0.0 return np.linalg.norm(diff) / np.sqrt(diff.size) - return _erosion_dilation_penalty + @classmethod + def from_params( + cls, + radius: Union[float, Tuple[float, ...]] = None, + dl: Union[float, Tuple[float, ...]] = None, + *, + size_px: Union[int, Tuple[int, ...]] = None, + beta: float = 100.0, + eta: float = 0.5, + delta_eta: float = 0.01, + padding: PaddingType = "reflect", + ) -> "ErosionDilationPenalty": + """Create an ErosionDilationPenalty instance from parameters. + + Parameters + ---------- + radius : Union[float, Tuple[float, ...]] = None + The radius of the kernel. Can be a scalar or a tuple. + dl : Union[float, Tuple[float, ...]] = None + The grid spacing. Can be a scalar or a tuple. + size_px : Union[int, Tuple[int, ...]] = None + The size of the kernel in pixels for each dimension. Can be a scalar or a tuple. + beta : float = 100.0 + Strength of the tanh projection. + eta : float = 0.5 + Midpoint of the tanh projection. + delta_eta : float = 0.01 + The binarization threshold for erosion and dilation operations. + padding : PaddingType = "reflect" + The padding type to use for the filter. + + Returns + ------- + ErosionDilationPenalty + An instance of ErosionDilationPenalty. + """ + filtproj = FilterAndProject.from_params( + radius, dl, size_px=size_px, beta=beta, eta=eta, padding=padding + ) + return cls(filtproj=filtproj, delta_eta=delta_eta) + + +def make_erosion_dilation_penalty( + radius: Union[float, Tuple[float, ...]] = None, + dl: Union[float, Tuple[float, ...]] = None, + *, + size_px: Union[int, Tuple[int, ...]] = None, + beta: float = 100.0, + eta: float = 0.5, + delta_eta: float = 0.01, + padding: PaddingType = "reflect", +) -> Callable: + """Computes a penalty for erosion/dilation of a parameter map not being unity. + + See Also + -------- + :func:`~penalties.ErosionDilationPenalty.from_params` for full method documentation. + """ + return ErosionDilationPenalty.from_params( + radius, dl, size_px=size_px, beta=beta, eta=eta, delta_eta=delta_eta, padding=padding + ) def curvature(dp: NDArray, ddp: NDArray) -> NDArray: @@ -141,7 +185,7 @@ def bezier_curvature(x: NDArray, y: NDArray, t: Union[NDArray, float] = 0.5) -> The x-coordinates of the control points. y : NDArray The y-coordinates of the control points. - t : Union[np.ndarray, float] = 0.5 + t : Union[NDArray, float] = 0.5 The parameter at which to evaluate the curvature. Returns diff --git a/tidy3d/plugins/autograd/utilities.py b/tidy3d/plugins/autograd/utilities.py index 5bda1cc96..c4bae3354 100644 --- a/tidy3d/plugins/autograd/utilities.py +++ b/tidy3d/plugins/autograd/utilities.py @@ -62,7 +62,7 @@ def make_kernel(kernel_type: KernelType, size: Iterable[int], normalize: bool = NDArray An n-dimensional array representing the specified type of kernel. """ - if not all(isinstance(dim, int) and dim > 0 for dim in size): + if not all(np.issubdtype(type(dim), int) and dim > 0 for dim in size): raise ValueError("'size' must be an iterable of positive integers.") if kernel_type == "circular":