zea 0.0.7__py3-none-any.whl → 0.0.9__py3-none-any.whl

zea/{tensor_ops.py → func/tensor.py}

@@ -1,6 +1,6 @@
 """Basic tensor operations implemented with the multi-backend ``keras.ops``."""
 
-from typing import Tuple, Union
+from typing import List, Tuple, Union
 
 import keras
 import numpy as np
@@ -9,7 +9,7 @@ from scipy.ndimage import _ni_support
 from scipy.ndimage._filters import _gaussian_kernel1d
 
 from zea import log
-from zea.utils import map_negative_indices
+from zea.utils import canonicalize_axis
 
 
 def split_seed(seed, n):
@@ -329,7 +329,7 @@ def _map(fun, in_axes=0, out_axes=0, map_fn=None, _use_torch_vmap=False):
     For jax, this uses the native vmap implementation.
     For other backends, this a wrapper that uses `ops.vectorized_map` under the hood.
 
-    Probably you want to use `zea.tensor_ops.vmap` instead, which uses this function
+    Probably you want to use `zea.func.vmap` instead, which uses this function
     with additional batching/chunking support.
 
     Args:
@@ -431,19 +431,20 @@ def _map(fun, in_axes=0, out_axes=0, map_fn=None, _use_torch_vmap=False):
 
 
 def vmap(
-    fun,
-    in_axes=0,
-    out_axes=0,
-    batch_size=None,
-    chunks=None,
-    fn_supports_batch=False,
-    disable_jit=False,
-    _use_torch_vmap=False,
+    fun: callable,
+    in_axes: List[Union[int, None]] | int = 0,
+    out_axes: List[Union[int, None]] | int = 0,
+    chunks: int | None = None,
+    batch_size: int | None = None,
+    fn_supports_batch: bool = False,
+    disable_jit: bool = False,
+    _use_torch_vmap: bool = False,
 ):
     """`vmap` with batching or chunking support to avoid memory issues.
 
     Basically a wrapper around `vmap` that splits the input into batches or chunks
-    to avoid memory issues with large inputs.
+    to avoid memory issues with large inputs. Choose the `batch_size` or `chunks` wisely, because
+    it pads the input to make it divisible, and then crop the output back to the original size.
 
     Args:
         fun: Function to be mapped.
@@ -648,7 +649,7 @@ def pad_array_to_divisible(arr, N, axis=0, mode="constant", pad_value=None):
     return padded_array
 
 
-def interpolate_data(subsampled_data, mask, order=1, axis=-1):
+def interpolate_data(subsampled_data, mask, order=1, axis=-1, fill_mode="nearest", fill_value=0):
     """Interpolate subsampled data along a specified axis using `map_coordinates`.
 
     Args:
@@ -658,6 +659,12 @@ def interpolate_data(subsampled_data, mask, order=1, axis=-1):
             `True` where data is known.
         order (int, optional): The order of the spline interpolation. Default is `1`.
         axis (int, optional): The axis along which the data is subsampled. Default is `-1`.
+        fill_mode (str, optional): Points outside the boundaries of the input are filled
+            according to the given mode. Default is 'nearest'. For more info see
+            `keras.ops.image.map_coordinates`.
+        fill_value (float, optional): Value to use for points outside the boundaries
+            of the input if `fill_mode` is 'constant'. Default is `0`. For more info see
+            `keras.ops.image.map_coordinates`.
 
     Returns:
         ndarray: The data interpolated back to the original grid.
@@ -722,6 +729,8 @@ def interpolate_data(subsampled_data, mask, order=1, axis=-1):
         subsampled_data,
         interp_coords,
         order=order,
+        fill_mode=fill_mode,
+        fill_value=fill_value,
     )
 
     interpolated_data = ops.reshape(interpolated_data, -1)
