Add more checks and add tests for missing functions in functions.py a…

…s well as tests for utilities

tests/test_plugins/autograd/test_functions.py

@@ -15,6 +15,8 @@
     morphological_gradient,
     morphological_gradient_internal,
     morphological_gradient_external,
+    rescale,
+    threshold,
 )
 from tidy3d.plugins.autograd.types import PaddingType
 
@@ -241,3 +243,61 @@ def test_morphology_val_structure_grad(
             """Test gradients of morphological operations for various kernel structures."""
             x, k = self._ary_and_kernel(rng, ary_size, kernel_size, full, square, flat)
             check_grads(op, modes=["rev"], order=1)(x, size=kernel_size, mode=mode)
+
+
+@pytest.mark.parametrize(
+    "array, out_min, out_max, in_min, in_max, expected",
+    [
+        (np.array([0, 0.5, 1]), 0, 10, 0, 1, np.array([0, 5, 10])),
+        (np.array([0, 0.5, 1]), -1, 1, 0, 1, np.array([-1, 0, 1])),
+        (np.array([0, 1, 2]), 0, 1, 0, 2, np.array([0, 0.5, 1])),
+        (np.array([-1, 0, 1]), -10, 10, -1, 1, np.array([-10, 0, 10])),
+        (np.array([-2, -1, 0]), -1, 1, -2, 0, np.array([-1, 0, 1])),
+    ],
+)
+def test_rescale(array, out_min, out_max, in_min, in_max, expected):
+    """Test rescale function for various input and output ranges."""
+    result = rescale(array, out_min, out_max, in_min, in_max)
+    npt.assert_allclose(result, expected)
+
+
+@pytest.mark.parametrize(
+    "array, out_min, out_max, in_min, in_max, expected_message",
+    [
+        (np.array([0, 0.5, 1]), 10, 0, 0, 1, "must be less than"),
+        (np.array([0, 0.5, 1]), 0, 10, 1, 1, "must not be equal"),
+        (np.array([0, 0.5, 1]), 0, 10, 1, 0, "must be less than"),
+    ],
+)
+def test_rescale_exceptions(array, out_min, out_max, in_min, in_max, expected_message):
+    """Test rescale function for expected exceptions."""
+    with pytest.raises(ValueError, match=expected_message):
+        rescale(array, out_min, out_max, in_min, in_max)
+
+
+@pytest.mark.parametrize(
+    "ary, vmin, vmax, level, expected",
+    [
+        (np.array([0, 0.5, 1]), 0, 1, 0.5, np.array([0, 1, 1])),
+        (np.array([0, 0.5, 1]), 0, 1, None, np.array([0, 1, 1])),
+        (np.array([0, 0.5, 1]), -1, 1, 0.5, np.array([-1, 1, 1])),
+    ],
+)
+def test_threshold(ary, vmin, vmax, level, expected):
+    """Test threshold function values for threshold levels and value ranges."""
+    result = threshold(ary, vmin, vmax, level)
+    npt.assert_allclose(result, expected)
+
+
+@pytest.mark.parametrize(
+    "array, vmin, vmax, level, expected_message",
+    [
+        (np.array([0, 0.5, 1]), 1, 0, None, "threshold range"),
+        (np.array([0, 0.5, 1]), 0, 1, -0.5, "threshold level"),
+        (np.array([0, 0.5, 1]), 0, 1, 1.5, "threshold level"),
+    ],
+)
+def test_threshold_exceptions(array, vmin, vmax, level, expected_message):
+    """Test threshold function for expected exceptions."""
+    with pytest.raises(ValueError, match=expected_message):
+        threshold(array, vmin, vmax, level)

tests/test_plugins/autograd/test_utilities.py

@@ -0,0 +1,85 @@
+import pytest
+import numpy as np
+import numpy.testing as npt
+from tidy3d.plugins.autograd.utilities import make_kernel, chain
+
+
+@pytest.mark.parametrize("size", [(3, 3), (4, 4), (5, 5)])
+@pytest.mark.parametrize("normalize", [True, False])
+class TestMakeKernel:
+    def test_make_kernel_circular(self, size, normalize):
+        """Test make_kernel function for circular kernel."""
+        kernel = make_kernel("circular", size, normalize=normalize)
+        assert kernel.shape == size
+        if normalize:
+            assert np.isclose(np.sum(kernel), 1.0)
+
+        # Check that the corners of the circular kernel are zero
+        assert all(kernel[i, j] == 0 for i in [0, -1] for j in [0, -1])
+
+    def test_make_kernel_conic(self, size, normalize):
+        """Test make_kernel function for conic kernel."""
+        kernel = make_kernel("conic", size, normalize=normalize)
+        assert kernel.shape == size
+        if normalize:
+            assert np.isclose(np.sum(kernel), 1.0)
+
+        # Check that the corners of the conic kernel are zero
+        assert all(kernel[i, j] == 0 for i in [0, -1] for j in [0, -1])
+
+
+class TestMakeKernelExceptions:
+    def test_make_kernel_invalid_type(self):
+        """Test make_kernel function for invalid kernel type."""
+        size = (5, 5)
+        with pytest.raises(ValueError, match="Unsupported kernel type"):
+            make_kernel("invalid_type", size)
+
+    def test_make_kernel_invalid_size(self):
+        """Test make_kernel function for invalid size."""
+        size = (5, -5)
+        with pytest.raises(ValueError, match="must be an iterable of positive integers"):
+            make_kernel("circular", size)
+
+
+class TestChain:
+    def test_chain_functions(self):
+        """Test chain function with multiple functions."""
+
+        def add_one(x):
+            return x + 1
+
+        def square(x):
+            return x**2
+
+        chained_func = chain(add_one, square)
+        array = np.array([1, 2, 3])
+        result = chained_func(array)
+        expected = np.array([4, 9, 16])
+        npt.assert_allclose(result, expected)
+
+    def test_chain_single_iterable(self):
+        """Test chain function with a single iterable of functions."""
+
+        def add_one(x):
+            return x + 1
+
+        def square(x):
+            return x**2
+
+        funcs = [add_one, square]
+        chained_func = chain(funcs)
+        array = np.array([1, 2, 3])
+        result = chained_func(array)
+        expected = np.array([4, 9, 16])
+        npt.assert_allclose(result, expected)
+
+    def test_chain_invalid_function(self):
+        """Test chain function with an invalid function in the list."""
+
+        def add_one(x):
+            return x + 1
+
+        funcs = [add_one, "not_a_function"]
+        with pytest.raises(TypeError, match="All elements in funcs must be callable"):
+            chain(funcs)

tidy3d/plugins/autograd/functions.py

@@ -608,6 +608,17 @@ def rescale(
     np.ndarray
         The rescaled array.
     """
+
+    if in_min == in_max:
+        raise ValueError(
+            f"'in_min' ({in_min}) must not be equal to 'in_max' ({in_max}) "
+            "to avoid division by zero."
+        )
+    if out_min >= out_max:
+        raise ValueError(f"'out_min' ({out_min}) must be less than 'out_max' ({out_max}).")
+    if in_min >= in_max:
+        raise ValueError(f"'in_min' ({in_min}) must be less than 'in_max' ({in_max}).")
+
     scaled = (array - in_min) / (in_max - in_min)
     return scaled * (out_max - out_min) + out_min
 
@@ -633,6 +644,17 @@ def threshold(
     np.ndarray
         The thresholded array.
     """
+    if vmin >= vmax:
+        raise ValueError(
+            f"Invalid threshold range: 'vmin' ({vmin}) must be smaller than 'vmax' ({vmax})."
+        )
+
     if level is None:
         level = (vmin + vmax) / 2
+    elif not (vmin <= level <= vmax):
+        raise ValueError(
+            f"Invalid threshold level: 'level' ({level}) must be "
+            f"between 'vmin' ({vmin}) and 'vmax' ({vmax})."
+        )
+
     return np.where(array < level, vmin, vmax)

tidy3d/plugins/autograd/utilities.py

@@ -1,17 +1,17 @@
 from functools import reduce
-from typing import Tuple, Callable, Iterable, Union
+from typing import Callable, Iterable, Union
 
 import numpy as np
 
 from .types import KernelType
 
 
-def _kernel_circular(size: Tuple[int, ...]) -> np.ndarray:
+def _kernel_circular(size: Iterable[int]) -> np.ndarray:
     """Create a circular kernel in n dimensions.
 
     Parameters
     ----------
-    size : Tuple[int, ...]
+    size : Iterable[int]
         The size of the circular kernel in pixels for each dimension.
 
     Returns
@@ -21,16 +21,16 @@ def _kernel_circular(size: Tuple[int, ...]) -> np.ndarray:
     """
     grids = np.ogrid[tuple(slice(-1, 1, 1j * s) for s in size)]
     squared_distances = sum(grid**2 for grid in grids)
-    kernel = squared_distances <= 1
+    kernel = np.array(squared_distances <= 1, dtype=np.float64)
     return kernel
 
 
-def _kernel_conic(size: Tuple[int, ...]) -> np.ndarray:
+def _kernel_conic(size: Iterable[int]) -> np.ndarray:
     """Create a conic kernel in n dimensions.
 
     Parameters
     ----------
-    size : Tuple[int, ...]
+    size : Iterable[int]
         The size of the conic kernel in pixels for each dimension.
 
     Returns
@@ -44,16 +44,14 @@ def _kernel_conic(size: Tuple[int, ...]) -> np.ndarray:
     return kernel
 
 
-def make_kernel(
-    kernel_type: KernelType, size: Tuple[int, ...], normalize: bool = True
-) -> np.ndarray:
+def make_kernel(kernel_type: KernelType, size: Iterable[int], normalize: bool = True) -> np.ndarray:
     """Create a kernel based on the specified type in n dimensions.
 
     Parameters
     ----------
     kernel_type : KernelType
         The type of kernel to create ('circular' or 'conic').
-    size : Tuple[int, ...]
+    size : Iterable[int]
         The size of the kernel in pixels for each dimension.
     normalize : bool, optional
         Whether to normalize the kernel so that it sums to 1. Default is True.
@@ -63,6 +61,9 @@ def make_kernel(
     np.ndarray
         An n-dimensional array representing the specified type of kernel.
     """
+    if not all(isinstance(dim, int) and dim > 0 for dim in size):
+        raise ValueError("'size' must be an iterable of positive integers.")
+
     if kernel_type == "circular":
         kernel = _kernel_circular(size)
     elif kernel_type == "conic":
@@ -111,6 +112,9 @@ def chain(*funcs: Union[Callable, Iterable[Callable]]):
     if len(funcs) == 1 and isinstance(funcs[0], Iterable):
         funcs = funcs[0]
 
+    if not all(callable(f) for f in funcs):
+        raise TypeError("All elements in funcs must be callable.")
+
     def chained(array: np.ndarray):
         return reduce(lambda x, y: y(x), funcs, array)