Source code for tidy3d.plugins.autograd.invdes.filters

from functools import partial
from typing import Tuple, Union

import numpy as np

from ..functions import convolve
from ..types import KernelType, PaddingType
from ..utilities import make_kernel




def make_filter(
    filter_type: KernelType,
    size: Union[int, Tuple[int, ...]],
    *,
    normalize: bool = True,
    padding: PaddingType = "reflect",
):
    """Create a filter function based on the specified kernel type and size.

    Parameters
    ----------
    filter_type : KernelType
        The type of kernel to create (`circular` or `conic`).
    size : Union[int, Tuple[int, ...]]
        The size of the kernel in pixels for each dimension. Can be a scalar or a tuple.
    normalize : bool, optional
        Whether to normalize the kernel so that it sums to 1. Default is True.
    padding : PadMode, optional
        The padding mode to use. Default is "reflect".

    Returns
    -------
    function
        A function that applies the created filter to an input array.
    """
    _kernel = {}

    def _filter(array):
        original_shape = array.shape
        squeezed_array = np.squeeze(array)

        if squeezed_array.ndim not in _kernel:
            if np.isscalar(size):
                kernel_size = (size,) * squeezed_array.ndim
            else:
                kernel_size = size
            _kernel[squeezed_array.ndim] = make_kernel(
                kernel_type=filter_type, size=kernel_size, normalize=normalize
            )

        convolved_array = convolve(squeezed_array, _kernel[squeezed_array.ndim], padding=padding)
        return np.reshape(convolved_array, original_shape)

    return _filter



make_conic_filter = partial(make_filter, filter_type="conic")
make_conic_filter.__doc__ = """make_filter() with a default filter_type value of `conic`.

See Also
--------
make_filter : Function to create a filter based on the specified kernel type and size.
"""

make_circular_filter = partial(make_filter, filter_type="circular")
make_circular_filter.__doc__ = """make_filter() with a default filter_type value of `circular`.

See Also
--------
make_filter : Function to create a filter based on the specified kernel type and size.
"""