@@ -836,7 +845,7 @@ def stack_volume_data_along_axis(data, batch_axis: int, stack_axis: int, number:
         .. doctest::
 
             >>> import keras
-            >>> from zea.tensor_ops import stack_volume_data_along_axis
+            >>> from zea.func import stack_volume_data_along_axis
 
             >>> data = keras.random.uniform((10, 20, 30))
             >>> # stacking along 1st axis with 2 frames per block
@@ -880,7 +889,7 @@ def split_volume_data_from_axis(data, batch_axis: int, stack_axis: int, number:
         .. doctest::
 
             >>> import keras
-            >>> from zea.tensor_ops import split_volume_data_from_axis
+            >>> from zea.func import split_volume_data_from_axis
 
             >>> data = keras.random.uniform((20, 10, 30))
             >>> split_data = split_volume_data_from_axis(data, 0, 1, 2, 2)
@@ -1003,7 +1012,7 @@ def check_patches_fit(
     Example:
         .. doctest::
 
-            >>> from zea.tensor_ops import check_patches_fit
+            >>> from zea.func import check_patches_fit
             >>> image_shape = (10, 10)
             >>> patch_shape = (4, 4)
             >>> overlap = (2, 2)
@@ -1071,7 +1080,7 @@ def images_to_patches(
         .. doctest::
 
             >>> import keras
-            >>> from zea.tensor_ops import images_to_patches
+            >>> from zea.func import images_to_patches
 
             >>> images = keras.random.uniform((2, 8, 8, 3))
             >>> patches = images_to_patches(images, patch_shape=(4, 4), overlap=(2, 2))
@@ -1157,7 +1166,7 @@ def patches_to_images(
         .. doctest::
 
             >>> import keras
-            >>> from zea.tensor_ops import patches_to_images
+            >>> from zea.func import patches_to_images
 
             >>> patches = keras.random.uniform((2, 3, 3, 4, 4, 3))
             >>> images = patches_to_images(patches, image_shape=(8, 8, 3), overlap=(2, 2))
@@ -1245,7 +1254,7 @@ def reshape_axis(data, newshape: tuple, axis: int):
         .. doctest::
 
             >>> import keras
-            >>> from zea.tensor_ops import reshape_axis
+            >>> from zea.func import reshape_axis
 
             >>> data = keras.random.uniform((3, 4, 5))
             >>> newshape = (2, 2)
@@ -1253,7 +1262,7 @@ def reshape_axis(data, newshape: tuple, axis: int):
             >>> reshaped_data.shape
             (3, 2, 2, 5)
     """
-    axis = map_negative_indices([axis], data.ndim)[0]
+    axis = canonicalize_axis(axis, data.ndim)
     shape = list(ops.shape(data))  # list
     shape = shape[:axis] + list(newshape) + shape[axis + 1 :]
     return ops.reshape(data, shape)
@@ -1512,7 +1521,8 @@ def sinc(x, eps=keras.config.epsilon()):
 def apply_along_axis(func1d, axis, arr, *args, **kwargs):
     """Apply a function to 1D array slices along an axis.
 
-    Keras implementation of numpy.apply_along_axis using keras.ops.vectorized_map.
+    Keras implementation of ``numpy.apply_along_axis``. Copies the ``jax`` implementation, which
+    uses ``vmap`` to vectorize the function application along the specified axis.
 
     Args:
         func1d: A callable function with signature ``func1d(arr, /, *args, **kwargs)``
@@ -1529,56 +1539,13 @@ def apply_along_axis(func1d, axis, arr, *args, **kwargs):
     # Convert to keras tensor
     arr = ops.convert_to_tensor(arr)
 
-    # Get array dimensions
-    num_dims = len(arr.shape)
-
-    # Canonicalize axis (handle negative indices)
-    if axis < 0:
-        axis = num_dims + axis
-
-    if axis < 0 or axis >= num_dims:
-        raise ValueError(f"axis {axis} is out of bounds for array of dimension {num_dims}")
-
-    # Create a wrapper function that applies func1d with the additional arguments
-    def func(slice_arr):
-        return func1d(slice_arr, *args, **kwargs)
-
-    # Recursively build up vectorized maps following the JAX pattern
-    # For dimensions after the target axis (right side)
+    num_dims = ops.ndim(arr)
+    axis = canonicalize_axis(axis, num_dims)
+    func = lambda arr: func1d(arr, *args, **kwargs)
     for i in range(1, num_dims - axis):
-        prev_func = func
-
-        def make_func(f, dim_offset):
-            def vectorized_func(x):
-                # Move the dimension we want to map over to the front
-                perm = list(range(len(x.shape)))
-                perm[0], perm[dim_offset] = perm[dim_offset], perm[0]
-                x_moved = ops.transpose(x, perm)
-                result = vectorized_map(f, x_moved)
-                # Move the result dimension back if needed
-                if len(result.shape) > 0:
-                    result_perm = list(range(len(result.shape)))
-                    if len(result_perm) > dim_offset:
-                        result_perm[0], result_perm[dim_offset] = (
-                            result_perm[dim_offset],
-                            result_perm[0],
-                        )
-                        result = ops.transpose(result, result_perm)
-                return result
-
-            return vectorized_func
-
-        func = make_func(prev_func, i)
-
-    # For dimensions before the target axis (left side)
+        func = vmap(func, in_axes=i, out_axes=-1, _use_torch_vmap=True)
     for i in range(axis):
-        prev_func = func
-
-        def make_func(f):
-            return lambda x: vectorized_map(f, x)
-
-        func = make_func(prev_func)
-
+        func = vmap(func, in_axes=0, out_axes=0, _use_torch_vmap=True)
     return func(arr)
 
 
@@ -1595,6 +1562,8 @@ def correlate(x, y, mode="full"):
         y: np.ndarray (complex or real)
         mode: "full", "valid", or "same"
     """
+    if keras.backend.backend() == "jax":
+        return ops.correlate(x, y, mode=mode)
     x = ops.convert_to_tensor(x)
     y = ops.convert_to_tensor(y)
 
@@ -1656,6 +1625,57 @@ def correlate(x, y, mode="full"):
         return ops.real(complex_tensor)
 
 
+def find_contour(binary_mask):
+    """Extract contour/boundary points from a binary mask using edge detection.
+
+    This function finds the boundary pixels of objects in a binary mask by detecting
+    pixels that have at least one neighbor with a different value (using 4-connectivity).
+
+    Args:
+        binary_mask: Binary mask tensor of shape (H, W) with values 0 or 1.
+
+    Returns:
+        Boundary points as tensor of shape (N, 2) in (row, col) format.
+        Returns empty tensor of shape (0, 2) if no boundaries are found.
+
+    Example:
+        .. doctest::
+
+            >>> from zea.func import find_contour
+            >>> import keras
+            >>> mask = keras.ops.zeros((10, 10))
+            >>> mask = keras.ops.scatter_update(
+            ...     mask, [[3, 3], [3, 4], [4, 3], [4, 4]], [1, 1, 1, 1]
+            ... )
+            >>> contour = find_contour(mask)
+            >>> contour.shape
+            (4, 2)
+    """
+    # Pad the mask to handle edges
+    padded = ops.pad(binary_mask, [[1, 1], [1, 1]], mode="constant", constant_values=0.0)
+
+    # Check 4-connectivity (up, down, left, right)
+    is_edge = (
+        (binary_mask != padded[:-2, 1:-1])  # top neighbor different
+        | (binary_mask != padded[2:, 1:-1])  # bottom neighbor different
+        | (binary_mask != padded[1:-1, :-2])  # left neighbor different
+        | (binary_mask != padded[1:-1, 2:])  # right neighbor different
+    )
+
+    # Only keep edges that are part of the foreground (binary_mask == 1)
+    is_boundary = is_edge & ops.cast(binary_mask, "bool")
+
+    boundary_indices = ops.where(is_boundary)
+
+    if ops.shape(boundary_indices[0])[0] > 0:
+        boundary_points = ops.stack(boundary_indices, axis=1)
+        boundary_points = ops.cast(boundary_points, "float32")
+    else:
+        boundary_points = ops.zeros((0, 2), dtype="float32")
+
+    return boundary_points
+
+
 def translate(array, range_from=None, range_to=(0, 255)):
     """Map values in array from one range to other.
 
@@ -1680,3 +1700,27 @@ def translate(array, range_from=None, range_to=(0, 255)):
 
     # Convert the 0-1 range into a value in the right range.
     return right_min + (value_scaled * (right_max - right_min))
+
+
+def normalize(data, output_range, input_range=None):
+    """Normalize data to a given range.
+
+    Equivalent to `translate` with clipping.
+
+    Args:
+        data (ops.Tensor): Input data to normalize.
+        output_range (tuple): Range to which data should be mapped, e.g., (0, 1).
+        input_range (tuple, optional): Range of input data.
+            If None, the range will be computed from the data.
+            Defaults to None.
+    """
+    if input_range is None:
+        input_range = (None, None)
+    minval, maxval = input_range
+    if minval is None:
+        minval = ops.min(data)
+    if maxval is None:
+        maxval = ops.max(data)
+    data = ops.clip(data, minval, maxval)
+    normalized_data = translate(data, (minval, maxval), output_range)
+    return normalized_data