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

zea/scan.py

@@ -153,7 +153,6 @@ class Scan(Parameters):
         apply_lens_correction (bool, optional): Whether to apply lens correction to
             delays. Defaults to False.
         lens_thickness (float, optional): Thickness of the lens in meters.
-            Defaults to None.
         f_number (float, optional): F-number of the transducer. Defaults to 1.0.
         theta_range (tuple, optional): Range of theta angles for 3D imaging.
         phi_range (tuple, optional): Range of phi angles for 3D imaging.
@@ -215,8 +214,8 @@ class Scan(Parameters):
         "initial_times": {"type": np.ndarray},
         "time_to_next_transmit": {"type": np.ndarray},
         "tgc_gain_curve": {"type": np.ndarray},
-        "waveforms_one_way": {"type": np.ndarray},
-        "waveforms_two_way": {"type": np.ndarray},
+        "waveforms_one_way": {"type": np.ndarray, "default": None},
+        "waveforms_two_way": {"type": np.ndarray, "default": None},
         "tx_waveform_indices": {"type": np.ndarray},
         "t_peak": {"type": np.ndarray},
         # scan conversion parameters
@@ -508,7 +507,7 @@ class Scan(Parameters):
         value = self._params.get("azimuth_angles")
         if value is None:
             log.warning("No azimuth angles provided, using zeros")
-            value = np.zeros(self.n_tx_selected)
+            return np.zeros(self.n_tx_selected)
 
         return value[self.selected_transmits]
 
@@ -529,7 +528,7 @@ class Scan(Parameters):
         value = self._params.get("tx_apodizations")
         if value is None:
             log.warning("No transmit apodizations provided, using ones")
-            value = np.ones((self.n_tx_selected, self.n_el))
+            return np.ones((self.n_tx_selected, self.n_el))
 
         return value[self.selected_transmits]
 
@@ -539,7 +538,7 @@ class Scan(Parameters):
         value = self._params.get("focus_distances")
         if value is None:
             log.warning("No focus distances provided, using zeros")
-            value = np.zeros(self.n_tx_selected)
+            return np.zeros(self.n_tx_selected)
 
         return value[self.selected_transmits]
 
@@ -549,7 +548,7 @@ class Scan(Parameters):
         value = self._params.get("initial_times")
         if value is None:
             log.warning("No initial times provided, using zeros")
-            value = np.zeros(self.n_tx_selected)
+            return np.zeros(self.n_tx_selected)
 
         return value[self.selected_transmits]
 
@@ -617,7 +616,7 @@ class Scan(Parameters):
 
     @cache_with_dependencies("pfield")
     def flat_pfield(self):
-        """Flattened pfield for weighting."""
+        """Flattened pfield for weighting of shape (n_pix, n_tx)."""
         return self.pfield.reshape(self.n_tx, -1).swapaxes(0, 1)
 
     @cache_with_dependencies("zlims")
@@ -674,6 +673,17 @@ class Scan(Parameters):
         otherwise 2D."""
         return self.coordinates_3d if getattr(self, "phi_range", None) else self.coordinates_2d
 
+    @property
+    def pulse_repetition_frequency(self):
+        """The pulse repetition frequency (PRF) [Hz]. Assumes a constant PRF."""
+        if self.time_to_next_transmit is None:
+            log.warning("Time to next transmit is not set, cannot compute PRF")
+            return None
+
+        pulse_repetition_interval = np.mean(self.time_to_next_transmit)
+
+        return 1 / pulse_repetition_interval
+
     @cache_with_dependencies("time_to_next_transmit")
     def frames_per_second(self):
         """The number of frames per second [Hz]. Assumes a constant frame rate.

zea/tensor_ops.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):
@@ -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)
@@ -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.tensor_ops 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.