Add class-based penalties and parametrizations

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",

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.

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",
 ]

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)

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
+    )